Separable Cost Matrix

[thungerland]

Hi everyone,

It is unclear to me how to set up the following cost matrix: Each point consists of multimodal information, i.e., part of the point vector corresponds to e.g. an actual position in Euclidean space whereas the remaining vector entries store e.g. visual information from a picture (and its corresponding embedding). Even though it is fine for me to use a Euclidean cost for the entire vector, the cost between each pair of points should be separated as a weighted cost such that I can tune the effects of the positional vs. visual data have. One way I have thought of doing this is that I could manually multiply all positions corresponding to one of the two features of my point vector by a certain factor. Yet, I feel like there must be a neater way to do this using the available functionalities. Moreover, it would be nice to really separate the two cost terms as this gives the option to combine two different costs for further analysis. Perhaps somebody else has already tried setting up a multi-feature cost function.

Thanks for your help,
Tiago Hungerland

[marcocuturi]

hi Tiago, apologies, your email fell on a holiday period here, and was a bit lost. I really apologize for not helping out earlier...

You question suggests that we should think about adding operators to add cost functions. This is a cool idea that's worth exploring.

In your case, I think the way to go could be either

• fall back on creating a geometry.Geometry(cost_matrix = C1+C2) by instantiating explicitly those matrices (a bit ugly)
• to define a custom cost function, as done here for instance,

@jax.tree_util.register_pytree_node_class
class Custom(costs.TICost):
    """Custom, translation invariant cost, sqrt of Euclidean norm."""

    def h(self, z):
        return jnp.sqrt(jnp.abs(jnp.linalg.norm(z)))

In your case, if you want to add a squared Euclidean to another Euclidean distance, I could imagine something of the form

@jax.tree_util.register_pytree_node_class
class Custom(costs.CostFn):
    """Mixture of costs"""
  def __init__(self, d1: int, lambda):
    super().__init__()
    self.d1 = d1
    self.lambda = lambda
    self._c1 = costs.SqEuclidean()
    self._c2 = costs.Euclidean()
    
  def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float:
    return self._c1(x[:d1], y[:d1]) + self.lambda * self._c2(x[d1:], y[d1:]) 

  def tree_flatten(self):  # noqa: D102
    return (lambda,), (d1,)

  @classmethod
  def tree_unflatten(cls, aux_data, children):  # noqa: D102
    return cls(*children, *aux_data)

essentially we've done these kind of things when defining elastic costs

[thungerland]

Hi Marco, thank you for this explanation. This helped me a lot for my problem.