Loss (stories.loss)

stories.loss.linear_loss(x: Array, a: Array, y: Array, b: Array, epsilon: float, debias: bool) → float

Compute the Sinkhorn loss (no quadratic component).

Parameters:

  • x – A pointcloud.

  • a – Histogram on x.

  • y – Another pointcloud.

  • b – Histogram on y.

  • epsilon – Entropic regularization parameter.

  • debias – Whether to debias the loss.

Returns:

The Sinkhorn loss.

stories.loss.loss_fn(params: Array | ndarray | bool_ | number | Iterable[ArrayTree] | Mapping[Any, ArrayTree], batch: Dict[str, Array], teacher_forcing: bool, quadratic: bool, proximal_step: ProximalStep, potential: Module, n_steps: int, epsilon: float, debias: bool, quadratic_weight: float, tau_diff: Array) → Array

The loss function

Parameters:

  • params – The parameters of the model.

  • batch – A batch of data.

  • teacher_forcing – Whether to use teacher forcing.

  • quadratic – Whether to use the quadratic (FGW) loss instead of the linear (W) one.

  • proximal_step – The proximal step, e.g. LinearExplicitStep.

  • potential – The potential function parametrized by a neural network.

  • n_steps – The number of steps to take.

  • epsilon – Entropic regularization parameter.

  • debias – Whether to debias the loss (see linear_loss or quadratic_loss).

  • quadratic_weight – The relative weight of the quadratic term, between 0 and 1.

  • tau_diff – The difference in time between each timepoint, e.g. [1., 1., 2.].

Returns:

The loss.

stories.loss.quadratic_loss(x: Array, a: Array, y: Array, b: Array, space_x: Array, space_y: Array, epsilon: float, quadratic_weight: float, debias: bool) → float

FGW loss The linear part of the loss operates on the gene coordinates. The quadratic part of the loss operates on the space coordinates.

The issue with entropic Fused Gromov-Wasserstein is that it is biased, ie FGW(x, y) != 0 when x=y. We can compute instead the following:

FGW(x, y) - 0.5 * FGW(x, x) - 0.5 * FGW(y, y)

As done in http://proceedings.mlr.press/v97/bunne19a/bunne19a.pdf

Parameters:

  • x – A pointcloud on the space of genes.

  • a – Histogram on x.

  • y – Another pointcloud on the space of genes.

  • b – Histogram on y.

  • space_x – A pointcloud on the space of spatial coordinates.

  • space_y – Another pointcloud on the space of spatial coordinates.

  • epsilon – Entropic regularization parameter.

  • quadratic_weight – The relative weight of the quadratic term, between 0 and 1.

  • debias – Whether to debias the loss.

Returns:

The FGW loss.