# Copyright OTT-JAX
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
from typing import Any, Callable, Literal, Optional, Tuple, Union
import jax
import jax.numpy as jnp
from ott import utils
from ott.geometry import costs, geometry, low_rank
from ott.math import utils as mu
__all__ = ["PointCloud"]
[docs]@jax.tree_util.register_pytree_node_class
class PointCloud(geometry.Geometry):
"""Defines geometry for 2 point clouds (possibly 1 vs itself).
Creates a geometry, specifying a cost function passed as CostFn type object.
When the number of points is large, setting the ``batch_size`` flag implies
that cost and kernel matrices used to update potentials or scalings
will be recomputed on the fly, rather than stored in memory. More precisely,
when setting ``batch_size``, the cost function will be partially cached by
storing norm values for each point in both point clouds, but the pairwise cost
function evaluations won't be.
Args:
x : n x d array of n d-dimensional vectors
y : m x d array of m d-dimensional vectors. If `None`, use ``x``.
cost_fn: a CostFn function between two points in dimension d.
batch_size: When ``None``, the cost matrix corresponding to that point cloud
is computed, stored and later re-used at each application of
:meth:`apply_lse_kernel`. When ``batch_size`` is a positive integer,
computations are done in an online fashion, namely the cost matrix is
recomputed at each call of the :meth:`apply_lse_kernel` step,
``batch_size`` lines at a time, used on a vector and discarded.
The online computation is particularly useful for big point clouds
whose cost matrix does not fit in memory.
scale_cost: option to rescale the cost matrix. Implemented scalings are
'median', 'mean', 'max_cost', 'max_norm' and 'max_bound'.
Alternatively, a float factor can be given to rescale the cost such
that ``cost_matrix /= scale_cost``. If `True`, use 'mean'.
kwargs: keyword arguments for :class:`~ott.geometry.geometry.Geometry`.
"""
def __init__(
self,
x: jnp.ndarray,
y: Optional[jnp.ndarray] = None,
cost_fn: Optional[costs.CostFn] = None,
batch_size: Optional[int] = None,
scale_cost: Union[bool, int, float,
Literal["mean", "max_norm", "max_bound", "max_cost",
"median"]] = 1.0,
**kwargs: Any
):
super().__init__(**kwargs)
self.x = x
self.y = self.x if y is None else y
self.cost_fn = costs.SqEuclidean() if cost_fn is None else cost_fn
self._axis_norm = 0 if callable(self.cost_fn.norm) else None
if batch_size is not None:
assert batch_size > 0, f"`batch_size={batch_size}` must be positive."
self._batch_size = batch_size
self._scale_cost = "mean" if scale_cost is True else scale_cost
@property
def _norm_x(self) -> Union[float, jnp.ndarray]:
if self._axis_norm == 0:
return self.cost_fn.norm(self.x)
return 0.
@property
def _norm_y(self) -> Union[float, jnp.ndarray]:
if self._axis_norm == 0:
return self.cost_fn.norm(self.y)
return 0.
@property
def can_LRC(self): # noqa: D102
return self.is_squared_euclidean and self._check_LRC_dim
@property
def _check_LRC_dim(self):
(n, m), d = self.shape, self.x.shape[1]
return n * m > (n + m) * d
@property
def cost_matrix(self) -> Optional[jnp.ndarray]: # noqa: D102
if self.is_online:
return None
cost_matrix = self._compute_cost_matrix()
return cost_matrix * self.inv_scale_cost
@property
def kernel_matrix(self) -> Optional[jnp.ndarray]: # noqa: D102
if self.is_online:
return None
return jnp.exp(-self.cost_matrix / self.epsilon)
@property
def shape(self) -> Tuple[int, int]: # noqa: D102
# in the process of flattening/unflattening in vmap, `__init__`
# can be called with dummy objects
# we optionally access `shape` in order to get the batch size
if self.x is None or self.y is None:
return 0, 0
return self.x.shape[0], self.y.shape[0]
@property
def is_symmetric(self) -> bool: # noqa: D102
return self.y is None or (
jnp.all(self.x.shape == self.y.shape) and jnp.all(self.x == self.y)
)
@property
def is_squared_euclidean(self) -> bool: # noqa: D102
return isinstance(self.cost_fn, costs.SqEuclidean)
@property
def is_online(self) -> bool:
"""Whether the cost/kernel is computed on-the-fly."""
return self.batch_size is not None
# TODO(michalk8): when refactoring, consider PC as a subclass of LR?
@property
def cost_rank(self) -> int: # noqa: D102
return self.x.shape[1]
@property
def inv_scale_cost(self) -> float: # noqa: D102
if isinstance(self._scale_cost,
(int, float)) or utils.is_jax_array(self._scale_cost):
return 1.0 / self._scale_cost
self = self._masked_geom()
if self._scale_cost == "max_cost":
if self.is_online:
return 1.0 / self._compute_summary_online(self._scale_cost)
return 1.0 / jnp.max(self._compute_cost_matrix())
if self._scale_cost == "mean":
if self.is_online:
return 1.0 / self._compute_summary_online(self._scale_cost)
if self.shape[0] > 0:
geom = self._masked_geom(mask_value=jnp.nan)._compute_cost_matrix()
return 1.0 / jnp.nanmean(geom)
return 1.0
if self._scale_cost == "median":
if not self.is_online:
geom = self._masked_geom(mask_value=jnp.nan)
return 1.0 / jnp.nanmedian(geom._compute_cost_matrix())
raise NotImplementedError(
"Using the median as scaling factor for "
"the cost matrix with the online mode is not implemented."
)
if self._scale_cost == "max_norm":
if self.cost_fn.norm is not None:
return 1.0 / jnp.maximum(self._norm_x.max(), self._norm_y.max())
return 1.0
if self._scale_cost == "max_bound":
if self.is_squared_euclidean:
x_argmax = jnp.argmax(self._norm_x)
y_argmax = jnp.argmax(self._norm_y)
max_bound = (
self._norm_x[x_argmax] + self._norm_y[y_argmax] +
2 * jnp.sqrt(self._norm_x[x_argmax] * self._norm_y[y_argmax])
)
return 1.0 / max_bound
raise NotImplementedError(
"Using max_bound as scaling factor for "
"the cost matrix when the cost is not squared euclidean "
"is not implemented."
)
raise ValueError(f"Scaling {self._scale_cost} not implemented.")
def _compute_cost_matrix(self) -> jnp.ndarray:
cost_matrix = self.cost_fn.all_pairs_pairwise(self.x, self.y)
if self._axis_norm is not None:
cost_matrix += self._norm_x[:, jnp.newaxis] + self._norm_y[jnp.newaxis, :]
return cost_matrix
[docs] def apply_lse_kernel( # noqa: D102
self,
f: jnp.ndarray,
g: jnp.ndarray,
eps: float,
vec: Optional[jnp.ndarray] = None,
axis: int = 0
) -> jnp.ndarray:
def body0(carry, i: int):
f, g, eps, vec = carry
y = jax.lax.dynamic_slice(
self.y, (i * self.batch_size, 0), (self.batch_size, self.y.shape[1])
)
g_ = jax.lax.dynamic_slice(g, (i * self.batch_size,), (self.batch_size,))
if self._axis_norm is None:
norm_y = self._norm_y
else:
norm_y = jax.lax.dynamic_slice(
self._norm_y, (i * self.batch_size,), (self.batch_size,)
)
h_res, h_sgn = app(
self.x, y, self._norm_x, norm_y, f, g_, eps, vec, self.cost_fn,
self.inv_scale_cost
)
return carry, (h_res, h_sgn)
def body1(carry, i: int):
f, g, eps, vec = carry
x = jax.lax.dynamic_slice(
self.x, (i * self.batch_size, 0), (self.batch_size, self.x.shape[1])
)
f_ = jax.lax.dynamic_slice(f, (i * self.batch_size,), (self.batch_size,))
if self._axis_norm is None:
norm_x = self._norm_x
else:
norm_x = jax.lax.dynamic_slice(
self._norm_x, (i * self.batch_size,), (self.batch_size,)
)
h_res, h_sgn = app(
self.y, x, self._norm_y, norm_x, g, f_, eps, vec, self.cost_fn,
self.inv_scale_cost
)
return carry, (h_res, h_sgn)
def finalize(i: int):
if axis == 0:
norm_y = self._norm_y if self._axis_norm is None else self._norm_y[i:]
return app(
self.x, self.y[i:], self._norm_x, norm_y, f, g[i:], eps, vec,
self.cost_fn, self.inv_scale_cost
)
norm_x = self._norm_x if self._axis_norm is None else self._norm_x[i:]
return app(
self.y, self.x[i:], self._norm_y, norm_x, g, f[i:], eps, vec,
self.cost_fn, self.inv_scale_cost
)
if not self.is_online:
return super().apply_lse_kernel(f, g, eps, vec, axis)
app = jax.vmap(
_apply_lse_kernel_xy,
in_axes=[
None, 0, None, self._axis_norm, None, 0, None, None, None, None
]
)
if axis == 0:
fun = body0
v, n = g, self._y_nsplit
elif axis == 1:
fun = body1
v, n = f, self._x_nsplit
else:
raise ValueError(axis)
_, (h_res, h_sign) = jax.lax.scan(
fun, init=(f, g, eps, vec), xs=jnp.arange(n)
)
h_res, h_sign = jnp.concatenate(h_res), jnp.concatenate(h_sign)
h_res_rest, h_sign_rest = finalize(n * self.batch_size)
h_res = jnp.concatenate([h_res, h_res_rest])
h_sign = jnp.concatenate([h_sign, h_sign_rest])
return eps * h_res - jnp.where(jnp.isfinite(v), v, 0), h_sign
[docs] def apply_kernel( # noqa: D102
self,
scaling: jnp.ndarray,
eps: Optional[float] = None,
axis: int = 0
) -> jnp.ndarray:
if eps is None:
eps = self.epsilon
if not self.is_online:
return super().apply_kernel(scaling, eps, axis)
app = jax.vmap(
_apply_kernel_xy,
in_axes=[None, 0, None, self._axis_norm, None, None, None, None]
)
if axis == 0:
return app(
self.x, self.y, self._norm_x, self._norm_y, scaling, eps,
self.cost_fn, self.inv_scale_cost
)
return app(
self.y, self.x, self._norm_y, self._norm_x, scaling, eps, self.cost_fn,
self.inv_scale_cost
)
[docs] def transport_from_potentials( # noqa: D102
self, f: jnp.ndarray, g: jnp.ndarray
) -> jnp.ndarray:
if not self.is_online:
return super().transport_from_potentials(f, g)
transport = jax.vmap(
_transport_from_potentials_xy,
in_axes=[None, 0, None, self._axis_norm, None, 0, None, None, None]
)
return transport(
self.y, self.x, self._norm_y, self._norm_x, g, f, self.epsilon,
self.cost_fn, self.inv_scale_cost
)
[docs] def transport_from_scalings( # noqa: D102
self, u: jnp.ndarray, v: jnp.ndarray
) -> jnp.ndarray:
if not self.is_online:
return super().transport_from_scalings(u, v)
transport = jax.vmap(
_transport_from_scalings_xy,
in_axes=[
None,
0,
None,
self._axis_norm,
None,
0,
None,
None,
None,
]
)
return transport(
self.y, self.x, self._norm_y, self._norm_x, v, u, self.epsilon,
self.cost_fn, self.inv_scale_cost
)
[docs] def apply_cost(
self,
arr: jnp.ndarray,
axis: int = 0,
fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None,
is_linear: bool = False,
) -> jnp.ndarray:
"""Apply cost matrix to array (vector or matrix).
This function applies the geometry's cost matrix, to perform either
output = C arr (if axis=1)
output = C' arr (if axis=0)
where C is [num_a, num_b] matrix resulting from the (optional) elementwise
application of fn to each entry of the :attr:`cost_matrix`.
Args:
arr: jnp.ndarray [num_a or num_b, batch], vector that will be multiplied
by the cost matrix.
axis: standard cost matrix if axis=1, transpose if 0.
fn: function optionally applied to cost matrix element-wise, before the
apply.
is_linear: Whether ``fn`` is a linear function.
If true and :attr:`is_squared_euclidean` is ``True``, efficient
implementation is used. See :func:`ott.geometry.geometry.is_linear`
for a heuristic to help determine if a function is linear.
Returns:
A jnp.ndarray, [num_b, batch] if axis=0 or [num_a, batch] if axis=1
"""
# switch to efficient computation for the squared euclidean case.
if self.is_squared_euclidean and (fn is None or is_linear):
return self.vec_apply_cost(arr, axis, fn=fn)
return self._apply_cost(arr, axis, fn=fn)
def _apply_cost(
self, arr: jnp.ndarray, axis: int = 0, fn=None
) -> jnp.ndarray:
"""See :meth:`apply_cost`."""
if not self.is_online:
return super().apply_cost(arr, axis, fn)
app = jax.vmap(
_apply_cost_xy,
in_axes=[None, 0, None, self._axis_norm, None, None, None, None]
)
if arr.ndim == 1:
arr = arr.reshape(-1, 1)
if axis == 0:
return app(
self.x, self.y, self._norm_x, self._norm_y, arr, self.cost_fn,
self.inv_scale_cost, fn
)
return app(
self.y, self.x, self._norm_y, self._norm_x, arr, self.cost_fn,
self.inv_scale_cost, fn
)
[docs] def vec_apply_cost(
self,
arr: jnp.ndarray,
axis: int = 0,
fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None
) -> jnp.ndarray:
"""Apply the geometry's cost matrix in a vectorized way.
This function can be used when the cost matrix is squared euclidean
and ``fn`` is a linear function.
Args:
arr: jnp.ndarray [num_a or num_b, p], vector that will be multiplied
by the cost matrix.
axis: standard cost matrix if axis=1, transport if 0.
fn: function optionally applied to cost matrix element-wise, before the
application.
Returns:
A jnp.ndarray, [num_b, p] if axis=0 or [num_a, p] if axis=1
"""
assert self.is_squared_euclidean, "Cost matrix is not a squared Euclidean."
rank = arr.ndim
x, y = (self.x, self.y) if axis == 0 else (self.y, self.x)
nx, ny = jnp.asarray(self._norm_x), jnp.asarray(self._norm_y)
nx, ny = (nx, ny) if axis == 0 else (ny, nx)
applied_cost = jnp.dot(nx, arr).reshape(1, -1)
applied_cost += ny.reshape(-1, 1) * jnp.sum(arr, axis=0).reshape(1, -1)
cross_term = -2.0 * jnp.dot(y, jnp.dot(x.T, arr))
applied_cost += cross_term[:, None] if rank == 1 else cross_term
if fn is not None:
applied_cost = fn(applied_cost)
return self.inv_scale_cost * applied_cost
def _leading_slice(self, t: jnp.ndarray, i: int) -> jnp.ndarray:
start_indices = [i * self.batch_size] + (t.ndim - 1) * [0]
slice_sizes = [self.batch_size] + list(t.shape[1:])
return jax.lax.dynamic_slice(t, start_indices, slice_sizes)
def _compute_summary_online(
self, summary: Literal["mean", "max_cost"]
) -> float:
"""Compute mean or max of cost matrix online, i.e. without instantiating it.
Args:
summary: can be 'mean' or 'max_cost'.
Returns:
summary statistics
"""
scale_cost = 1.0
def body0(carry, i: int):
vec, = carry
y = self._leading_slice(self.y, i)
if self._axis_norm is None:
norm_y = self._norm_y
else:
norm_y = self._leading_slice(self._norm_y, i)
h_res = app(
self.x, y, self._norm_x, norm_y, vec, self.cost_fn, scale_cost
)
return carry, h_res
def body1(carry, i: int):
vec, = carry
x = self._leading_slice(self.x, i)
if self._axis_norm is None:
norm_x = self._norm_x
else:
norm_x = self._leading_slice(self._norm_x, i)
h_res = app(
self.y, x, self._norm_y, norm_x, vec, self.cost_fn, scale_cost
)
return carry, h_res
def finalize(i: int):
if batch_for_y:
norm_y = self._norm_y if self._axis_norm is None else self._norm_y[i:]
return app(
self.x, self.y[i:], self._norm_x, norm_y, vec, self.cost_fn,
scale_cost
)
norm_x = self._norm_x if self._axis_norm is None else self._norm_x[i:]
return app(
self.y, self.x[i:], self._norm_y, norm_x, vec, self.cost_fn,
scale_cost
)
if summary == "mean":
fn = _apply_cost_xy
elif summary == "max_cost":
fn = _apply_max_xy
else:
raise ValueError(
f"Scaling method {summary} does not exist for online mode."
)
app = jax.vmap(
fn, in_axes=[None, 0, None, self._axis_norm, None, None, None]
)
batch_for_y = self.shape[0] < self.shape[1]
if batch_for_y:
fun = body0
n = self._y_nsplit
vec, other = self._n_normed_ones, self._m_normed_ones
else:
fun = body1
n = self._x_nsplit
vec, other = self._m_normed_ones, self._n_normed_ones
_, val = jax.lax.scan(fun, init=(vec,), xs=jnp.arange(n))
val = jnp.concatenate(val).squeeze()
val_rest = finalize(n * self.batch_size)
val_res = jnp.concatenate([val, val_rest])
if summary == "mean":
return jnp.sum(val_res * other)
if summary == "max_cost":
# TODO(michalk8): explain why scaling is not needed
return jnp.max(val_res)
raise ValueError(
f"Scaling method {summary} does not exist for online mode."
)
[docs] def barycenter(self, weights: jnp.ndarray) -> jnp.ndarray:
"""Compute barycenter of points in self.x using weights."""
return self.cost_fn.barycenter(self.x, weights)[0]
[docs] @classmethod
def prepare_divergences(
cls,
x: jnp.ndarray,
y: jnp.ndarray,
static_b: bool = False,
src_mask: Optional[jnp.ndarray] = None,
tgt_mask: Optional[jnp.ndarray] = None,
**kwargs: Any
) -> Tuple["PointCloud", ...]:
"""Instantiate the geometries used for a divergence computation."""
couples = [(x, y), (x, x)]
masks = [(src_mask, tgt_mask), (src_mask, src_mask)]
if not static_b:
couples += [(y, y)]
masks += [(tgt_mask, tgt_mask)]
return tuple(
cls(x, y, src_mask=x_mask, tgt_mask=y_mask, **kwargs)
for ((x, y), (x_mask, y_mask)) in zip(couples, masks)
)
def tree_flatten(self): # noqa: D102
return (
self.x,
self.y,
self._src_mask,
self._tgt_mask,
self._epsilon_init,
self.cost_fn,
), {
"batch_size": self._batch_size,
"scale_cost": self._scale_cost
}
@classmethod
def tree_unflatten(cls, aux_data, children): # noqa: D102
x, y, src_mask, tgt_mask, epsilon, cost_fn = children
return cls(
x,
y,
cost_fn=cost_fn,
src_mask=src_mask,
tgt_mask=tgt_mask,
epsilon=epsilon,
**aux_data
)
def _cosine_to_sqeucl(self) -> "PointCloud":
assert isinstance(self.cost_fn, costs.Cosine), type(self.cost_fn)
(x, y, *args, _), aux_data = self.tree_flatten()
x = x / jnp.linalg.norm(x, axis=-1, keepdims=True)
y = y / jnp.linalg.norm(y, axis=-1, keepdims=True)
# TODO(michalk8): find a better way
aux_data["scale_cost"] = 2. / self.inv_scale_cost
cost_fn = costs.SqEuclidean()
return type(self).tree_unflatten(aux_data, [x, y] + args + [cost_fn])
[docs] def to_LRCGeometry(
self,
scale: float = 1.0,
**kwargs: Any,
) -> Union[low_rank.LRCGeometry, "PointCloud"]:
r"""Convert point cloud to low-rank geometry.
Args:
scale: Value used to rescale the factors of the low-rank geometry.
Useful when this geometry is used in the linear term of fused GW.
kwargs: Keyword arguments, such as ``rank``, to
:meth:`~ott.geometry.geometry.Geometry.to_LRCGeometry` used when
the point cloud does not have squared Euclidean cost.
Returns:
Returns the unmodified point cloud if :math:`n m \ge (n + m) d`, where
:math:`n, m` is the shape and :math:`d` is the dimension of the point
cloud with squared Euclidean cost.
Otherwise, returns the re-scaled low-rank geometry.
"""
if self.is_squared_euclidean:
if self._check_LRC_dim:
return self._sqeucl_to_lr(scale)
# we don't update the `scale_factor` because in GW, the linear cost
# is first materialized and then scaled by `fused_penalty` afterwards
# TODO(michalk8): in the future, consider defining point cloud as a
# subclass of LRCGeometry
return self
return super().to_LRCGeometry(scale=scale, **kwargs)
def _sqeucl_to_lr(self, scale: float = 1.0) -> low_rank.LRCGeometry:
assert self.is_squared_euclidean, "Geometry must be squared Euclidean."
n, m = self.shape
nx = jnp.sum(self.x ** 2, axis=1, keepdims=True)
ny = jnp.sum(self.y ** 2, axis=1, keepdims=True)
cost_1 = jnp.concatenate((nx, jnp.ones((n, 1)), -jnp.sqrt(2.0) * self.x),
axis=1)
cost_2 = jnp.concatenate((jnp.ones((m, 1)), ny, jnp.sqrt(2.0) * self.y),
axis=1)
return low_rank.LRCGeometry(
cost_1=cost_1,
cost_2=cost_2,
scale_factor=scale,
epsilon=self._epsilon_init,
relative_epsilon=self._relative_epsilon,
scale_cost=self._scale_cost,
src_mask=self.src_mask,
tgt_mask=self.tgt_mask,
)
[docs] def subset( # noqa: D102
self, src_ixs: Optional[jnp.ndarray], tgt_ixs: Optional[jnp.ndarray],
**kwargs: Any
) -> "PointCloud":
def subset_fn(
arr: Optional[jnp.ndarray],
ixs: Optional[jnp.ndarray],
) -> jnp.ndarray:
return arr if arr is None or ixs is None else arr[jnp.atleast_1d(ixs)]
return self._mask_subset_helper(
src_ixs, tgt_ixs, fn=subset_fn, propagate_mask=True, **kwargs
)
[docs] def mask( # noqa: D102
self,
src_mask: Optional[jnp.ndarray],
tgt_mask: Optional[jnp.ndarray],
mask_value: float = 0.,
) -> "PointCloud":
def mask_fn(
arr: Optional[jnp.ndarray],
mask: Optional[jnp.ndarray],
) -> Optional[jnp.ndarray]:
if arr is None or mask is None:
return arr
return jnp.where(mask[:, None], arr, mask_value)
src_mask = self._normalize_mask(src_mask, self.shape[0])
tgt_mask = self._normalize_mask(tgt_mask, self.shape[1])
return self._mask_subset_helper(
src_mask, tgt_mask, fn=mask_fn, propagate_mask=False
)
def _mask_subset_helper(
self,
src_ixs: Optional[jnp.ndarray],
tgt_ixs: Optional[jnp.ndarray],
*,
fn: Callable[[Optional[jnp.ndarray], Optional[jnp.ndarray]],
Optional[jnp.ndarray]],
propagate_mask: bool,
**kwargs: Any,
) -> "PointCloud":
(x, y, src_mask, tgt_mask, *children), aux_data = self.tree_flatten()
x = fn(x, src_ixs)
y = fn(y, tgt_ixs)
if propagate_mask:
src_mask = self._normalize_mask(src_mask, self.shape[0])
tgt_mask = self._normalize_mask(tgt_mask, self.shape[1])
src_mask = fn(src_mask, src_ixs)
tgt_mask = fn(tgt_mask, tgt_ixs)
aux_data = {**aux_data, **kwargs}
return type(self).tree_unflatten(
aux_data, [x, y, src_mask, tgt_mask] + children
)
@property
def dtype(self) -> jnp.dtype: # noqa: D102
return self.x.dtype
@property
def batch_size(self) -> Optional[int]:
"""Batch size for online mode."""
if self._batch_size is None:
return None
n, m = self.shape
return min(n, m, self._batch_size)
@property
def _x_nsplit(self) -> Optional[int]:
if self.batch_size is None:
return None
n, _ = self.shape
return int(math.floor(n / self.batch_size))
@property
def _y_nsplit(self) -> Optional[int]:
if self.batch_size is None:
return None
_, m = self.shape
return int(math.floor(m / self.batch_size))
def _apply_lse_kernel_xy(
x, y, norm_x, norm_y, f, g, eps, vec, cost_fn, scale_cost
):
c = _cost(x, y, norm_x, norm_y, cost_fn, scale_cost)
return mu.logsumexp((f + g - c) / eps, b=vec, return_sign=True, axis=-1)
def _transport_from_potentials_xy(
x, y, norm_x, norm_y, f, g, eps, cost_fn, scale_cost
):
return jnp.exp(
(f + g - _cost(x, y, norm_x, norm_y, cost_fn, scale_cost)) / eps
)
def _apply_kernel_xy(x, y, norm_x, norm_y, vec, eps, cost_fn, scale_cost):
c = _cost(x, y, norm_x, norm_y, cost_fn, scale_cost)
return jnp.dot(jnp.exp(-c / eps), vec)
def _transport_from_scalings_xy(
x, y, norm_x, norm_y, u, v, eps, cost_fn, scale_cost
):
return jnp.exp(
-_cost(x, y, norm_x, norm_y, cost_fn, scale_cost) * scale_cost / eps
) * u * v
def _cost(x, y, norm_x, norm_y, cost_fn, scale_cost):
one_line_pairwise = jax.vmap(cost_fn.pairwise, in_axes=[0, None])
cost = norm_x + norm_y + one_line_pairwise(x, y)
return cost * scale_cost
def _apply_cost_xy(x, y, norm_x, norm_y, vec, cost_fn, scale_cost, fn=None):
"""Apply [num_b, num_a] fn(cost) matrix (or transpose) to vector.
Applies [num_b, num_a] ([num_a, num_b] if axis=1 from `apply_cost`)
fn(cost) matrix (or transpose) to vector.
Args:
x: jnp.ndarray [num_a, d], first pointcloud
y: jnp.ndarray [num_b, d], second pointcloud
norm_x: jnp.ndarray [num_a,], (squared) norm as defined in by cost_fn
norm_y: jnp.ndarray [num_b,], (squared) norm as defined in by cost_fn
vec: jnp.ndarray [num_a,] ([num_b,] if axis=1 from `apply_cost`) vector
cost_fn: a CostFn function between two points in dimension d.
scale_cost: scaling factor of the cost matrix.
fn: function optionally applied to cost matrix element-wise, before the
apply.
Returns:
A jnp.ndarray corresponding to cost x vector
"""
c = _cost(x, y, norm_x, norm_y, cost_fn, scale_cost)
return jnp.dot(c, vec) if fn is None else jnp.dot(fn(c), vec)
def _apply_max_xy(x, y, norm_x, norm_y, vec, cost_fn, scale_cost):
del vec
c = _cost(x, y, norm_x, norm_y, cost_fn, scale_cost)
return jnp.max(jnp.abs(c))
