Custom cost function?

[prabhant]

Hi,

is there a way to define custom cost function in the PointCloud geometry?

https://ott-jax.readthedocs.io/en/latest/_modules/ott/geometry/costs.html#CostFn

Here I can't find a clean way to define a cost function that is not there.

[marcocuturi]

Hi @prabhant!

To define a custom cost function, you should create a class that inherits from CostFn, and, at the very least, define a pairwise function that computes the difference between two vectors (at this moment we only handle distance functions between vectors).

The simple example is the Euclidean distance:

class Euclidean(CostFn):
  """Euclidean distance."""
  def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float:
    """Compute Euclidean norm."""
    return jnp.linalg.norm(x - y)

so for instance you could define your own distance,

from ott.geometry import costs
class MyCost(costs.CostFn):
  """My weird cost."""
  def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float:
    """Compute a weird cost between two vectors."""
    return jnp.sum(jnp.abs(jnp.sin(x) - jnp.cos(x)))

and then instantiate a point cloud with that distance should work (in principle!)
geom = pointcloud.PointCloud(x, y, cost_fn = MyDist())