From 1dea1bb6190e4a11913331a4b722b8b385ba87d7 Mon Sep 17 00:00:00 2001 From: Aelrach Date: Sat, 22 Nov 2025 23:20:00 +0100 Subject: [PATCH 1/7] Applied depedency patch --- .python-version | 1 + ott/build/lib/ott/__init__.py | 33 + ott/build/lib/ott/_version.py | 21 + ott/build/lib/ott/datasets.py | 154 +++ ott/build/lib/ott/geometry/__init__.py | 24 + ott/build/lib/ott/geometry/costs.py | 1141 +++++++++++++++ ott/build/lib/ott/geometry/distrib_costs.py | 89 ++ .../lib/ott/geometry/epsilon_scheduler.py | 102 ++ ott/build/lib/ott/geometry/geodesic.py | 277 ++++ ott/build/lib/ott/geometry/geometry.py | 921 ++++++++++++ ott/build/lib/ott/geometry/graph.py | 272 ++++ ott/build/lib/ott/geometry/grid.py | 415 ++++++ ott/build/lib/ott/geometry/low_rank.py | 508 +++++++ ott/build/lib/ott/geometry/pointcloud.py | 792 +++++++++++ ott/build/lib/ott/geometry/segment.py | 188 +++ ott/build/lib/ott/initializers/__init__.py | 14 + .../lib/ott/initializers/linear/__init__.py | 14 + .../ott/initializers/linear/initializers.py | 408 ++++++ .../initializers/linear/initializers_lr.py | 654 +++++++++ .../ott/initializers/quadratic/__init__.py | 14 + .../initializers/quadratic/initializers.py | 193 +++ ott/build/lib/ott/math/__init__.py | 19 + ott/build/lib/ott/math/fixed_point_loop.py | 239 ++++ ott/build/lib/ott/math/matrix_square_root.py | 337 +++++ .../lib/ott/math/unbalanced_functions.py | 90 ++ ott/build/lib/ott/math/utils.py | 298 ++++ ott/build/lib/ott/neural/__init__.py | 14 + ott/build/lib/ott/neural/layers.py | 213 +++ ott/build/lib/ott/neural/losses.py | 309 +++++ ott/build/lib/ott/neural/models.py | 408 ++++++ ott/build/lib/ott/neural/solvers/__init__.py | 14 + ott/build/lib/ott/neural/solvers/conjugate.py | 121 ++ .../lib/ott/neural/solvers/map_estimator.py | 289 ++++ .../lib/ott/neural/solvers/neuraldual.py | 700 ++++++++++ ott/build/lib/ott/problems/__init__.py | 14 + ott/build/lib/ott/problems/linear/__init__.py | 14 + .../ott/problems/linear/barycenter_problem.py | 200 +++ .../lib/ott/problems/linear/linear_problem.py | 132 ++ .../lib/ott/problems/linear/potentials.py | 437 ++++++ .../lib/ott/problems/quadratic/__init__.py | 14 + .../ott/problems/quadratic/gw_barycenter.py | 325 +++++ .../ott/problems/quadratic/quadratic_costs.py | 78 ++ .../problems/quadratic/quadratic_problem.py | 588 ++++++++ ott/build/lib/ott/py.typed | 0 ott/build/lib/ott/solvers/__init__.py | 14 + ott/build/lib/ott/solvers/linear/__init__.py | 30 + ott/build/lib/ott/solvers/linear/_solve.py | 60 + .../lib/ott/solvers/linear/acceleration.py | 172 +++ .../solvers/linear/continuous_barycenter.py | 233 ++++ .../ott/solvers/linear/discrete_barycenter.py | 251 ++++ .../linear/implicit_differentiation.py | 324 +++++ .../lib/ott/solvers/linear/lineax_implicit.py | 98 ++ ott/build/lib/ott/solvers/linear/lr_utils.py | 371 +++++ ott/build/lib/ott/solvers/linear/sinkhorn.py | 1230 +++++++++++++++++ .../lib/ott/solvers/linear/sinkhorn_lr.py | 822 +++++++++++ .../lib/ott/solvers/linear/univariate.py | 346 +++++ .../lib/ott/solvers/quadratic/__init__.py | 20 + ott/build/lib/ott/solvers/quadratic/_solve.py | 91 ++ .../solvers/quadratic/gromov_wasserstein.py | 461 ++++++ .../quadratic/gromov_wasserstein_lr.py | 897 ++++++++++++ .../ott/solvers/quadratic/gw_barycenter.py | 339 +++++ .../lib/ott/solvers/quadratic/lower_bound.py | 96 ++ ott/build/lib/ott/solvers/was_solver.py | 131 ++ ott/build/lib/ott/tools/__init__.py | 21 + .../ott/tools/gaussian_mixture/__init__.py | 14 + .../lib/ott/tools/gaussian_mixture/fit_gmm.py | 303 ++++ .../tools/gaussian_mixture/fit_gmm_pair.py | 393 ++++++ .../ott/tools/gaussian_mixture/gaussian.py | 217 +++ .../gaussian_mixture/gaussian_mixture.py | 329 +++++ .../gaussian_mixture/gaussian_mixture_pair.py | 222 +++ .../lib/ott/tools/gaussian_mixture/linalg.py | 138 ++ .../tools/gaussian_mixture/probabilities.py | 100 ++ .../ott/tools/gaussian_mixture/scale_tril.py | 214 +++ ott/build/lib/ott/tools/k_means.py | 412 ++++++ ott/build/lib/ott/tools/plot.py | 247 ++++ ott/build/lib/ott/tools/segment_sinkhorn.py | 143 ++ .../lib/ott/tools/sinkhorn_divergence.py | 354 +++++ ott/build/lib/ott/tools/soft_sort.py | 706 ++++++++++ ott/build/lib/ott/types.py | 42 + ott/build/lib/ott/utils.py | 210 +++ ott/src/ott_jax.egg-info/PKG-INFO | 356 +++++ ott/src/ott_jax.egg-info/SOURCES.txt | 92 ++ ott/src/ott_jax.egg-info/dependency_links.txt | 1 + ott/src/ott_jax.egg-info/requires.txt | 41 + ott/src/ott_jax.egg-info/top_level.txt | 1 + perturbot/build/lib/perturbot/__init__.py | 0 .../build/lib/perturbot/eval/.run_loo.py | 126 ++ .../build/lib/perturbot/eval/__init__.py | 0 perturbot/build/lib/perturbot/eval/all.py | 190 +++ perturbot/build/lib/perturbot/eval/cv.py | 251 ++++ .../build/lib/perturbot/eval/cv_inner_loop.py | 673 +++++++++ .../build/lib/perturbot/eval/cv_outer_loop.py | 325 +++++ .../lib/perturbot/eval/feature_matching.py | 160 +++ perturbot/build/lib/perturbot/eval/loo.py | 432 ++++++ perturbot/build/lib/perturbot/eval/match.py | 242 ++++ .../build/lib/perturbot/eval/prediction.py | 210 +++ perturbot/build/lib/perturbot/eval/utils.py | 105 ++ .../build/lib/perturbot/match/__init__.py | 27 + perturbot/build/lib/perturbot/match/cot.py | 389 ++++++ .../build/lib/perturbot/match/cot_labels.py | 340 +++++ perturbot/build/lib/perturbot/match/fot.py | 220 +++ perturbot/build/lib/perturbot/match/gw.py | 137 ++ .../build/lib/perturbot/match/gw_labels.py | 148 ++ .../build/lib/perturbot/match/ot_labels.py | 78 ++ .../build/lib/perturbot/match/ott_egwl.py | 450 ++++++ perturbot/build/lib/perturbot/match/utils.py | 184 +++ .../build/lib/perturbot/predict/__init__.py | 3 + .../perturbot/predict/linear_regression.py | 155 +++ perturbot/build/lib/perturbot/predict/mlp.py | 250 ++++ .../build/lib/perturbot/predict/scvi_vae.py | 177 +++ perturbot/build/lib/perturbot/predict/vae.py | 493 +++++++ .../lib/perturbot/preprocess/__init__.py | 0 .../build/lib/perturbot/preprocess/vae.py | 65 + perturbot/build/lib/perturbot/utils.py | 10 + perturbot/perturbot.egg-info/PKG-INFO | 17 + perturbot/perturbot.egg-info/SOURCES.txt | 8 + .../perturbot.egg-info/dependency_links.txt | 1 + perturbot/perturbot.egg-info/requires.txt | 5 + perturbot/perturbot.egg-info/top_level.txt | 1 + pyproject.toml | 7 + 120 files changed, 27509 insertions(+) create mode 100644 .python-version create mode 100644 ott/build/lib/ott/__init__.py create mode 100644 ott/build/lib/ott/_version.py create mode 100644 ott/build/lib/ott/datasets.py create mode 100644 ott/build/lib/ott/geometry/__init__.py create mode 100644 ott/build/lib/ott/geometry/costs.py create mode 100644 ott/build/lib/ott/geometry/distrib_costs.py create mode 100644 ott/build/lib/ott/geometry/epsilon_scheduler.py create mode 100644 ott/build/lib/ott/geometry/geodesic.py create mode 100644 ott/build/lib/ott/geometry/geometry.py create mode 100644 ott/build/lib/ott/geometry/graph.py create mode 100644 ott/build/lib/ott/geometry/grid.py create mode 100644 ott/build/lib/ott/geometry/low_rank.py create mode 100644 ott/build/lib/ott/geometry/pointcloud.py create mode 100644 ott/build/lib/ott/geometry/segment.py create mode 100644 ott/build/lib/ott/initializers/__init__.py create mode 100644 ott/build/lib/ott/initializers/linear/__init__.py create mode 100644 ott/build/lib/ott/initializers/linear/initializers.py create mode 100644 ott/build/lib/ott/initializers/linear/initializers_lr.py create mode 100644 ott/build/lib/ott/initializers/quadratic/__init__.py create mode 100644 ott/build/lib/ott/initializers/quadratic/initializers.py create mode 100644 ott/build/lib/ott/math/__init__.py create mode 100644 ott/build/lib/ott/math/fixed_point_loop.py create mode 100644 ott/build/lib/ott/math/matrix_square_root.py create mode 100644 ott/build/lib/ott/math/unbalanced_functions.py create mode 100644 ott/build/lib/ott/math/utils.py create mode 100644 ott/build/lib/ott/neural/__init__.py create mode 100644 ott/build/lib/ott/neural/layers.py create mode 100644 ott/build/lib/ott/neural/losses.py create mode 100644 ott/build/lib/ott/neural/models.py create mode 100644 ott/build/lib/ott/neural/solvers/__init__.py create mode 100644 ott/build/lib/ott/neural/solvers/conjugate.py create mode 100644 ott/build/lib/ott/neural/solvers/map_estimator.py create mode 100644 ott/build/lib/ott/neural/solvers/neuraldual.py create mode 100644 ott/build/lib/ott/problems/__init__.py create mode 100644 ott/build/lib/ott/problems/linear/__init__.py create mode 100644 ott/build/lib/ott/problems/linear/barycenter_problem.py create mode 100644 ott/build/lib/ott/problems/linear/linear_problem.py create mode 100644 ott/build/lib/ott/problems/linear/potentials.py create mode 100644 ott/build/lib/ott/problems/quadratic/__init__.py create mode 100644 ott/build/lib/ott/problems/quadratic/gw_barycenter.py create mode 100644 ott/build/lib/ott/problems/quadratic/quadratic_costs.py create mode 100644 ott/build/lib/ott/problems/quadratic/quadratic_problem.py create mode 100755 ott/build/lib/ott/py.typed create mode 100644 ott/build/lib/ott/solvers/__init__.py create mode 100644 ott/build/lib/ott/solvers/linear/__init__.py create mode 100644 ott/build/lib/ott/solvers/linear/_solve.py create mode 100644 ott/build/lib/ott/solvers/linear/acceleration.py create mode 100644 ott/build/lib/ott/solvers/linear/continuous_barycenter.py create mode 100644 ott/build/lib/ott/solvers/linear/discrete_barycenter.py create mode 100644 ott/build/lib/ott/solvers/linear/implicit_differentiation.py create mode 100644 ott/build/lib/ott/solvers/linear/lineax_implicit.py create mode 100644 ott/build/lib/ott/solvers/linear/lr_utils.py create mode 100644 ott/build/lib/ott/solvers/linear/sinkhorn.py create mode 100644 ott/build/lib/ott/solvers/linear/sinkhorn_lr.py create mode 100644 ott/build/lib/ott/solvers/linear/univariate.py create mode 100644 ott/build/lib/ott/solvers/quadratic/__init__.py create mode 100644 ott/build/lib/ott/solvers/quadratic/_solve.py create mode 100644 ott/build/lib/ott/solvers/quadratic/gromov_wasserstein.py create mode 100644 ott/build/lib/ott/solvers/quadratic/gromov_wasserstein_lr.py create mode 100644 ott/build/lib/ott/solvers/quadratic/gw_barycenter.py create mode 100644 ott/build/lib/ott/solvers/quadratic/lower_bound.py create mode 100644 ott/build/lib/ott/solvers/was_solver.py create mode 100644 ott/build/lib/ott/tools/__init__.py create mode 100644 ott/build/lib/ott/tools/gaussian_mixture/__init__.py create mode 100644 ott/build/lib/ott/tools/gaussian_mixture/fit_gmm.py create mode 100644 ott/build/lib/ott/tools/gaussian_mixture/fit_gmm_pair.py create mode 100644 ott/build/lib/ott/tools/gaussian_mixture/gaussian.py create mode 100644 ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture.py create mode 100644 ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture_pair.py create mode 100644 ott/build/lib/ott/tools/gaussian_mixture/linalg.py create mode 100644 ott/build/lib/ott/tools/gaussian_mixture/probabilities.py create mode 100644 ott/build/lib/ott/tools/gaussian_mixture/scale_tril.py create mode 100644 ott/build/lib/ott/tools/k_means.py create mode 100644 ott/build/lib/ott/tools/plot.py create mode 100644 ott/build/lib/ott/tools/segment_sinkhorn.py create mode 100644 ott/build/lib/ott/tools/sinkhorn_divergence.py create mode 100644 ott/build/lib/ott/tools/soft_sort.py create mode 100644 ott/build/lib/ott/types.py create mode 100644 ott/build/lib/ott/utils.py create mode 100644 ott/src/ott_jax.egg-info/PKG-INFO create mode 100644 ott/src/ott_jax.egg-info/SOURCES.txt create mode 100644 ott/src/ott_jax.egg-info/dependency_links.txt create mode 100644 ott/src/ott_jax.egg-info/requires.txt create mode 100644 ott/src/ott_jax.egg-info/top_level.txt create mode 100644 perturbot/build/lib/perturbot/__init__.py create mode 100755 perturbot/build/lib/perturbot/eval/.run_loo.py create mode 100755 perturbot/build/lib/perturbot/eval/__init__.py create mode 100755 perturbot/build/lib/perturbot/eval/all.py create mode 100755 perturbot/build/lib/perturbot/eval/cv.py create mode 100755 perturbot/build/lib/perturbot/eval/cv_inner_loop.py create mode 100755 perturbot/build/lib/perturbot/eval/cv_outer_loop.py create mode 100755 perturbot/build/lib/perturbot/eval/feature_matching.py create mode 100755 perturbot/build/lib/perturbot/eval/loo.py create mode 100755 perturbot/build/lib/perturbot/eval/match.py create mode 100755 perturbot/build/lib/perturbot/eval/prediction.py create mode 100755 perturbot/build/lib/perturbot/eval/utils.py create mode 100755 perturbot/build/lib/perturbot/match/__init__.py create mode 100755 perturbot/build/lib/perturbot/match/cot.py create mode 100755 perturbot/build/lib/perturbot/match/cot_labels.py create mode 100755 perturbot/build/lib/perturbot/match/fot.py create mode 100755 perturbot/build/lib/perturbot/match/gw.py create mode 100755 perturbot/build/lib/perturbot/match/gw_labels.py create mode 100755 perturbot/build/lib/perturbot/match/ot_labels.py create mode 100755 perturbot/build/lib/perturbot/match/ott_egwl.py create mode 100755 perturbot/build/lib/perturbot/match/utils.py create mode 100755 perturbot/build/lib/perturbot/predict/__init__.py create mode 100755 perturbot/build/lib/perturbot/predict/linear_regression.py create mode 100755 perturbot/build/lib/perturbot/predict/mlp.py create mode 100755 perturbot/build/lib/perturbot/predict/scvi_vae.py create mode 100755 perturbot/build/lib/perturbot/predict/vae.py create mode 100755 perturbot/build/lib/perturbot/preprocess/__init__.py create mode 100755 perturbot/build/lib/perturbot/preprocess/vae.py create mode 100644 perturbot/build/lib/perturbot/utils.py create mode 100644 perturbot/perturbot.egg-info/PKG-INFO create mode 100644 perturbot/perturbot.egg-info/SOURCES.txt create mode 100644 perturbot/perturbot.egg-info/dependency_links.txt create mode 100644 perturbot/perturbot.egg-info/requires.txt create mode 100644 perturbot/perturbot.egg-info/top_level.txt create mode 100644 pyproject.toml diff --git a/.python-version b/.python-version new file mode 100644 index 0000000..e4fba21 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.12 diff --git a/ott/build/lib/ott/__init__.py b/ott/build/lib/ott/__init__.py new file mode 100644 index 0000000..dac0eb8 --- /dev/null +++ b/ott/build/lib/ott/__init__.py @@ -0,0 +1,33 @@ +# 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 contextlib + +from . import ( + datasets, + geometry, + initializers, + math, + problems, + solvers, + tools, + utils, +) + +with contextlib.suppress(ImportError): + # TODO(michalk8): add warning that neural module is not imported + from . import neural + +from ._version import __version__ + +del contextlib diff --git a/ott/build/lib/ott/_version.py b/ott/build/lib/ott/_version.py new file mode 100644 index 0000000..972f66c --- /dev/null +++ b/ott/build/lib/ott/_version.py @@ -0,0 +1,21 @@ +# 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. +from importlib.metadata import PackageNotFoundError, version + +try: + __version__ = version("ott-jax") +except PackageNotFoundError: + __version__ = "" + +del version, PackageNotFoundError diff --git a/ott/build/lib/ott/datasets.py b/ott/build/lib/ott/datasets.py new file mode 100644 index 0000000..2a8d353 --- /dev/null +++ b/ott/build/lib/ott/datasets.py @@ -0,0 +1,154 @@ +# 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 dataclasses +from typing import Iterator, Literal, NamedTuple, Optional, Tuple + +import jax +import jax.numpy as jnp +import numpy as np + +__all__ = ["create_gaussian_mixture_samplers", "Dataset", "GaussianMixture"] + +from ott import utils + +Name_t = Literal["simple", "circle", "square_five", "square_four"] + + +class Dataset(NamedTuple): + r"""Samplers from source and target measures. + + Args: + source_iter: loader for the source measure + target_iter: loader for the target measure + """ + source_iter: Iterator[jnp.ndarray] + target_iter: Iterator[jnp.ndarray] + + +@dataclasses.dataclass +class GaussianMixture: + """A mixture of Gaussians. + + Args: + name: the name specifying the centers of the mixture components: + + - ``simple`` - data clustered in one center, + - ``circle`` - two-dimensional Gaussians arranged on a circle, + - ``square_five`` - two-dimensional Gaussians on a square with + one Gaussian in the center, and + - ``square_four`` - two-dimensional Gaussians in the corners of a + rectangle + + batch_size: batch size of the samples + init_rng: initial PRNG key + scale: scale of the Gaussian means + std: the standard deviation of the individual Gaussian samples + """ + name: Name_t + batch_size: int + init_rng: jax.Array + scale: float = 5.0 + std: float = 0.5 + + def __post_init__(self): + gaussian_centers = { + "simple": + np.array([[0, 0]]), + "circle": + np.array([ + (1, 0), + (-1, 0), + (0, 1), + (0, -1), + (1.0 / np.sqrt(2), 1.0 / np.sqrt(2)), + (1.0 / np.sqrt(2), -1.0 / np.sqrt(2)), + (-1.0 / np.sqrt(2), 1.0 / np.sqrt(2)), + (-1.0 / np.sqrt(2), -1.0 / np.sqrt(2)), + ]), + "square_five": + np.array([[0, 0], [1, 1], [-1, 1], [-1, -1], [1, -1]]), + "square_four": + np.array([[1, 0], [0, 1], [-1, 0], [0, -1]]), + } + if self.name not in gaussian_centers: + raise ValueError( + f"{self.name} is not a valid dataset for GaussianMixture" + ) + self.centers = gaussian_centers[self.name] + + def __iter__(self) -> Iterator[jnp.array]: + """Random sample generator from Gaussian mixture. + + Returns: + A generator of samples from the Gaussian mixture. + """ + return self._create_sample_generators() + + def _create_sample_generators(self) -> Iterator[jnp.array]: + rng = self.init_rng + while True: + rng1, rng2, rng = jax.random.split(rng, 3) + means = jax.random.choice(rng1, self.centers, (self.batch_size,)) + normal_samples = jax.random.normal(rng2, (self.batch_size, 2)) + samples = self.scale * means + (self.std ** 2) * normal_samples + yield samples + + +def create_gaussian_mixture_samplers( + name_source: Name_t, + name_target: Name_t, + train_batch_size: int = 2048, + valid_batch_size: int = 2048, + rng: Optional[jax.Array] = None, +) -> Tuple[Dataset, Dataset, int]: + """Gaussian samplers. + + Args: + name_source: name of the source sampler + name_target: name of the target sampler + train_batch_size: the training batch size + valid_batch_size: the validation batch size + rng: initial PRNG key + + Returns: + The dataset and dimension of the data. + """ + rng = utils.default_prng_key(rng) + rng1, rng2, rng3, rng4 = jax.random.split(rng, 4) + train_dataset = Dataset( + source_iter=iter( + GaussianMixture( + name_source, batch_size=train_batch_size, init_rng=rng1 + ) + ), + target_iter=iter( + GaussianMixture( + name_target, batch_size=train_batch_size, init_rng=rng2 + ) + ) + ) + valid_dataset = Dataset( + source_iter=iter( + GaussianMixture( + name_source, batch_size=valid_batch_size, init_rng=rng3 + ) + ), + target_iter=iter( + GaussianMixture( + name_target, batch_size=valid_batch_size, init_rng=rng4 + ) + ) + ) + dim_data = 2 + return train_dataset, valid_dataset, dim_data diff --git a/ott/build/lib/ott/geometry/__init__.py b/ott/build/lib/ott/geometry/__init__.py new file mode 100644 index 0000000..32d9d96 --- /dev/null +++ b/ott/build/lib/ott/geometry/__init__.py @@ -0,0 +1,24 @@ +# 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. +from . import ( + costs, + distrib_costs, + epsilon_scheduler, + geodesic, + geometry, + graph, + grid, + pointcloud, + segment, +) diff --git a/ott/build/lib/ott/geometry/costs.py b/ott/build/lib/ott/geometry/costs.py new file mode 100644 index 0000000..de76be3 --- /dev/null +++ b/ott/build/lib/ott/geometry/costs.py @@ -0,0 +1,1141 @@ +# 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 abc +import functools +import math +from typing import Any, Callable, Dict, Optional, Tuple, Union + +import jax +import jax.numpy as jnp +import jaxopt +import numpy as np + +from ott.math import fixed_point_loop, matrix_square_root +from ott.math import utils as mu + +__all__ = [ + "PNormP", + "SqPNorm", + "Euclidean", + "SqEuclidean", + "Cosine", + "Arccos", + "ElasticL1", + "ElasticL2", + "ElasticSTVS", + "ElasticSqKOverlap", + "Bures", + "UnbalancedBures", + "SoftDTW", +] + + +@jax.tree_util.register_pytree_node_class +class CostFn(abc.ABC): + """Base class for all costs. + + Cost functions evaluate a function on a pair of inputs. For convenience, + that function is split into two norms -- evaluated on each input separately -- + followed by a pairwise cost that involves both inputs, as in: + + .. math:: + c(x, y) = norm(x) + norm(y) + pairwise(x, y) + + If the :attr:`norm` function is not implemented, that value is handled as + :math:`0`, and only :func:`pairwise` is used. + """ + + # no norm function created by default. + norm: Optional[Callable[[jnp.ndarray], Union[float, jnp.ndarray]]] = None + + @abc.abstractmethod + def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + """Compute cost between :math:`x` and :math:`y`. + + Args: + x: Array. + y: Array. + + Returns: + The cost. + """ + + def barycenter(self, weights: jnp.ndarray, + xs: jnp.ndarray) -> Tuple[jnp.ndarray, Any]: + """Barycentric operator. + + Args: + weights: Convex set of weights. + xs: Points. + + Returns: + A list, whose first element is the barycenter of `xs` using `weights` + coefficients, followed by auxiliary information on the convergence of + the algorithm. + """ + raise NotImplementedError("Barycenter is not implemented.") + + @classmethod + def _padder(cls, dim: int) -> jnp.ndarray: + """Create a padding vector of adequate dimension, well-suited to a cost. + + Args: + dim: Dimensionality of the data. + + Returns: + The padding vector. + """ + return jnp.zeros((1, dim)) + + def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + """Compute cost between :math:`x` and :math:`y`. + + Args: + x: Array. + y: Array. + + Returns: + The cost, optionally including the :attr:`norms ` of + :math:`x`/:math:`y`. + """ + cost = self.pairwise(x, y) + if self.norm is None: + return cost + return cost + self.norm(x) + self.norm(y) + + def all_pairs(self, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: + """Compute matrix of all pairwise costs, including the :attr:`norms `. + + Args: + x: Array of shape ``[n, ...]``. + y: Array of shape ``[m, ...]``. + + Returns: + Array of shape ``[n, m]`` of cost evaluations. + """ + return jax.vmap(lambda x_: jax.vmap(lambda y_: self(x_, y_))(y))(x) + + def all_pairs_pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: + """Compute matrix of all pairwise costs, excluding the :attr:`norms `. + + Args: + x: Array of shape ``[n, ...]``. + y: Array of shape ``[m, ...]``. + + Returns: + Array of shape ``[n, m]`` of cost evaluations. + """ + return jax.vmap(lambda x_: jax.vmap(lambda y_: self.pairwise(x_, y_))(y))(x) + + def tree_flatten(self): # noqa: D102 + return (), None + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + del aux_data + return cls(*children) + + +@jax.tree_util.register_pytree_node_class +class TICost(CostFn): + """Base class for translation invariant (TI) costs. + + Such costs are defined using a function :math:`h`, mapping vectors to + real-values, to be used as: + + .. math:: + c(x, y) = h(z), z := x - y. + + If that cost function is used to form an Entropic map using the + :cite:`brenier:91` theorem, then the user should ensure :math:`h` is + strictly convex, as well as provide the Legendre transform of :math:`h`, + whose gradient is necessarily the inverse of the gradient of :math:`h`. + """ + + @abc.abstractmethod + def h(self, z: jnp.ndarray) -> float: + """TI function acting on difference of :math:`x-y` to output cost. + + Args: + z: Array of shape ``[d,]``. + + Returns: + The cost. + """ + + def h_legendre(self, z: jnp.ndarray) -> float: + """Legendre transform of :func:`h` when it is convex.""" + raise NotImplementedError("Legendre transform of `h` is not implemented.") + + def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + """Compute cost as evaluation of :func:`h` on :math:`x-y`.""" + return self.h(x - y) + + +@jax.tree_util.register_pytree_node_class +class SqPNorm(TICost): + r"""Squared p-norm of the difference of two vectors. + + Uses custom implementation of `norm` to avoid `NaN` values when + differentiating the norm of `x-x`. + + Args: + p: Power of the p-norm, :math:`\ge 1`. + """ + + def __init__(self, p: float): + super().__init__() + self.p = p + self.q = 1.0 / (1.0 - (1.0 / p)) if p > 1.0 else jnp.inf + + def h(self, z: jnp.ndarray) -> float: # noqa: D102 + return 0.5 * mu.norm(z, self.p) ** 2 + + def h_legendre(self, z: jnp.ndarray) -> float: + """Legendre transform of :func:`h`. + + For details on the derivation, see e.g., :cite:`boyd:04`, p. 93/94. + """ + return 0.5 * mu.norm(z, self.q) ** 2 + + def tree_flatten(self): # noqa: D102 + return (), (self.p,) + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + del children + return cls(*aux_data) + + +@jax.tree_util.register_pytree_node_class +class PNormP(TICost): + r"""p-norm to the power p (and divided by p) of the difference of two vectors. + + Uses custom implementation of `norm` to avoid `NaN` values when + differentiating the norm of `x-x`. + + Args: + p: Power of the p-norm in :math:`[1, +\infty)`. + Note that :func:`h_legendre` is not defined for ``p = 1``. + """ + + def __init__(self, p: float): + super().__init__() + self.p = p + self.q = 1.0 / (1.0 - (1.0 / p)) if p > 1.0 else jnp.inf + + def h(self, z: jnp.ndarray) -> float: # noqa: D102 + return mu.norm(z, self.p) ** self.p / self.p + + def h_legendre(self, z: jnp.ndarray) -> float: # noqa: D102 + # not defined for `p=1` + return mu.norm(z, self.q) ** self.q / self.q + + def tree_flatten(self): # noqa: D102 + return (), (self.p,) + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + del children + return cls(*aux_data) + + +@jax.tree_util.register_pytree_node_class +class Euclidean(CostFn): + """Euclidean distance. + + Note that the Euclidean distance is not cast as a + :class:`~ott.geometry.costs.TICost`, since this would correspond to :math:`h` + being :func:`jax.numpy.linalg.norm`, whose gradient is not invertible, + because the function is not strictly convex (it is linear on rays). + """ + + def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + """Compute Euclidean norm using custom jvp implementation. + + Here we use a custom jvp implementation for the norm that does not yield + `NaN` gradients when differentiating the norm of `(x-x)`, but defaults + instead to zero, using a `custom_jvp` rule. + """ + return mu.norm(x - y) + + +@jax.tree_util.register_pytree_node_class +class SqEuclidean(TICost): + r"""Squared Euclidean distance. + + Implemented as a translation invariant cost, :math:`h(z) = \|z\|^2`. + """ + + def norm(self, x: jnp.ndarray) -> Union[float, jnp.ndarray]: + """Compute squared Euclidean norm for vector.""" + return jnp.sum(x ** 2, axis=-1) + + def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + """Compute minus twice the dot-product between vectors.""" + return -2.0 * jnp.vdot(x, y) + + def h(self, z: jnp.ndarray) -> float: # noqa: D102 + return jnp.sum(z ** 2) + + def h_legendre(self, z: jnp.ndarray) -> float: # noqa: D102 + return 0.25 * jnp.sum(z ** 2) + + def barycenter(self, weights: jnp.ndarray, + xs: jnp.ndarray) -> Tuple[jnp.ndarray, Any]: + """Output barycenter of vectors when using squared-Euclidean distance.""" + return jnp.average(xs, weights=weights, axis=0), None + + +@jax.tree_util.register_pytree_node_class +class Cosine(CostFn): + """Cosine distance cost function. + + Args: + ridge: Ridge regularization. + """ + + def __init__(self, ridge: float = 1e-8): + super().__init__() + self._ridge = ridge + + def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + """Cosine distance between vectors, denominator regularized with ridge.""" + x_norm = jnp.linalg.norm(x, axis=-1) + y_norm = jnp.linalg.norm(y, axis=-1) + cosine_similarity = jnp.vdot(x, y) / (x_norm * y_norm + self._ridge) + return 1.0 - cosine_similarity + + @classmethod + def _padder(cls, dim: int) -> jnp.ndarray: + return jnp.ones((1, dim)) + + +@jax.tree_util.register_pytree_node_class +class Arccos(CostFn): + r"""Arc-cosine cost function :cite:`cho:09`. + + The cost is implemented as: + + .. math:: + c_n(x, y) = -\log(\frac{1}{\pi} \|x\|^n \|y\|^n J_n(\theta)) + + where :math:`\theta := \arccos(\frac{x \cdot y}{\|x\| \|y\|})` and + :math:`J_n(\theta) := (-1)^n (\sin \theta)^{2n + 1} + (\frac{1}{\sin \theta}\frac{\partial}{\partial \theta})^n + (\frac{\pi - \theta}{\sin \theta})`. + + Args: + n: Order of the kernel. For :math:`n > 2`, successive applications of + :func:`~jax.grad` are used to compute the :math:`J_n(\theta)`. + ridge: Ridge regularization. + """ + + def __init__(self, n: int, ridge: float = 1e-8): + self.n = n + self._ridge = ridge + + def pairwise(self, x: jnp.ndarray, y: jnp.ndarray): # noqa: D102 + x_norm = jnp.linalg.norm(x, axis=-1) + y_norm = jnp.linalg.norm(y, axis=-1) + cosine_similarity = jnp.vdot(x, y) / (x_norm * y_norm + self._ridge) + theta = jnp.arccos(cosine_similarity) + + if self.n == 0: + m = 1.0 - theta / jnp.pi + elif self.n == 1: + j = jnp.sin(theta) + (jnp.pi - theta) * jnp.cos(theta) + m = (x_norm * y_norm) * (j / jnp.pi) + elif self.n == 2: + j = 3.0 * jnp.sin(theta) * jnp.cos(theta) + (jnp.pi - theta) * ( + 1.0 + 2.0 * jnp.cos(theta) ** 2 + ) + m = (x_norm * y_norm) ** 2 * (j / jnp.pi) + else: + j = self._j(theta) # less optimized version using autodiff + m = (x_norm * y_norm) ** self.n * (j / jnp.pi) + + return -jnp.log(m + self._ridge) + + @jax.jit + def _j(self, theta: float) -> float: + + def f(t: float, i: int) -> float: + if i == 0: + return (jnp.pi - t) / jnp.sin(t) + return jax.grad(f)(t, i - 1) / jnp.sin(t) + + n = self.n + return (-1) ** n * jnp.sin(theta) ** (2.0 * n + 1.0) * f(theta, n) + + def tree_flatten(self): # noqa: D102 + return [], {"n": self.n, "ridge": self._ridge} + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + del children + return cls(**aux_data) + + +class RegTICost(TICost, abc.ABC): + r"""Base class for regularized translation-invariant costs. + + .. math:: + \frac{1}{2} \|\cdot\|_2^2 + \text{scaling_reg} reg\left(matrix \cdot\right) + + where :func:`reg` is the regularization function. + + Args: + scaling_reg: Strength of the :meth:`regularization `. + matrix: :math:`p \times d` projection matrix in the Stiefel manifold, + namely with **orthonormalized rows**. + orthogonal: Whether to regularize in the orthogonal complement + to promote displacements in the span of ``matrix``. + """ + + def __init__( + self, + scaling_reg: float = 1.0, + matrix: Optional[jnp.ndarray] = None, + orthogonal: bool = False, + ): + super().__init__() + self.scaling_reg = scaling_reg + self.matrix = matrix + self.orthogonal = orthogonal + + @abc.abstractmethod + def _reg(self, z: jnp.ndarray) -> float: + """Regularization function.""" + + def _reg_stiefel_orth(self, z: jnp.ndarray) -> float: + raise NotImplementedError( + "Regularization in the orthogonal " + "subspace is not implemented." + ) + + def reg(self, z: jnp.ndarray) -> float: + """Regularization function. + + Args: + z: Array of shape ``[d,]``. + + Returns: + The regularization value. + """ + if self.matrix is None: + return self._reg(z) + if self.orthogonal: + return self._reg_stiefel_orth(z) + return self._reg(self.matrix @ z) + + def prox_reg(self, z: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: + """Proximal operator of :meth:`reg`. + + Args: + z: Array of shape ``[d,]``. + tau: Positive weight. + + Returns: + The prox of ``z``. + """ + if self.matrix is None: + return self._prox_reg(z, tau) + if self.orthogonal: + # regularization in the orthogonal subspace + return self._prox_reg_stiefel_orth(z, tau) + return self._prox_reg_stiefel(z, tau) + + def _prox_reg(self, z: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: + raise NotImplementedError("Proximal operator is not implemented.") + + def _prox_reg_stiefel_orth( + self, z: jnp.ndarray, tau: float = 1.0 + ) -> jnp.ndarray: + + def orth(x: jnp.ndarray) -> jnp.ndarray: + return x - self.matrix.T @ (self.matrix @ x) + + # assumes `matrix` has orthogonal rows + tmp = orth(z) + return z - orth(tmp - self._prox_reg(tmp, tau)) + + def _prox_reg_stiefel(self, z: jnp.ndarray, tau: float) -> jnp.ndarray: + # assumes `matrix` has orthogonal rows + tmp = self.matrix @ z + return z - self.matrix.T @ (tmp - self._prox_reg(tmp, tau)) + + def prox_legendre_reg(self, z: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: + r"""Proximal operator of the Legendre transform of :meth:`reg`. + + Uses Moreau's decomposition: + + .. math:: + x = \text{prox}_{\tau f} \left(x\right) + + \tau \text{prox}_{\frac{1}{\tau} f^*} \left(\frac{x}{\tau}\right) + + Args: + z: Array of shape ``[d,]``. + tau: Positive weight. + + Returns: + The prox of ``z``. + """ + return z - tau * self.prox_reg(z / tau, 1.0 / tau) + + def h(self, z: jnp.ndarray) -> float: # noqa: D102 + out = 0.5 * jnp.sum(z ** 2) + return out + self.scaling_reg * self.reg(z) + + def h_legendre(self, z: jnp.ndarray) -> float: # noqa: D102 + q = jax.lax.stop_gradient(self.prox_reg(z)) + return jnp.sum(q * z) - self.h(q) + + def h_transform(self, f: Callable[[jnp.ndarray], float], + **kwargs: Any) -> Callable[[jnp.ndarray], float]: + r"""Compute the h-transform of a concave function. + + Return a callable :math:`f_h` defined as: + + .. math:: + f_h(x) = \min_y h(x - y) - f(y) + + This is equivalent, up to a change of variables, :math:`z = x - y`, to + define + + .. math:: + \min_z h(z) - f(x - z). \\ + \min_z h(z) + \tilde{f}(z, x). + + where :math:`\tilde{f}(z, x) := -f(x - z)`. + + This is solved using proximal gradient descent, which requires having + access to the prox of :math:`\text{scaling_h} \cdot h` and not only to that + of :meth:`h`. Given the properties of :meth:`h`, the prox is obtained by + rescaling the output of the prox of a suitable scaling of :meth:`prox_reg`. + + Args: + f: Concave function. + kwargs: Keyword arguments for :class:`~jaxopt.ProximalGradient`. + + Returns: + The h-transform of ``f``. + """ + + def minus_f(z: jnp.ndarray, x: jnp.ndarray) -> float: + return -f(x - z) + + def prox( + x: jnp.ndarray, scaling_reg: float, scaling_h: float + ) -> jnp.ndarray: + # https://web.stanford.edu/~boyd/papers/pdf/prox_algs.pdf 2.2. + tmp = 1.0 / (1.0 + scaling_h) + tau = scaling_reg * scaling_h * tmp + return self.prox_reg(x * tmp, tau) + + def f_h(x: jnp.ndarray) -> float: + pg = jaxopt.ProximalGradient(fun=minus_f, prox=prox, **kwargs) + pg_run = pg.run(x, self.scaling_reg, x=x) + pg_sol = jax.lax.stop_gradient(pg_run.params) + return self.h(pg_sol) + minus_f(pg_sol, x) + + return f_h + + def tree_flatten(self): # noqa: D102 + return (self.scaling_reg, self.matrix), {"orthogonal": self.orthogonal} + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + return cls(*children, **aux_data) + + +@jax.tree_util.register_pytree_node_class +class ElasticL1(RegTICost): + r"""Cost inspired by elastic net :cite:`zou:05` regularization. + + .. math:: + \frac{1}{2} \|\cdot\|_2^2 + \text{scaling_reg} \|\text{matrix} \cdot\|_1 + + Args: + scaling_reg: Strength of the :meth:`regularization `. + matrix: :math:`p \times d` projection matrix with **orthogonal rows**. + orthogonal: Whether to regularize in the orthogonal complement + to promote displacements in the span of ``matrix``. + """ + + def _reg(self, z: jnp.ndarray) -> float: # noqa: D102 + return jnp.linalg.norm(z, ord=1) + + def _prox_reg(self, z: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: + return jnp.sign(z) * jax.nn.relu(jnp.abs(z) - tau * self.scaling_reg) + + +@jax.tree_util.register_pytree_node_class +class ElasticL2(RegTICost): + r"""Cost with L2 regularization. + + .. math:: + \frac{1}{2} \|\cdot\|_2^2 + \text{scaling_reg} \|\text{matrix} \cdot\|_2^2 + + Args: + scaling_reg: Strength of the :meth:`regularization `. + matrix: :math:`p \times d` projection matrix with **orthogonal rows**. + orthogonal: Whether to regularize in the orthogonal complement + to promote displacements in the span of ``matrix``. + """ + + def _reg(self, z: jnp.ndarray) -> float: # noqa: D102 + return 0.5 * jnp.sum(z ** 2) + + def _reg_stiefel_orth(self, z: jnp.ndarray) -> float: + # Pythagorean identity + return self._reg(z) - self._reg(self.matrix @ z) + + def _prox_reg(self, z: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: + return z / (1.0 + tau * self.scaling_reg) + + def _prox_reg_stiefel_orth( + self, z: jnp.ndarray, tau: float = 1.0 + ) -> jnp.ndarray: + out = z + tau * self.scaling_reg * self.matrix.T @ (self.matrix @ z) + return self._prox_reg(out, tau) + + +@jax.tree_util.register_pytree_node_class +class ElasticSTVS(RegTICost): + r"""Cost with soft thresholding operator with vanishing shrinkage (STVS) + :cite:`schreck:15` regularization. + + .. math:: + \frac{1}{2} \|\cdot\|_2^2 + \text{scaling_reg}^2\mathbf{1}_d^T\left(\sigma(\cdot) - + \frac{1}{2} \exp\left(-2\sigma(\cdot)\right) + \frac{1}{2}\right) + + where :math:`\sigma(\cdot) := \text{asinh}\left(\frac{\cdot} + {2\text{scaling_reg}}\right)` + + Args: + scaling_reg: Strength of the :meth:`regularization `. + matrix: :math:`p \times d` projection matrix with **orthogonal rows**. + orthogonal: Whether to regularize in the orthogonal complement + to promote displacements in the span of ``matrix``. + """ # noqa: D205,E501 + + def _reg(self, z: jnp.ndarray) -> float: # noqa: D102 + u = jnp.arcsinh(jnp.abs(z) / (2 * self.scaling_reg)) + out = u - 0.5 * jnp.exp(-2.0 * u) + # Lemma 2.1 of `schreck:15`; + # don't use `self.scaling_reg ** 2` because it's included in `h` + return self.scaling_reg * jnp.sum(out + 0.5) # make positive + + def _prox_reg( # noqa: D102 + self, z: jnp.ndarray, tau: float = 1.0 + ) -> jnp.ndarray: + tmp = 1.0 - (self.scaling_reg * tau / (jnp.abs(z) + 1e-12)) ** 2 + return jax.nn.relu(tmp) * z + + +@jax.tree_util.register_pytree_node_class +class ElasticSqKOverlap(RegTICost): + r"""Cost with squared k-overlap norm regularization :cite:`argyriou:12`. + + .. math:: + \frac{1}{2} \|\cdot\|_2^2 + \frac{1}{2} \text{scaling_reg} \|\cdot\|_{ovk}^2 + + where :math:`\|\cdot\|_{ovk}^2` is the squared k-overlap norm, + see def. 2.1 of :cite:`argyriou:12`. + + Args: + k: Number of groups. Must be in ``[0, d)`` where :math:`d` is the + dimensionality of the data. + args: Positional arguments for :class:`~ott.geometry.costs.RegTICost`. + kwargs: Keyword arguments for :class:`~ott.geometry.costs.RegTICost`. + """ + + def __init__(self, k: int, *args, **kwargs: Any): + super().__init__(*args, **kwargs) + self.k = k + + def _reg(self, z: jnp.ndarray) -> float: # noqa: D102 + # Prop 2.1 in :cite:`argyriou:12` + k = self.k + top_w = jax.lax.top_k(jnp.abs(z), k)[0] # Fetch largest k values + top_w = jnp.flip(top_w) # Sort k-largest from smallest to largest + # sum (dim - k) smallest values + sum_bottom = jnp.sum(jnp.abs(z)) - jnp.sum(top_w) + cumsum_top = jnp.cumsum(top_w) + # Cesaro mean of top_w (each term offset with sum_bottom). + cesaro = sum_bottom + cumsum_top + cesaro /= jnp.arange(k) + 1 + # Choose first index satisfying constraint in Prop 2.1 + lower_bound = cesaro - top_w >= 0 + # Last upper bound is always True. + upper_bound = jnp.concatenate(((top_w[1:] - cesaro[:-1] + > 0), jnp.array((True,)))) + r = jnp.argmax(lower_bound * upper_bound) + s = jnp.sum(jnp.where(jnp.arange(k) < k - r - 1, jnp.flip(top_w) ** 2, 0)) + + return 0.5 * (s + (r + 1) * cesaro[r] ** 2) + + def prox_reg(self, z: jnp.ndarray, tau: float = 1.0) -> float: # noqa: D102 + + @functools.partial(jax.vmap, in_axes=[0, None, None]) + def find_indices(r: int, l: jnp.ndarray, + z: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]: + + @functools.partial(jax.vmap, in_axes=[None, 0, None]) + def inner(r: int, l: int, + z: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]: + i = k - r - 1 + res = jnp.sum(z * ((i <= ixs) & (ixs < l))) + res /= l - k + (beta + 1) * r + beta + 1 + + cond1_left = jnp.logical_or(i == 0, (z[i - 1] / beta + 1) > res) + cond1_right = res >= (z[i] / (beta + 1)) + cond1 = jnp.logical_and(cond1_left, cond1_right) + + cond2_left = z[l - 1] > res + cond2_right = jnp.logical_or(l == d, res >= z[l]) + cond2 = jnp.logical_and(cond2_left, cond2_right) + + return res, cond1 & cond2 + + return inner(r, l, z) + + del tau # this case is not handled and currently not needed + # Alg. 1 of :cite:`argyriou:12` + k, d, beta = self.k, z.shape[-1], 1.0 / self.scaling_reg + + ixs = jnp.arange(d) + z, sgn = jnp.abs(z), jnp.sign(z) + z_ixs = jnp.argsort(z)[::-1] + z_sorted = z[z_ixs] + + # (k, d - k + 1) + T, mask = find_indices(jnp.arange(k), jnp.arange(k, d + 1), z_sorted) + (r,), (l,) = jnp.where(mask, size=1) # size=1 for jitting + T = T[r, l] + + q1 = (beta / (beta + 1)) * z_sorted * (ixs < (k - r - 1)) + q2 = (z_sorted - T) * jnp.logical_and((k - r - 1) <= ixs, ixs < (l + k)) + q = q1 + q2 + + # change sign and reorder + return sgn * q[jnp.argsort(z_ixs.astype(float))] + + def tree_flatten(self): # noqa: D102 + children, aux_data = super().tree_flatten() + return children, (self.k, aux_data) + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + k, aux_data = aux_data + return cls(k, *children, **aux_data) + + +@jax.tree_util.register_pytree_node_class +class Bures(CostFn): + """Bures distance between a pair of (mean, covariance matrix). + + Args: + dimension: Dimensionality of the data. + sqrtm_kw: Dictionary of keyword arguments to control the + behavior of inner calls to :func:`~ott.math.matrix_square_root.sqrtm`. + """ + + def __init__(self, dimension: int, sqrtm_kw: Optional[Dict[str, Any]] = None): + super().__init__() + self._dimension = dimension + self._sqrtm_kw = {} if sqrtm_kw is None else sqrtm_kw + + def norm(self, x: jnp.ndarray) -> jnp.ndarray: + """Compute norm of Gaussian, sq. 2-norm of mean + trace of covariance.""" + mean, cov = x_to_means_and_covs(x, self._dimension) + norm = jnp.sum(mean ** 2, axis=-1) + norm += jnp.trace(cov, axis1=-2, axis2=-1) + return norm + + def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + """Compute - 2 x Bures dot-product.""" + mean_x, cov_x = x_to_means_and_covs(x, self._dimension) + mean_y, cov_y = x_to_means_and_covs(y, self._dimension) + mean_dot_prod = jnp.vdot(mean_x, mean_y) + sq_x = matrix_square_root.sqrtm(cov_x, self._dimension, **self._sqrtm_kw)[0] + sq_x_y_sq_x = jnp.matmul(sq_x, jnp.matmul(cov_y, sq_x)) + sq__sq_x_y_sq_x = matrix_square_root.sqrtm( + sq_x_y_sq_x, self._dimension, **self._sqrtm_kw + )[0] + return -2 * (mean_dot_prod + jnp.trace(sq__sq_x_y_sq_x, axis1=-2, axis2=-1)) + + def covariance_fixpoint_iter( + self, + covs: jnp.ndarray, + weights: jnp.ndarray, + tolerance: float = 1e-4, + sqrtm_kw: Optional[Dict[str, Any]] = None, + **kwargs: Any + ) -> jnp.ndarray: + """Iterate fix-point updates to compute barycenter of Gaussians. + + Args: + covs: [batch, d^2] covariance matrices + weights: simplicial weights (non-negative, sum to 1) + tolerance: tolerance of the fixed-point procedure. That tolerance is + applied to the Frobenius norm (normalized by total size) + of two successive iterations of the algorithm + sqrtm_kw: keyword arguments for :func:`~ott.math.matrix_square_root.sqrtm` + kwargs: keyword arguments for the outer fixed-point iteration + + Returns: + List containing Weighted Bures average of the covariance matrices, and + vector of (normalized) 2-norms of successive differences between iterates, + to monitor convergence. + """ + sqrtm_kw = {} if sqrtm_kw is None else sqrtm_kw + # Pop values or set defaults for fixed-point loop. + min_iterations = kwargs.pop("min_iterations", 1) + max_iterations = kwargs.pop("max_iterations", 100) + inner_iterations = kwargs.pop("inner_iterations", 5) + + @functools.partial(jax.vmap, in_axes=[None, 0, 0]) + def scale_covariances( + cov_sqrt: jnp.ndarray, cov: jnp.ndarray, weight: jnp.ndarray + ) -> jnp.ndarray: + """Rescale covariance in barycenter step.""" + return weight * matrix_square_root.sqrtm_only((cov_sqrt @ cov) @ cov_sqrt, + **sqrtm_kw) + + def cond_fn(iteration: int, constants: Tuple[Any, ...], state) -> bool: + del constants + _, diffs = state + return diffs[iteration // inner_iterations] > tolerance + + def body_fn( + iteration: int, constants: Tuple[Any, ...], + state: Tuple[jnp.ndarray, float], compute_error: bool + ) -> Tuple[jnp.ndarray, float]: + del constants, compute_error + cov, diffs = state + cov_sqrt, cov_inv_sqrt, _ = matrix_square_root.sqrtm(cov, **sqrtm_kw) + scaled_cov = jnp.linalg.matrix_power( + jnp.sum(scale_covariances(cov_sqrt, covs, weights), axis=0), 2 + ) + next_cov = (cov_inv_sqrt @ scaled_cov) @ cov_inv_sqrt + diff = jnp.sum((next_cov - cov) ** 2) / jnp.prod(jnp.array(cov.shape)) + diffs = diffs.at[iteration // inner_iterations].set(diff) + return next_cov, diffs + + def init_state() -> Tuple[jnp.ndarray, float]: + cov_init = jnp.eye(self._dimension) + diffs = -jnp.ones(math.ceil(max_iterations / inner_iterations)) + return cov_init, diffs + + cov, diffs = fixed_point_loop.fixpoint_iter( + cond_fn=cond_fn, + body_fn=body_fn, + min_iterations=min_iterations, + max_iterations=max_iterations, + inner_iterations=inner_iterations, + constants=(), + state=init_state(), + ) + return cov, diffs + + def barycenter( + self, + weights: jnp.ndarray, + xs: jnp.ndarray, + tolerance: float = 1e-4, + sqrtm_kw: Optional[Dict[str, Any]] = None, + **kwargs: Any + ) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Compute the Bures barycenter of weighted Gaussian distributions. + + Implements the fixed point approach proposed in :cite:`alvarez-esteban:16` + for the computation of the mean and the covariance of the barycenter of + weighted Gaussian distributions. + + Args: + weights: The barycentric weights. + xs: The points to be used in the computation of the barycenter, where + each point is described by a concatenation of the mean and the + covariance (raveled). + tolerance: convergence tolerance to control the termination of the + algorithm. + sqrtm_kw: Arguments passed on to the + :func:`~ott.math.matrix_square_root.sqrtm` function used within + :meth:`covariance_fixpoint_iter`. This defines the precision + (in terms of convergence threshold, and number of iterations) of the + matrix square root calls that are used at each outer iteration of + the computation of Gaussian barycenters. These values are, by default, + the same as those used to define the Bures cost object itself. + kwargs: Passed on to :meth:`covariance_fixpoint_iter`, to specify the + number of iterations and tolerance of the fixed-point iteration of the + barycenter routine, by parameterizing `tolerance` and other relevant + arguments passed on to :func:`~ott.math.fixed_point_loop.fixpoint_iter`, + namely `min_iterations`, `max_iterations` and `inner_iterations`. + + Returns: + A list holding a concatenation of the mean and the raveled covariance + of the barycenter as its first element, followed by a vector of + norms of successive differences in iterates. + """ + # Ensure that barycentric weights sum to 1. + weights = weights / jnp.sum(weights) + mus, covs = x_to_means_and_covs(xs, self._dimension) + mu_bary = jnp.sum(weights[:, None] * mus, axis=0) + cov_bary, diffs = self.covariance_fixpoint_iter( + covs=covs, + weights=weights, + tolerance=tolerance, + sqrtm_kw=sqrtm_kw if sqrtm_kw is not None else self._sqrtm_kw, + **kwargs + ) + return mean_and_cov_to_x(mu_bary, cov_bary, self._dimension), diffs + + @classmethod + def _padder(cls, dim: int) -> jnp.ndarray: + dimension = int((-1 + math.sqrt(1 + 4 * dim)) / 2) + padding = mean_and_cov_to_x( + jnp.zeros((dimension,)), jnp.eye(dimension), dimension + ) + return padding[jnp.newaxis, :] + + def tree_flatten(self): # noqa: D102 + return (), (self._dimension, self._sqrtm_kw) + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + del children + return cls(*aux_data) + + +@jax.tree_util.register_pytree_node_class +class UnbalancedBures(CostFn): + """Unbalanced Bures distance between two triplets of `(mass, mean, cov)`. + + This cost uses the notation defined in :cite:`janati:20`, eq. 37, 39, 40. + + Args: + dimension: Dimensionality of the data. + sigma: Entropic regularization. + gamma: KL-divergence regularization for the marginals. + kwargs: Keyword arguments for :func:`~ott.math.matrix_square_root.sqrtm`. + """ + + def __init__( + self, + dimension: int, + *, + sigma: float = 1.0, + gamma: float = 1.0, + **kwargs: Any, + ): + super().__init__() + self._dimension = dimension + self._sigma = sigma + self._gamma = gamma + self._sqrtm_kw = kwargs + + def norm(self, x: jnp.ndarray) -> jnp.ndarray: + """Compute norm of Gaussian for unbalanced Bures. + + Args: + x: Array of shape ``[n_points + n_points + n_dim ** 2,]``, potentially + batched, corresponding to the raveled mass, means and the covariance + matrix. + + Returns: + The norm, array of shape ``[]`` or ``[batch,]`` in the batched case. + """ + return self._gamma * x[..., 0] + + def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + """Compute dot-product for unbalanced Bures. + + Args: + x: Array of shape ``[n_points + n_points + n_dim ** 2,]`` + corresponding to the raveled mass, means and the covariance matrix. + y: Array of shape ``[n_points + n_points + n_dim ** 2,]`` + corresponding to the raveled mass, means and the covariance matrix. + + Returns: + The cost. + """ + # Sets a few constants + gam = self._gamma + sig2 = self._sigma ** 2 + lam = sig2 + gam / 2.0 + tau = gam / (2.0 * lam) + + # Extracts mass, mean vector, covariance matrices + mass_x, mass_y = x[0], y[0] + mean_x, cov_x = x_to_means_and_covs(x[1:], self._dimension) + mean_y, cov_y = x_to_means_and_covs(y[1:], self._dimension) + + diff_means = mean_x - mean_y + + # Identity matrix of suitable size + iden = jnp.eye(self._dimension) + + # Creates matrices needed in the computation + tilde_a = 0.5 * gam * (iden - lam * jnp.linalg.inv(cov_x + lam * iden)) + tilde_b = 0.5 * gam * (iden - lam * jnp.linalg.inv(cov_y + lam * iden)) + + tilde_a_b = jnp.matmul(tilde_a, tilde_b) + c_mat = matrix_square_root.sqrtm( + 1 / tau * tilde_a_b + 0.25 * (sig2 ** 2) * iden, **self._sqrtm_kw + )[0] + c_mat -= 0.5 * sig2 * iden + + # Computes log determinants (their sign should be >0). + sldet_c, ldet_c = jnp.linalg.slogdet(c_mat) + sldet_t_ab, ldet_t_ab = jnp.linalg.slogdet(tilde_a_b) + sldet_ab, ldet_ab = jnp.linalg.slogdet(jnp.matmul(cov_x, cov_y)) + sldet_c_ab, ldet_c_ab = jnp.linalg.slogdet(c_mat - 2.0 * tilde_a_b / gam) + + # Gathers all these results to compute log total mass of transport + log_m_pi = (0.5 * self._dimension * sig2 / (gam + sig2)) * jnp.log(sig2) + log_m_pi += (1.0 / (tau + 1.0)) * ( + jnp.log(mass_x) + jnp.log(mass_y) + ldet_c + 0.5 * + (tau * ldet_t_ab - ldet_ab) + ) + log_m_pi += -jnp.sum( + diff_means * jnp.linalg.solve(cov_x + cov_y + lam * iden, diff_means) + ) / (2.0 * (tau + 1.0)) + log_m_pi += -0.5 * ldet_c_ab + + # if all logdet signs are 1, output value, nan otherwise + pos_signs = (sldet_c + sldet_c_ab + sldet_t_ab + sldet_t_ab) == 4 + + return jax.lax.cond( + pos_signs, lambda: 2 * sig2 * mass_x * mass_y - 2 * + (sig2 + gam) * jnp.exp(log_m_pi), lambda: jnp.nan + ) + + def tree_flatten(self): # noqa: D102 + return (), (self._dimension, self._sigma, self._gamma, self._sqrtm_kw) + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + del children + dim, sigma, gamma, kwargs = aux_data + return cls(dim, sigma=sigma, gamma=gamma, **kwargs) + + +@jax.tree_util.register_pytree_node_class +class SoftDTW(CostFn): + """Soft dynamic time warping (DTW) cost :cite:`cuturi:17`. + + Args: + gamma: Smoothing parameter :math:`> 0` for the soft-min operator. + ground_cost: Ground cost function. If ``None``, + use :class:`~ott.geometry.costs.SqEuclidean`. + debiased: Whether to compute the debiased soft-DTW :cite:`blondel:21`. + """ + + def __init__( + self, + gamma: float, + ground_cost: Optional[CostFn] = None, + debiased: bool = False + ): + self.gamma = gamma + self.ground_cost = SqEuclidean() if ground_cost is None else ground_cost + self.debiased = debiased + + def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: # noqa: D102 + c_xy = self._soft_dtw(x, y) + if self.debiased: + return c_xy - 0.5 * (self._soft_dtw(x, x) + self._soft_dtw(y, y)) + return c_xy + + def _soft_dtw(self, t1: jnp.ndarray, t2: jnp.ndarray) -> float: + + def body( + carry: Tuple[jnp.ndarray, jnp.ndarray], + current_antidiagonal: jnp.ndarray + ) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], jnp.ndarray]: + # modified from: https://github.com/khdlr/softdtw_jax + two_ago, one_ago = carry + + diagonal, right, down = two_ago[:-1], one_ago[:-1], one_ago[1:] + best = mu.softmin( + jnp.stack([diagonal, right, down], axis=-1), self.gamma, axis=-1 + ) + + next_row = best + current_antidiagonal + next_row = jnp.pad(next_row, (1, 0), constant_values=jnp.inf) + + return (one_ago, next_row), next_row + + t1 = t1[:, None] if t1.ndim == 1 else t1 + t2 = t2[:, None] if t2.ndim == 1 else t2 + dist = self.ground_cost.all_pairs(t1, t2) + + n, m = dist.shape + if n < m: + dist = dist.T + n, m = m, n + + model_matrix = jnp.full((n + m - 1, n), fill_value=jnp.inf) + mask = np.tri(n + m - 1, n, k=0, dtype=bool) + mask = mask & mask[::-1, ::-1] + model_matrix = model_matrix.T.at[mask.T].set(dist.ravel()).T + + init = ( + jnp.pad(model_matrix[0], (1, 0), constant_values=jnp.inf), + jnp.pad( + model_matrix[1] + model_matrix[0, 0], (1, 0), + constant_values=jnp.inf + ) + ) + + (_, carry), _ = jax.lax.scan(body, init, model_matrix[2:]) + return carry[-1] + + def tree_flatten(self): # noqa: D102 + return (self.gamma, self.ground_cost), {"debiased": self.debiased} + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + return cls(*children, **aux_data) + + +def x_to_means_and_covs(x: jnp.ndarray, + dimension: int) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Extract means and covariance matrices of Gaussians from raveled vector. + + Args: + x: [num_gaussians, dimension, (1 + dimension)] array of concatenated means + and covariances (raveled) dimension: the dimension of the Gaussians. + dimension: Dimensionality of the Gaussians. + + Returns: + Means and covariances of shape ``[num_gaussian, dimension]``. + """ + x = jnp.atleast_2d(x) + means = x[:, :dimension] + covariances = jnp.reshape( + x[:, dimension:dimension + dimension ** 2], (-1, dimension, dimension) + ) + return jnp.squeeze(means), jnp.squeeze(covariances) + + +def mean_and_cov_to_x( + mean: jnp.ndarray, covariance: jnp.ndarray, dimension: int +) -> jnp.ndarray: + """Ravel a Gaussian's mean and covariance matrix to d(1 + d) vector.""" + return jnp.concatenate( + (mean, jnp.reshape(covariance, (dimension * dimension))) + ) diff --git a/ott/build/lib/ott/geometry/distrib_costs.py b/ott/build/lib/ott/geometry/distrib_costs.py new file mode 100644 index 0000000..3ddc01a --- /dev/null +++ b/ott/build/lib/ott/geometry/distrib_costs.py @@ -0,0 +1,89 @@ +# 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. +from typing import Any, Optional + +import jax.numpy as jnp +import jax.tree_util as jtu + +from ott.geometry import costs, pointcloud +from ott.problems.linear import linear_problem +from ott.solvers.linear import univariate + +__all__ = [ + "UnivariateWasserstein", +] + + +@jtu.register_pytree_node_class +class UnivariateWasserstein(costs.CostFn): + """1D Wasserstein cost for two 1D distributions. + + This ground cost between considers vectors as a family of values. + The Wasserstein distance between them is the 1D OT cost, using a user-defined + ground cost. + + Args: + ground_cost: Cost used to compute the 1D optimal transport between vector, + should be a translation-invariant (TI) cost for correctness. + If :obj:`None`, defaults to :class:`~ott.geometry.costs.SqEuclidean`. + solver: 1D optimal transport solver. + kwargs: Arguments passed on when calling the + :class:`~ott.solvers.linear.univariate.UnivariateSolver`. May include + random key, or specific instructions to subsample or compute using + quantiles. + """ + + def __init__( + self, + ground_cost: Optional[costs.TICost] = None, + solver: Optional[univariate.UnivariateSolver] = None, + **kwargs: Any + ): + super().__init__() + + self.ground_cost = ( + costs.SqEuclidean() if ground_cost is None else ground_cost + ) + self._solver = univariate.UnivariateSolver() if solver is None else solver + self._kwargs_solve = kwargs + # ensure transport solutions are neither computed nor stored + self._kwargs_solve["return_transport"] = False + + def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + """Wasserstein distance between :math:`x` and :math:`y` seen as a 1D dist. + + Args: + x: Array of shape ``[n,]``. + y: Array of shape ``[m,]``. + + Returns: + The transport cost. + """ + out = self._solver( + linear_problem.LinearProblem( + pointcloud.PointCloud( + x[:, None], y[:, None], cost_fn=self.ground_cost + ) + ), **self._kwargs_solve + ) + return jnp.squeeze(out.ot_costs) + + def tree_flatten(self): # noqa: D102 + return (self.ground_cost,), (self._solver, self._kwargs_solve) + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + ground_cost, = children + solver, solve_kwargs = aux_data + return cls(ground_cost, solver, **solve_kwargs) diff --git a/ott/build/lib/ott/geometry/epsilon_scheduler.py b/ott/build/lib/ott/geometry/epsilon_scheduler.py new file mode 100644 index 0000000..a5cca8f --- /dev/null +++ b/ott/build/lib/ott/geometry/epsilon_scheduler.py @@ -0,0 +1,102 @@ +# 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. +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +__all__ = ["Epsilon"] + + +@jax.tree_util.register_pytree_node_class +class Epsilon: + """Scheduler class for the regularization parameter epsilon. + + An epsilon scheduler outputs a regularization strength, to be used by in a + Sinkhorn-type algorithm, at any iteration count. That value is either the + final, targeted regularization, or one that is larger, obtained by + geometric decay of an initial value that is larger than the intended target. + Concretely, the value returned by such a scheduler will consider first + the max between ``target`` and ``init * target * decay ** iteration``. + If the ``scale_epsilon`` parameter is provided, that value is used to + multiply the max computed previously by ``scale_epsilon``. + + Args: + target: the epsilon regularizer that is targeted. + If ``None``, use :math:`0.05`. + scale_epsilon: if passed, used to multiply the regularizer, to rescale it. + If ``None``, use :math:`1`. + init: initial value when using epsilon scheduling, understood as multiple + of target value. if passed, ``int * decay ** iteration`` will be used + to rescale target. + decay: geometric decay factor, :math:`<1`. + """ + + def __init__( + self, + target: Optional[float] = None, + scale_epsilon: Optional[float] = None, + init: float = 1.0, + decay: float = 1.0 + ): + self._target_init = target + self._scale_epsilon = scale_epsilon + self._init = init + self._decay = decay + + @property + def target(self) -> float: + """Return the final regularizer value of scheduler.""" + target = 5e-2 if self._target_init is None else self._target_init + scale = 1.0 if self._scale_epsilon is None else self._scale_epsilon + return scale * target + + def at(self, iteration: Optional[int] = 1) -> float: + """Return (intermediate) regularizer value at a given iteration.""" + if iteration is None: + return self.target + # check the decay is smaller than 1.0. + decay = jnp.minimum(self._decay, 1.0) + # the multiple is either 1.0 or a larger init value that is decayed. + multiple = jnp.maximum(self._init * (decay ** iteration), 1.0) + return multiple * self.target + + def done(self, eps: float) -> bool: + """Return whether the scheduler is done at a given value.""" + return eps == self.target + + def done_at(self, iteration: Optional[int]) -> bool: + """Return whether the scheduler is done at a given iteration.""" + return self.done(self.at(iteration)) + + def set(self, **kwargs: Any) -> "Epsilon": + """Return a copy of self, with potential overwrites.""" + kwargs = { + "target": self._target_init, + "scale_epsilon": self._scale_epsilon, + "init": self._init, + "decay": self._decay, + **kwargs + } + return Epsilon(**kwargs) + + def tree_flatten(self): # noqa: D102 + return ( + self._target_init, self._scale_epsilon, self._init, self._decay + ), None + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + del aux_data + return cls(*children) diff --git a/ott/build/lib/ott/geometry/geodesic.py b/ott/build/lib/ott/geometry/geodesic.py new file mode 100644 index 0000000..9375b23 --- /dev/null +++ b/ott/build/lib/ott/geometry/geodesic.py @@ -0,0 +1,277 @@ +# 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. +from typing import Any, Dict, Optional, Sequence, Tuple + +import jax +import jax.experimental.sparse as jesp +import jax.numpy as jnp +import numpy as np +from scipy.special import ive + +from ott import utils +from ott.geometry import geometry +from ott.math import utils as mu +from ott.types import Array_g + +__all__ = ["Geodesic"] + + +@jax.tree_util.register_pytree_node_class +class Geodesic(geometry.Geometry): + r"""Graph distance approximation using heat kernel :cite:`huguet:2023`. + + .. note:: + This constructor is not meant to be called by the user, + please use the :meth:`from_graph` method instead. + + Approximates the heat kernel using + `Chebyshev polynomials `_ + of the first kind of max order ``order``, which for small ``t`` + approximates the geodesic exponential kernel :math:`e^{\frac{-d(x, y)^2}{t}}`. + + Args: + scaled_laplacian: The Laplacian scaled by the largest eigenvalue. + eigval: The largest eigenvalue of the Laplacian. + chebyshev_coeffs: Coefficients of the Chebyshev polynomials. + t: Time parameter for the heat kernel. + kwargs: Keyword arguments for :class:`~ott.geometry.geometry.Geometry`. + """ + + def __init__( + self, + scaled_laplacian: Array_g, + eigval: jnp.ndarray, + chebyshev_coeffs: jnp.ndarray, + t: float = 1e-3, + **kwargs: Any + ): + super().__init__(epsilon=1.0, **kwargs) + self.scaled_laplacian = scaled_laplacian + self.eigval = eigval + self.chebyshev_coeffs = chebyshev_coeffs + self.t = t + + @classmethod + def from_graph( + cls, + G: Array_g, + t: Optional[float] = 1e-3, + eigval: Optional[jnp.ndarray] = None, + order: int = 100, + directed: bool = False, + normalize: bool = False, + rng: Optional[jax.Array] = None, + **kwargs: Any + ) -> "Geodesic": + r"""Construct a Geodesic geometry from an adjacency matrix. + + Args: + G: Adjacency matrix. + t: Time parameter for approximating the geodesic exponential kernel. + If `None`, it defaults to :math:`\frac{1}{|E|} \sum_{(u, v) \in E} + \text{weight}(u, v)` :cite:`crane:13`. In this case, the ``graph`` + must be specified and the edge weights are assumed to be positive. + eigval: Largest eigenvalue of the Laplacian. If :obj:`None`, it's + computed using :func:`jax.experimental.sparse.linalg.lobpcg_standard`. + order: Max order of Chebyshev polynomials. + directed: Whether the ``graph`` is directed. If :obj:`True`, it's made + undirected as :math:`G + G^T`. This parameter is ignored when passing + the Laplacian directly, assumed to be symmetric. + normalize: Whether to normalize the Laplacian as + :math:`L^{sym} = \left(D^+\right)^{\frac{1}{2}} L + \left(D^+\right)^{\frac{1}{2}}`, where :math:`L` is the + non-normalized Laplacian and :math:`D` is the degree matrix. + rng: Random key used when computing the largest eigenvalue. + kwargs: Keyword arguments for :class:`~ott.geometry.geodesic.Geodesic`. + + Returns: + The Geodesic geometry. + """ + assert G.shape[0] == G.shape[1], G.shape + rng = utils.default_prng_key(rng) + + if directed: + G = G + G.T + if t is None: + t = (jnp.sum(G) / jnp.sum(G > 0.0)) ** 2 + + degree = jnp.sum(G, axis=1) + laplacian = jnp.diag(degree) - G + if normalize: + inv_sqrt_deg = jnp.diag( + jnp.where(degree > 0.0, 1.0 / jnp.sqrt(degree), 0.0) + ) + laplacian = inv_sqrt_deg @ laplacian @ inv_sqrt_deg + + if eigval is None: + eigval = compute_largest_eigenvalue(laplacian, rng) + + scaled_laplacian, eigval = jax.lax.cond((eigval > 2.0), lambda l: + (2.0 * l / eigval, 2.0), lambda l: + (l, eigval), laplacian) + + # compute the coeffs of the Chebyshev pols approx using Bessel funcs + chebyshev_coeffs = compute_chebychev_coeff_all( + 0.5 * eigval, t, order, laplacian.dtype + ) + + return cls( + scaled_laplacian=scaled_laplacian, + eigval=eigval, + chebyshev_coeffs=chebyshev_coeffs, + t=t, + **kwargs + ) + + def apply_kernel( + self, + scaling: jnp.ndarray, + eps: Optional[float] = None, + axis: int = 0, + ) -> jnp.ndarray: + r"""Apply :attr:`kernel_matrix` on positive scaling vector. + + Args: + scaling: Scaling to apply the kernel to. + eps: passed for consistency, not used yet. + axis: passed for consistency, not used yet. + + Returns: + Kernel applied to ``scaling``. + """ + return expm_multiply( + self.scaled_laplacian, scaling, self.chebyshev_coeffs, 0.5 * self.eigval + ) + + @property + def kernel_matrix(self) -> jnp.ndarray: # noqa: D102 + n, _ = self.shape + kernel = self.apply_kernel(jnp.eye(n)) + return jax.lax.cond( + jnp.allclose(kernel, kernel.T, atol=1e-8, rtol=1e-8), lambda x: x, + lambda x: (x + x.T) / 2.0, kernel + ) + + @property + def cost_matrix(self) -> jnp.ndarray: # noqa: D102 + # Calculate the cost matrix using the formula (5) from the main reference + return -4.0 * self.t * mu.safe_log(self.kernel_matrix) + + @property + def shape(self) -> Tuple[int, int]: # noqa: D102 + return self.scaled_laplacian.shape + + @property + def is_symmetric(self) -> bool: # noqa: D102 + return True + + @property + def dtype(self) -> jnp.dtype: # noqa: D102 + return self.scaled_laplacian.dtype + + def transport_from_potentials( + self, f: jnp.ndarray, g: jnp.ndarray + ) -> jnp.ndarray: + """Not implemented.""" + raise ValueError("Not implemented.") + + def apply_transport_from_potentials( + self, + f: jnp.ndarray, + g: jnp.ndarray, + vec: jnp.ndarray, + axis: int = 0 + ) -> jnp.ndarray: + """Not implemented.""" + raise ValueError("Not implemented.") + + def marginal_from_potentials( + self, + f: jnp.ndarray, + g: jnp.ndarray, + axis: int = 0, + ) -> jnp.ndarray: + """Not implemented.""" + raise ValueError("Not implemented.") + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return [ + self.scaled_laplacian, + self.eigval, + self.chebyshev_coeffs, + self.t, + ], {} + + @classmethod + def tree_unflatten( # noqa: D102 + cls, aux_data: Dict[str, Any], children: Sequence[Any] + ) -> "Geodesic": + return cls(*children, **aux_data) + + +def compute_largest_eigenvalue( + laplacian_matrix: jnp.ndarray, + rng: jax.Array, +) -> float: + # Compute the largest eigenvalue of the Laplacian matrix. + n = laplacian_matrix.shape[0] + # Generate random initial directions for eigenvalue computation + initial_dirs = jax.random.normal(rng, (n, 1)) + + # Create a sparse matrix-vector product function using sparsify + # This function multiplies the sparse laplacian_matrix with a vector + lapl_vector_product = jesp.sparsify(lambda v: laplacian_matrix @ v) + + # Compute eigenvalues using the sparse matrix-vector product + eigvals, _, _ = jesp.linalg.lobpcg_standard( + lapl_vector_product, + initial_dirs, + ) + return eigvals[0] + + +def expm_multiply( + L: jnp.ndarray, X: jnp.ndarray, coeff: jnp.ndarray, eigval: float +) -> jnp.ndarray: + + def body(carry, c): + T0, T1, Y = carry + T2 = (2.0 / eigval) * L @ T1 - 2.0 * T1 - T0 + Y = Y + c * T2 + return (T1, T2, Y), None + + T0 = X + Y = 0.5 * coeff[0] * T0 + T1 = (1.0 / eigval) * L @ X - T0 + Y = Y + coeff[1] * T1 + + initial_state = (T0, T1, Y) + (_, _, Y), _ = jax.lax.scan(body, initial_state, coeff[2:]) + return Y + + +def compute_chebychev_coeff_all( + eigval: float, tau: float, K: int, dtype: np.dtype +) -> jnp.ndarray: + """Jax wrapper to compute the K+1 Chebychev coefficients.""" + result_shape_dtype = jax.ShapeDtypeStruct( + shape=(K + 1,), + dtype=dtype, + ) + + chebychev_coeff = lambda eigval, tau, K: ( + 2.0 * ive(np.arange(0, K + 1), -tau * eigval) + ).astype(dtype) + + return jax.pure_callback(chebychev_coeff, result_shape_dtype, eigval, tau, K) diff --git a/ott/build/lib/ott/geometry/geometry.py b/ott/build/lib/ott/geometry/geometry.py new file mode 100644 index 0000000..bcbc724 --- /dev/null +++ b/ott/build/lib/ott/geometry/geometry.py @@ -0,0 +1,921 @@ +# 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 functools +from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Tuple, Union + +if TYPE_CHECKING: + from ott.geometry import low_rank + +import jax +import jax.numpy as jnp +import jax.scipy as jsp +import numpy as np + +from ott import utils +from ott.geometry import epsilon_scheduler +from ott.math import utils as mu + +__all__ = ["Geometry", "is_linear", "is_affine"] + + +@jax.tree_util.register_pytree_node_class +class Geometry: + r"""Base class to define ground costs/kernels used in optimal transport. + + Optimal transport problems are intrinsically geometric: they compute an + optimal way to transport mass from one configuration onto another. To define + what is meant by optimality of transport requires defining a cost, of moving + mass from one among several sources, towards one out of multiple targets. + These sources and targets can be provided as points in vectors spaces, grids, + or more generally exclusively described through a (dissimilarity) cost matrix, + or almost equivalently, a (similarity) kernel matrix. + + Once that cost or kernel matrix is set, the ``Geometry`` class provides a + basic operations to be run with the Sinkhorn algorithm. + + Args: + cost_matrix: Cost matrix of shape ``[n, m]``. + kernel_matrix: Kernel matrix of shape ``[n, m]``. + epsilon: Regularization parameter. If ``None`` and either + ``relative_epsilon = True`` or ``relative_epsilon = None``, this defaults + to the value computed in :attr:`mean_cost_matrix` / 20. If passed as a + ``float``, then the regularizer that is ultimately used is either that + ``float`` value (if ``relative_epsilon = False`` or ``None``) or that + ``float`` times the :attr:`mean_cost_matrix` + (if ``relative_epsilon = True``). Look for + :class:`~ott.geometry.epsilon_scheduler.Epsilon` when passed as a + scheduler. + relative_epsilon: when `False`, the parameter ``epsilon`` specifies the + value of the entropic regularization parameter. When `True`, ``epsilon`` + refers to a fraction of the :attr:`mean_cost_matrix`, which is computed + adaptively from data. + scale_cost: option to rescale the cost matrix. Implemented scalings are + 'median', 'mean' and 'max_cost'. Alternatively, a float factor can be + given to rescale the cost such that ``cost_matrix /= scale_cost``. + If `True`, use 'mean'. + src_mask: Mask specifying valid rows when computing some statistics of + :attr:`cost_matrix`, see :attr:`src_mask`. + tgt_mask: Mask specifying valid columns when computing some statistics of + :attr:`cost_matrix`, see :attr:`tgt_mask`. + + Note: + When defining a :class:`~ott.geometry.geometry.Geometry` through a + ``cost_matrix``, it is important to select an ``epsilon`` regularization + parameter that is meaningful. That parameter can be provided by the user, + or assigned a default value through a simple rule, + using the :attr:`mean_cost_matrix`. + """ + + def __init__( + self, + cost_matrix: Optional[jnp.ndarray] = None, + kernel_matrix: Optional[jnp.ndarray] = None, + epsilon: Optional[Union[float, epsilon_scheduler.Epsilon]] = None, + relative_epsilon: Optional[bool] = None, + scale_cost: Union[bool, int, float, Literal["mean", "max_cost", + "median"]] = 1.0, + src_mask: Optional[jnp.ndarray] = None, + tgt_mask: Optional[jnp.ndarray] = None, + ): + self._cost_matrix = cost_matrix + self._kernel_matrix = kernel_matrix + + # needed for `copy_epsilon`, because of the `isinstance` check + self._epsilon_init = epsilon if isinstance( + epsilon, epsilon_scheduler.Epsilon + ) else epsilon_scheduler.Epsilon(epsilon) + self._relative_epsilon = relative_epsilon + + self._scale_cost = "mean" if scale_cost is True else scale_cost + + self._src_mask = src_mask + self._tgt_mask = tgt_mask + + @property + def cost_rank(self) -> Optional[int]: + """Output rank of cost matrix, if any was provided.""" + + @property + def cost_matrix(self) -> jnp.ndarray: + """Cost matrix, recomputed from kernel if only kernel was specified.""" + if self._cost_matrix is None: + # If no epsilon was passed on to the geometry, then assume it is one by + # default. + eps = jnp.finfo(self._kernel_matrix.dtype).tiny + cost = -jnp.log(self._kernel_matrix + eps) + cost *= self.inv_scale_cost + return cost if self._epsilon_init is None else self.epsilon * cost + return self._cost_matrix * self.inv_scale_cost + + @property + def median_cost_matrix(self) -> float: + """Median of the :attr:`cost_matrix`.""" + geom = self._masked_geom(mask_value=jnp.nan) + return jnp.nanmedian(geom.cost_matrix) # will fail for online PC + + @property + def mean_cost_matrix(self) -> float: + """Mean of the :attr:`cost_matrix`.""" + tmp = self._masked_geom().apply_cost(self._n_normed_ones).squeeze() + return jnp.sum(tmp * self._m_normed_ones) + + @property + def kernel_matrix(self) -> jnp.ndarray: + """Kernel matrix. + + Either provided by user or recomputed from :attr:`cost_matrix`. + """ + if self._kernel_matrix is None: + return jnp.exp(-(self._cost_matrix * self.inv_scale_cost / self.epsilon)) + return self._kernel_matrix ** self.inv_scale_cost + + @property + def _epsilon(self) -> epsilon_scheduler.Epsilon: + (target, scale_eps, _, _), _ = self._epsilon_init.tree_flatten() + rel = self._relative_epsilon + + use_mean_scale = rel is True or (rel is None and target is None) + if scale_eps is None and use_mean_scale: + scale_eps = jax.lax.stop_gradient(self.mean_cost_matrix) + + if isinstance(self._epsilon_init, epsilon_scheduler.Epsilon): + return self._epsilon_init.set(scale_epsilon=scale_eps) + + return epsilon_scheduler.Epsilon( + target=5e-2 if target is None else target, scale_epsilon=scale_eps + ) + + @property + def epsilon(self) -> float: + """Epsilon regularization value.""" + return self._epsilon.target + + @property + def shape(self) -> Tuple[int, int]: + """Shape of the geometry.""" + mat = ( + self._kernel_matrix if self._cost_matrix is None else self._cost_matrix + ) + if mat is not None: + return mat.shape + return 0, 0 + + @property + def can_LRC(self) -> bool: + """Check quickly if casting geometry as LRC makes sense. + + This check is only carried out using basic considerations from the geometry, + not using a rigorous check involving, e.g., SVD. + """ + return False + + @property + def is_squared_euclidean(self) -> bool: + """Whether cost is computed by taking squared Euclidean distance.""" + return False + + @property + def is_online(self) -> bool: + """Whether geometry cost/kernel should be recomputed on the fly.""" + return False + + @property + def is_symmetric(self) -> bool: + """Whether geometry cost/kernel is a symmetric matrix.""" + mat = self.kernel_matrix if self.cost_matrix is None else self.cost_matrix + return ( + mat.shape[0] == mat.shape[1] and jnp.all(mat == mat.T) + ) if mat is not None else False + + @property + def inv_scale_cost(self) -> float: + """Compute and return inverse of scaling factor for cost matrix.""" + if isinstance(self._scale_cost, (int, float, np.number, jax.Array)): + return 1.0 / self._scale_cost + self = self._masked_geom(mask_value=jnp.nan) + if self._scale_cost == "max_cost": + return 1.0 / jnp.nanmax(self._cost_matrix) + if self._scale_cost == "mean": + return 1.0 / jnp.nanmean(self._cost_matrix) + if self._scale_cost == "median": + return 1.0 / jnp.nanmedian(self._cost_matrix) + raise ValueError(f"Scaling {self._scale_cost} not implemented.") + + def set_scale_cost(self, scale_cost: Union[bool, float, str]) -> "Geometry": + """Modify how to rescale of the :attr:`cost_matrix`.""" + # case when `geom` doesn't have `scale_cost` or doesn't need to be modified + # `False` retains the original scale + if scale_cost is False or scale_cost == self._scale_cost: + return self + children, aux_data = self.tree_flatten() + aux_data["scale_cost"] = scale_cost + return type(self).tree_unflatten(aux_data, children) + + def copy_epsilon(self, other: "Geometry") -> "Geometry": + """Copy the epsilon parameters from another geometry.""" + other_epsilon = other._epsilon + children, aux_data = self.tree_flatten() + + new_children = [] + for child in children: + if isinstance(child, epsilon_scheduler.Epsilon): + child = child.set( + target=other_epsilon._target_init, + scale_epsilon=other_epsilon._scale_epsilon + ) + new_children.append(child) + + aux_data["relative_epsilon"] = False + return type(self).tree_unflatten(aux_data, new_children) + + # The functions below are at the core of Sinkhorn iterations, they + # are implemented here in their default form, either in lse (using directly + # cost matrices in stabilized form) or kernel mode (using kernel matrices). + + def apply_lse_kernel( + self, + f: jnp.ndarray, + g: jnp.ndarray, + eps: float, + vec: jnp.ndarray = None, + axis: int = 0 + ) -> jnp.ndarray: + r"""Apply :attr:`kernel_matrix` in log domain. + + This function applies the ground geometry's kernel in log domain, using + a stabilized formulation. At a high level, this iteration performs either: + + - output = eps * log (K (exp(g / eps) * vec)) (1) + - output = eps * log (K'(exp(f / eps) * vec)) (2) + + K is implicitly exp(-cost_matrix/eps). + + To carry this out in a stabilized way, we take advantage of the fact that + the entries of the matrix ``f[:,*] + g[*,:] - C`` are all negative, and + therefore their exponential never overflows, to add (and subtract after) + f and g in iterations 1 & 2 respectively. + + Args: + f: jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix + g: jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix + eps: float, regularization strength + vec: jnp.ndarray [num_a or num_b,] , when not None, this has the effect of + doing log-Kernel computations with an addition elementwise + multiplication of exp(g / eps) by a vector. This is carried out by + adding weights to the log-sum-exp function, and needs to handle signs + separately. + axis: summing over axis 0 when doing (2), or over axis 1 when doing (1) + + Returns: + A jnp.ndarray corresponding to output above, depending on axis. + """ + w_res, w_sgn = self._softmax(f, g, eps, vec, axis) + remove = f if axis == 1 else g + return w_res - jnp.where(jnp.isfinite(remove), remove, 0), w_sgn + + def apply_kernel( + self, + scaling: jnp.ndarray, + eps: Optional[float] = None, + axis: int = 0, + ) -> jnp.ndarray: + """Apply :attr:`kernel_matrix` on positive scaling vector. + + Args: + scaling: jnp.ndarray [num_a or num_b] , scaling of size num_rows or + num_cols of kernel_matrix + eps: passed for consistency, not used yet. + axis: standard kernel product if axis is 1, transpose if 0. + + Returns: + a jnp.ndarray corresponding to output above, depending on axis. + """ + if eps is None: + kernel = self.kernel_matrix + else: + kernel = self.kernel_matrix ** (self.epsilon / eps) + kernel = kernel if axis == 1 else kernel.T + + return jnp.dot(kernel, scaling) + + def marginal_from_potentials( + self, + f: jnp.ndarray, + g: jnp.ndarray, + axis: int = 0, + ) -> jnp.ndarray: + """Output marginal of transportation matrix from potentials. + + This applies first lse kernel in the standard way, removes the + correction used to stabilize computations, and lifts this with an exp to + recover either of the marginals corresponding to the transport map induced + by potentials. + + Args: + f: jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix + g: jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix + axis: axis along which to integrate, returns marginal on other axis. + + Returns: + a vector of marginals of the transport matrix. + """ + h = (f if axis == 1 else g) + z = self.apply_lse_kernel(f, g, self.epsilon, axis=axis)[0] + return jnp.exp((z + h) / self.epsilon) + + def marginal_from_scalings( + self, + u: jnp.ndarray, + v: jnp.ndarray, + axis: int = 0, + ) -> jnp.ndarray: + """Output marginal of transportation matrix from scalings.""" + u, v = (v, u) if axis == 0 else (u, v) + return u * self.apply_kernel(v, eps=self.epsilon, axis=axis) + + def transport_from_potentials( + self, f: jnp.ndarray, g: jnp.ndarray + ) -> jnp.ndarray: + """Output transport matrix from potentials.""" + return jnp.exp(self._center(f, g) / self.epsilon) + + def transport_from_scalings( + self, u: jnp.ndarray, v: jnp.ndarray + ) -> jnp.ndarray: + """Output transport matrix from pair of scalings.""" + return self.kernel_matrix * u[:, jnp.newaxis] * v[jnp.newaxis, :] + + # Functions that are not supposed to be changed by inherited classes. + # These are the point of entry for Sinkhorn's algorithm to use a geometry. + + def update_potential( + self, + f: jnp.ndarray, + g: jnp.ndarray, + log_marginal: jnp.ndarray, + iteration: Optional[int] = None, + axis: int = 0, + ) -> jnp.ndarray: + """Carry out one Sinkhorn update for potentials, i.e. in log space. + + Args: + f: jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix + g: jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix + log_marginal: targeted marginal + iteration: used to compute epsilon from schedule, if provided. + axis: axis along which the update should be carried out. + + Returns: + new potential value, g if axis=0, f if axis is 1. + """ + eps = self._epsilon.at(iteration) + app_lse = self.apply_lse_kernel(f, g, eps, axis=axis)[0] + return eps * log_marginal - jnp.where(jnp.isfinite(app_lse), app_lse, 0) + + def update_scaling( + self, + scaling: jnp.ndarray, + marginal: jnp.ndarray, + iteration: Optional[int] = None, + axis: int = 0, + ) -> jnp.ndarray: + """Carry out one Sinkhorn update for scalings, using kernel directly. + + Args: + scaling: jnp.ndarray of num_a or num_b positive values. + marginal: targeted marginal + iteration: used to compute epsilon from schedule, if provided. + axis: axis along which the update should be carried out. + + Returns: + new scaling vector, of size num_b if axis=0, num_a if axis is 1. + """ + eps = self._epsilon.at(iteration) + app_kernel = self.apply_kernel(scaling, eps, axis=axis) + return marginal / jnp.where(app_kernel > 0, app_kernel, 1.0) + + # Helper functions + def _center(self, f: jnp.ndarray, g: jnp.ndarray) -> jnp.ndarray: + return f[:, jnp.newaxis] + g[jnp.newaxis, :] - self.cost_matrix + + def _softmax( + self, f: jnp.ndarray, g: jnp.ndarray, eps: float, + vec: Optional[jnp.ndarray], axis: int + ) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Apply softmax row or column wise, weighted by vec.""" + if vec is not None: + if axis == 0: + vec = vec.reshape((-1, 1)) + lse_output = mu.logsumexp( + self._center(f, g) / eps, b=vec, axis=axis, return_sign=True + ) + return eps * lse_output[0], lse_output[1] + lse_output = mu.logsumexp( + self._center(f, g) / eps, axis=axis, return_sign=False + ) + return eps * lse_output, jnp.array([1.0]) + + @functools.partial(jax.vmap, in_axes=[None, None, None, 0, None]) + def _apply_transport_from_potentials( + self, f: jnp.ndarray, g: jnp.ndarray, vec: jnp.ndarray, axis: int + ) -> jnp.ndarray: + """Apply lse_kernel to arbitrary vector while keeping track of signs.""" + lse_res, lse_sgn = self.apply_lse_kernel( + f, g, self.epsilon, vec=vec, axis=axis + ) + lse_res += f if axis == 1 else g + return lse_sgn * jnp.exp(lse_res / self.epsilon) + + # wrapper to allow default option for axis. + def apply_transport_from_potentials( + self, + f: jnp.ndarray, + g: jnp.ndarray, + vec: jnp.ndarray, + axis: int = 0 + ) -> jnp.ndarray: + """Apply transport matrix computed from potentials to a (batched) vec. + + This approach does not instantiate the transport matrix itself, but uses + instead potentials to apply the transport using apply_lse_kernel, therefore + guaranteeing stability and lower memory footprint. + + Computations are done in log space, and take advantage of the + (b=..., return_sign=True) optional parameters of logsumexp. + + Args: + f: jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix + g: jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix + vec: jnp.ndarray [batch, num_a or num_b], vector that will be multiplied + by transport matrix corresponding to potentials f, g, and geom. + axis: axis to differentiate left (0) or right (1) multiply. + + Returns: + ndarray of the size of vec. + """ + if vec.ndim == 1: + return self._apply_transport_from_potentials( + f, g, vec[jnp.newaxis, :], axis + )[0, :] + return self._apply_transport_from_potentials(f, g, vec, axis) + + @functools.partial(jax.vmap, in_axes=[None, None, None, 0, None]) + def _apply_transport_from_scalings( + self, u: jnp.ndarray, v: jnp.ndarray, vec: jnp.ndarray, axis: int + ): + u, v = (u, v * vec) if axis == 1 else (v, u * vec) + return u * self.apply_kernel(v, eps=self.epsilon, axis=axis) + + # wrapper to allow default option for axis + def apply_transport_from_scalings( + self, + u: jnp.ndarray, + v: jnp.ndarray, + vec: jnp.ndarray, + axis: int = 0 + ) -> jnp.ndarray: + """Apply transport matrix computed from scalings to a (batched) vec. + + This approach does not instantiate the transport matrix itself, but + relies instead on the apply_kernel function. + + Args: + u: jnp.ndarray [num_a,] , scaling of size num_rows of cost_matrix + v: jnp.ndarray [num_b,] , scaling of size num_cols of cost_matrix + vec: jnp.ndarray [batch, num_a or num_b], vector that will be multiplied + by transport matrix corresponding to scalings u, v, and geom. + axis: axis to differentiate left (0) or right (1) multiply. + + Returns: + ndarray of the size of vec. + """ + if vec.ndim == 1: + return self._apply_transport_from_scalings( + u, v, vec[jnp.newaxis, :], axis + )[0, :] + return self._apply_transport_from_scalings(u, v, vec, axis) + + def potential_from_scaling(self, scaling: jnp.ndarray) -> jnp.ndarray: + """Compute dual potential vector from scaling vector. + + Args: + scaling: vector. + + Returns: + a vector of the same size. + """ + return self.epsilon * jnp.log(scaling) + + def scaling_from_potential(self, potential: jnp.ndarray) -> jnp.ndarray: + """Compute scaling vector from dual potential. + + Args: + potential: vector. + + Returns: + a vector of the same size. + """ + finite = jnp.isfinite(potential) + return jnp.where( + finite, jnp.exp(jnp.where(finite, potential / self.epsilon, 0.0)), 0.0 + ) + + def apply_square_cost(self, arr: jnp.ndarray, axis: int = 0, sub_idx = None) -> jnp.ndarray: + """Apply elementwise-square of cost matrix to array (vector or matrix). + + This function applies the ground 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], when the cost matrix itself is computed as a + squared-Euclidean distance between vectors, and therefore admits an + explicit low-rank factorization. + + Args: + arr: array. + axis: axis of the array on which the cost matrix should be applied. + + Returns: + An array, [num_b, p] if axis=0 or [num_a, p] if axis=1. + """ + return self.apply_cost(arr, axis=axis, fn=lambda x: x ** 2, sub_idx = sub_idx) + + def apply_cost( + self, + arr: jnp.ndarray, + axis: int = 0, + fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, + sub_idx = None, + **kwargs: Any + ) -> jnp.ndarray: + """Apply :attr:`cost_matrix` to array (vector or matrix). + + This function applies the ground 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] + + 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, transpose if 0 + fn: function to apply to cost matrix element-wise before the dot product + kwargs: Keyword arguments for :meth:`_apply_cost_to_vec`. + + Returns: + An array, [num_b, p] if axis=0 or [num_a, p] if axis=1 + """ + if arr.ndim == 1: + arr = arr.reshape(-1, 1) + + app = functools.partial(self._apply_cost_to_vec, axis=axis, fn=fn, sub_idx=sub_idx,**kwargs) + return jax.vmap(app, in_axes=1, out_axes=1)(arr) + + def _apply_cost_to_vec( + self, + vec: jnp.ndarray, + axis: int = 0, + fn=None, + sub_idx = None, + **_: Any, + ) -> jnp.ndarray: + """Apply ``[num_a, num_b]`` fn(cost) (or transpose) to vector. + + Args: + vec: jnp.ndarray [num_a,] ([num_b,] if axis=1) vector + axis: axis on which the reduction is done. + fn: function optionally applied to cost matrix element-wise, before the + doc product + + Returns: + A jnp.ndarray corresponding to cost x vector + """ + if sub_idx is None: + matrix = self.cost_matrix.T if axis == 0 else self.cost_matrix + matrix = fn(matrix) if fn is not None else matrix + return jnp.dot(matrix, vec) + matrix = self.cost_matrix[sub_idx, :].T if axis == 0 else self.cost_matrix[:, sub_idx] + matrix = fn(matrix) if fn is not None else matrix + return jnp.dot(matrix, vec[sub_idx]) + + @classmethod + def prepare_divergences( + cls, + *args: Any, + static_b: bool = False, + **kwargs: Any + ) -> Tuple["Geometry", ...]: + """Instantiate 2 (or 3) geometries to compute a Sinkhorn divergence.""" + size = 2 if static_b else 3 + nones = [None, None, None] + cost_matrices = kwargs.pop("cost_matrix", args) + kernel_matrices = kwargs.pop("kernel_matrix", nones) + cost_matrices = cost_matrices if cost_matrices is not None else nones + return tuple( + cls(cost_matrix=arg1, kernel_matrix=arg2, **kwargs) + for arg1, arg2, _ in zip(cost_matrices, kernel_matrices, range(size)) + ) + + def to_LRCGeometry( + self, + rank: int = 0, + tol: float = 1e-2, + rng: Optional[jax.Array] = None, + scale: float = 1.0 + ) -> "low_rank.LRCGeometry": + r"""Factorize the cost matrix using either SVD (full) or :cite:`indyk:19`. + + When `rank=min(n,m)` or `0` (by default), use :func:`jax.numpy.linalg.svd`. + + For other values, use the routine in sublinear time :cite:`indyk:19`. + Uses the implementation of :cite:`scetbon:21`, algorithm 4. + + It holds that with probability *0.99*, + :math:`||A - UV||_F^2 \leq || A - A_k ||_F^2 + tol \cdot ||A||_F^2`, + where :math:`A` is ``n x m`` cost matrix, :math:`UV` the factorization + computed in sublinear time and :math:`A_k` the best rank-k approximation. + + Args: + rank: Target rank of the :attr:`cost_matrix`. + tol: Tolerance of the error. The total number of sampled points is + :math:`min(n, m,\frac{rank}{tol})`. + rng: The PRNG key to use for initializing the model. + 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. + + Returns: + Low-rank geometry. + """ + from ott.geometry import low_rank + assert rank >= 0, f"Rank must be non-negative, got {rank}." + n, m = self.shape + + if rank == 0 or rank >= min(n, m): + # TODO(marcocuturi): add hermitian=self.is_symmetric, currently bugging. + u, s, vh = jnp.linalg.svd( + self.cost_matrix, + full_matrices=False, + compute_uv=True, + ) + + cost_1 = u + cost_2 = (s[:, None] * vh).T + else: + rng = utils.default_prng_key(rng) + rng1, rng2, rng3, rng4, rng5 = jax.random.split(rng, 5) + n_subset = min(int(rank / tol), n, m) + + i_star = jax.random.randint(rng1, shape=(), minval=0, maxval=n) + j_star = jax.random.randint(rng2, shape=(), minval=0, maxval=m) + + ci_star = self.subset([i_star], None).cost_matrix.ravel() ** 2 # (m,) + cj_star = self.subset(None, [j_star]).cost_matrix.ravel() ** 2 # (n,) + + p_row = cj_star + ci_star[j_star] + jnp.mean(ci_star) # (n,) + p_row /= jnp.sum(p_row) + row_ixs = jax.random.choice(rng3, n, shape=(n_subset,), p=p_row) + # (n_subset, m) + s = self.subset(row_ixs, None).cost_matrix + s /= jnp.sqrt(n_subset * p_row[row_ixs][:, None]) + + p_col = jnp.sum(s ** 2, axis=0) # (m,) + p_col /= jnp.sum(p_col) + # (n_subset,) + col_ixs = jax.random.choice(rng4, m, shape=(n_subset,), p=p_col) + # (n_subset, n_subset) + w = s[:, col_ixs] / jnp.sqrt(n_subset * p_col[col_ixs][None, :]) + + U, _, V = jsp.linalg.svd(w) + U = U[:, :rank] # (n_subset, rank) + U = (s.T @ U) / jnp.linalg.norm(w.T @ U, axis=0) # (m, rank) + + _, d, v = jnp.linalg.svd(U.T @ U) # (k,), (k, k) + v = v.T / jnp.sqrt(d)[None, :] + + inv_scale = (1.0 / jnp.sqrt(n_subset)) + col_ixs = jax.random.choice(rng5, m, shape=(n_subset,)) # (n_subset,) + + # (n, n_subset) + A_trans = self.subset(None, col_ixs).cost_matrix * inv_scale + B = (U[col_ixs, :] @ v * inv_scale) # (n_subset, k) + M = jnp.linalg.inv(B.T @ B) # (k, k) + V = jnp.linalg.multi_dot([A_trans, B, M.T, v.T]) # (n, k) + cost_1 = V + cost_2 = U + + return low_rank.LRCGeometry( + cost_1=cost_1, + cost_2=cost_2, + epsilon=self._epsilon_init, + relative_epsilon=self._relative_epsilon, + scale_cost=self._scale_cost, + scale_factor=scale, + ) + + def subset( + self, src_ixs: Optional[jnp.ndarray], tgt_ixs: Optional[jnp.ndarray], + **kwargs: Any + ) -> "Geometry": + """Subset rows or columns of a geometry. + + Args: + src_ixs: Row indices. If ``None``, use all rows. + tgt_ixs: Column indices. If ``None``, use all columns. + kwargs: Keyword arguments to override the initialization. + + Returns: + The modified geometry. + """ + + def subset_fn( + arr: Optional[jnp.ndarray], + src_ixs: Optional[jnp.ndarray], + tgt_ixs: Optional[jnp.ndarray], + ) -> Optional[jnp.ndarray]: + if arr is None: + return None + if src_ixs is not None: + arr = arr[src_ixs, ...] + if tgt_ixs is not None: + arr = arr[:, tgt_ixs] + return arr # noqa: RET504 + + return self._mask_subset_helper( + src_ixs, + tgt_ixs, + fn=subset_fn, + propagate_mask=True, + **kwargs, + ) + + def mask( + self, + src_mask: Optional[jnp.ndarray], + tgt_mask: Optional[jnp.ndarray], + mask_value: float = 0.0, + ) -> "Geometry": + """Mask rows or columns of a geometry. + + The mask is used only when computing some statistics of the + :attr:`cost_matrix`. + + - :attr:`mean_cost_matrix` + - :attr:`median_cost_matrix` + - :attr:`inv_scale_cost` + + Args: + src_mask: Row mask. Can be specified either as a boolean array of shape + ``[num_a,]`` or as an array of indices. If ``None``, no mask is applied. + tgt_mask: Column mask. Can be specified either as a boolean array of shape + ``[num_b,]`` or as an array of indices. If ``None``, no mask is applied. + mask_value: Value to use for masking. + + Returns: + The masked geometry. + """ + + def mask_fn( + arr: Optional[jnp.ndarray], + src_mask: Optional[jnp.ndarray], + tgt_mask: Optional[jnp.ndarray], + ) -> Optional[jnp.ndarray]: + if arr is None: + return arr + assert arr.ndim == 2, arr.ndim + if src_mask is not None: + arr = jnp.where(src_mask[:, None], arr, mask_value) + if tgt_mask is not None: + arr = jnp.where(tgt_mask[None, :], arr, mask_value) + return arr # noqa: RET504 + + 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]], + Optional[jnp.ndarray]], + propagate_mask: bool, + **kwargs: Any, + ) -> "Geometry": + (cost, kernel, eps, src_mask, tgt_mask), aux_data = self.tree_flatten() + cost = fn(cost, src_ixs, tgt_ixs) + kernel = fn(kernel, src_ixs, 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, None) + tgt_mask = fn(tgt_mask, tgt_ixs, None) + + aux_data = {**aux_data, **kwargs} + return type(self).tree_unflatten( + aux_data, [cost, kernel, eps, src_mask, tgt_mask] + ) + + @property + def src_mask(self) -> Optional[jnp.ndarray]: + """Mask of shape ``[num_a,]`` to compute :attr:`cost_matrix` statistics. + + Specifically, it is used when computing: + + - :attr:`mean_cost_matrix` + - :attr:`median_cost_matrix` + - :attr:`inv_scale_cost` + """ + return self._normalize_mask(self._src_mask, self.shape[0]) + + @property + def tgt_mask(self) -> Optional[jnp.ndarray]: + """Mask of shape ``[num_b,]`` to compute :attr:`cost_matrix` statistics. + + Specifically, it is used when computing: + + - :attr:`mean_cost_matrix` + - :attr:`median_cost_matrix` + - :attr:`inv_scale_cost` + """ + return self._normalize_mask(self._tgt_mask, self.shape[1]) + + @property + def dtype(self) -> jnp.dtype: + """The data type.""" + return ( + self._kernel_matrix if self._cost_matrix is None else self._cost_matrix + ).dtype + + def _masked_geom(self, mask_value: float = 0.0) -> "Geometry": + """Mask geometry based on :attr:`src_mask` and :attr:`tgt_mask`.""" + src_mask, tgt_mask = self.src_mask, self.tgt_mask + if src_mask is None and tgt_mask is None: + return self + return self.mask(src_mask, tgt_mask, mask_value=mask_value) + + @property + def _n_normed_ones(self) -> jnp.ndarray: + """Normalized array of shape ``[num_a,]``.""" + mask = self.src_mask + arr = jnp.ones(self.shape[0]) if mask is None else mask + return arr / jnp.sum(arr) + + @property + def _m_normed_ones(self) -> jnp.ndarray: + """Normalized array of shape ``[num_b,]``.""" + mask = self.tgt_mask + arr = jnp.ones(self.shape[1]) if mask is None else mask + return arr / jnp.sum(arr) + + @staticmethod + def _normalize_mask(mask: Optional[jnp.ndarray], + size: int) -> Optional[jnp.ndarray]: + """Convert array of indices to a boolean mask.""" + if mask is None: + return None + if not jnp.issubdtype(mask, (bool, jnp.bool_)): + mask = jnp.isin(jnp.arange(size), mask) + assert mask.shape == (size,) + return mask + + def tree_flatten(self): # noqa: D102 + return ( + self._cost_matrix, self._kernel_matrix, self._epsilon_init, + self._src_mask, self._tgt_mask + ), { + "scale_cost": self._scale_cost, + "relative_epsilon": self._relative_epsilon + } + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + cost, kernel, eps, src_mask, tgt_mask = children + return cls( + cost, kernel, eps, src_mask=src_mask, tgt_mask=tgt_mask, **aux_data + ) + + +def is_affine(fn) -> bool: + """Test heuristically if a function is affine.""" + x = jnp.arange(10.0) + out = jax.vmap(jax.grad(fn))(x) + return jnp.sum(jnp.diff(jnp.abs(out))) == 0.0 + + +def is_linear(fn) -> bool: + """Test heuristically if a function is linear.""" + return jnp.logical_and(fn(0.0) == 0.0, is_affine(fn)) diff --git a/ott/build/lib/ott/geometry/graph.py b/ott/build/lib/ott/geometry/graph.py new file mode 100644 index 0000000..4a97707 --- /dev/null +++ b/ott/build/lib/ott/geometry/graph.py @@ -0,0 +1,272 @@ +# 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. +from typing import Any, Dict, Literal, Optional, Sequence, Tuple + +import jax +import jax.numpy as jnp +import jax.scipy as jsp + +from ott.geometry import geometry +from ott.math import fixed_point_loop +from ott.math import utils as mu + +__all__ = ["Graph"] + + +@jax.tree_util.register_pytree_node_class +class Graph(geometry.Geometry): + r"""Graph distance approximation using heat kernel :cite:`heitz:21,crane:13`. + + Approximates the heat kernel for large ``n_steps``, which for small ``t`` + approximates the geodesic exponential kernel :math:`e^{\frac{-d(x, y)^2}{t}}`. + + Args: + laplacian: Symmetric graph Laplacian. The check for symmetry is **NOT** + performed. See also :meth:`from_graph`. + n_steps: Maximum number of steps used to approximate the heat kernel. + numerical_scheme: Numerical scheme used to solve the heat diffusion. + normalize: Whether to normalize the Laplacian as + :math:`L^{sym} = \left(D^+\right)^{\frac{1}{2}} L + \left(D^+\right)^{\frac{1}{2}}`, where :math:`L` is the + non-normalized Laplacian and :math:`D` is the degree matrix. + tol: Relative tolerance with respect to the Hilbert metric, see + :cite:`peyre:19`, Remark 4.12. Used when iteratively updating scalings. + If negative, this option is ignored and only ``n_steps`` is used. + kwargs: Keyword arguments for :class:`~ott.geometry.geometry.Geometry`. + """ + + def __init__( + self, + laplacian: jnp.ndarray, + t: float = 1e-3, + n_steps: int = 100, + numerical_scheme: Literal["backward_euler", + "crank_nicolson"] = "backward_euler", + tol: float = -1.0, + **kwargs: Any + ): + super().__init__(epsilon=1.0, **kwargs) + self.laplacian = laplacian + self.t = t + self.n_steps = n_steps + self.numerical_scheme = numerical_scheme + self.tol = tol + + @classmethod + def from_graph( + cls, + G: jnp.ndarray, + t: Optional[float] = 1e-3, + directed: bool = False, + normalize: bool = False, + **kwargs: Any + ) -> "Graph": + r"""Construct :class:`~ott.geometry.graph.Graph` from an adjacency matrix. + + Args: + G: Adjacency matrix. + t: Constant used when approximating the geodesic exponential kernel. + If `None`, use :math:`\frac{1}{|E|} \sum_{(u, v) \in E} weight(u, v)` + :cite:`crane:13`. In this case, the ``graph`` must be specified + and the edge weights are all assumed to be positive. + directed: Whether the ``graph`` is directed. If not, it will be made + undirected as :math:`G + G^T`. This parameter is ignored when directly + passing the Laplacian, which is assumed to be symmetric. + normalize: Whether to normalize the Laplacian as + :math:`L^{sym} = \left(D^+\right)^{\frac{1}{2}} L + \left(D^+\right)^{\frac{1}{2}}`, where :math:`L` is the + non-normalized Laplacian and :math:`D` is the degree matrix. + kwargs: Keyword arguments for :class:`~ott.geometry.graph.Graph`. + + Returns: + The graph geometry. + """ + assert G.shape[0] == G.shape[1], G.shape + + if directed: + G = G + G.T + + degree = jnp.sum(G, axis=1) + laplacian = jnp.diag(degree) - G + + if normalize: + inv_sqrt_deg = jnp.diag( + jnp.where(degree > 0.0, 1.0 / jnp.sqrt(degree), 0.0) + ) + laplacian = inv_sqrt_deg @ laplacian @ inv_sqrt_deg + + if t is None: + t = (jnp.sum(G) / jnp.sum(G > 0.0)) ** 2 + + return cls(laplacian, t=t, **kwargs) + + def apply_kernel( + self, + scaling: jnp.ndarray, + eps: Optional[float] = None, + axis: int = 0, + ) -> jnp.ndarray: + r"""Apply :attr:`kernel_matrix` on positive scaling vector. + + Args: + scaling: Scaling to apply the kernel to. + eps: passed for consistency, not used yet. + axis: passed for consistency, not used yet. + + Returns: + Kernel applied to ``scaling``. + """ + + def conf_fn( + iteration: int, consts: Tuple[jnp.ndarray, Optional[jnp.ndarray]], + old_new: Tuple[jnp.ndarray, jnp.ndarray] + ) -> bool: + del iteration, consts + + x_old, x_new = old_new + x_old, x_new = mu.safe_log(x_old), mu.safe_log(x_new) + # center + x_old, x_new = x_old - jnp.nanmax(x_old), x_new - jnp.nanmax(x_new) + # Hilbert metric, see Remark 4.12 in `Computational Optimal Transport` + f = x_new - x_old + return (jnp.nanmax(f) - jnp.nanmin(f)) > self.tol + + def body_fn( + iteration: int, consts: Tuple[jnp.ndarray, Optional[jnp.ndarray]], + old_new: Tuple[jnp.ndarray, jnp.ndarray], compute_errors: bool + ) -> Tuple[jnp.ndarray, jnp.ndarray]: + del iteration, compute_errors + + L, scaled_lap = consts + _, b = old_new + + if self.numerical_scheme == "crank_nicolson": + # below is a preferred way of specifying the update (albeit more FLOPS), + # as CSR/CSC/COO matrices don't support adding a diagonal matrix now: + # b' = (2 * I - M) @ b = (2 * I - (I + c * L)) @ b = (I - c * L) @ b + b = b - scaled_lap @ b + return b, jsp.linalg.solve_triangular(L, b, lower=True) + + # eps we cannot use since it would require a re-solve + # axis we can ignore since the matrix is symmetric + del eps, axis + + force_scan = self.tol < 0.0 + fixpoint_fn = ( + fixed_point_loop.fixpoint_iter + if force_scan else fixed_point_loop.fixpoint_iter_backprop + ) + + state = (jnp.full_like(scaling, jnp.nan), scaling) + L = jsp.linalg.cholesky(self._M, lower=True) + if self.numerical_scheme == "crank_nicolson": + constants = L, self._scaled_laplacian + else: + constants = L, None + + return fixpoint_fn( + cond_fn=(lambda *_, **__: True) if force_scan else conf_fn, + body_fn=body_fn, + min_iterations=self.n_steps if force_scan else 1, + max_iterations=self.n_steps, + inner_iterations=1, + constants=constants, + state=state, + )[1] + + @property + def kernel_matrix(self) -> jnp.ndarray: # noqa: D102 + n, _ = self.shape + kernel = self.apply_kernel(jnp.eye(n)) + # Symmetrize the kernel if needed. Numerical imprecision + # happens when `numerical_scheme='backward_euler'` and small `t` + return jax.lax.cond( + jnp.allclose(kernel, kernel.T, atol=1e-8, rtol=1e-8), lambda x: x, + lambda x: (x + x.T) / 2.0, kernel + ) + + @property + def cost_matrix(self) -> jnp.ndarray: # noqa: D102 + return -self.t * mu.safe_log(self.kernel_matrix) + + @property + def _scale(self) -> float: + """Constant used to scale the Laplacian.""" + if self.numerical_scheme == "backward_euler": + return self.t / (4.0 * self.n_steps) + if self.numerical_scheme == "crank_nicolson": + return self.t / (2.0 * self.n_steps) + raise NotImplementedError( + f"Numerical scheme `{self.numerical_scheme}` is not implemented." + ) + + @property + def _scaled_laplacian(self) -> jnp.ndarray: + """Laplacian scaled by a constant, depending on the numerical scheme.""" + return self._scale * self.laplacian + + @property + def _M(self) -> jnp.ndarray: + n, _ = self.shape + return self._scaled_laplacian + jnp.eye(n) + + @property + def shape(self) -> Tuple[int, int]: # noqa: D102 + return self.laplacian.shape + + @property + def is_symmetric(self) -> bool: # noqa: D102 + return True + + @property + def dtype(self) -> jnp.dtype: # noqa: D102 + return self.laplacian.dtype + + def transport_from_potentials( + self, f: jnp.ndarray, g: jnp.ndarray + ) -> jnp.ndarray: + """Not implemented.""" + raise ValueError("Not implemented.") + + def apply_transport_from_potentials( + self, + f: jnp.ndarray, + g: jnp.ndarray, + vec: jnp.ndarray, + axis: int = 0 + ) -> jnp.ndarray: + """Not implemented.""" + raise ValueError("Not implemented.") + + def marginal_from_potentials( + self, + f: jnp.ndarray, + g: jnp.ndarray, + axis: int = 0, + ) -> jnp.ndarray: + """Not implemented.""" + raise ValueError("Not implemented.") + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return [self.laplacian, self.t], { + "n_steps": self.n_steps, + "numerical_scheme": self.numerical_scheme, + "tol": self.tol, + } + + @classmethod + def tree_unflatten( # noqa: D102 + cls, aux_data: Dict[str, Any], children: Sequence[Any] + ) -> "Graph": + return cls(*children, **aux_data) diff --git a/ott/build/lib/ott/geometry/grid.py b/ott/build/lib/ott/geometry/grid.py new file mode 100644 index 0000000..708753d --- /dev/null +++ b/ott/build/lib/ott/geometry/grid.py @@ -0,0 +1,415 @@ +# 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 itertools +from typing import Any, List, NoReturn, Optional, Sequence, Tuple + +import jax +import jax.numpy as jnp +import numpy as np + +from ott.geometry import costs, geometry, low_rank, pointcloud +from ott.math import utils + +__all__ = ["Grid"] + + +@jax.tree_util.register_pytree_node_class +class Grid(geometry.Geometry): + r"""Class describing the geometry of points taken in a Cartesian product. + + This class implements a geometry in which probability measures are supported + on a :math:`d`-dimensional Cartesian grid, a Cartesian product of :math:`d` + lists of values, each list being itself of size :math:`n_i`. + + The transportation cost between points in the grid is assumed to be separable, + namely a sum of coordinate-wise cost functions, as in: + + .. math:: + + cost(x,y) = \sum_{i=1}^d cost_i(x_i, y_i) + + where :math:`cost_i`: R x R → R. + + In such a regime, and despite the fact that the total number :math:`n_{total}` + of points in the grid is exponential :math:`d` (namely :math:`\prod_i n_i`), + applying a kernel in the context of regularized optimal transport can be + carried out in time that is of the order of :math:`n_{total}^{(1+1/d)}` using + convolutions, either in the original domain or log-space domain. This class + precomputes :math:`d` :math:`n_i` x :math:`n_i` cost matrices (one per + dimension) and implements these two operations by carrying out these + convolutions one dimension at a time. + + Args: + x: list of arrays of varying sizes, describing the locations of the grid. + Locations are provided as a list of arrays, that is :math:`d` + vectors of (possibly varying) size :math:`n_i`. The resulting grid + is the Cartesian product of these vectors. + grid_size: tuple of integers describing grid sizes, namely + :math:`(n_1,...,n_d)`. This will only be used if x is None. + In that case the grid will be assumed to lie in the hypercube + :math:`[0,1]^d`, with the :math:`d` dimensions, described as points + regularly sampled in :math:`[0,1]`. + cost_fns: a sequence of :math:`d` cost functions, each being a cost taking + two reals as inputs to output a real number. + num_a: total size of grid. This parameters will be computed from other + inputs. + grid_dimension: dimension of grid. This parameters will be computed from + other inputs. + kwargs: keyword arguments for :class:`~ott.geometry.geometry.Geometry`. + """ + + def __init__( + self, + x: Optional[Sequence[jnp.ndarray]] = None, + grid_size: Optional[Sequence[int]] = None, + cost_fns: Optional[Sequence[costs.CostFn]] = None, + num_a: Optional[int] = None, + grid_dimension: Optional[int] = None, + **kwargs: Any, + ): + if ( + grid_size is not None and x is not None and num_a is not None and + grid_dimension is not None + ): + self.grid_size = tuple(map(int, grid_size)) + self.x = x + self.num_a = num_a + self.grid_dimension = grid_dimension + elif x is not None: + self.x = x + self.grid_size = tuple(xs.shape[0] for xs in x) + self.num_a = np.prod(np.array(self.grid_size)) + self.grid_dimension = len(self.x) + elif grid_size is not None: + self.grid_size = tuple(map(int, grid_size)) + self.x = tuple(jnp.linspace(0, 1, n) for n in self.grid_size) + self.num_a = np.prod(np.array(grid_size)) + self.grid_dimension = len(self.grid_size) + else: + raise ValueError("Input either grid_size tuple or grid locations x.") + + if cost_fns is None: + cost_fns = [costs.SqEuclidean()] + self.cost_fns = cost_fns + self.kwargs = { + "num_a": self.num_a, + "grid_size": self.grid_size, + "grid_dimension": self.grid_dimension + } + + super().__init__(**kwargs) + + @property + def geometries(self) -> List[geometry.Geometry]: + """Cost matrices along each dimension of the grid.""" + geometries = [] + for dimension, cost_fn in itertools.zip_longest( + range(self.grid_dimension), self.cost_fns, fillvalue=self.cost_fns[-1] + ): + x_values = self.x[dimension][:, jnp.newaxis] + geom = pointcloud.PointCloud( + x_values, + cost_fn=cost_fn, + epsilon=self._epsilon_init, + ) + geometries.append(geom) + return geometries + + @property + def median_cost_matrix(self) -> NoReturn: + """Not implemented.""" + raise NotImplementedError("Median cost not implemented for grids.") + + @property + def can_LRC(self) -> bool: # noqa: D102 + return True + + @property + def shape(self) -> Tuple[int, int]: # noqa: D102 + return self.num_a, self.num_a + + @property + def is_symmetric(self) -> bool: # noqa: D102 + return True + + # Reimplemented functions to be used in regularized OT + def apply_lse_kernel( + self, + f: jnp.ndarray, + g: jnp.ndarray, + eps: float, + vec: Optional[jnp.ndarray] = None, + axis: int = 0 + ) -> jnp.ndarray: + """Apply grid kernel in log space. See notes in parent class for use case. + + Reshapes vector inputs below as grids, applies kernels onto each slice, and + then expands the outputs as vectors. + + More implementation details in :cite:`schmitz:18`. + + Args: + f: jnp.ndarray, a vector of potentials + g: jnp.ndarray, a vector of potentials + eps: float, regularization strength + vec: jnp.ndarray, if needed, a vector onto which apply the kernel weighted + by f and g. + axis: axis (0 or 1) along which summation should be carried out. + + Returns: + a vector, the result of kernel applied in lse space onto vec. + """ + f, g = jnp.reshape(f, self.grid_size), jnp.reshape(g, self.grid_size) + + if vec is not None: + vec = jnp.reshape(vec, self.grid_size) + + if axis == 0: + f, g = g, f + + for dimension in range(self.grid_dimension): + g, vec = self._apply_lse_kernel_one_dimension(dimension, f, g, eps, vec) + g -= jnp.where(jnp.isfinite(f), f, 0) + + if vec is None: + vec = jnp.array(1.0) + return g.ravel(), vec.ravel() + + def _apply_lse_kernel_one_dimension(self, dimension, f, g, eps, vec=None): + """Permute axis & apply the kernel on a single slice.""" + indices = np.arange(self.grid_dimension) + indices[dimension], indices[0] = 0, dimension + f, g = jnp.transpose(f, indices), jnp.transpose(g, indices) + centered_cost = ( + f[:, jnp.newaxis, ...] + g[jnp.newaxis, ...] - jnp.expand_dims( + self.geometries[dimension].cost_matrix, + axis=tuple(range(2, 1 + self.grid_dimension)) + ) + ) / eps + + if vec is not None: + vec = jnp.transpose(vec, indices) + softmax_res, softmax_sgn = utils.logsumexp( + centered_cost, b=vec, axis=1, return_sign=True + ) + return eps * jnp.transpose(softmax_res, + indices), jnp.transpose(softmax_sgn, indices) + softmax_res = eps * utils.logsumexp(centered_cost, axis=1) + return jnp.transpose(softmax_res, indices), None + + def _apply_cost_to_vec( + self, vec: jnp.ndarray, axis: int = 0, fn=None + ) -> jnp.ndarray: + r"""Apply grid's cost matrix (without instantiating it) to a vector. + + The `apply_cost` operation on grids rests on the following identity. + If it were to be cast as a [num_a, num_a] matrix, the corresponding cost + matrix :math:`C` would be a sum of `grid_dimension` matrices, each of the + form (here for the j-th slice) + :math:`\tilde{C}_j : = 1_{n_1} \otimes \dots \otimes C_j \otimes 1_{n_d}` + where each :math:`1_{n}` is the :math:`n\times n` square matrix full of 1's. + + Applying :math:`\tilde{C}_j` on a vector grid consists in carrying a tensor + multiplication on the dimension of that vector reshaped as a grid, followed + by a summation on all other axis while keeping dimensions. That identity is + a generalization of the formula + :math:`(1_{n} \otimes A) vec(X) = vec( A X 1_n)` + where the last multiplication by the matrix of ones is equivalent to + summation while keeping dimensions. + + Args: + vec: jnp.ndarray, flat vector of total size prod(grid_size). + axis: axis 0 if applying transpose costs, 1 if using the original cost. + fn: function optionally applied to cost matrix element-wise, before the + dot product. + + Returns: + A jnp.ndarray corresponding to cost x matrix + """ + vec = jnp.reshape(vec, self.grid_size) + accum_vec = jnp.zeros_like(vec) + indices = list(range(1, self.grid_dimension)) + for dimension, geom in enumerate(self.geometries): + cost = geom.cost_matrix + ind = indices.copy() + ind.insert(dimension, 0) + if axis == 0: + cost = cost.T + accum_vec += jnp.sum( + jnp.tensordot(cost, vec, axes=([0], [dimension])), + axis=indices, + keepdims=True + ).transpose(ind) + return accum_vec.ravel() + + def apply_kernel( + self, + scaling: jnp.ndarray, + eps: Optional[float] = None, + axis: Optional[int] = None + ) -> jnp.ndarray: + """Apply grid kernel on scaling vector. + + See notes in parent class for use. + + Reshapes scaling vector as a grid, applies kernels onto each slice, and + then ravels backs the output as a vector. + + More implementation details in :cite:`schmitz:18`, + + Args: + scaling: jnp.ndarray, a vector of scaling (>0) values. + eps: float, regularization strength + axis: axis (0 or 1) along which summation should be carried out. + + Returns: + a vector, the result of kernel applied onto scaling. + """ + scaling = jnp.reshape(scaling, self.grid_size) + indices = list(range(1, self.grid_dimension)) + for dimension, geom in enumerate(self.geometries): + kernel = geom.kernel_matrix + kernel = kernel if eps is None else kernel ** (self.epsilon / eps) + ind = indices.copy() + ind.insert(dimension, 0) + scaling = jnp.tensordot( + kernel, scaling, axes=([0], [dimension]) + ).transpose(ind) + return scaling.ravel() + + def transport_from_potentials( + self, f: jnp.ndarray, g: jnp.ndarray, axis: int = 0 + ) -> NoReturn: + """Not implemented, use :meth:`apply_transport_from_potentials` instead.""" + raise ValueError( + "Grid geometry cannot instantiate a transport matrix, use", + " apply_transport_from_potentials(...) if you wish to ", + " apply the transport matrix to a vector, or use a point " + " cloud geometry instead" + ) + + def transport_from_scalings( + self, f: jnp.ndarray, g: jnp.ndarray, axis: int = 0 + ) -> NoReturn: + """Not implemented, use :meth:`apply_transport_from_scalings` instead.""" + raise ValueError( + "Grid geometry cannot instantiate a transport matrix, use ", + "apply_transport_from_scalings(...) if you wish to ", + "apply the transport matrix to a vector, or use a point " + "cloud geometry instead." + ) + + def subset( + self, src_ixs: Optional[jnp.ndarray], tgt_ixs: Optional[jnp.ndarray] + ) -> NoReturn: + """Not implemented.""" + raise NotImplementedError("Subsetting is not implemented for grids.") + + def mask( + self, + src_mask: Optional[jnp.ndarray], + tgt_mask: Optional[jnp.ndarray], + mask_value: float = 0.0, + ) -> NoReturn: + """Not implemented.""" + raise NotImplementedError("Masking is not implemented for grids.") + + @classmethod + def prepare_divergences( + cls, + *args: Any, + static_b: bool = False, + **kwargs: Any + ) -> Tuple["Grid", ...]: + """Instantiate the geometries used for a divergence computation.""" + grid_size = kwargs.pop("grid_size", None) + x = kwargs.pop("x", args) + + sep_grid = cls(x=x, grid_size=grid_size, **kwargs) + size = 2 if static_b else 3 + return tuple(sep_grid for _ in range(size)) + + @property + def dtype(self) -> jnp.dtype: # noqa: D102 + return self.x[0].dtype + + def tree_flatten(self): # noqa: D102 + return (self.x, self.cost_fns, self._epsilon_init), self.kwargs + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + return cls( + x=children[0], cost_fns=children[1], epsilon=children[2], **aux_data + ) + + def to_LRCGeometry( + self, + scale: float = 1.0, + **kwargs: Any, + ) -> low_rank.LRCGeometry: + """Converts grid to low-rank geometry. + + Conversion is carried out by taking advantage of the fact that the true cost + matrix of a grid geometry is a sum of Kronecker products of local cost + matrices (for each dimension) with matrices of 1's (both on left and right + sides) of varying dimension. Each of the matrices in that sum can be + factorized if each of these cost matrices can be factorized, which we do + by forcing a conversion to a low rank geometry object. + + 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 + geometries on each slice are not low-rank. + + Returns: + :class:`~ott.geometry.low_rank.LRCGeometry` object. + """ + cost_1, cost_2 = [], [] + for dimension, geom in enumerate(self.geometries): + # An overall low-rank conversion of the cost matrix on a grid, to an + # object of :class:`~ott.geometry.low_rank.LRCGeometry`, necesitates an + # exact low-rank matrix decompisition of the cost matrix of each slice + # of that grid, even if costs on such slices are not low-rank. + # The idea here is that even if the cost matrix on slice `i` is full rank + # `n_i`, we are better off doing 2 redundant `n_i x n_i` matrix products, + # because this is the only way to access to an overall low-rank + # factorization for the entire cost matrix. To get such an exact + # decomposition, the parameter `rank` is set to `0`, triggering a full + # singular value decomposition if needed. + geom = geom.to_LRCGeometry(rank=0, scale=scale, **kwargs) + c_1, c_2 = geom.cost_1, geom.cost_2 + l, r = self.grid_size[:dimension], self.grid_size[dimension + 1:] + l = int(np.prod(np.array(l))) + r = int(np.prod(np.array(r))) + cost_1.append( + jnp.kron(jnp.ones((l, 1)), jnp.kron(c_1, jnp.ones((r, 1),))) + ) + cost_2.append( + jnp.kron(jnp.ones((l, 1)), jnp.kron(c_2, jnp.ones((r, 1),))) + ) + cost_1 = jnp.concatenate(cost_1, axis=-1) + cost_2 = jnp.concatenate(cost_2, 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, + ) diff --git a/ott/build/lib/ott/geometry/low_rank.py b/ott/build/lib/ott/geometry/low_rank.py new file mode 100644 index 0000000..c28c314 --- /dev/null +++ b/ott/build/lib/ott/geometry/low_rank.py @@ -0,0 +1,508 @@ +# 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. +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 +from ott.math import utils as mu + +__all__ = ["LRCGeometry", "LRKGeometry"] + + +@jax.tree_util.register_pytree_node_class +class LRCGeometry(geometry.Geometry): + """Geometry whose cost is defined by product of two low-rank matrices. + + Implements geometries that are defined as low rank products, i.e. for which + there exists two matrices :math:`A` and :math:`B` of :math:`r` columns such + that the cost of the geometry equals :math:`AB^T`. Apart from being faster to + apply to a vector, these geometries are characterized by the fact that adding + two such geometries should be carried out by concatenating factors, i.e. + if :math:`C = AB^T` and :math:`D = EF^T` then :math:`C + D = [A,E][B,F]^T` + + Args: + cost_1: Array of shape ``[num_a, r]``. + cost_2: Array of shape ``[num_b, r]``. + bias: constant added to entire cost matrix. + scale: Value used to rescale the factors of the low-rank geometry. + scale_cost: option to rescale the cost matrix. Implemented scalings are + 'max_bound', 'mean' and 'max_cost'. Alternatively, a float + factor can be given to rescale the cost such that + ``cost_matrix /= scale_cost``. If `True`, use 'mean'. + batch_size: optional size of the batch to compute online (without + instantiating the matrix) the scale factor ``scale_cost`` of the + :attr:`cost_matrix` when ``scale_cost = 'max_cost'``. If `None`, the batch + size is set to `1024` or to the largest number of samples between + :attr:`cost_1` and :attr:`cost_2` if smaller than `1024`. + kwargs: keyword arguments for :class:`~ott.geometry.geometry.Geometry`. + """ + + def __init__( + self, + cost_1: jnp.ndarray, + cost_2: jnp.ndarray, + bias: float = 0.0, + scale_factor: float = 1.0, + scale_cost: Union[bool, int, float, Literal["mean", "max_bound", + "max_cost"]] = 1.0, + batch_size: Optional[int] = None, + **kwargs: Any, + ): + super().__init__(**kwargs) + self._cost_1 = cost_1 + self._cost_2 = cost_2 + self._bias = bias + self._scale_factor = scale_factor + self._scale_cost = "mean" if scale_cost is True else scale_cost + self.batch_size = batch_size + + @property + def cost_1(self) -> jnp.ndarray: + """First factor of the :attr:`cost_matrix`.""" + scale_factor = jnp.sqrt(self._scale_factor * self.inv_scale_cost) + return scale_factor * self._cost_1 + + @property + def cost_2(self) -> jnp.ndarray: + """Second factor of the :attr:`cost_matrix`.""" + scale_factor = jnp.sqrt(self._scale_factor * self.inv_scale_cost) + return scale_factor * self._cost_2 + + @property + def bias(self) -> float: + """Constant offset added to the entire :attr:`cost_matrix`.""" + return self._bias * self.inv_scale_cost + + @property + def cost_rank(self) -> int: # noqa: D102 + return self._cost_1.shape[1] + + @property + def cost_matrix(self) -> jnp.ndarray: + """Materialize the cost matrix.""" + return jnp.matmul(self.cost_1, self.cost_2.T) + self.bias + + @property + def shape(self) -> Tuple[int, int]: # noqa: D102 + return self._cost_1.shape[0], self._cost_2.shape[0] + + @property + def is_symmetric(self) -> bool: # noqa: D102 + return ( + self._cost_1.shape[0] == self._cost_2.shape[0] and + jnp.all(self._cost_1 == self._cost_2) + ) + + @property + def inv_scale_cost(self) -> float: # noqa: D102 + if isinstance(self._scale_cost, (int, float, jax.Array)): + return 1.0 / self._scale_cost + self = self._masked_geom() + if self._scale_cost == "max_bound": + x_norm = self._cost_1[:, 0].max() + y_norm = self._cost_2[:, 1].max() + max_bound = x_norm + y_norm + 2 * jnp.sqrt(x_norm * y_norm) + return 1.0 / (max_bound + self._bias) + if self._scale_cost == "mean": + factor1 = jnp.dot(self._n_normed_ones, self._cost_1) + factor2 = jnp.dot(self._cost_2.T, self._m_normed_ones) + mean = jnp.dot(factor1, factor2) + self._bias + return 1.0 / mean + if self._scale_cost == "max_cost": + return 1.0 / self.compute_max_cost() + raise ValueError(f"Scaling {self._scale_cost} not implemented.") + + def apply_square_cost(self, arr: jnp.ndarray, axis: int = 0) -> jnp.ndarray: + """Apply elementwise-square of cost matrix to array (vector or matrix).""" + (n, m), r = self.shape, self.cost_rank + # When applying square of a LRCGeometry, one can either elementwise square + # the cost matrix, or instantiate an augmented (rank^2) LRCGeometry + # and apply it. First is O(nm), the other is O((n+m)r^2). + if n * m < (n + m) * r ** 2: # better use regular apply + return super().apply_square_cost(arr, axis) + + new_cost_1 = self.cost_1[:, :, None] * self.cost_1[:, None, :] + new_cost_2 = self.cost_2[:, :, None] * self.cost_2[:, None, :] + return LRCGeometry( + cost_1=new_cost_1.reshape((n, r ** 2)), + cost_2=new_cost_2.reshape((m, r ** 2)) + ).apply_cost(arr, axis) + + def _apply_cost_to_vec( + self, + vec: jnp.ndarray, + axis: int = 0, + fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, + is_linear: bool = False, + ) -> jnp.ndarray: + """Apply [num_a, num_b] fn(cost) (or transpose) to vector. + + Args: + vec: jnp.ndarray [num_a,] ([num_b,] if axis=1) vector + axis: axis on which the reduction is done. + fn: function optionally applied to cost matrix element-wise, before the + doc product + is_linear: Whether ``fn`` is a linear function to enable efficient + implementation. See :func:`ott.geometry.geometry.is_linear` + for a heuristic to help determine if a function is linear. + + Returns: + A jnp.ndarray corresponding to cost x vector + """ + + def linear_apply( + vec: jnp.ndarray, axis: int, fn: Callable[[jnp.ndarray], jnp.ndarray] + ) -> jnp.ndarray: + c1 = self.cost_1 if axis == 1 else self.cost_2 + c2 = self.cost_2 if axis == 1 else self.cost_1 + c2 = fn(c2) if fn is not None else c2 + bias = fn(self.bias) if fn is not None else self.bias + out = jnp.dot(c1, jnp.dot(c2.T, vec)) + return out + bias * jnp.sum(vec) * jnp.ones_like(out) + + if fn is None or is_linear: + return linear_apply(vec, axis, fn=fn) + return super()._apply_cost_to_vec(vec, axis, fn=fn) + + def compute_max_cost(self) -> float: + """Compute the maximum of the :attr:`cost_matrix`. + + Three cases are taken into account: + + - If the number of samples of ``cost_1`` and ``cost_2`` are both smaller + than 1024 and if ``batch_size`` is `None`, the ``cost_matrix`` is + computed to obtain its maximum entry. + - If one of the number of samples of ``cost_1`` or ``cost_2`` is larger + than 1024 and if ``batch_size`` is `None`, then the maximum of the + cost matrix is calculated by batch. The batches are created on the + longest axis of the cost matrix and their size is fixed to 1024. + - If ``batch_size`` is provided as a float, then the maximum of the cost + matrix is calculated by batch. The batches are created on the longest + axis of the cost matrix and their size if fixed by ``batch_size``. + + Returns: + Maximum of the cost matrix. + """ + batch_for_y = self.shape[1] > self.shape[0] + + n = self.shape[1] if batch_for_y else self.shape[0] + p = self._cost_2.shape[1] if batch_for_y else self._cost_1.shape[1] + carry = ((self._cost_1, self._cost_2) if batch_for_y else + (self._cost_2, self._cost_1)) + + if self.batch_size: + batch_size = min(self.batch_size, n) + else: + batch_size = min(1024, max(self.shape[0], self.shape[1])) + n_batch = n // batch_size + + def body(carry, slice_idx): + cost1, cost2 = carry + cost2_slice = jax.lax.dynamic_slice( + cost2, (slice_idx * batch_size, 0), (batch_size, p) + ) + out_slice = jnp.max(jnp.dot(cost2_slice, cost1.T)) + return carry, out_slice + + def finalize(carry): + cost1, cost2 = carry + return jnp.dot(cost2[n_batch * batch_size:], cost1.T) + + _, out = jax.lax.scan(body, carry, jnp.arange(n_batch)) + last_slice = finalize(carry) + max_value = jnp.max(jnp.concatenate((out, last_slice.reshape(-1)))) + return max_value + self._bias + + def to_LRCGeometry( + self, + rank: int = 0, + tol: float = 1e-2, + rng: Optional[jax.Array] = None, + scale: float = 1.0, + ) -> "LRCGeometry": + """Return self.""" + del rank, tol, rng, scale + return self + + @property + def can_LRC(self): # noqa: D102 + return True + + def subset( # noqa: D102 + self, src_ixs: Optional[jnp.ndarray], tgt_ixs: Optional[jnp.ndarray], + **kwargs: Any + ) -> "LRCGeometry": + + 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[ixs, ...] + + return self._mask_subset_helper( + src_ixs, tgt_ixs, fn=subset_fn, propagate_mask=True, **kwargs + ) + + def mask( # noqa: D102 + self, + src_mask: Optional[jnp.ndarray], + tgt_mask: Optional[jnp.ndarray], + mask_value: float = 0.0, + ) -> "LRCGeometry": + + 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, + ) -> "LRCGeometry": + (c1, c2, src_mask, tgt_mask, *children), aux_data = self.tree_flatten() + c1 = fn(c1, src_ixs) + c2 = fn(c2, 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, [c1, c2, src_mask, tgt_mask] + children + ) + + def __add__(self, other: "LRCGeometry") -> "LRCGeometry": + if not isinstance(other, LRCGeometry): + return NotImplemented + return LRCGeometry( + cost_1=jnp.concatenate((self.cost_1, other.cost_1), axis=1), + cost_2=jnp.concatenate((self.cost_2, other.cost_2), axis=1), + bias=self._bias + other._bias, + # already included in `cost_{1,2}` + scale_factor=1.0, + scale_cost=1.0, + ) + + @property + def dtype(self) -> jnp.dtype: # noqa: D102 + return self._cost_1.dtype + + def tree_flatten(self): # noqa: D102 + return ( + self._cost_1, + self._cost_2, + self._src_mask, + self._tgt_mask, + self._epsilon_init, + self._bias, + self._scale_factor, + ), { + "scale_cost": self._scale_cost, + "batch_size": self.batch_size + } + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + c1, c2, src_mask, tgt_mask, epsilon, bias, scale_factor = children + return cls( + c1, + c2, + bias=bias, + scale_factor=scale_factor, + epsilon=epsilon, + src_mask=src_mask, + tgt_mask=tgt_mask, + **aux_data + ) + + +@jax.tree_util.register_pytree_node_class +class LRKGeometry(geometry.Geometry): + """Low-rank kernel geometry. + + .. note:: + This constructor is not meant to be called by the user, + please use the :meth:`from_pointcloud` method instead. + + Args: + k1: Array of shape ``[num_a, r]`` with positive features. + k2: Array of shape ``[num_b, r]`` with positive features. + epsilon: Epsilon regularization. + kwargs: Keyword arguments for :class:`~ott.geometry.geometry.Geometry`. + """ + + def __init__( + self, + k1: jnp.ndarray, + k2: jnp.ndarray, + epsilon: Optional[float] = None, + **kwargs: Any + ): + super().__init__(epsilon=epsilon, relative_epsilon=False, **kwargs) + self.k1 = k1 + self.k2 = k2 + + @classmethod + def from_pointcloud( + cls, + x: jnp.ndarray, + y: jnp.ndarray, + *, + kernel: Literal["gaussian", "arccos"], + rank: int = 100, + std: float = 1.0, + n: int = 1, + rng: Optional[jax.Array] = None + ) -> "LRKGeometry": + r"""Low-rank kernel approximation :cite:`scetbon:20`. + + Args: + x: Array of shape ``[n, d]``. + y: Array of shape ``[m, d]``. + kernel: Type of the kernel to approximate. + rank: Rank of the approximation. + std: Depending on the ``kernel`` approximation: + + - ``'gaussian'`` - scale of the Gibbs kernel. + - ``'arccos'`` - standard deviation of the random projections. + n: Order of the arc-cosine kernel, see :cite:`cho:09` for reference. + rng: Random key used for seeding. + + Returns: + Low-rank kernel geometry. + """ + rng = utils.default_prng_key(rng) + if kernel == "gaussian": + r = jnp.maximum( + jnp.linalg.norm(x, axis=-1).max(), + jnp.linalg.norm(y, axis=-1).max() + ) + k1 = _gaussian_kernel(rng, x, rank, eps=std, R=r) + k2 = _gaussian_kernel(rng, y, rank, eps=std, R=r) + eps = std + elif kernel == "arccos": + k1 = _arccos_kernel(rng, x, rank, n=n, std=std) + k2 = _arccos_kernel(rng, y, rank, n=n, std=std) + eps = 1.0 + else: + raise NotImplementedError(kernel) + + return cls(k1, k2, epsilon=eps) + + def apply_kernel( # noqa: D102 + self, + scaling: jnp.ndarray, + eps: Optional[float] = None, + axis: int = 0, + ) -> jnp.ndarray: + if axis == 0: + return self.k2 @ (self.k1.T @ scaling) + return self.k1 @ (self.k2.T @ scaling) + + @property + def kernel_matrix(self) -> jnp.ndarray: # noqa: D102 + return self.k1 @ self.k2.T + + @property + def cost_matrix(self) -> jnp.ndarray: # noqa: D102 + eps = jnp.finfo(self.dtype).tiny + return -self.epsilon * jnp.log(self.kernel_matrix + eps) + + @property + def rank(self) -> int: # noqa: D102 + return self.k1.shape[1] + + @property + def shape(self) -> Tuple[int, int]: # noqa: D102 + return self.k1.shape[0], self.k2.shape[0] + + @property + def dtype(self) -> jnp.dtype: # noqa: D102 + return self.k1.dtype + + def transport_from_potentials( + self, f: jnp.ndarray, g: jnp.ndarray + ) -> jnp.ndarray: + """Not implemented.""" + raise ValueError("Not implemented.") + + def tree_flatten(self): # noqa: D102 + return [self.k1, self.k2, self._epsilon_init], {} + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + return cls(*children, **aux_data) + + +def _gaussian_kernel( + rng: jax.Array, + x: jnp.ndarray, + n_features: int, + eps: float, + R: jnp.ndarray, +) -> jnp.ndarray: + _, d = x.shape + cost_fn = costs.SqEuclidean() + + y = (R ** 2) / (eps * d) + q = y / (2.0 * mu.lambertw(y)) + sigma = jnp.sqrt(q * eps * 0.25) + + u = jax.random.normal(rng, shape=(n_features, d)) * sigma + cost = cost_fn.all_pairs(x, u) + norm_u = cost_fn.norm(u) + + tmp = -2.0 * (cost / eps) + (norm_u / (eps * q)) + phi = (2 * q) ** (d / 4) * jnp.exp(tmp) + + return (1.0 / jnp.sqrt(n_features)) * phi + + +def _arccos_kernel( + rng: jax.Array, + x: jnp.ndarray, + n_features: int, + n: int, + std: float = 1.0, + kappa: float = 1e-6, +) -> jnp.ndarray: + n_points, d = x.shape + c = jnp.sqrt(2) * (std ** (d / 2)) + + u = jax.random.normal(rng, shape=(n_features, d)) * std + tmp = -(1 / 4) * jnp.sum(u ** 2, axis=-1) * (1.0 - (1.0 / (std ** 2))) + phi = c * (jnp.maximum(0.0, (x @ u.T)) ** n) * jnp.exp(tmp) + + return jnp.c_[(1.0 / jnp.sqrt(n_features)) * phi, + jnp.full((n_points,), fill_value=kappa)] diff --git a/ott/build/lib/ott/geometry/pointcloud.py b/ott/build/lib/ott/geometry/pointcloud.py new file mode 100644 index 0000000..313ecb3 --- /dev/null +++ b/ott/build/lib/ott/geometry/pointcloud.py @@ -0,0 +1,792 @@ +# 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.geometry import costs, geometry, low_rank +from ott.math import utils as mu + +__all__ = ["PointCloud"] + + +@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.0 + + @property + def _norm_y(self) -> Union[float, jnp.ndarray]: + if self._axis_norm == 0: + return self.cost_fn.norm(self.y) + return 0.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, jax.Array)): + 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 + + 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 + + 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 + ) + + 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 + ) + + 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 + ) + + def apply_cost( + self, + arr: jnp.ndarray, + axis: int = 0, + fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, + is_linear: bool = False, + sub_idx=None, + ) -> 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, sub_idx=None) + + def _apply_cost( + self, arr: jnp.ndarray, axis: int = 0, fn=None, sub_idx=None + ) -> jnp.ndarray: + """See :meth:`apply_cost`.""" + if not self.is_online or sub_idx is not None: + return super().apply_cost(arr, axis, fn, sub_idx=sub_idx) + + 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 + ) + + def vec_apply_cost( + self, + arr: jnp.ndarray, + axis: int = 0, + fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, + sub_idx=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." + ) + + 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] + + @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.0 / self.inv_scale_cost + cost_fn = costs.SqEuclidean() + return type(self).tree_unflatten(aux_data, [x, y] + args + [cost_fn]) + + 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, + ) + + 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[ixs, ...] + + return self._mask_subset_helper( + src_ixs, tgt_ixs, fn=subset_fn, propagate_mask=True, **kwargs + ) + + def mask( # noqa: D102 + self, + src_mask: Optional[jnp.ndarray], + tgt_mask: Optional[jnp.ndarray], + mask_value: float = 0.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)) diff --git a/ott/build/lib/ott/geometry/segment.py b/ott/build/lib/ott/geometry/segment.py new file mode 100644 index 0000000..20a1ee9 --- /dev/null +++ b/ott/build/lib/ott/geometry/segment.py @@ -0,0 +1,188 @@ +# +# 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. +from typing import Callable, Optional, Tuple + +import jax +import jax.numpy as jnp + +__all__ = ["segment_point_cloud"] + + +def segment_point_cloud( + x: jnp.ndarray, + a: Optional[jnp.ndarray] = None, + num_segments: Optional[int] = None, + max_measure_size: Optional[int] = None, + segment_ids: Optional[jnp.ndarray] = None, + indices_are_sorted: bool = False, + num_per_segment: Optional[Tuple[int, ...]] = None, + padding_vector: Optional[jnp.ndarray] = None +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Segment and pad as needed the entries of a point cloud. + + There are two interfaces: + + #. use ``segment_ids``, and optionally ``indices_are_sorted`` to describe + for each data point in the matrix to which segment it belongs to. + #. use ``num_per_segment`` which describes contiguous segments. + + If using the first interface, ``num_segments`` is required for jitting. + Assumes ``range(0, num_segments)`` are the segment ids. + + In both cases, jitting requires defining a ``max_measure_size``, the + upper bound on the maximal size of measures, which will be used for padding. + + Args: + x: Array of input points, of shape ``[num_x, ndim]``. + Multiple segments are held in this single array. + a: Array of shape ``[num_x,]`` containing the weights (within each measure) + of all the points. + num_segments: Number of segments. Required for jitting. + If `None` and using the second interface, it will be computed as + ``len(num_per_segment)``. + max_measure_size: Overall size of padding. Required for jitting. + If `None` and using the second interface, it will be computed as + ``max(num_per_segment)``. + segment_ids: **1st interface** The segment ids for which each row of ``x`` + belongs. This is a similar interface to :func:`jax.ops.segment_sum`. + indices_are_sorted: **1st interface** Whether ``segment_ids`` are sorted. + num_per_segment: **2nd interface** Number of points in each segment. + For example, `[100, 20, 30]` would imply that ``x`` is segmented into 3 + arrays of length `[100]`, `[20]`, and `[30]`, respectively. + Must be a tuple and not a :class:`jax.numpy.ndarray` to allow jitting. + This means changes in ``num_per_segment`` will re-trigger compilation. + padding_vector: vector to be used to pad point cloud matrices. Most likely + to be zero, but can be adjusted to be other values to avoid errors or + over/underflow in cost matrix that could be problematic (even these values + are not supposed to be taken given their corresponding masses are 0). + See also :func:`~ott.geometry.costs.CostFn._padder`. + If ``None``, vector of 0s of shape ``[1, ndim]`` is used. + + Returns: + Segmented ``x`` as an array of shape + ``[num_measures, max_measure_size, ndim]`` and ``a`` as an array of shape + ``[num_measures, max_measure_size]``. + """ + num, dim = x.shape + use_segment_ids = segment_ids is not None + if use_segment_ids: + assert num_segments is not None, "Please specify `num_segments`." + assert max_measure_size is not None, "Please specify `max_measure_size`." + num_per_segment = jax.ops.segment_sum( + jnp.ones_like(segment_ids), + segment_ids, + num_segments=num_segments, + indices_are_sorted=indices_are_sorted + ) + else: + assert num_per_segment is not None, "Please specify `num_per_segment`." + if max_measure_size is None: + max_measure_size = max(num_per_segment) + if num_segments is None: + num_segments = len(num_per_segment) + else: + assert num_segments == len(num_per_segment) + # conversion to facilitate computation of default weight below. + num_per_segment = jnp.array(num_per_segment) + segment_ids = jnp.arange(num_segments).repeat( + num_per_segment, total_repeat_length=num + ) + + if a is None: + a = jnp.array( + (1.0 / + num_per_segment).repeat(num_per_segment, total_repeat_length=num) + ) + + if padding_vector is None: + padding_vector = jnp.zeros((1, dim)) + + x = jnp.concatenate((x, padding_vector)) + a = jnp.concatenate((a, jnp.zeros((1,)))) + segmented_a, segmented_x = [], [] + + for i in range(num_segments): + idx = jnp.where(segment_ids == i, jnp.arange(num), num + 1) + idx = jax.lax.dynamic_slice(jnp.sort(idx), (0,), (max_measure_size,)) + + # segment the weights + segmented_a.append(a.at[idx].get()) + # segment the positions + segmented_x.append(x.at[idx].get()) + + segmented_a = jnp.stack(segmented_a) + segmented_x = jnp.stack(segmented_x) + + return segmented_x, segmented_a + + +def _segment_interface( + x: jnp.ndarray, + y: jnp.ndarray, + eval_fn: Callable[[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray], + jnp.ndarray], + num_segments: Optional[int] = None, + max_measure_size: Optional[int] = None, + segment_ids_x: Optional[jnp.ndarray] = None, + segment_ids_y: Optional[jnp.ndarray] = None, + indices_are_sorted: bool = False, + num_per_segment_x: Optional[jnp.ndarray] = None, + num_per_segment_y: Optional[jnp.ndarray] = None, + weights_x: Optional[jnp.ndarray] = None, + weights_y: Optional[jnp.ndarray] = None, + padding_vector: Optional[jnp.ndarray] = None, +) -> jnp.ndarray: + """Wrapper to segment two point clouds and return parallel evaluations. + + Utility function that segments two point clouds using the approach outlined + in `segment_point_cloud` and evaluates `eval_fn` on pairs of segmented point + clouds. + """ + use_segment_ids = segment_ids_x is not None + if use_segment_ids: + assert segment_ids_y is not None + else: + assert num_per_segment_x is not None + assert num_per_segment_y is not None + + segmented_x, segmented_weights_x = segment_point_cloud( + x, + a=weights_x, + num_segments=num_segments, + max_measure_size=max_measure_size, + segment_ids=segment_ids_x, + indices_are_sorted=indices_are_sorted, + num_per_segment=num_per_segment_x, + padding_vector=padding_vector + ) + + segmented_y, segmented_weights_y = segment_point_cloud( + y, + a=weights_y, + num_segments=num_segments, + max_measure_size=max_measure_size, + segment_ids=segment_ids_y, + indices_are_sorted=indices_are_sorted, + num_per_segment=num_per_segment_y, + padding_vector=padding_vector + ) + + v_eval = jax.vmap(eval_fn, in_axes=[0] * 4) + return v_eval( + segmented_x, + segmented_y, + segmented_weights_x, + segmented_weights_y, + ) diff --git a/ott/build/lib/ott/initializers/__init__.py b/ott/build/lib/ott/initializers/__init__.py new file mode 100644 index 0000000..5406247 --- /dev/null +++ b/ott/build/lib/ott/initializers/__init__.py @@ -0,0 +1,14 @@ +# 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. +from . import linear, quadratic diff --git a/ott/build/lib/ott/initializers/linear/__init__.py b/ott/build/lib/ott/initializers/linear/__init__.py new file mode 100644 index 0000000..c7a7dc4 --- /dev/null +++ b/ott/build/lib/ott/initializers/linear/__init__.py @@ -0,0 +1,14 @@ +# 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. +from . import initializers, initializers_lr diff --git a/ott/build/lib/ott/initializers/linear/initializers.py b/ott/build/lib/ott/initializers/linear/initializers.py new file mode 100644 index 0000000..e486349 --- /dev/null +++ b/ott/build/lib/ott/initializers/linear/initializers.py @@ -0,0 +1,408 @@ +# 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 abc +from typing import Any, Dict, Optional, Sequence, Tuple + +import jax +import jax.numpy as jnp + +from ott import utils +from ott.geometry import pointcloud +from ott.problems.linear import linear_problem + +__all__ = [ + "DefaultInitializer", "GaussianInitializer", "SortingInitializer", + "SubsampleInitializer" +] + + +@jax.tree_util.register_pytree_node_class +class SinkhornInitializer(abc.ABC): + """Base class for Sinkhorn initializers.""" + + @abc.abstractmethod + def init_dual_a( + self, + ot_prob: linear_problem.LinearProblem, + lse_mode: bool, + rng: Optional[jax.Array] = None, + ) -> jnp.ndarray: + """Initialize Sinkhorn potential/scaling f_u. + + Args: + ot_prob: Linear OT problem. + lse_mode: Return potential if ``True``, scaling if ``False``. + rng: Random number generator for stochastic initializers. + + Returns: + potential/scaling, array of size ``[n,]``. + """ + + @abc.abstractmethod + def init_dual_b( + self, + ot_prob: linear_problem.LinearProblem, + lse_mode: bool, + rng: Optional[jax.Array] = None, + ) -> jnp.ndarray: + """Initialize Sinkhorn potential/scaling g_v. + + Args: + ot_prob: Linear OT problem. + lse_mode: Return potential if ``True``, scaling if ``False``. + rng: Random number generator for stochastic initializers. + + Returns: + potential/scaling, array of size ``[m,]``. + """ + + def __call__( + self, + ot_prob: linear_problem.LinearProblem, + a: Optional[jnp.ndarray], + b: Optional[jnp.ndarray], + lse_mode: bool, + rng: Optional[jax.Array] = None, + ) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Initialize Sinkhorn potentials/scalings f_u and g_v. + + Args: + ot_prob: Linear OT problem. + a: Initial potential/scaling f_u. + If ``None``, it will be initialized using :meth:`init_dual_a`. + b: Initial potential/scaling g_v. + If ``None``, it will be initialized using :meth:`init_dual_b`. + lse_mode: Return potentials if ``True``, scalings if ``False``. + rng: Random number generator for stochastic initializers. + + Returns: + The initial potentials/scalings. + """ + rng = utils.default_prng_key(rng) + rng_x, rng_y = jax.random.split(rng, 2) + n, m = ot_prob.geom.shape + if a is None: + a = self.init_dual_a(ot_prob, lse_mode=lse_mode, rng=rng_x) + if b is None: + b = self.init_dual_b(ot_prob, lse_mode=lse_mode, rng=rng_y) + + assert a.shape == ( + n, + ), f"Expected `f_u` to have shape `{n,}`, found `{a.shape}`." + assert b.shape == ( + m, + ), f"Expected `g_v` to have shape `{m,}`, found `{b.shape}`." + + # cancel dual variables for zero weights + a = jnp.where(ot_prob.a > 0.0, a, -jnp.inf if lse_mode else 0.0) + b = jnp.where(ot_prob.b > 0.0, b, -jnp.inf if lse_mode else 0.0) + + return a, b + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return [], {} + + @classmethod + def tree_unflatten( # noqa: D102 + cls, aux_data: Dict[str, Any], children: Sequence[Any] + ) -> "SinkhornInitializer": + return cls(*children, **aux_data) + + +@jax.tree_util.register_pytree_node_class +class DefaultInitializer(SinkhornInitializer): + """Default initialization of Sinkhorn dual potentials/primal scalings.""" + + def init_dual_a( # noqa: D102 + self, + ot_prob: linear_problem.LinearProblem, + lse_mode: bool, + rng: Optional[jax.Array] = None, + ) -> jnp.ndarray: + del rng + return jnp.zeros_like(ot_prob.a) if lse_mode else jnp.ones_like(ot_prob.a) + + def init_dual_b( # noqa: D102 + self, + ot_prob: linear_problem.LinearProblem, + lse_mode: bool, + rng: Optional[jax.Array] = None, + ) -> jnp.ndarray: + del rng + return jnp.zeros_like(ot_prob.b) if lse_mode else jnp.ones_like(ot_prob.b) + + +@jax.tree_util.register_pytree_node_class +class GaussianInitializer(DefaultInitializer): + """Gaussian initializer :cite:`thornton2022rethinking:22`. + + Compute Gaussian approximations of each + :class:`~ott.geometry.pointcloud.PointCloud`, then compute closed from + Kantorovich potential between Gaussian approximations using Brenier's theorem + (adapt convex/Brenier potential to Kantorovich). Use this Gaussian potential + to initialize Sinkhorn potentials/scalings. + """ + + def init_dual_a( # noqa: D102 + self, + ot_prob: linear_problem.LinearProblem, + lse_mode: bool, + rng: Optional[jax.Array] = None, + ) -> jnp.ndarray: + # import Gaussian here due to circular imports + from ott.tools.gaussian_mixture import gaussian + + del rng + assert isinstance( + ot_prob.geom, pointcloud.PointCloud + ), "Gaussian initializer valid only for pointcloud geoms." + + x, y = ot_prob.geom.x, ot_prob.geom.y + a, b = ot_prob.a, ot_prob.b + + gaussian_a = gaussian.Gaussian.from_samples(x, weights=a) + gaussian_b = gaussian.Gaussian.from_samples(y, weights=b) + # Brenier potential for cost ||x-y||^2/2, multiply by two for ||x-y||^2 + f_potential = 2 * gaussian_a.f_potential(dest=gaussian_b, points=x) + f_potential = f_potential - jnp.mean(f_potential) + return f_potential if lse_mode else ot_prob.geom.scaling_from_potential( + f_potential + ) + + +@jax.tree_util.register_pytree_node_class +class SortingInitializer(DefaultInitializer): + """Sorting initializer :cite:`thornton2022rethinking:22`. + + Solve non-regularized OT problem via sorting, then compute potential through + iterated minimum on C-transform and use this potential to initialize + regularized potential. + + Args: + vectorized_update: Whether to use vectorized loop. + tolerance: DualSort convergence threshold. + max_iter: Max DualSort steps. + """ + + def __init__( + self, + vectorized_update: bool = True, + tolerance: float = 1e-2, + max_iter: int = 100 + ): + super().__init__() + self.tolerance = tolerance + self.max_iter = max_iter + self.vectorized_update = vectorized_update + + def _init_sorting_dual( + self, modified_cost: jnp.ndarray, init_f: jnp.ndarray + ) -> jnp.ndarray: + """Run DualSort algorithm. + + Args: + modified_cost: cost matrix minus diagonal column-wise. + init_f: potential f, array of size n. This is the starting potential, + which is then updated to make the init potential, so an init of an init. + + Returns: + potential f, array of size n. + """ + + def body_fn( + state: Tuple[jnp.ndarray, float, int] + ) -> Tuple[jnp.ndarray, float, int]: + prev_f, _, it = state + new_f = fn(prev_f, modified_cost) + diff = jnp.sum((new_f - prev_f) ** 2) + it += 1 + return new_f, diff, it + + def cond_fn(state: Tuple[jnp.ndarray, float, int]) -> bool: + _, diff, it = state + return jnp.logical_and(diff > self.tolerance, it < self.max_iter) + + fn = _vectorized_update if self.vectorized_update else _coordinate_update + state = (init_f, jnp.inf, 0) # init, error, iter + f_potential, _, _ = jax.lax.while_loop( + cond_fun=cond_fn, body_fun=body_fn, init_val=state + ) + + return f_potential + + def init_dual_a( + self, + ot_prob: linear_problem.LinearProblem, + lse_mode: bool, + rng: Optional[jax.Array] = None, + init_f: Optional[jnp.ndarray] = None, + ) -> jnp.ndarray: + """Apply DualSort algorithm. + + Args: + ot_prob: OT problem between discrete distributions. + lse_mode: Return potential if ``True``, scaling if ``False``. + rng: Random number generator for stochastic initializers, unused. + init_f: potential f, array of size ``[n,]``. This is the starting + potential, which is then updated to make the init potential, + so an init of an init. + + Returns: + potential/scaling f_u, array of size ``[n,]``. + """ + del rng + assert not ot_prob.geom.is_online, \ + "Sorting initializer does not work for online geometry." + # check for sorted x, y requires point cloud and could slow initializer + cost_matrix = ot_prob.geom.cost_matrix + + assert cost_matrix.shape[0] == cost_matrix.shape[ + 1], "Requires square cost matrix." + + modified_cost = cost_matrix - jnp.diag(cost_matrix)[None, :] + + n = cost_matrix.shape[0] + init_f = jnp.zeros(n) if init_f is None else init_f + + f_potential = self._init_sorting_dual(modified_cost, init_f) + f_potential = f_potential - jnp.mean(f_potential) + + return f_potential if lse_mode else ot_prob.geom.scaling_from_potential( + f_potential + ) + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return ([], { + "tolerance": self.tolerance, + "max_iter": self.max_iter, + "vectorized_update": self.vectorized_update + }) + + +@jax.tree_util.register_pytree_node_class +class SubsampleInitializer(DefaultInitializer): + """Subsample initializer :cite:`thornton2022rethinking:22`. + + Subsample each :class:`~ott.geometry.pointcloud.PointCloud`, then compute + :class:`Sinkhorn potential ` + from the subsampled approximations and use this potential to initialize + Sinkhorn potentials/scalings for the original problem. + + Args: + subsample_n_x: number of points to subsample from the first measure in + :class:`~ott.geometry.pointcloud.PointCloud`. + subsample_n_y: number of points to subsample from the second measure in + :class:`~ott.geometry.pointcloud.PointCloud`. + If ``None``, use ``subsample_n_x``. + kwargs: Keyword arguments for + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. + """ + + def __init__( + self, + subsample_n_x: int, + subsample_n_y: Optional[int] = None, + **kwargs: Any, + ): + super().__init__() + self.subsample_n_x = subsample_n_x + self.subsample_n_y = subsample_n_y or subsample_n_x + self.sinkhorn_kwargs = kwargs + + def init_dual_a( # noqa: D102 + self, + ot_prob: linear_problem.LinearProblem, + lse_mode: bool, + rng: Optional[jax.Array] = None, + ) -> jnp.ndarray: + from ott.solvers import linear + + assert isinstance( + ot_prob.geom, pointcloud.PointCloud + ), "Subsample initializer valid only for pointcloud geom." + rng = utils.default_prng_key(rng) + rng_x, rng_y = jax.random.split(rng, 2) + + x, y = ot_prob.geom.x, ot_prob.geom.y + a, b = ot_prob.a, ot_prob.b + + # subsample + sub_x = jax.random.choice( + rng_x, a=x, shape=(self.subsample_n_x,), replace=True, p=a, axis=0 + ) + sub_y = jax.random.choice( + rng_y, a=y, shape=(self.subsample_n_y,), replace=True, p=b, axis=0 + ) + + # create subsampled point cloud geometry + sub_geom = pointcloud.PointCloud( + sub_x, + sub_y, + epsilon=ot_prob.geom.epsilon, + scale_cost=ot_prob.geom._scale_cost, + cost_fn=ot_prob.geom.cost_fn + ) + + # run sinkhorn + subsample_sink_out = linear.solve(sub_geom, **self.sinkhorn_kwargs) + + # interpolate potentials + dual_potentials = subsample_sink_out.to_dual_potentials() + f_potential = jax.vmap(dual_potentials.f)(x) + + return f_potential if lse_mode else ot_prob.geom.scaling_from_potential( + f_potential + ) + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return ([], { + "subsample_n_x": self.subsample_n_x, + "subsample_n_y": self.subsample_n_y, + **self.sinkhorn_kwargs + }) + + +def _vectorized_update( + f: jnp.ndarray, modified_cost: jnp.ndarray +) -> jnp.ndarray: + """Inner loop DualSort Update. + + Args: + f: potential f, array of size n. + modified_cost: cost matrix minus diagonal column-wise. + + Returns: + updated potential vector, f. + """ + return jnp.min(modified_cost + f[None, :], axis=1) + + +def _coordinate_update( + f: jnp.ndarray, modified_cost: jnp.ndarray +) -> jnp.ndarray: + """Coordinate-wise updates within inner loop. + + Args: + f: potential f, array of size n. + modified_cost: cost matrix minus diagonal column-wise. + + Returns: + updated potential vector, f. + """ + + def body_fn(i: int, f: jnp.ndarray) -> jnp.ndarray: + new_f = jnp.min(modified_cost[i, :] + f) + return f.at[i].set(new_f) + + return jax.lax.fori_loop(0, len(f), body_fn, f) diff --git a/ott/build/lib/ott/initializers/linear/initializers_lr.py b/ott/build/lib/ott/initializers/linear/initializers_lr.py new file mode 100644 index 0000000..8206f29 --- /dev/null +++ b/ott/build/lib/ott/initializers/linear/initializers_lr.py @@ -0,0 +1,654 @@ +# 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 abc +import functools +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Literal, + Mapping, + NamedTuple, + Optional, + Sequence, + Tuple, + Union, +) + +import jax +import jax.numpy as jnp +import numpy as np + +from ott import utils +from ott.geometry import geometry, low_rank, pointcloud +from ott.math import fixed_point_loop +from ott.math import utils as mu + +if TYPE_CHECKING: + from ott.problems.linear import linear_problem + from ott.problems.quadratic import quadratic_problem + from ott.solvers.linear import sinkhorn, sinkhorn_lr + from ott.solvers.quadratic import gromov_wasserstein_lr + +Problem_t = Union["linear_problem.LinearProblem", + "quadratic_problem.QuadraticProblem"] + +__all__ = [ + "RandomInitializer", "Rank2Initializer", "KMeansInitializer", + "GeneralizedKMeansInitializer" +] + + +@jax.tree_util.register_pytree_node_class +class LRInitializer(abc.ABC): + """Base class for low-rank initializers. + + Args: + rank: Rank of the factorization. + kwargs: Additional keyword arguments. + """ + + def __init__(self, rank: int, **kwargs: Any): + self._rank = rank + self._kwargs = kwargs + + @abc.abstractmethod + def init_q( + self, + ot_prob: Problem_t, + rng: jax.Array, + *, + init_g: jnp.ndarray, + **kwargs: Any, + ) -> jnp.ndarray: + """Initialize the low-rank factor :math:`Q`. + + Args: + ot_prob: OT problem. + rng: Random key for seeding. + init_g: Initial value for :math:`g` factor. + kwargs: Additional keyword arguments. + + Returns: + Array of shape ``[n, rank]``. + """ + + @abc.abstractmethod + def init_r( + self, + ot_prob: Problem_t, + rng: jax.Array, + *, + init_g: jnp.ndarray, + **kwargs: Any, + ) -> jnp.ndarray: + """Initialize the low-rank factor :math:`R`. + + Args: + ot_prob: Linear OT problem. + rng: Random key for seeding. + init_g: Initial value for :math:`g` factor. + kwargs: Additional keyword arguments. + + Returns: + Array of shape ``[m, rank]``. + """ + + @abc.abstractmethod + def init_g( + self, + ot_prob: Problem_t, + rng: jax.Array, + **kwargs: Any, + ) -> jnp.ndarray: + """Initialize the low-rank factor :math:`g`. + + Args: + ot_prob: OT problem. + rng: Random key for seeding. + kwargs: Additional keyword arguments. + + Returns: + Array of shape ``[rank,]``. + """ + + @classmethod + def from_solver( + cls, + solver: Union["sinkhorn_lr.LRSinkhorn", + "gromov_wasserstein_lr.LRGromovWasserstein"], + *, + kind: Literal["random", "rank2", "k-means", "generalized-k-means"], + **kwargs: Any, + ) -> "LRInitializer": + """Create a low-rank initializer from a linear or quadratic solver. + + Args: + solver: Low-rank linear or quadratic solver. + kind: Which initializer to instantiate. + kwargs: Keyword arguments when creating the initializer. + + Returns: + Low-rank initializer. + """ + rank = solver.rank + sinkhorn_kwargs = { + "norm_error": solver._norm_error, + "lse_mode": solver.lse_mode, + "implicit_diff": solver.implicit_diff, + "use_danskin": solver.use_danskin + } + + if kind == "random": + return RandomInitializer(rank, **kwargs) + if kind == "rank2": + return Rank2Initializer(rank, **kwargs) + if kind == "k-means": + return KMeansInitializer(rank, sinkhorn_kwargs=sinkhorn_kwargs, **kwargs) + if kind == "generalized-k-means": + return GeneralizedKMeansInitializer( + rank, sinkhorn_kwargs=sinkhorn_kwargs, **kwargs + ) + raise NotImplementedError(f"Initializer `{kind}` is not implemented.") + + def __call__( + self, + ot_prob: Problem_t, + q: Optional[jnp.ndarray] = None, + r: Optional[jnp.ndarray] = None, + g: Optional[jnp.ndarray] = None, + *, + rng: Optional[jax.Array] = None, + **kwargs: Any + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Initialize the factors :math:`Q`, :math:`R` and :math:`g`. + + Args: + ot_prob: OT problem. + q: Factor of shape ``[n, rank]``. If `None`, it will be initialized + using :meth:`init_q`. + r: Factor of shape ``[m, rank]``. If `None`, it will be initialized + using :meth:`init_r`. + g: Factor of shape ``[rank,]``. If `None`, it will be initialized + using :meth:`init_g`. + rng: Random key for seeding. + kwargs: Additional keyword arguments for :meth:`init_q`, :meth:`init_r` + and :meth:`init_g`. + + Returns: + The factors :math:`Q`, :math:`R` and :math:`g`, respectively. + """ + rng = utils.default_prng_key(rng) + rng1, rng2, rng3 = jax.random.split(rng, 3) + + if g is None: + g = self.init_g(ot_prob, rng1, **kwargs) + if q is None: + q = self.init_q(ot_prob, rng2, init_g=g, **kwargs) + if r is None: + r = self.init_r(ot_prob, rng3, init_g=g, **kwargs) + + assert g.shape == (self.rank,) + assert q.shape == (ot_prob.a.shape[0], self.rank) + assert r.shape == (ot_prob.b.shape[0], self.rank) + + return q, r, g + + @property + def rank(self) -> int: + """Rank of the transport matrix factorization.""" + return self._rank + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return [], {**self._kwargs, "rank": self.rank} + + @classmethod + def tree_unflatten( # noqa: D102 + cls, aux_data: Dict[str, Any], children: Sequence[Any] + ) -> "LRInitializer": + return cls(*children, **aux_data) + + +@jax.tree_util.register_pytree_node_class +class RandomInitializer(LRInitializer): + """Low-rank Sinkhorn factorization using random factors. + + Args: + rank: Rank of the factorization. + kwargs: Additional keyword arguments. + """ + + def init_q( # noqa: D102 + self, + ot_prob: Problem_t, + rng: jax.Array, + *, + init_g: jnp.ndarray, + **kwargs: Any, + ) -> jnp.ndarray: + del kwargs, init_g + a = ot_prob.a + init_q = jnp.abs(jax.random.normal(rng, (a.shape[0], self.rank))) + return a[:, None] * (init_q / jnp.sum(init_q, axis=1, keepdims=True)) + + def init_r( # noqa: D102 + self, + ot_prob: Problem_t, + rng: jax.Array, + *, + init_g: jnp.ndarray, + **kwargs: Any, + ) -> jnp.ndarray: + del kwargs, init_g + b = ot_prob.b + init_r = jnp.abs(jax.random.normal(rng, (b.shape[0], self.rank))) + return b[:, None] * (init_r / jnp.sum(init_r, axis=1, keepdims=True)) + + def init_g( # noqa: D102 + self, + ot_prob: Problem_t, + rng: jax.Array, + **kwargs: Any, + ) -> jnp.ndarray: + del kwargs + init_g = jnp.abs(jax.random.uniform(rng, (self.rank,))) + 1.0 + return init_g / jnp.sum(init_g) + + +@jax.tree_util.register_pytree_node_class +class Rank2Initializer(LRInitializer): + """Low-rank Sinkhorn factorization using rank-2 factors :cite:`scetbon:21`. + + Args: + rank: Rank of the factorization. + kwargs: Additional keyword arguments. + """ + + def _compute_factor( + self, + ot_prob: Problem_t, + init_g: jnp.ndarray, + *, + which: Literal["q", "r"], + ) -> jnp.ndarray: + a, b = ot_prob.a, ot_prob.b + marginal = a if which == "q" else b + n, r = marginal.shape[0], self.rank + + lambda_1 = jnp.min( + jnp.array([jnp.min(a), jnp.min(init_g), + jnp.min(b)]) + ) * 0.5 + + g1 = jnp.arange(1, r + 1) + g1 /= g1.astype(float).sum() + g2 = (init_g - lambda_1 * g1) / (1.0 - lambda_1) + + x = jnp.arange(1, n + 1) + x /= x.astype(float).sum() + y = (marginal - lambda_1 * x) / (1.0 - lambda_1) + + return ((lambda_1 * x[:, None] @ g1.reshape(1, -1)) + + ((1.0 - lambda_1) * y[:, None] @ g2.reshape(1, -1))) + + def init_q( # noqa: D102 + self, + ot_prob: Problem_t, + rng: jax.Array, + *, + init_g: jnp.ndarray, + **kwargs: Any, + ) -> jnp.ndarray: + del rng, kwargs + return self._compute_factor(ot_prob, init_g, which="q") + + def init_r( # noqa: D102 + self, + ot_prob: Problem_t, + rng: jax.Array, + *, + init_g: jnp.ndarray, + **kwargs: Any, + ) -> jnp.ndarray: + del rng, kwargs + return self._compute_factor(ot_prob, init_g, which="r") + + def init_g( # noqa: D102 + self, + ot_prob: Problem_t, + rng: jax.Array, + **kwargs: Any, + ) -> jnp.ndarray: + del rng, kwargs + return jnp.ones((self.rank,)) / self.rank + + +@jax.tree_util.register_pytree_node_class +class KMeansInitializer(LRInitializer): + """K-means initializer for low-rank Sinkhorn :cite:`scetbon:22b`. + + Applicable for :class:`~ott.geometry.pointcloud.PointCloud` and + :class:`~ott.geometry.low_rank.LRCGeometry`. + + Args: + rank: Rank of the factorization. + min_iterations: Minimum number of k-means iterations. + max_iterations: Maximum number of k-means iterations. + sinkhorn_kwargs: Keyword arguments for + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. + kwargs: Keyword arguments for :func:`~ott.tools.k_means.k_means`. + """ + + def __init__( + self, + rank: int, + min_iterations: int = 100, + max_iterations: int = 100, + sinkhorn_kwargs: Optional[Mapping[str, Any]] = None, + **kwargs: Any + ): + super().__init__(rank, **kwargs) + self._min_iter = min_iterations + self._max_iter = max_iterations + self._sinkhorn_kwargs = {} if sinkhorn_kwargs is None else sinkhorn_kwargs + + @staticmethod + def _extract_array(geom: geometry.Geometry, *, first: bool) -> jnp.ndarray: + if isinstance(geom, pointcloud.PointCloud): + return geom.x if first else geom.y + if isinstance(geom, low_rank.LRCGeometry): + return geom.cost_1 if first else geom.cost_2 + raise TypeError( + f"k-means initializer not implemented for `{type(geom).__name__}`." + ) + + def _compute_factor( + self, + ot_prob: Problem_t, + rng: jax.Array, + *, + init_g: jnp.ndarray, + which: Literal["q", "r"], + **kwargs: Any, + ) -> jnp.ndarray: + from ott.problems.linear import linear_problem + from ott.problems.quadratic import quadratic_problem + from ott.solvers.linear import sinkhorn + from ott.tools import k_means + + del kwargs + fn = functools.partial( + k_means.k_means, + min_iterations=self._min_iter, + max_iterations=self._max_iter, + **self._kwargs + ) + + if isinstance(ot_prob, quadratic_problem.QuadraticProblem): + if ot_prob.geom_xy is not None and ot_prob.fused_penalty >= 1.0: + # prefer the linear term if it has a higher weight + geom = ot_prob.geom_xy + else: + geom = ot_prob.geom_xx if which == "q" else ot_prob.geom_yy + else: + geom = ot_prob.geom + arr = self._extract_array(geom, first=which == "q") + marginals = ot_prob.a if which == "q" else ot_prob.b + + centroids = fn(arr, self.rank, rng=rng).centroids + geom = pointcloud.PointCloud( + arr, centroids, epsilon=1e-1, scale_cost="max_cost" + ) + + prob = linear_problem.LinearProblem(geom, marginals, init_g) + solver = sinkhorn.Sinkhorn(**self._sinkhorn_kwargs) + return solver(prob).matrix + + def init_q( # noqa: D102 + self, + ot_prob: Problem_t, + rng: jax.Array, + *, + init_g: jnp.ndarray, + **kwargs: Any, + ) -> jnp.ndarray: + return self._compute_factor( + ot_prob, rng, init_g=init_g, which="q", **kwargs + ) + + def init_r( # noqa: D102 + self, + ot_prob: Problem_t, + rng: jax.Array, + *, + init_g: jnp.ndarray, + **kwargs: Any, + ) -> jnp.ndarray: + return self._compute_factor( + ot_prob, rng, init_g=init_g, which="r", **kwargs + ) + + def init_g( # noqa: D102 + self, + ot_prob: Problem_t, + rng: jax.Array, + **kwargs: Any, + ) -> jnp.ndarray: + del rng, kwargs + return jnp.ones((self.rank,)) / self.rank + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + children, aux_data = super().tree_flatten() + aux_data["sinkhorn_kwargs"] = self._sinkhorn_kwargs + aux_data["min_iterations"] = self._min_iter + aux_data["max_iterations"] = self._max_iter + return children, aux_data + + +class GeneralizedKMeansInitializer(KMeansInitializer): + """Generalized k-means initializer :cite:`scetbon:22b`. + + Applicable for any :class:`~ott.geometry.geometry.Geometry` with a + square shape. + + Args: + rank: Rank of the factorization. + gamma: The (inverse of) gradient step size used by mirror descent. + min_iterations: Minimum number of iterations. + max_iterations: Maximum number of iterations. + inner_iterations: Number of iterations used by the algorithm before + re-evaluating progress. + threshold: Convergence threshold. + sinkhorn_kwargs: Keyword arguments for + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. + """ + + def __init__( + self, + rank: int, + gamma: float = 10.0, + min_iterations: int = 0, + max_iterations: int = 100, + inner_iterations: int = 10, + threshold: float = 1e-6, + sinkhorn_kwargs: Optional[Mapping[str, Any]] = None, + ): + super().__init__( + rank, + sinkhorn_kwargs=sinkhorn_kwargs, + # below argument are stored in `_kwargs` + gamma=gamma, + min_iterations=min_iterations, + max_iterations=max_iterations, + inner_iterations=inner_iterations, + threshold=threshold, + ) + + class Constants(NamedTuple): # noqa: D106 + solver: "sinkhorn.Sinkhorn" + geom: geometry.Geometry # (n, n) + marginal: jnp.ndarray # (n,) + g: jnp.ndarray # (r,) + gamma: float + threshold: float + + class State(NamedTuple): # noqa: D106 + factor: jnp.ndarray + criterions: jnp.ndarray + crossed_threshold: bool + + def _compute_factor( + self, + ot_prob: Problem_t, + rng: jax.Array, + *, + init_g: jnp.ndarray, + which: Literal["q", "r"], + **kwargs: Any, + ) -> jnp.ndarray: + from ott.problems.linear import linear_problem + from ott.problems.quadratic import quadratic_problem + from ott.solvers.linear import sinkhorn + + def init_fn() -> GeneralizedKMeansInitializer.State: + n = geom.shape[0] + factor = jnp.abs(jax.random.normal(rng, (n, self.rank))) + 1.0 # (n, r) + factor *= consts.marginal[:, None] / jnp.sum( + factor, axis=1, keepdims=True + ) + + return self.State( + factor, + criterions=-jnp.ones(outer_iterations), + crossed_threshold=False + ) + + # see the explanation in `ott.solvers.linear.sinkhorn_lr` + def converged( + state: GeneralizedKMeansInitializer.State, + consts: GeneralizedKMeansInitializer.Constants, iteration: int + ) -> bool: + + def conv_crossed(prev_err: float, curr_err: float) -> bool: + return jnp.logical_and( + prev_err < consts.threshold, curr_err < consts.threshold + ) + + def conv_not_crossed(prev_err: float, curr_err: float) -> bool: + return jnp.logical_and(curr_err < prev_err, curr_err < consts.threshold) + + it = iteration // inner_iterations + return jax.lax.cond( + state.crossed_threshold, conv_crossed, conv_not_crossed, + state.criterions[it - 2], state.criterions[it - 1] + ) + + def diverged( + state: GeneralizedKMeansInitializer.State, iteration: int + ) -> bool: + it = iteration // inner_iterations + return jnp.logical_not(jnp.isfinite(state.criterions[it - 1])) + + def cond_fn( + iteration: int, + consts: GeneralizedKMeansInitializer.Constants, + state: GeneralizedKMeansInitializer.State, + ) -> bool: + return jnp.logical_or( + iteration <= 2, + jnp.logical_and( + jnp.logical_not(diverged(state, iteration)), + jnp.logical_not(converged(state, consts, iteration)) + ) + ) + + def body_fn( + iteration: int, consts: GeneralizedKMeansInitializer.Constants, + state: GeneralizedKMeansInitializer.State, compute_error: bool + ) -> GeneralizedKMeansInitializer.State: + del compute_error + it = iteration // inner_iterations + + grad = consts.geom.apply_cost(state.factor, axis=1) # (n, r) + grad = grad + consts.geom.apply_cost(state.factor, axis=0) # (n, r) + grad = grad / consts.g + + norm = jnp.max(jnp.abs(grad)) ** 2 + gamma = consts.gamma / norm + eps = 1.0 / gamma + + cost = grad - eps * mu.safe_log(state.factor) # (n, r) + cost = geometry.Geometry( + cost_matrix=cost, + epsilon=eps, + ) + problem = linear_problem.LinearProblem( + cost, a=consts.marginal, b=consts.g + ) + + out = consts.solver(problem) + new_factor = out.matrix + + criterion = ((1 / gamma) ** 2) * ( + mu.kl(new_factor, state.factor) + mu.kl(state.factor, new_factor) + ) + crossed_threshold = jnp.logical_or( + state.crossed_threshold, + jnp.logical_and( + state.criterions[it - 1] >= consts.threshold, criterion + < consts.threshold + ) + ) + + return self.State( + factor=new_factor, + criterions=state.criterions.at[it].set(criterion), + crossed_threshold=crossed_threshold + ) + + del kwargs + + if isinstance(ot_prob, quadratic_problem.QuadraticProblem): + geom = ot_prob.geom_xx if which == "q" else ot_prob.geom_yy + else: + geom = ot_prob.geom + assert geom.shape[0] == geom.shape[ + 1], f"Expected the shape to be square, found `{geom.shape}`." + + inner_iterations = self._kwargs["inner_iterations"] + outer_iterations = np.ceil(self._max_iter / inner_iterations).astype(int) + force_scan = self._min_iter == self._max_iter + fixpoint_fn = ( + fixed_point_loop.fixpoint_iter + if force_scan else fixed_point_loop.fixpoint_iter_backprop + ) + + consts = self.Constants( + solver=sinkhorn.Sinkhorn(**self._sinkhorn_kwargs), + geom=geom.set_scale_cost("max_cost"), + marginal=ot_prob.a if which == "q" else ot_prob.b, + g=init_g, + gamma=self._kwargs["gamma"], + threshold=self._kwargs["threshold"], + ) + + return fixpoint_fn( + cond_fn, + body_fn, + min_iterations=self._min_iter, + max_iterations=self._max_iter, + inner_iterations=inner_iterations, + constants=consts, + state=init_fn(), + ).factor diff --git a/ott/build/lib/ott/initializers/quadratic/__init__.py b/ott/build/lib/ott/initializers/quadratic/__init__.py new file mode 100644 index 0000000..f188cae --- /dev/null +++ b/ott/build/lib/ott/initializers/quadratic/__init__.py @@ -0,0 +1,14 @@ +# 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. +from . import initializers diff --git a/ott/build/lib/ott/initializers/quadratic/initializers.py b/ott/build/lib/ott/initializers/quadratic/initializers.py new file mode 100644 index 0000000..795e81c --- /dev/null +++ b/ott/build/lib/ott/initializers/quadratic/initializers.py @@ -0,0 +1,193 @@ +# 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 abc +from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple + +import jax +import jax.numpy as jnp + +from ott.geometry import geometry + +if TYPE_CHECKING: + from ott.problems.linear import linear_problem + from ott.problems.quadratic import quadratic_problem + +__all__ = ["BaseQuadraticInitializer", "QuadraticInitializer"] + + +@jax.tree_util.register_pytree_node_class +class BaseQuadraticInitializer(abc.ABC): + """Base class for quadratic initializers. + + Args: + kwargs: Keyword arguments. + """ + + def __init__(self, **kwargs: Any): + self._kwargs = kwargs + + def __call__( + self, quad_prob: "quadratic_problem.QuadraticProblem", **kwargs: Any + ) -> "linear_problem.LinearProblem": + """Compute the initial linearization of a quadratic problem. + + Args: + quad_prob: Quadratic problem to linearize. + kwargs: Additional keyword arguments. + + Returns: + Linear problem. + """ + from ott.problems.linear import linear_problem + + n, m = quad_prob.geom_xx.shape[0], quad_prob.geom_yy.shape[0] + geom = self._create_geometry(quad_prob, **kwargs) + assert geom.shape == (n, m), ( + f"Expected geometry of shape `{n, m}`, " + f"found `{geom.shape}`." + ) + return linear_problem.LinearProblem( + geom, + a=quad_prob.a, + b=quad_prob.b, + tau_a=quad_prob.tau_a, + tau_b=quad_prob.tau_b, + ) + + @abc.abstractmethod + def _create_geometry( + self, quad_prob: "quadratic_problem.QuadraticProblem", **kwargs: Any + ) -> geometry.Geometry: + """Compute initial geometry for linearization. + + Args: + quad_prob: Quadratic problem. + kwargs: Additional keyword arguments. + + Returns: + Geometry used to initialize the linearized problem. + """ + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return [], self._kwargs + + @classmethod + def tree_unflatten( # noqa: D102 + cls, aux_data: Dict[str, Any], children: Sequence[Any] + ) -> "BaseQuadraticInitializer": + return cls(*children, **aux_data) + + +class QuadraticInitializer(BaseQuadraticInitializer): + r"""Initialize a linear problem locally around a selected coupling. + + If the problem is balanced (``tau_a = 1`` and ``tau_b = 1``), + the equation of the cost follows eq. 6, p. 1 of :cite:`peyre:16`. + + If the problem is unbalanced (``tau_a < 1`` or ``tau_b < 1``), there are two + possible cases. A first possibility is to introduce a quadratic KL + divergence on the marginals in the objective as done in :cite:`sejourne:21` + (``gw_unbalanced_correction = True``), which in turns modifies the + local cost matrix. + + Alternatively, it could be possible to leave the formulation of the + local cost unchanged, i.e. follow eq. 6, p. 1 of :cite:`peyre:16` + (``gw_unbalanced_correction = False``) and include the unbalanced terms + at the level of the linear problem only. + + Let :math:`P` [num_a, num_b] be the transport matrix, `cost_xx` is the + cost matrix of `geom_xx` and `cost_yy` is the cost matrix of `geom_yy`. + `left_x` and `right_y` depend on the loss chosen for GW. + `gw_unbalanced_correction` is flag indicating whether the unbalanced + correction applies. The equation of the local cost can be written as: + + .. math:: + + \text{marginal_dep_term} + \text{left}_x(\text{cost_xx}) P + \text{right}_y(\text{cost_yy}) + \text{unbalanced_correction} + + When working with the fused problem, a linear term is added to the cost + matrix: `cost_matrix` += `fused_penalty` * `geom_xy.cost_matrix` + + Args: + init_coupling: The coupling to use for initialization. If :obj:`None`, + defaults to the product coupling :math:`ab^T`. + """ + + def __init__( + self, init_coupling: Optional[jnp.ndarray] = None, **kwargs: Any + ): + super().__init__(**kwargs) + self.init_coupling = init_coupling + + def _create_geometry( + self, + quad_prob: "quadratic_problem.QuadraticProblem", + *, + epsilon: float, + relative_epsilon: Optional[bool] = None, + **kwargs: Any, + ) -> geometry.Geometry: + """Compute initial geometry for linearization. + + Args: + quad_prob: Quadratic OT problem. + epsilon: Epsilon regularization. + relative_epsilon: Flag, use `relative_epsilon` or not in geometry. + kwargs: Keyword arguments for :class:`~ott.geometry.geometry.Geometry`. + + Returns: + The initial geometry used to initialize the linearized problem. + """ + from ott.problems.quadratic import quadratic_problem + + del kwargs + + marginal_cost = quad_prob.marginal_dependent_cost(quad_prob.a, quad_prob.b) + geom_xx, geom_yy = quad_prob.geom_xx, quad_prob.geom_yy + + h1, h2 = quad_prob.quad_loss + if self.init_coupling is None: + tmp1 = quadratic_problem.apply_cost(geom_xx, quad_prob.a, axis=1, fn=h1) + tmp2 = quadratic_problem.apply_cost(geom_yy, quad_prob.b, axis=1, fn=h2) + tmp = jnp.outer(tmp1, tmp2) + else: + tmp1 = h1.func(geom_xx.cost_matrix) + tmp2 = h2.func(geom_yy.cost_matrix) + tmp = tmp1 @ self.init_coupling @ tmp2.T + + if quad_prob.is_balanced: + cost_matrix = marginal_cost.cost_matrix - tmp + else: + # initialize epsilon for Unbalanced GW according to Sejourne et. al (2021) + init_transport = jnp.outer(quad_prob.a, quad_prob.b) + marginal_1, marginal_2 = init_transport.sum(1), init_transport.sum(0) + + epsilon = quadratic_problem.update_epsilon_unbalanced( + epsilon=epsilon, transport_mass=marginal_1.sum() + ) + unbalanced_correction = quad_prob.cost_unbalanced_correction( + init_transport, marginal_1, marginal_2, epsilon=epsilon + ) + cost_matrix = marginal_cost.cost_matrix - tmp + unbalanced_correction + + cost_matrix += quad_prob.fused_penalty * quad_prob._fused_cost_matrix + return geometry.Geometry( + cost_matrix=cost_matrix, + epsilon=epsilon, + relative_epsilon=relative_epsilon + ) + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return [self.init_coupling], self._kwargs diff --git a/ott/build/lib/ott/math/__init__.py b/ott/build/lib/ott/math/__init__.py new file mode 100644 index 0000000..64bc1c0 --- /dev/null +++ b/ott/build/lib/ott/math/__init__.py @@ -0,0 +1,19 @@ +# 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. +from . import ( + fixed_point_loop, + matrix_square_root, + unbalanced_functions, + utils, +) diff --git a/ott/build/lib/ott/math/fixed_point_loop.py b/ott/build/lib/ott/math/fixed_point_loop.py new file mode 100644 index 0000000..9034eba --- /dev/null +++ b/ott/build/lib/ott/math/fixed_point_loop.py @@ -0,0 +1,239 @@ +# 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. +from typing import Any, Callable + +import jax +import jax.numpy as jnp +import numpy as np + +__all__ = ["fixpoint_iter", "fixpoint_iter_backprop"] + + +def fixpoint_iter( + cond_fn: Callable[[int, Any, Any], bool], + body_fn: Callable[[Any, Any, Any, Any], Any], min_iterations: int, + max_iterations: int, inner_iterations: int, constants: Any, state: Any +): + """Implementation of a fixed point loop. + + This fixed point loop iterator applies ``body_fn`` to a tuple + ``(iteration, constants, state, compute_error)`` to output a new state, using + context provided in iteration and constants. + + ``body_fn`` is iterated (inner_iterations -1) times, and one last time with + the ``compute_error`` flag to ``True``, indicating that additional + computational effort can be spent on recalculating the latest error + (``errors`` are stored as the first element of the state tuple). + + upon termination of these ``inner_iterations``, the loop is continued if + iteration is smaller than ``min_iterations``, stopped if equal/larger than + ``max_iterations``, and interrupted if ``cond_fn`` returns False. + + Args: + cond_fn : termination condition function + body_fn : body loop instructions + min_iterations : lower bound on the total amount of fixed point iterations + max_iterations : upper bound on the total amount of fixed point iterations + inner_iterations : number of iterations ``body_fn`` will be executed + successively before calling ``cond_fn``. + constants : constant (during loop) parameters passed on to body + state : state variable + + Returns: + outputs state returned by ``body_fn`` upon termination. + """ # noqa: D401 + # If number of minimal iterations matches maximal number, force a scan instead + # of a while loop. + + force_scan = (min_iterations == max_iterations) + + compute_error_flags = jnp.arange(inner_iterations) == inner_iterations - 1 + + def max_cond_fn(iteration_state): + iteration, state = iteration_state + return jnp.logical_and( + iteration < max_iterations, + jnp.logical_or( + iteration < min_iterations, cond_fn(iteration, constants, state) + ) + ) + + def unrolled_body_fn(iteration_state): + + def one_iteration(iteration_state, compute_error): + iteration, state = iteration_state + state = body_fn(iteration, constants, state, compute_error) + iteration += 1 + return (iteration, state), None + + iteration_state, _ = jax.lax.scan( + one_iteration, iteration_state, compute_error_flags + ) + return (iteration_state, None) if force_scan else iteration_state + + if force_scan: + (_, state), _ = jax.lax.scan( + lambda carry, x: unrolled_body_fn(carry), (0, state), + None, + length=max_iterations // inner_iterations + ) + else: + _, state = jax.lax.while_loop(max_cond_fn, unrolled_body_fn, (0, state)) + return state + + +def fixpoint_iter_fwd( + cond_fn, body_fn, min_iterations, max_iterations, inner_iterations, + constants, state +): + """Forward iteration of fixed point iteration to handle backpropagation. + + The main difference with fixpoint_iter is the checkpointing, in variable + states, of the state variables as they are recorded through iterations, every + inner_iterations. This sequence of states will be used in the backward loop. + + Args: + cond_fn : termination condition function + body_fn : body loop instructions + min_iterations : lower bound on the total amount of fixed point iterations + max_iterations : upper bound on the total amount of fixed point iterations + inner_iterations : number of iterations body_fn will be executed + successively before calling cond_fn. + constants : constant (during loop) parameters passed on to body + state : state variable + + Returns: + outputs state returned by body_fn upon termination. + """ + force_scan = min_iterations == max_iterations + compute_error_flags = jnp.arange(inner_iterations) == inner_iterations - 1 + states = jax.tree_util.tree_map( + lambda x: jnp.zeros( + (max_iterations // inner_iterations + 1,) + jnp.shape(x), + dtype=jax.dtypes.result_type(x) + ), state + ) + + def max_cond_fn(iteration_states_state): + iteration, _, state = iteration_states_state + return jnp.logical_and( + iteration < max_iterations, + jnp.logical_or( + iteration < min_iterations, cond_fn(iteration, constants, state) + ) + ) + + def unrolled_body_fn(iteration_states_state): + iteration, states, state = iteration_states_state + states = jax.tree_util.tree_map( + lambda states, state: jax.lax.dynamic_update_index_in_dim( + states, state, iteration // inner_iterations, 0 + ), states, state + ) + + def one_iteration(iteration_state, compute_error): + iteration, state = iteration_state + state = body_fn(iteration, constants, state, compute_error) + iteration += 1 + return (iteration, state), None + + iteration_state, _ = jax.lax.scan( + one_iteration, (iteration, state), compute_error_flags + ) + iteration, state = iteration_state + out = (iteration, states, state) + return (out, None) if force_scan else out + + if force_scan: + (iteration, states, state), _ = jax.lax.scan( + lambda carry, x: unrolled_body_fn(carry), (0, states, state), + None, + length=max_iterations // inner_iterations + ) + else: + iteration, states, state = jax.lax.while_loop( + max_cond_fn, unrolled_body_fn, (0, states, state) + ) + + return state, (constants, iteration, states) + + +def fixpoint_iter_bwd( + cond_fn, body_fn, min_iterations, max_iterations, inner_iterations, res, g +): + """Backward iteration of fixed point iteration, using checkpointed states.""" + del cond_fn + force_scan = (min_iterations == max_iterations) + constants, iteration, states = res + # The tree may contain some python floats + g_constants = jax.tree_util.tree_map( + lambda x: jnp.zeros_like(x, dtype=x.dtype) + if isinstance(x, (np.ndarray, jnp.ndarray)) else 0, constants + ) + + def bwd_cond_fn(iteration_g_gconst): + iteration, _, _ = iteration_g_gconst + return iteration >= 0 + + def unrolled_body_fn_no_errors(iteration, constants, state): + compute_error_flags = jnp.zeros((inner_iterations,), dtype=bool) + + def one_iteration(iteration_state, compute_error): + iteration, state = iteration_state + state = body_fn(iteration, constants, state, compute_error) + iteration += 1 + return (iteration, state), None + + iteration_state, _ = jax.lax.scan( + one_iteration, (iteration, state), compute_error_flags + ) + _, state = iteration_state + return state + + def unrolled_body_fn(iteration_g_gconst): + iteration, g, g_constants = iteration_g_gconst + state = jax.tree_util.tree_map( + lambda x: x[iteration // inner_iterations], states + ) + _, pullback = jax.vjp( + unrolled_body_fn_no_errors, iteration, constants, state + ) + _, gi_constants, g_state = pullback(g) + g_constants = jax.tree_util.tree_map( + lambda x, y: x + y, g_constants, gi_constants + ) + out = (iteration - inner_iterations, g_state, g_constants) + return (out, None) if force_scan else out + + if force_scan: + (_, g_state, g_constants), _ = jax.lax.scan( + lambda carry, x: unrolled_body_fn(carry), (0, g, g_constants), + None, + length=max_iterations // inner_iterations + ) + else: + _, g_state, g_constants = jax.lax.while_loop( + bwd_cond_fn, unrolled_body_fn, + (iteration - inner_iterations, g, g_constants) + ) + + return g_constants, g_state + + +# definition of backprop friendly variant of fixpoint_iter. +fixpoint_iter_backprop = jax.custom_vjp( + fixpoint_iter, nondiff_argnums=(0, 1, 2, 3, 4) +) + +fixpoint_iter_backprop.defvjp(fixpoint_iter_fwd, fixpoint_iter_bwd) diff --git a/ott/build/lib/ott/math/matrix_square_root.py b/ott/build/lib/ott/math/matrix_square_root.py new file mode 100644 index 0000000..324a5ea --- /dev/null +++ b/ott/build/lib/ott/math/matrix_square_root.py @@ -0,0 +1,337 @@ +# 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 functools +import math +from typing import Tuple + +import jax +import jax.numpy as jnp + +from ott.math import fixed_point_loop + +__all__ = ["sqrtm", "sqrtm_only", "inv_sqrtm_only"] + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2, 3, 4, 5)) +def sqrtm( + x: jnp.ndarray, + threshold: float = 1e-6, + min_iterations: int = 0, + inner_iterations: int = 10, + max_iterations: int = 1000, + regularization: float = 1e-6 +) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Higham algorithm to compute matrix square root of p.d. matrix. + + See :cite:`higham:97`, eq. 2.6b + + Args: + x: a (batch of) square p.s.d. matrices of the same size. + threshold: convergence tolerance threshold for Newton-Schulz iterations. + min_iterations: min number of iterations after which error is computed. + inner_iterations: error is re-evaluated every inner_iterations iterations. + max_iterations: max number of iterations. + regularization: small regularizer added to norm of x, before normalization. + + Returns: + Square root matrix of x (or x's if batch), its inverse, + errors along iterates. + """ + dimension = x.shape[-1] + norm_x = jnp.linalg.norm(x, axis=(-2, -1)) * (1 + regularization) + + if jnp.ndim(x) > 2: + norm_x = norm_x[..., jnp.newaxis, jnp.newaxis] + + def cond_fn(iteration, const, state): + """Stopping criterion. Checking decrease of objective is needed here.""" + _, threshold = const + errors, _, _ = state + err = errors[iteration // inner_iterations - 1] + + return jnp.logical_or( + iteration == 0, + jnp.logical_and( + jnp.logical_and(jnp.isfinite(err), err > threshold), + jnp.all(jnp.diff(errors) <= 0) + ) + ) # check decreasing obj, else stop + + def body_fn(iteration, const, state, compute_error): + """Carry out matrix updates on y and z, stores error if requested. + + Args: + iteration: iteration number + const: tuple of constant parameters that do not change throughout the + loop. + state: state variables currently updated in the loop. + compute_error: flag to indicate this iteration computes/stores an error + + Returns: + state variables. + """ + x, _ = const + errors, y, z = state + w = 0.5 * jnp.matmul(z, y) + y = 1.5 * y - jnp.matmul(y, w) + z = 1.5 * z - jnp.matmul(w, z) + + err = jnp.where(compute_error, new_err(x, norm_x, y), jnp.inf) + + errors = errors.at[iteration // inner_iterations].set(err) + + return errors, y, z + + def new_err(x, norm_x, y): + res = x - norm_x * jnp.matmul(y, y) + norm_fn = functools.partial(jnp.linalg.norm, axis=(-2, -1)) + return jnp.max(norm_fn(res) / norm_fn(x)) + + y = x / norm_x + z = jnp.eye(dimension) + if jnp.ndim(x) > 2: + z = jnp.tile(z, list(x.shape[:-2]) + [1, 1]) + errors = -jnp.ones(math.ceil(max_iterations / inner_iterations)) + state = (errors, y, z) + const = (x, threshold) + errors, y, z = fixed_point_loop.fixpoint_iter_backprop( + cond_fn, body_fn, min_iterations, max_iterations, inner_iterations, const, + state + ) + sqrt_x = jnp.sqrt(norm_x) * y + inv_sqrt_x = z / jnp.sqrt(norm_x) + + return sqrt_x, inv_sqrt_x, errors + + +def solve_sylvester_bartels_stewart( + a: jnp.ndarray, + b: jnp.ndarray, + c: jnp.ndarray, +) -> jnp.ndarray: + """Solve the real Sylvester equation AX - XB = C using Bartels-Stewart.""" + # See https://nhigham.com/2020/09/01/what-is-the-sylvester-equation/ for + # discussion of the algorithm (but note that in the derivation, the sign on + # the right hand side is flipped in the equation in which the columns are set + # to be equal). + m = a.shape[-1] + n = b.shape[-1] + # Cast a and b to complex to ensure we get the complex Schur decomposition + # (the real Schur decomposition may not give an upper triangular solution). + # For the decomposition below, a = u r u* and b = v s v* + r, u = jax.lax.linalg.schur(a + 0j) + s, v = jax.lax.linalg.schur(b + 0j) + d = jnp.matmul( + jnp.conjugate(jnp.swapaxes(u, axis1=-2, axis2=-1)), jnp.matmul(c, v) + ) + # The solution in the transformed space will in general be complex, too. + y = jnp.zeros(a.shape[:-2] + (m, n)) + 0j + idx = jnp.arange(m) + for j in range(n): + lhs = r.at[..., idx, idx].add(-s[..., j:j + 1, j]) + rhs = d[..., j] + jnp.matmul(y[..., :j], s[..., :j, j:j + 1])[..., 0] + y = y.at[..., j].set(jax.scipy.linalg.solve_triangular(lhs, rhs)) + + x = jnp.matmul( + u, jnp.matmul(y, jnp.conjugate(jnp.swapaxes(v, axis1=-2, axis2=-1))) + ) + # The end result should be real; remove the imaginary part of the solution. + return jnp.real(x) + + +def sqrtm_fwd( + x: jnp.ndarray, + threshold: float, + min_iterations: int, + inner_iterations: int, + max_iterations: int, + regularization: float, +) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], Tuple[jnp.ndarray, + jnp.ndarray]]: + """Forward pass of custom VJP.""" + sqrt_x, inv_sqrt_x, errors = sqrtm( + x=x, + threshold=threshold, + min_iterations=min_iterations, + inner_iterations=inner_iterations, + max_iterations=max_iterations, + regularization=regularization, + ) + return (sqrt_x, inv_sqrt_x, errors), (sqrt_x, inv_sqrt_x) + + +def sqrtm_bwd( + threshold: float, + min_iterations: int, + inner_iterations: int, + max_iterations: int, + regularization: float, + residual: Tuple[jnp.ndarray, jnp.ndarray], + cotangent: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], +) -> Tuple[jnp.ndarray]: + """Compute the derivative by solving a Sylvester equation.""" + del threshold, min_iterations, inner_iterations, \ + max_iterations, regularization + sqrt_x, inv_sqrt_x = residual + # ignores cotangent associated with errors + cot_sqrt, cot_inv_sqrt, _ = cotangent + + # Solve for d(X^{1/2}): + # Start with X^{1/2} X^{1/2} = X + # Differentiate to obtain + # d(X^{1/2}) X^{1/2} + X^{1/2} d(X^{1/2}) = dX + # The above is a Sylvester equation that we can solve using Bartels-Stewart. + # Below think of cot_sqrt as (dX)^T and vjp_cot_sqrt as d(X^{1/2})^T. + # See https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html + vjp_cot_sqrt = jnp.swapaxes( + solve_sylvester_bartels_stewart( + a=sqrt_x, b=-sqrt_x, c=jnp.swapaxes(cot_sqrt, axis1=-1, axis2=-2) + ), + axis1=-1, + axis2=-2 + ) + + # Now solve for d(X^{-1/2}): + # Start with X^{-1/2} X^{-1/2} = X^{-1} + # Use the product rule and the fact that d(X^{-1}) = -X^{-1} dX X^{-1} + # to obtain + # (See The Matrix Cookbook section on derivatives of an inverse + # https://www.math.uwaterloo.ca/~hwolkowi/matrixcookbook.pdf ) + # d(X^{-1/2}) X^{-1/2} + X^{-1/2} d(X^{-1/2}) = -X^{-1} dX X^{-1} + # Again we have a Sylvester equation that we solve as above, and again we + # think of cot_inv_sqrt as (dX)^T and vjp_cot_inv_sqrt as d(X^{-1/2})^T + inv_x = jnp.matmul(inv_sqrt_x, inv_sqrt_x) + vjp_cot_inv_sqrt = jnp.swapaxes( + solve_sylvester_bartels_stewart( + a=inv_sqrt_x, + b=-inv_sqrt_x, + c=-jnp.matmul( + inv_x, + jnp.matmul(jnp.swapaxes(cot_inv_sqrt, axis1=-2, axis2=-1), inv_x) + ) + ), + axis1=-1, + axis2=-2 + ) + return vjp_cot_sqrt + vjp_cot_inv_sqrt, + + +sqrtm.defvjp(sqrtm_fwd, sqrtm_bwd) + +# Specialized versions of sqrtm that compute only the square root or inverse. +# These functions have lower complexity gradients than sqrtm. + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2, 3, 4, 5)) +def sqrtm_only( # noqa: D103 + x: jnp.ndarray, + threshold: float = 1e-6, + min_iterations: int = 0, + inner_iterations: int = 10, + max_iterations: int = 1000, + regularization: float = 1e-6 +) -> jnp.ndarray: + return sqrtm( + x, threshold, min_iterations, inner_iterations, max_iterations, + regularization + )[0] + + +def sqrtm_only_fwd( # noqa: D103 + x: jnp.ndarray, threshold: float, min_iterations: int, + inner_iterations: int, max_iterations: int, regularization: float +) -> Tuple[jnp.ndarray, jnp.ndarray]: + sqrt_x = sqrtm( + x, threshold, min_iterations, inner_iterations, max_iterations, + regularization + )[0] + return sqrt_x, sqrt_x + + +def sqrtm_only_bwd( # noqa: D103 + threshold: float, min_iterations: int, inner_iterations: int, + max_iterations: int, regularization: float, sqrt_x: jnp.ndarray, + cotangent: jnp.ndarray +) -> Tuple[jnp.ndarray]: + del threshold, min_iterations, inner_iterations, \ + max_iterations, regularization + vjp = jnp.swapaxes( + solve_sylvester_bartels_stewart( + a=sqrt_x, b=-sqrt_x, c=jnp.swapaxes(cotangent, axis1=-2, axis2=-1) + ), + axis1=-2, + axis2=-1 + ) + return vjp, + + +sqrtm_only.defvjp(sqrtm_only_fwd, sqrtm_only_bwd) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2, 3, 4, 5)) +def inv_sqrtm_only( # noqa: D103 + x: jnp.ndarray, + threshold: float = 1e-6, + min_iterations: int = 0, + inner_iterations: int = 10, + max_iterations: int = 1000, + regularization: float = 1e-6 +) -> jnp.ndarray: + return sqrtm( + x, threshold, min_iterations, inner_iterations, max_iterations, + regularization + )[1] + + +def inv_sqrtm_only_fwd( # noqa: D103 + x: jnp.ndarray, + threshold: float, + min_iterations: int, + inner_iterations: int, + max_iterations: int, + regularization: float, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + inv_sqrt_x = sqrtm( + x, threshold, min_iterations, inner_iterations, max_iterations, + regularization + )[1] + return inv_sqrt_x, inv_sqrt_x + + +def inv_sqrtm_only_bwd( # noqa: D103 + threshold: float, min_iterations: int, inner_iterations: int, + max_iterations: int, regularization: float, residual: jnp.ndarray, + cotangent: jnp.ndarray +) -> Tuple[jnp.ndarray]: + del threshold, min_iterations, inner_iterations, \ + max_iterations, regularization + + inv_sqrt_x = residual + inv_x = jnp.matmul(inv_sqrt_x, inv_sqrt_x) + vjp = jnp.swapaxes( + solve_sylvester_bartels_stewart( + a=inv_sqrt_x, + b=-inv_sqrt_x, + c=-jnp.matmul( + inv_x, + jnp.matmul(jnp.swapaxes(cotangent, axis1=-2, axis2=-1), inv_x) + ) + ), + axis1=-1, + axis2=-2 + ) + return vjp, + + +inv_sqrtm_only.defvjp(inv_sqrtm_only_fwd, inv_sqrtm_only_bwd) diff --git a/ott/build/lib/ott/math/unbalanced_functions.py b/ott/build/lib/ott/math/unbalanced_functions.py new file mode 100644 index 0000000..9ba20ba --- /dev/null +++ b/ott/build/lib/ott/math/unbalanced_functions.py @@ -0,0 +1,90 @@ +# 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. +from typing import Callable + +import jax.numpy as jnp + + +def phi_star(h: jnp.ndarray, rho: float) -> jnp.ndarray: + """Legendre transform of KL, :cite:`sejourne:19`, p. 9.""" + return rho * (jnp.exp(h / rho) - 1) + + +def derivative_phi_star(f: jnp.ndarray, rho: float) -> jnp.ndarray: + """Derivative of Legendre transform of phi_starKL, see phi_star.""" + # TODO(cuturi): use jax.grad directly. + return jnp.exp(f / rho) + + +def grad_of_marginal_fit( + c: jnp.ndarray, h: jnp.ndarray, tau: float, epsilon: float +) -> jnp.ndarray: + """Compute grad of terms linked to marginals in objective. + + Computes gradient w.r.t. f ( or g) of terms in :cite:`sejourne:19`, + left-hand-side of eq. 15 terms involving phi_star). + + Args: + c: jnp.ndarray, first target marginal (either a or b in practice) + h: jnp.ndarray, potential (either f or g in practice) + tau: float, strength (in ]0,1]) of regularizer w.r.t. marginal + epsilon: regularization + + Returns: + a vector of the same size as c or h + """ + if tau == 1.0: + return c + r = rho(epsilon, tau) + return jnp.where(c > 0, c * derivative_phi_star(-h, r), 0.0) + + +def second_derivative_phi_star(f: jnp.ndarray, rho: float) -> jnp.ndarray: + """Second Derivative of Legendre transform of KL, see phi_star.""" + return jnp.exp(f / rho) / rho + + +def diag_jacobian_of_marginal_fit( + c: jnp.ndarray, h: jnp.ndarray, tau: float, epsilon: float, + derivative: Callable[[jnp.ndarray, float], jnp.ndarray] +): + """Compute grad of terms linked to marginals in objective. + + Computes second derivative w.r.t. f ( or g) of terms in :cite:`sejourne:19`, + left-hand-side of eq. 32 (terms involving phi_star) + + Args: + c: jnp.ndarray, first target marginal (either a or b in practice) + h: jnp.ndarray, potential (either f or g in practice) + tau: float, strength (in ]0,1]) of regularizer w.r.t. marginal + epsilon: regularization + derivative: Callable + + Returns: + a vector of the same size as c or h. + """ + if tau == 1.0: + return 0.0 + + r = rho(epsilon, tau) + # here no minus sign because we are taking derivative w.r.t -h + return jnp.where( + c > 0, + c * second_derivative_phi_star(-h, r) * + derivative(c * derivative_phi_star(-h, r)), 0.0 + ) + + +def rho(epsilon: float, tau: float) -> float: # noqa: D103 + return (epsilon * tau) / (1.0 - tau) diff --git a/ott/build/lib/ott/math/utils.py b/ott/build/lib/ott/math/utils.py new file mode 100644 index 0000000..afad18d --- /dev/null +++ b/ott/build/lib/ott/math/utils.py @@ -0,0 +1,298 @@ +# 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 functools +from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Union + +import jax +import jax.numpy as jnp +import jax.scipy as jsp + +if TYPE_CHECKING: + from ott.geometry import costs + +__all__ = [ + "safe_log", + "norm", + "kl", + "gen_kl", + "gen_js", + "logsumexp", + "softmin", + "barycentric_projection", + "sort_and_argsort", + "lambertw", +] + + +def safe_log( # noqa: D103 + x: jnp.ndarray, + *, + eps: Optional[float] = None +) -> jnp.ndarray: + if eps is None: + eps = jnp.finfo(x.dtype).tiny + return jnp.where(x > 0.0, jnp.log(x), jnp.log(eps)) + + +@functools.partial(jax.custom_jvp, nondiff_argnums=[1, 2, 3]) +@functools.partial(jax.jit, static_argnames=("ord", "axis", "keepdims")) +def norm( + x: jnp.ndarray, + ord: Union[int, str, None] = None, + axis: Union[None, Sequence[int], int] = None, + keepdims: bool = False +) -> jnp.ndarray: + """Computes order ord norm of vector, using `jnp.linalg` in forward pass. + + Evaluations of distances between a vector and itself using translation + invariant costs, typically norms, result in functions of the form + ``lambda x : jnp.linalg.norm(x-x)``. Such functions output `NaN` gradients, + because they involve computing the derivative of a negative exponent of 0 + (e.g. when differentiating the Euclidean norm, one gets a 0-denominator in the + expression, see e.g. https://github.com/google/jax/issues/6484 for context). + + While this makes sense mathematically, in the context of optimal transport + such distances between a point and itself can be safely ignored when they + contribute to an OT cost (when, for instance, computing Sinkhorn divergences, + involving computing the OT cost of a point cloud with itself). + + To avoid such `NaN` values, this custom norm implementation uses the + double-where trick, to avoid having branches that output any `NaN`, and + safely output a 0 instead. + + Args: + x: Input array. If `axis` is None, `x` must be 1-D or 2-D, unless `ord` + is None. If both `axis` and `ord` are None, the 2-norm of ``x.ravel`` + will be returned. + ord: `{non-zero int, jnp.inf, -jnp.inf, 'fro', 'nuc'}`, Order of the norm. + The default is `None`, which is equivalent to `2.0` for vectors. + axis: `{None, int, 2-tuple of ints}`, optional. If `axis` is an integer, it + specifies the axis of `x` along which to compute the vector norms. + If `axis` is a 2-tuple, it specifies the axes that hold 2-D matrices, and + the matrix norms of these matrices are computed. If `axis` is None then + either a vector norm (when `x` is 1-D) or a matrix norm (when `x` is 2-D) + is returned. The default is None. + keepdims: If set to True, the axes which are normed over are left in the + result as dimensions with size one. With this option the result will + broadcast correctly against the original `x`. + + Returns: + float or ndarray, Norm of the matrix or vector(s). + """ + return jnp.linalg.norm(x, ord=ord, axis=axis, keepdims=keepdims) + + +@norm.defjvp +def norm_jvp(ord, axis, keepdims, primals, tangents): + """Custom_jvp for norm, that returns 0.0 when evaluated at 0.""" + x, = primals + x_is_zero = jnp.all(jnp.logical_not(x)) + clean_x = jnp.where(x_is_zero, jnp.ones_like(x), x) + primals, tangents = jax.jvp( + functools.partial(jnp.linalg.norm, ord=ord, axis=axis, keepdims=keepdims), + (clean_x,), tangents + ) + return primals, jnp.where(x_is_zero, 0.0, tangents) + + +# TODO(michalk8): add axis argument +def kl(p: jnp.ndarray, q: jnp.ndarray) -> float: + """Kullback-Leibler divergence.""" + return jnp.vdot(p, (safe_log(p) - safe_log(q))) + + +def gen_kl(p: jnp.ndarray, q: jnp.ndarray) -> float: + """Generalized Kullback-Leibler divergence.""" + return jnp.vdot(p, (safe_log(p) - safe_log(q))) + jnp.sum(q) - jnp.sum(p) + + +# TODO(michalk8): add axis argument +def gen_js(p: jnp.ndarray, q: jnp.ndarray, c: float = 0.5) -> float: + """Jensen-Shannon divergence.""" + return c * (gen_kl(p, q) + gen_kl(q, p)) + + +@functools.partial(jax.custom_jvp, nondiff_argnums=(1, 2, 4)) +def logsumexp( # noqa: D103 + mat, axis=None, keepdims=False, b=None, return_sign=False +): + return jax.scipy.special.logsumexp( + mat, axis=axis, keepdims=keepdims, b=b, return_sign=return_sign + ) + + +@logsumexp.defjvp +def logsumexp_jvp(axis, keepdims, return_sign, primals, tangents): + """Custom derivative rule for lse that does not blow up with -inf. + + This logsumexp implementation uses the standard jax one in forward mode but + implements a custom rule to differentiate. Given the preference of jax for + jvp over vjp, and the fact that this is a simple linear rule, jvp is used. + This custom differentiation address issues when the output of lse is + -inf (which corresponds to the case where all inputs in a slice are -inf, + which happens typically when ``a`` or ``b`` weight vectors have zeros.) + + Although both exp(lse) and its derivative should be 0, automatic + differentiation returns a NaN derivative because of a -inf - (-inf) operation + appearing in the definition of centered_exp below. This is corrected in the + implementation below. + + Args: + axis: argument from original logsumexp + keepdims: argument from original logsumexp + return_sign: argument from original logsumexp + primals: mat and b, the two arguments against which we differentiate. + tangents: of same size as mat and b. + + Returns: + original primal outputs + their tangent. + """ # noqa: D401 + mat, b = primals + tan_mat, tan_b = tangents + lse = logsumexp(mat, axis, keepdims, b, return_sign) + if return_sign: + lse, sign = lse + lse = jnp.where(jnp.isfinite(lse), lse, 0.0) + centered_exp = jnp.exp(mat - jnp.expand_dims(lse, axis=axis)) + + if b is None: + res = jnp.sum(centered_exp * tan_mat, axis=axis, keepdims=keepdims) + else: + res = jnp.sum(b * centered_exp * tan_mat, axis=axis, keepdims=keepdims) + res += jnp.sum(tan_b * centered_exp, axis=axis, keepdims=keepdims) + if return_sign: + return (lse, sign), (sign * res, jnp.zeros_like(sign)) + return lse, res + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(2,)) +def softmin( + x: jnp.ndarray, gamma: float, axis: Optional[int] = None +) -> jnp.ndarray: + r"""Soft-min operator. + + Args: + x: Input data. + gamma: Smoothing parameter :math:`> 0`. + axis: Axis or axes over which to operate. If ``None``, use flattened input. + + Returns: + The soft minimum. + """ + return -gamma * jsp.special.logsumexp(x / -gamma, axis=axis) + + +softmin.defvjp( + lambda x, gamma, axis: (softmin(x, gamma, axis), (x / -gamma, axis)), + lambda axis, res, g: ( + jnp.where( + jnp.isinf(res[0]), 0.0, + jax.nn.softmax(res[0], axis=axis) * + (g if axis is None else jnp.expand_dims(g, axis=axis)) + ), None + ) +) + + +@functools.partial(jax.vmap, in_axes=[0, 0, None]) +def barycentric_projection( + matrix: jnp.ndarray, y: jnp.ndarray, cost_fn: "costs.CostFn" +) -> jnp.ndarray: + """Compute the barycentric projection of a matrix. + + Args: + matrix: a matrix of shape (n, m) + y: a vector of shape (m,) + cost_fn: a CostFn instance. + + Returns: + a vector of shape (n,) containing the barycentric projection of matrix. + """ + return jax.vmap( + lambda m, y: cost_fn.barycenter(m, y)[0], in_axes=[0, None] + )(matrix, y) + + +def sort_and_argsort( + x: jnp.array, + *, + argsort: bool = False +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: + """Unified function that returns both sort and argsort, if latter needed.""" + if argsort: + i_x = jnp.argsort(x) + return x[i_x], i_x + return jnp.sort(x), None + + +@functools.partial(jax.custom_jvp, nondiff_argnums=(1, 2)) +def lambertw( + z: jnp.ndarray, tol: float = 1e-8, max_iter: int = 100 +) -> jnp.ndarray: + """Principal branch of the + `Lambert W function `_. + + This implementation uses Halley's iteration and the global initialization + proposed in :cite:`iacono:17`, Eq. 20 . + + Args: + z: Array. + tol: Tolerance threshold. + max_iter: Maximum number of iterations. + + Returns: + The Lambert W evaluated at ``z``. + """ # noqa: D205 + + def initial_iacono(x: jnp.ndarray) -> jnp.ndarray: + y = jnp.sqrt(1.0 + jnp.e * x) + num = 1.0 + 1.14956131 * y + denom = 1.0 + 0.45495740 * jnp.log1p(y) + return -1.0 + 2.036 * jnp.log(num / denom) + + def cond_fun(container): + it, converged, _ = container + return jnp.logical_and(jnp.any(~converged), it < max_iter) + + def halley_iteration(container): + it, _, w = container + + # modified from `tensorflow_probability` + f = w - z * jnp.exp(-w) + delta = f / (w + 1.0 - 0.5 * (w + 2.0) * f / (w + 1.0)) + + w_next = w - delta + + not_converged = jnp.abs(delta) <= tol * jnp.abs(w_next) + return it + 1, not_converged, w_next + + w0 = initial_iacono(z) + converged = jnp.zeros_like(w0, dtype=bool) + + _, _, w = jax.lax.while_loop( + cond_fun=cond_fun, body_fun=halley_iteration, init_val=(0, converged, w0) + ) + return w + + +@lambertw.defjvp +def _lambertw_jvp( + tol: float, max_iter: int, primals: Tuple[jnp.ndarray, ...], + tangents: Tuple[jnp.ndarray, ...] +) -> Tuple[jnp.ndarray, jnp.ndarray]: + z, = primals + dz, = tangents + w = lambertw(z, tol=tol, max_iter=max_iter) + pz = jnp.where(z == 0.0, 1.0, w / ((1.0 + w) * z)) + return w, pz * dz diff --git a/ott/build/lib/ott/neural/__init__.py b/ott/build/lib/ott/neural/__init__.py new file mode 100644 index 0000000..aa1ca23 --- /dev/null +++ b/ott/build/lib/ott/neural/__init__.py @@ -0,0 +1,14 @@ +# 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. +from . import layers, losses, models, solvers diff --git a/ott/build/lib/ott/neural/layers.py b/ott/build/lib/ott/neural/layers.py new file mode 100644 index 0000000..78c2ef3 --- /dev/null +++ b/ott/build/lib/ott/neural/layers.py @@ -0,0 +1,213 @@ +# 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. +from typing import Any, Callable, Optional, Tuple + +import jax +import jax.numpy as jnp +from flax import linen as nn + +__all__ = ["PositiveDense", "PosDefPotentials"] + +PRNGKey = jax.Array +Shape = Tuple[int, ...] +Dtype = Any +Array = jnp.ndarray + +# wrap to silence docs linter +DEFAULT_KERNEL_INIT = lambda *a, **k: nn.initializers.lecun_normal()(*a, **k) +DEFAULT_BIAS_INIT = nn.initializers.zeros +DEFAULT_RECTIFIER = nn.activation.relu + + +class PositiveDense(nn.Module): + """A linear transformation using a matrix with all entries non-negative. + + Args: + dim_hidden: Number of output dimensions. + rectifier_fn: Rectifier function. The default is + :func:`~flax.linen.activation.relu`. + use_bias: Whether to add bias to the output. + kernel_init: Initializer for the matrix. The default is + :func:`~flax.linen.initializers.lecun_normal`. + bias_init: Initializer for the bias. The default is + :func:`~flax.linen.initializers.zeros`. + precision: Numerical precision of the computation. + """ + + dim_hidden: int + rectifier_fn: Callable[[Array], Array] = DEFAULT_RECTIFIER + use_bias: bool = True + kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = DEFAULT_KERNEL_INIT + bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = DEFAULT_BIAS_INIT + precision: Optional[jax.lax.Precision] = None + + @nn.compact + def __call__(self, x: jnp.ndarray) -> jnp.ndarray: + """Applies a linear transformation to x along the last dimension. + + Args: + x: Array of shape ``[batch, ..., features]``. + + Returns: + Array of shape ``[batch, ..., dim_hidden]``. + """ + # TODO(michalk8): update when refactoring neuraldual + # assert x.ndim > 1, x.ndim + + kernel = self.param( + "kernel", self.kernel_init, (x.shape[-1], self.dim_hidden) + ) + kernel = self.rectifier_fn(kernel) + + x = jnp.tensordot(x, kernel, axes=(-1, 0), precision=self.precision) + if self.use_bias: + x = x + self.param("bias", self.bias_init, (self.dim_hidden,)) + + return x + + +class PosDefPotentials(nn.Module): + r""":math:`\frac{1}{2} x^T (A_i A_i^T + \text{Diag}(d_i)) x + b_i^T x^2 + c_i` potentials. + + This class implements a layer that takes (batched) ``d``-dimensional vectors + ``x`` in, to output a ``num_potentials``-dimensional vector. Each of the + entries in that output is a positive definite quadratic form evaluated at + ``x``; each of these quadratic terms is parameterized as a low-rank plus + diagonal matrix. The low-rank term is parameterized as :math:`A_i A_i^T`, + where each of these matrices is of size ``(rank, d)``. Taken together, + these matrices form a tensor ``(num_potentials, rank, d)``. + The diagonal terms :math:`d_i` form a ``(num_potentials, d)`` matrix of + positive values; the linear terms :math:`b_i` form a ``(num_potentials, d)`` + matrix. Finally, the :math:`c_i` are contained in a vector of size + ``(num_potentials,)``. + + Args: + num_potentials: Dimension of the output. + rank: Rank of the matrices :math:`A_i` used as low-rank factors + for the quadratic potentials. + rectifier_fn: Rectifier function to ensure non-negativity of the diagonals + :math:`d_i`. The default is :func:`~flax.linen.activation.relu`. + use_linear: Whether to add a linear layers :math:`b_i` to the outputs. + use_bias: Whether to add biases :math:`c_i` to the outputs. + kernel_lr_init: Initializer for the matrices :math:`A_i` + of the quadratic potentials when ``rank > 0``. + The default is :func:`~flax.linen.initializers.lecun_normal`. + kernel_diag_init: Initializer for the diagonals :math:`d_i`. + The default is :func:`~flax.linen.initializers.ones`. + kernel_linear_init: Initializer for the linear layers :math:`b_i`. + The default is :func:`~flax.linen.initializers.lecun_normal`. + bias_init: Initializer for the bias. The default is + :func:`~flax.linen.initializers.zeros`. + precision: Numerical precision of the computation. + """ # noqa: E501 + + num_potentials: int + rank: int = 0 + rectifier_fn: Callable[[Array], Array] = DEFAULT_RECTIFIER + use_linear: bool = True + use_bias: bool = True + kernel_lr_init: Callable[[PRNGKey, Shape, Dtype], Array] = DEFAULT_KERNEL_INIT + kernel_diag_init: Callable[[PRNGKey, Shape, Dtype], + Array] = nn.initializers.ones + kernel_linear_init: Callable[[PRNGKey, Shape, Dtype], + Array] = DEFAULT_KERNEL_INIT + bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = DEFAULT_BIAS_INIT + precision: Optional[jax.lax.Precision] = None + + @nn.compact + def __call__(self, x: jnp.ndarray) -> jnp.ndarray: + """Compute quadratic forms of the input. + + Args: + x: Array of shape ``[batch, ..., features]``. + + Returns: + Array of shape ``[batch, ..., num_potentials]``. + """ + # TODO(michalk8): update when refactoring neuraldual + # assert x.ndim > 1, x.ndim + + dim_data = x.shape[-1] + x = x.reshape((-1, dim_data)) + + diag_kernel = self.param( + "diag_kernel", self.kernel_diag_init, (dim_data, self.num_potentials) + ) + # ensures the diag_kernel parameter stays non negative + diag_kernel = self.rectifier_fn(diag_kernel) + + # (batch, dim_data, 1), (1, dim_data, num_potentials) + y = 0.5 * jnp.sum(((x ** 2)[..., None] * diag_kernel[None]), axis=1) + + if self.rank > 0: + quad_kernel = self.param( + "quad_kernel", self.kernel_lr_init, + (self.num_potentials, dim_data, self.rank) + ) + # (batch, num_potentials, rank) + quad = 0.5 * jnp.tensordot( + x, quad_kernel, axes=(-1, 1), precision=self.precision + ) ** 2 + y = y + jnp.sum(quad, axis=-1) + + if self.use_linear: + linear_kernel = self.param( + "lin_kernel", self.kernel_linear_init, + (dim_data, self.num_potentials) + ) + y = y + jnp.dot(x, linear_kernel, precision=self.precision) + + if self.use_bias: + y = y + self.param("bias", self.bias_init, (self.num_potentials,)) + + return y + + @classmethod + def init_from_samples( + cls, source: jnp.ndarray, target: jnp.ndarray, **kwargs: Any + ) -> "PosDefPotentials": + """Initialize the layer using Gaussian approximation :cite:`bunne:22`. + + Args: + source: Samples from the source distribution, array of shape ``[n, d]``. + target: Samples from the target distribution, array of shape ``[m, d]``. + kwargs: Keyword arguments for initialization. Note that ``use_linear`` + will be always set to :obj:`True`. + + Returns: + The layer with fixed linear and quadratic initialization. + """ + factor, mean = _compute_gaussian_map_params(source, target) + + kwargs["use_linear"] = True + return cls( + kernel_lr_init=lambda *_, **__: factor, + kernel_linear_init=lambda *_, **__: mean.T, + **kwargs, + ) + + +def _compute_gaussian_map_params( + source: jnp.ndarray, target: jnp.ndarray +) -> Tuple[jnp.ndarray, jnp.ndarray]: + from ott.math import matrix_square_root + from ott.tools.gaussian_mixture import gaussian + + g_s = gaussian.Gaussian.from_samples(source) + g_t = gaussian.Gaussian.from_samples(target) + lin_op = g_s.scale.gaussian_map(g_t.scale) + b = jnp.squeeze(g_t.loc) - lin_op @ jnp.squeeze(g_s.loc) + lin_op = matrix_square_root.sqrtm_only(lin_op) + + return jnp.expand_dims(lin_op, 0), jnp.expand_dims(b, 0) diff --git a/ott/build/lib/ott/neural/losses.py b/ott/build/lib/ott/neural/losses.py new file mode 100644 index 0000000..ca5abc6 --- /dev/null +++ b/ott/build/lib/ott/neural/losses.py @@ -0,0 +1,309 @@ +# 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. +from typing import Any, Callable, Literal, Optional, Tuple, Union + +import jax +import jax.numpy as jnp +import numpy as np + +from ott.geometry import costs, pointcloud +from ott.solvers import linear +from ott.solvers.linear import sinkhorn +from ott.problems.quadratic import quadratic_problem +from ott.solvers.quadratic import gromov_wasserstein + +__all__ = ["monge_gap", "monge_gap_from_samples"] + + +def monge_gap( + map_fn: Callable[[jnp.ndarray], jnp.ndarray], + reference_points: jnp.ndarray, + cost_fn: Optional[costs.CostFn] = None, + epsilon: Optional[float] = None, + relative_epsilon: Optional[bool] = None, + scale_cost: Union[ + bool, int, float, Literal["mean", "max_cost", "median"] + ] = 1.0, + return_output: bool = False, + **kwargs: Any, +) -> Union[float, Tuple[float, sinkhorn.SinkhornOutput]]: + r"""Monge gap regularizer :cite:`uscidda:23`. + + For a cost function :math:`c` and empirical reference measure + :math:`\hat{\rho}_n=\frac{1}{n}\sum_{i=1}^n \delta_{x_i}`, the + (entropic) Monge gap of a map function + :math:`T:\mathbb{R}^d\rightarrow\mathbb{R}^d` is defined as: + + .. math:: + \mathcal{M}^c_{\hat{\rho}_n, \varepsilon} (T) + = \frac{1}{n} \sum_{i=1}^n c(x_i, T(x_i)) - + W_{c, \varepsilon}(\hat{\rho}_n, T \sharp \hat{\rho}_n) + + See :cite:`uscidda:23` Eq. (8). This function is a thin wrapper that calls + :func:`~ott.neural.losses.monge_gap_from_samples`. + + Args: + map_fn: Callable corresponding to map :math:`T` in definition above. The + callable should be vectorized (e.g. using :func:`jax.vmap`), i.e, + able to process a *batch* of vectors of size `d`, namely + ``map_fn`` applied to an array returns an array of the same shape. + reference_points: Array of `[n,d]` points, :math:`\hat\rho_n` in paper + cost_fn: An object of class :class:`~ott.geometry.costs.CostFn`. + epsilon: Regularization parameter. See + :class:`~ott.geometry.pointcloud.PointCloud` + relative_epsilon: when `False`, the parameter ``epsilon`` specifies the + value of the entropic regularization parameter. When `True`, ``epsilon`` + refers to a fraction of the + :attr:`~ott.geometry.pointcloud.PointCloud.mean_cost_matrix`, which is + computed adaptively using ``source`` and ``target`` points. + scale_cost: option to rescale the cost matrix. Implemented scalings are + 'median', 'mean' and 'max_cost'. Alternatively, a float factor can be + given to rescale the cost such that ``cost_matrix /= scale_cost``. + If `True`, use 'mean'. + return_output: boolean to also return the + :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput`. + kwargs: holds the kwargs to instantiate the or + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver to + compute the regularized OT cost. + + Returns: + The Monge gap value and optionally the + :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput` + """ + target = map_fn(reference_points) + return monge_gap_from_samples( + source=reference_points, + target=target, + cost_fn=cost_fn, + epsilon=epsilon, + relative_epsilon=relative_epsilon, + scale_cost=scale_cost, + return_output=return_output, + **kwargs, + ) + + +def monge_gap_from_samples( + source: jnp.ndarray, + target: jnp.ndarray, + cost_fn: Optional[costs.CostFn] = None, + epsilon: Optional[float] = None, + relative_epsilon: Optional[bool] = None, + scale_cost: Union[ + bool, int, float, Literal["mean", "max_cost", "median"] + ] = 1.0, + return_output: bool = False, + **kwargs: Any, +) -> Union[float, Tuple[float, sinkhorn.SinkhornOutput]]: + r"""Monge gap, instantiated in terms of samples before / after applying map. + + .. math:: + \frac{1}{n} \sum_{i=1}^n c(x_i, y_i)) - + W_{c, \varepsilon}(\frac{1}{n}\sum_i \delta_{x_i}, + \frac{1}{n}\sum_i \delta_{y_i}) + + where :math:`W_{c, \varepsilon}` is an entropy-regularized optimal transport + cost, the :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.ent_reg_cost`. + + Args: + source: samples from first measure, array of shape ``[n, d]``. + target: samples from second measure, array of shape ``[n, d]``. + cost_fn: a cost function between two points in dimension :math:`d`. + If :obj:`None`, :class:`~ott.geometry.costs.SqEuclidean` is used. + epsilon: Regularization parameter. See + :class:`~ott.geometry.pointcloud.PointCloud` + relative_epsilon: when `False`, the parameter ``epsilon`` specifies the + value of the entropic regularization parameter. When `True`, ``epsilon`` + refers to a fraction of the + :attr:`~ott.geometry.pointcloud.PointCloud.mean_cost_matrix`, which is + computed adaptively using ``source`` and ``target`` points. + scale_cost: option to rescale the cost matrix. Implemented scalings are + 'median', 'mean' and 'max_cost'. Alternatively, a float factor can be + given to rescale the cost such that ``cost_matrix /= scale_cost``. + If `True`, use 'mean'. + return_output: boolean to also return the + :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput`. + kwargs: holds the kwargs to instantiate the or + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver to + compute the regularized OT cost. + + Returns: + The Monge gap value and optionally the + :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput` + """ + cost_fn = costs.SqEuclidean() if cost_fn is None else cost_fn + geom = pointcloud.PointCloud( + x=source, + y=target, + cost_fn=cost_fn, + epsilon=epsilon, + relative_epsilon=relative_epsilon, + scale_cost=scale_cost, + ) + gt_displacement_cost = jnp.mean(jax.vmap(cost_fn)(source, target)) + out = linear.solve(geom=geom, **kwargs) + loss = gt_displacement_cost - out.ent_reg_cost + return (loss, out) if return_output else loss + + +def gw_monge_gap_from_samples( + source: jnp.ndarray, + target: jnp.ndarray, + cost_fn: Optional[costs.CostFn] = None, + epsilon: Optional[float] = None, + scale_cost: Union[ + bool, int, float, Literal["mean", "max_cost", "median"] + ] = 1.0, + return_output: bool = False, + **kwargs: Any, +) -> Union[float, Tuple[float, sinkhorn.SinkhornOutput]]: + r"""Monge gap, instantiated in terms of samples before / after applying map. + + .. math:: + \frac{1}{n} \sum_{i=1}^n c(x_i, y_i)) - + W_{c, \varepsilon}(\frac{1}{n}\sum_i \delta_{x_i}, + \frac{1}{n}\sum_i \delta_{y_i}) + + where :math:`W_{c, \varepsilon}` is an entropy-regularized optimal transport + cost, the :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.ent_reg_cost`. + + Args: + source: samples from first measure, array of shape ``[n, d]``. + target: samples from second measure, array of shape ``[n, d]``. + cost_fn: a cost function between two points in dimension :math:`d`. + If :obj:`None`, :class:`~ott.geometry.costs.SqEuclidean` is used. + epsilon: Regularization parameter. See + :class:`~ott.geometry.pointcloud.PointCloud` + relative_epsilon: when `False`, the parameter ``epsilon`` specifies the + value of the entropic regularization parameter. When `True`, ``epsilon`` + refers to a fraction of the + :attr:`~ott.geometry.pointcloud.PointCloud.mean_cost_matrix`, which is + computed adaptively using ``source`` and ``target`` points. + scale_cost: option to rescale the cost matrix. Implemented scalings are + 'median', 'mean' and 'max_cost'. Alternatively, a float factor can be + given to rescale the cost such that ``cost_matrix /= scale_cost``. + If `True`, use 'mean'. + return_output: boolean to also return the + :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput`. + kwargs: holds the kwargs to instantiate the or + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver to + compute the regularized OT cost. + + Returns: + The Monge gap value and optionally the + :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput` + """ + cost_fn = costs.SqEuclidean() if cost_fn is None else cost_fn + geom_xx = pointcloud.PointCloud( + x=source, y=source, cost_fn=cost_fn, scale_cost=scale_cost + ) + geom_yy = pointcloud.PointCloud( + x=target, y=target, cost_fn=cost_fn, scale_cost=scale_cost + ) + prob = quadratic_problem.QuadraticProblem(geom_xx, geom_yy) + solver = gromov_wasserstein.GromovWasserstein(epsilon=epsilon) + out = solver(prob) + + gt_displacement_cost = ( + (geom_xx.cost_matrix - geom_yy.cost_matrix) ** 2 + ).mean() + + loss = gt_displacement_cost - out.reg_gw_cost + return (loss, out) if return_output else loss + + +def labeled_gw_monge_gap_from_samples( + source: jnp.ndarray, + target: jnp.ndarray, + labels: jnp.array, + n_labels: int, + cost_fn: Optional[costs.CostFn] = None, + epsilon: Optional[float] = None, + scale_cost: Union[ + bool, int, float, Literal["mean", "max_cost", "median"] + ] = 1.0, + return_output: bool = False, + **kwargs: Any, +) -> Union[float, Tuple[float, sinkhorn.SinkhornOutput]]: + r"""Monge gap, instantiated in terms of samples before / after applying map. + + .. math:: + \frac{1}{n} \sum_{i=1}^n c(x_i, y_i)) - + W_{c, \varepsilon}(\frac{1}{n}\sum_i \delta_{x_i}, + \frac{1}{n}\sum_i \delta_{y_i}) + + where :math:`W_{c, \varepsilon}` is an entropy-regularized optimal transport + cost, the :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.ent_reg_cost`. + + Args: + source: samples from first measure, array of shape ``[n, d]``. + target: samples from second measure, array of shape ``[n, d]``. + cost_fn: a cost function between two points in dimension :math:`d`. + If :obj:`None`, :class:`~ott.geometry.costs.SqEuclidean` is used. + epsilon: Regularization parameter. See + :class:`~ott.geometry.pointcloud.PointCloud` + relative_epsilon: when `False`, the parameter ``epsilon`` specifies the + value of the entropic regularization parameter. When `True`, ``epsilon`` + refers to a fraction of the + :attr:`~ott.geometry.pointcloud.PointCloud.mean_cost_matrix`, which is + computed adaptively using ``source`` and ``target`` points. + scale_cost: option to rescale the cost matrix. Implemented scalings are + 'median', 'mean' and 'max_cost'. Alternatively, a float factor can be + given to rescale the cost such that ``cost_matrix /= scale_cost``. + If `True`, use 'mean'. + return_output: boolean to also return the + :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput`. + kwargs: holds the kwargs to instantiate the or + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver to + compute the regularized OT cost. + + Returns: + The Monge gap value and optionally the + :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput` + """ + cost_fn = costs.SqEuclidean() if cost_fn is None else cost_fn + geom_xx = pointcloud.PointCloud( + x=source, y=source, cost_fn=cost_fn, scale_cost=scale_cost + ) + geom_yy = pointcloud.PointCloud( + x=target, y=target, cost_fn=cost_fn, scale_cost=scale_cost + ) + bdm = create_block_diag_mat(labels, labels) + prob = quadratic_problem.QuadraticProblem( + geom_xx, + geom_yy, + labels_a=labels, + labels_b=labels, + n_labels=n_labels, + block_diag_mat=bdm, + ) + solver = gromov_wasserstein.GromovWasserstein(epsilon=epsilon) + out = solver(prob) + + gt_displacement_cost = ( + (geom_xx.cost_matrix - geom_yy.cost_matrix) ** 2 * bdm + ).mean() + + loss = gt_displacement_cost - out.reg_gw_cost + return (loss, out) if return_output else loss + + +def create_block_diag_mat(labels_a, labels_b): + """Creates block diagonal matrix that has 1 entry when label_a == label_b otherwise 0.""" + block_diag_mat = np.zeros((len(labels_a), len(labels_b))) + for l in np.unique(labels_a): + block_diag_mat[ + np.ix_(np.where(labels_a == l)[0], np.where(labels_b == l)[0]) + ] = 1.0 + return jnp.array(block_diag_mat) diff --git a/ott/build/lib/ott/neural/models.py b/ott/build/lib/ott/neural/models.py new file mode 100644 index 0000000..c8fc502 --- /dev/null +++ b/ott/build/lib/ott/neural/models.py @@ -0,0 +1,408 @@ +# 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 functools +from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union + +import jax +import jax.numpy as jnp +import optax +from flax import linen as nn +from flax.core import frozen_dict +from flax.training import train_state + +from ott import utils +from ott.geometry import geometry +from ott.initializers.linear import initializers as lin_init +from ott.neural import layers +from ott.neural.solvers import neuraldual +from ott.problems.linear import linear_problem + +__all__ = ["ICNN", "MLP", "MetaInitializer"] + +# wrap to silence docs linter +DEFAULT_KERNEL_INIT = lambda *a, **k: nn.initializers.normal()(*a, **k) +DEFAULT_RECTIFIER = nn.activation.relu +DEFAULT_ACTIVATION = nn.activation.relu + + +class ICNN(neuraldual.BaseW2NeuralDual): + """Input convex neural network (ICNN). + + Implementation of input convex neural networks as introduced in + :cite:`amos:17` with initialization schemes proposed by :cite:`bunne:22`. + + Args: + dim_data: data dimensionality. + dim_hidden: sequence specifying size of hidden dimensions. The + output dimension of the last layer is 1 by default. + ranks: ranks of the matrices :math:`A_i` used as low-rank factors + for the quadratic potentials. If a sequence is passed, it must contain + ``len(dim_hidden) + 2`` elements, where the last 2 elements correspond + to the ranks of the final layer with dimension 1 and the potentials, + respectively. + init_fn: Initializer for the kernel weight matrices. + The default is :func:`~flax.linen.initializers.normal`. + act_fn: choice of activation function used in network architecture, + needs to be convex. The default is :func:`~flax.linen.activation.relu`. + pos_weights: Enforce positive weights with a projection. + If :obj:`False`, the positive weights should be enforced with clipping + or regularization in the loss. + rectifier_fn: function to ensure the non negativity of the weights. + The default is :func:`~flax.linen.activation.relu`. + gaussian_map_samples: Tuple of source and target points, used to initialize + the ICNN to mimic the linear Bures map that morphs the (Gaussian + approximation) of the input measure to that of the target measure. If + :obj:`None`, the identity initialization is used, and ICNN mimics half the + squared Euclidean norm. + """ + + dim_data: int + dim_hidden: Sequence[int] + ranks: Union[int, Tuple[int, ...]] = 1 + init_fn: Callable[[jax.Array, Tuple[int, ...], Any], jnp.ndarray] = ( + DEFAULT_KERNEL_INIT + ) + act_fn: Callable[[jnp.ndarray], jnp.ndarray] = DEFAULT_ACTIVATION + pos_weights: bool = False + rectifier_fn: Callable[[jnp.ndarray], jnp.ndarray] = DEFAULT_RECTIFIER + gaussian_map_samples: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None + + def setup(self) -> None: # noqa: D102 + dim_hidden = list(self.dim_hidden) + [1] + *ranks, pos_def_rank = self._normalize_ranks() + + # final layer computes average, still with normalized rescaling + self.w_zs = [self._get_wz(dim) for dim in dim_hidden[1:]] + # subsequent layers re-injected into convex functions + self.w_xs = [ + self._get_wx(dim, rank) for dim, rank in zip(dim_hidden, ranks) + ] + self.pos_def_potentials = self._get_pos_def_potentials(pos_def_rank) + + @nn.compact + def __call__(self, x: jnp.ndarray) -> float: # noqa: D102 + w_x, *w_xs = self.w_xs + assert len(self.w_zs) == len(w_xs), (len(self.w_zs), len(w_xs)) + + z = self.act_fn(w_x(x)) + for w_z, w_x in zip(self.w_zs, w_xs): + z = self.act_fn(w_z(z) + w_x(x)) + z = z + self.pos_def_potentials(x) + + return z.squeeze() + + def _get_wz(self, dim: int) -> nn.Module: + if self.pos_weights: + return layers.PositiveDense( + dim, + kernel_init=self.init_fn, + use_bias=False, + rectifier_fn=self.rectifier_fn, + ) + + return nn.Dense( + dim, + kernel_init=self.init_fn, + use_bias=False, + ) + + def _get_wx(self, dim: int, rank: int) -> nn.Module: + return layers.PosDefPotentials( + rank=rank, + num_potentials=dim, + use_linear=True, + use_bias=True, + kernel_diag_init=nn.initializers.zeros, + kernel_lr_init=self.init_fn, + kernel_linear_init=self.init_fn, + bias_init=nn.initializers.zeros, + ) + + def _get_pos_def_potentials(self, rank: int) -> layers.PosDefPotentials: + kwargs = { + "num_potentials": 1, + "use_linear": True, + "use_bias": True, + "bias_init": nn.initializers.zeros, + } + + if self.gaussian_map_samples is None: + return layers.PosDefPotentials( + rank=rank, + kernel_diag_init=nn.initializers.ones, + kernel_lr_init=nn.initializers.zeros, + kernel_linear_init=nn.initializers.zeros, + **kwargs, + ) + + source, target = self.gaussian_map_samples + return layers.PosDefPotentials.init_from_samples( + source, + target, + rank=self.dim_data, + kernel_diag_init=nn.initializers.zeros, + **kwargs, + ) + + def _normalize_ranks(self) -> Tuple[int, ...]: + # +2 for the newly added layer with 1 + the final potentials + n_ranks = len(self.dim_hidden) + 2 + if isinstance(self.ranks, int): + return (self.ranks,) * n_ranks + + assert len(self.ranks) == n_ranks, (len(self.ranks), n_ranks) + return tuple(self.ranks) + + @property + def is_potential(self) -> bool: # noqa: D102 + return True + + +class MLP(neuraldual.BaseW2NeuralDual): + """A generic, not-convex MLP. + + Args: + dim_hidden: sequence specifying size of hidden dimensions. The output + dimension of the last layer is automatically set to 1 if + :attr:`is_potential` is ``True``, or the dimension of the input otherwise + is_potential: Model the potential if ``True``, otherwise + model the gradient of the potential + act_fn: Activation function + """ + + dim_hidden: Sequence[int] + dim_out: int + is_potential: bool = True + act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.leaky_relu + + @nn.compact + def __call__(self, x: jnp.ndarray) -> jnp.ndarray: # noqa: D102 + squeeze = x.ndim == 1 + if squeeze: + x = jnp.expand_dims(x, 0) + assert x.ndim == 2, x.ndim + n_input = x.shape[-1] + + z = x + for n_hidden in self.dim_hidden: + Wx = nn.Dense(n_hidden, use_bias=True) + z = self.act_fn(Wx(z)) + + if self.is_potential: + Wx = nn.Dense(1, use_bias=True) + z = Wx(z).squeeze(-1) + + quad_term = 0.5 * jax.vmap(jnp.dot)(x, x) + z += quad_term + elif self.dim_out is None: + Wx = nn.Dense(n_input, use_bias=True) + z = x + Wx(z) + else: + Wx = nn.Dense(self.dim_out, use_bias=True) + z = Wx(z) + + return z.squeeze(0) if squeeze else z + + +@jax.tree_util.register_pytree_node_class +class MetaInitializer(lin_init.DefaultInitializer): + """Meta OT Initializer with a fixed geometry :cite:`amos:22`. + + This initializer consists of a predictive model that outputs the + :math:`f` duals to solve the entropy-regularized OT problem given + input probability weights ``a`` and ``b``, and a given (assumed to be + fixed) geometry ``geom``. + + The model's parameters are learned using a training set of OT + instances (multiple pairs of probability weights), that assume the + **same** geometry ``geom`` is used throughout, both for training and + evaluation. + + Args: + geom: The fixed geometry of the problem instances. + meta_model: The model to predict the potential :math:`f` from the measures. + opt: The optimizer to update the parameters. If :obj:`None`, use + :func:`optax.adam` with :math:`0.001` learning rate. + rng: The PRNG key to use for initializing the model. + state: The training state of the model to start from. + + Examples: + The following code shows a simple + example of using ``update`` to train the model, where + ``a`` and ``b`` are the weights of the measures and + ``geom`` is the fixed geometry. + + .. code-block:: python + + meta_initializer = init_lib.MetaInitializer(geom) + while training(): + a, b = sample_batch() + loss, init_f, meta_initializer.state = meta_initializer.update( + meta_initializer.state, a=a, b=b + ) + """ + + def __init__( + self, + geom: geometry.Geometry, + meta_model: nn.Module, + opt: Optional[optax.GradientTransformation] = optax.adam( + learning_rate=1e-3 + ), # noqa: B008 + rng: Optional[jax.Array] = None, + state: Optional[train_state.TrainState] = None, + ): + self.geom = geom + self.opt = opt + self.rng = utils.default_prng_key(rng) + + na, nb = geom.shape + self.meta_model = meta_model + + if state is None: + # Initialize the model's training state. + placeholder = jnp.concatenate( + [jnp.zeros(na), jnp.zeros(nb)], axis=-1 + ) + params = self.meta_model.init(self.rng, placeholder)["params"] + self.state = train_state.TrainState.create( + apply_fn=self.meta_model.apply, params=params, tx=opt + ) + else: + self.state = state + + self.update_impl = self._get_update_fn() + + def update( + self, state: train_state.TrainState, a: jnp.ndarray, b: jnp.ndarray + ) -> Tuple[jnp.ndarray, jnp.ndarray, train_state.TrainState]: + r"""Update the meta model with the dual objective. + + The goal is for the model to match the optimal duals, i.e., + :math:`\hat f_\theta \approx f^\star`. + This can be done by training the predictions of :math:`\hat f_\theta` + to optimize the dual objective, which :math:`f^\star` also optimizes for. + The overall learning setup can thus be written as: + + .. math:: + \min_\theta\; {\mathbb E}_{(\alpha,\beta)\sim{\mathcal{D}}}\; + J(\hat f_\theta(a, b); \alpha, \beta), + + where :math:`a,b` are the probabilities of the measures :math:`\alpha,\beta` + ,:math:`\mathcal{D}` is a meta distribution of optimal transport problems, + + .. math:: + -J(f; \alpha, \beta, c) := \langle f, a\rangle + \langle g, b \rangle - + \varepsilon\left\langle \exp\{f/\varepsilon\}, K\exp\{g/\varepsilon\} + \right\rangle + + is the entropic dual objective, + and :math:`K_{i,j} := -C_{i,j}/\varepsilon` is the *Gibbs kernel*. + + Args: + state: Optimizer state of the meta model. + a: Probabilities of the :math:`\alpha` measure's atoms. + b: Probabilities of the :math:`\beta` measure's atoms. + + Returns: + The training loss, :math:`f`, and updated state. + """ + return self.update_impl(state, a, b) + + def init_dual_a( # noqa: D102 + self, + ot_prob: "linear_problem.LinearProblem", + lse_mode: bool, + rng: Optional[jax.Array] = None, + ) -> jnp.ndarray: + del rng + # Detect if the problem is batched. + assert ot_prob.a.ndim in (1, 2) + assert ot_prob.b.ndim in (1, 2) + vmap_a_val = 0 if ot_prob.a.ndim == 2 else None + vmap_b_val = 0 if ot_prob.b.ndim == 2 else None + + if vmap_a_val is not None or vmap_b_val is not None: + compute_f_maybe_batch = jax.vmap( + self._compute_f, in_axes=(vmap_a_val, vmap_b_val, None) + ) + else: + compute_f_maybe_batch = self._compute_f + + init_f = compute_f_maybe_batch(ot_prob.a, ot_prob.b, self.state.params) + return ( + init_f if lse_mode else ot_prob.geom.scaling_from_potential(init_f) + ) + + def _get_update_fn(self): + """Return the implementation (and jitted) update function.""" + from ott.problems.linear import linear_problem + from ott.solvers.linear import sinkhorn + + def dual_obj_loss_single(params, a, b): + f_pred = self._compute_f(a, b, params) + g_pred = self.geom.update_potential( + f_pred, jnp.zeros_like(b), jnp.log(b), 0, axis=0 + ) + g_pred = jnp.where(jnp.isfinite(g_pred), g_pred, 0.0) + + ot_prob = linear_problem.LinearProblem(geom=self.geom, a=a, b=b) + dual_obj = sinkhorn.compute_kl_reg_cost( + f_pred, g_pred, ot_prob, lse_mode=True + ) + loss = -dual_obj + return loss, f_pred + + def loss_batch(params, a, b): + loss_fn = functools.partial(dual_obj_loss_single, params=params) + loss, f_pred = jax.vmap(loss_fn)(a=a, b=b) + return jnp.mean(loss), f_pred + + @jax.jit + def update(state, a, b): + a = jnp.atleast_2d(a) + b = jnp.atleast_2d(b) + grad_fn = jax.value_and_grad(loss_batch, has_aux=True) + (loss, init_f), grads = grad_fn(state.params, a, b) + return loss, init_f, state.apply_gradients(grads=grads) + + return update + + def _compute_f( + self, + a: jnp.ndarray, + b: jnp.ndarray, + params: frozen_dict.FrozenDict[str, jnp.ndarray], + ) -> jnp.ndarray: + r"""Predict the optimal :math:`f` potential. + + Args: + a: Probabilities of the :math:`\alpha` measure's atoms. + b: Probabilities of the :math:`\beta` measure's atoms. + params: The parameters of the Meta model. + + Returns: + The :math:`f` potential. + """ + return self.meta_model.apply( + {"params": params}, jnp.concatenate([a, b], axis=-1) + ) + + def tree_flatten( + self, + ) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return [self.geom, self.meta_model, self.opt], { + "rng": self.rng, + "state": self.state, + } diff --git a/ott/build/lib/ott/neural/solvers/__init__.py b/ott/build/lib/ott/neural/solvers/__init__.py new file mode 100644 index 0000000..b09d8c6 --- /dev/null +++ b/ott/build/lib/ott/neural/solvers/__init__.py @@ -0,0 +1,14 @@ +# 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. +from . import conjugate, map_estimator, neuraldual diff --git a/ott/build/lib/ott/neural/solvers/conjugate.py b/ott/build/lib/ott/neural/solvers/conjugate.py new file mode 100644 index 0000000..7d0098f --- /dev/null +++ b/ott/build/lib/ott/neural/solvers/conjugate.py @@ -0,0 +1,121 @@ +# 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 abc +from typing import Callable, Literal, NamedTuple, Optional + +import jax.numpy as jnp +from jaxopt import LBFGS + +from ott import utils + +__all__ = [ + "ConjugateResults", + "FenchelConjugateSolver", + "FenchelConjugateLBFGS", + "DEFAULT_CONJUGATE_SOLVER", +] + + +class ConjugateResults(NamedTuple): + r"""Holds the results of numerically conjugating a function. + + Args: + val: the conjugate value, i.e., :math:`f^\star(y)` + grad: the gradient, i.e., :math:`\nabla f^\star(y)` + num_iter: the number of iterations taken by the solver + """ + val: float + grad: jnp.ndarray + num_iter: int + + +class FenchelConjugateSolver(abc.ABC): + r"""Abstract conjugate solver class. + + Given a function :math:`f`, numerically estimate the Fenchel conjugate + :math:`f^\star(y) := -\inf_{x\in\mathbb{R}^n} f(x)-\langle x, y\rangle`. + """ + + @abc.abstractmethod + def solve( + self, + f: Callable[[jnp.ndarray], jnp.ndarray], + y: jnp.ndarray, + x_init: Optional[jnp.ndarray] = None + ) -> ConjugateResults: + """Solve for the conjugate. + + Args: + f: function to conjugate + y: point to conjugate + x_init: initial point to search over + + Returns: + The solution to the conjugation. + """ + + +@utils.register_pytree_node +class FenchelConjugateLBFGS(FenchelConjugateSolver): + """Solve for the conjugate using :class:`~jaxopt.LBFGS`. + + Args: + gtol: gradient tolerance + max_iter: maximum number of iterations + max_linesearch_iter: maximum number of line search iterations + linesearch_type: type of line search + linesearch_init: strategy for line search initialization + increase_factor: factor by which to increase the step size during + the line search + """ + + gtol: float = 1e-3 + max_iter: int = 10 + max_linesearch_iter: int = 10 + linesearch_type: Literal["zoom", "backtracking", + "hager-zhang"] = "backtracking" + linesearch_init: Literal["increase", "max", "current"] = "increase" + increase_factor: float = 1.5 + + def solve( # noqa: D102 + self, + f: Callable[[jnp.ndarray], jnp.ndarray], + y: jnp.ndarray, + x_init: Optional[jnp.array] = None + ) -> ConjugateResults: + assert y.ndim == 1, y.ndim + + solver = LBFGS( + fun=lambda x: f(x) - x.dot(y), + tol=self.gtol, + maxiter=self.max_iter, + linesearch=self.linesearch_type, + linesearch_init=self.linesearch_init, + increase_factor=self.increase_factor, + implicit_diff=False, + unroll=False + ) + + out = solver.run(y if x_init is None else x_init) + return ConjugateResults( + val=-out.state.value, grad=out.params, num_iter=out.state.iter_num + ) + + +DEFAULT_CONJUGATE_SOLVER = FenchelConjugateLBFGS( + gtol=1e-5, + max_iter=20, + max_linesearch_iter=20, + linesearch_type="backtracking", +) diff --git a/ott/build/lib/ott/neural/solvers/map_estimator.py b/ott/build/lib/ott/neural/solvers/map_estimator.py new file mode 100644 index 0000000..f5389f5 --- /dev/null +++ b/ott/build/lib/ott/neural/solvers/map_estimator.py @@ -0,0 +1,289 @@ +# 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 collections +import functools +from typing import ( + Any, + Callable, + Dict, + Iterator, + Optional, + Sequence, + Tuple, + Union, +) + +import jax +import jax.numpy as jnp +import optax +from flax.core import frozen_dict +from flax.training import train_state + +from ott import utils +from ott.neural.solvers import neuraldual + +__all__ = ["MapEstimator"] + + +class MapEstimator: + r"""Mapping estimator between probability measures. + + It estimates a map :math:`T` by minimizing the loss: + + .. math:: + \text{min}_{\theta}\; \Delta(T_\theta \sharp \mu, \theta) + + \lambda R(T_\theta \sharp \rho, \rho) + + where :math:`\Delta` is a fitting loss and :math:`R` is a regularizer. + :math:`\Delta` allows to fit the marginal constraint, i.e. transport + :math:`\mu` to :math:`\nu` via :math:`T`, while :math:`R` + is a regularizer imposing an inductive bias on the learned map. The + regularizer in this case is a function used to compute a metric between two + sets of points. + + For instance, :math:`\Delta` can be the + :func:`~ott.tools.sinkhorn_divergence.sinkhorn_divergence` + and :math:`R` the :func:`~ott.neural.losses.monge_gap_from_samples` + :cite:`uscidda:23` for a given cost function :math:`c`. + In that case, it estimates a :math:`c`-OT map, i.e. a map :math:`T` + optimal for the Monge problem induced by :math:`c`. + + Args: + dim_data: input dimensionality of data required for network init. + model: network architecture for map :math:`T`. + optimizer: optimizer function for map :math:`T`. + fitting_loss: function that outputs a fitting loss :math:`\Delta` between + two families of points, as well as any log object. + regularizer: function that outputs a score from two families of points, + here assumed to be of the same size, as well as any log object. + regularizer_strength: strength of the :attr:`regularizer`. + num_train_iters: number of total training iterations. + logging: option to return logs. + valid_freq: frequency with training and validation are logged. + rng: random key used for seeding for network initializations. + """ + + def __init__( + self, + dim_data: int, + model: neuraldual.BaseW2NeuralDual, + optimizer: Optional[optax.OptState] = None, + fitting_loss: Optional[Callable[[jnp.ndarray, jnp.ndarray], + Tuple[float, Optional[Any]]]] = None, + regularizer: Optional[Callable[[jnp.ndarray, jnp.ndarray], + Tuple[float, Optional[Any]]]] = None, + regularizer_strength: Union[float, Sequence[float]] = 1.0, + num_train_iters: int = 10_000, + logging: bool = False, + valid_freq: int = 500, + rng: Optional[jax.Array] = None, + ): + self._fitting_loss = fitting_loss + self._regularizer = regularizer + # Can use either a fixed strength, or generalize to a schedule. + self.regularizer_strength = jnp.repeat( + jnp.atleast_2d(regularizer_strength), + num_train_iters, + total_repeat_length=num_train_iters, + axis=0 + ).ravel() + self.num_train_iters = num_train_iters + self.logging = logging + self.valid_freq = valid_freq + self.rng = utils.default_prng_key(rng) + + # set default optimizer + if optimizer is None: + optimizer = optax.adam(learning_rate=0.001) + + # setup training + self.setup(dim_data, model, optimizer) + + def setup( + self, + dim_data: int, + neural_net: neuraldual.BaseW2NeuralDual, + optimizer: optax.OptState, + ): + """Setup all components required to train the network.""" + # neural network + self.state_neural_net = neural_net.create_train_state( + self.rng, optimizer, dim_data + ) + + # step function + self.step_fn = self._get_step_fn() + + @property + def regularizer(self) -> Callable[[jnp.ndarray, jnp.ndarray], float]: + """Regularizer added to the fitting loss. + + Can be, e.g. the :func:`~ott.neural.losses.monge_gap_from_samples`. + If no regularizer is passed for solver instantiation, + or regularization weight :attr:`regularizer_strength` is 0, + return 0 by default along with an empty set of log values. + """ + if self._regularizer is not None: + return self._regularizer + return lambda *_, **__: (0.0, None) + + @property + def fitting_loss(self) -> Callable[[jnp.ndarray, jnp.ndarray], float]: + """Fitting loss to fit the marginal constraint. + + Can be, e.g. :func:`~ott.tools.sinkhorn_divergence.sinkhorn_divergence`. + If no fitting_loss is passed for solver instantiation, return 0 by default, + and no log values. + """ + if self._fitting_loss is not None: + return self._fitting_loss + return lambda *_, **__: (0.0, None) + + @staticmethod + def _generate_batch( + loader_source: Iterator[jnp.ndarray], + loader_target: Iterator[jnp.ndarray], + ) -> Dict[str, jnp.ndarray]: + """Generate batches a batch of samples. + + ``loader_source`` and ``loader_target`` can be training or + validation dataloaders. + """ + return { + "source": next(loader_source), + "target": next(loader_target), + } + + def train_map_estimator( + self, + trainloader_source: Iterator[jnp.ndarray], + trainloader_target: Iterator[jnp.ndarray], + validloader_source: Iterator[jnp.ndarray], + validloader_target: Iterator[jnp.ndarray], + ) -> Tuple[train_state.TrainState, Dict[str, Any]]: + """Training loop.""" + # define logs + logs = collections.defaultdict(lambda: collections.defaultdict(list)) + + # try to display training progress with tqdm + try: + from tqdm import trange + tbar = trange(self.num_train_iters, leave=True) + except ImportError: + tbar = range(self.num_train_iters) + + for step in tbar: + # update step + is_logging_step = ( + self.logging and ((step % self.valid_freq == 0) or + (step == self.num_train_iters - 1)) + ) + train_batch = self._generate_batch( + loader_source=trainloader_source, + loader_target=trainloader_target, + ) + valid_batch = ( + None if not is_logging_step else self._generate_batch( + loader_source=validloader_source, + loader_target=validloader_target, + ) + ) + self.state_neural_net, current_logs = self.step_fn( + self.state_neural_net, train_batch, valid_batch, is_logging_step, step + ) + + # store and print metrics if logging step + if is_logging_step: + for log_key in current_logs: + for metric_key in current_logs[log_key]: + logs[log_key][metric_key].append(current_logs[log_key][metric_key]) + + # update the tqdm bar if tqdm is available + if not isinstance(tbar, range): + reg_msg = ( + "NA" if current_logs["eval"]["regularizer"] == 0.0 else + f"{current_logs['eval']['regularizer']:.4f}" + ) + postfix_str = ( + f"fitting_loss: {current_logs['eval']['fitting_loss']:.4f}, " + f"regularizer: {reg_msg} ," + f"total: {current_logs['eval']['total_loss']:.4f}" + ) + tbar.set_postfix_str(postfix_str) + + return self.state_neural_net, logs + + def _get_step_fn(self) -> Callable: + """Create a one step training and evaluation function.""" + + def loss_fn( + params: frozen_dict.FrozenDict, apply_fn: Callable, + batch: Dict[str, jnp.ndarray], step: int + ) -> Tuple[float, Dict[str, float]]: + """Loss function.""" + # map samples with the fitted map + mapped_samples = apply_fn({"params": params}, batch["source"]) + + # compute the loss + val_fitting_loss, log_fitting_loss = self.fitting_loss( + mapped_samples, batch["target"] + ) + val_regularizer, log_regularizer = self.regularizer( + batch["source"], mapped_samples + ) + val_tot_loss = ( + val_fitting_loss + self.regularizer_strength[step] * val_regularizer + ) + + # store training logs + loss_logs = { + "total_loss": val_tot_loss, + "fitting_loss": val_fitting_loss, + "regularizer": val_regularizer, + "log_regularizer": log_regularizer, + "log_fitting": log_fitting_loss, + } + + return val_tot_loss, loss_logs + + @functools.partial(jax.jit, static_argnums=3) + def step_fn( + state_neural_net: train_state.TrainState, + train_batch: Dict[str, jnp.ndarray], + valid_batch: Optional[Dict[str, jnp.ndarray]] = None, + is_logging_step: bool = False, + step: int = 0 + ) -> Tuple[train_state.TrainState, Dict[str, float]]: + """One step function.""" + # compute loss and gradients + grad_fn = jax.value_and_grad(loss_fn, argnums=0, has_aux=True) + (_, current_train_logs), grads = grad_fn( + state_neural_net.params, state_neural_net.apply_fn, train_batch, step + ) + + # logging step + current_logs = {"train": current_train_logs, "eval": {}} + if is_logging_step: + _, current_eval_logs = loss_fn( + params=state_neural_net.params, + apply_fn=state_neural_net.apply_fn, + batch=valid_batch, + step=step + ) + current_logs["eval"] = current_eval_logs + + # update state + return state_neural_net.apply_gradients(grads=grads), current_logs + + return step_fn diff --git a/ott/build/lib/ott/neural/solvers/neuraldual.py b/ott/build/lib/ott/neural/solvers/neuraldual.py new file mode 100644 index 0000000..fffa927 --- /dev/null +++ b/ott/build/lib/ott/neural/solvers/neuraldual.py @@ -0,0 +1,700 @@ +# 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 abc +import warnings +from typing import ( + Any, + Callable, + Dict, + Iterator, + List, + Literal, + Optional, + Tuple, + Union, +) + +import jax +import jax.numpy as jnp +import optax +from flax import linen as nn +from flax import struct +from flax.core import frozen_dict +from flax.training import train_state + +from ott import utils +from ott.geometry import costs +from ott.neural import models +from ott.neural.solvers import conjugate +from ott.problems.linear import potentials + +__all__ = ["W2NeuralTrainState", "BaseW2NeuralDual", "W2NeuralDual"] + +Train_t = Dict[Literal["train_logs", "valid_logs"], Dict[str, List[float]]] +Callback_t = Callable[[int, potentials.DualPotentials], None] + +PotentialValueFn_t = Callable[[jnp.ndarray], jnp.ndarray] +PotentialGradientFn_t = Callable[[jnp.ndarray], jnp.ndarray] + + +class W2NeuralTrainState(train_state.TrainState): + """Adds information about the model's value and gradient to the state. + + This extends :class:`~flax.training.train_state.TrainState` to include + the potential methods from the + :class:`~ott.neural.solvers.neuraldual.BaseW2NeuralDual` used during training. + + Args: + potential_value_fn: the potential's value function + potential_gradient_fn: the potential's gradient function + """ + potential_value_fn: Callable[ + [frozen_dict.FrozenDict[str, jnp.ndarray], Optional[PotentialValueFn_t]], + PotentialValueFn_t] = struct.field(pytree_node=False) + potential_gradient_fn: Callable[[frozen_dict.FrozenDict[str, jnp.ndarray]], + PotentialGradientFn_t] = struct.field( + pytree_node=False + ) + + +class BaseW2NeuralDual(abc.ABC, nn.Module): + """Base class for the neural solver models.""" + + @property + @abc.abstractmethod + def is_potential(self) -> bool: + """Indicates if the module implements a potential value or a vector field. + + Returns: + ``True`` if the module defines a potential, ``False`` if it defines a + vector field. + """ + + def potential_value_fn( + self, + params: frozen_dict.FrozenDict[str, jnp.ndarray], + other_potential_value_fn: Optional[PotentialValueFn_t] = None, + ) -> PotentialValueFn_t: + r"""Return a function giving the value of the potential. + + Applies the module if :attr:`is_potential` is ``True``, otherwise + constructs the value of the potential from the gradient with + + .. math:: + + g(y) = -f(\nabla_y g(y)) + y^T \nabla_y g(y) + + where :math:`\nabla_y g(y)` is detached for the envelope theorem + :cite:`danskin:67,bertsekas:71` + to give the appropriate first derivatives of this construction. + + Args: + params: parameters of the module + other_potential_value_fn: function giving the value of the other + potential. Only needed when :attr:`is_potential` is ``False``. + + Returns: + A function that can be evaluated to obtain a potential value, or a linear + interpolation of a potential. + """ + if self.is_potential: + return lambda x: self.apply({"params": params}, x) + + assert other_potential_value_fn is not None, \ + "The value of the gradient-based potential depends " \ + "on the value of the other potential." + + def value_fn(x: jnp.ndarray) -> jnp.ndarray: + squeeze = x.ndim == 1 + if squeeze: + x = jnp.expand_dims(x, 0) + grad_g_x = jax.lax.stop_gradient(self.apply({"params": params}, x)) + value = -other_potential_value_fn(grad_g_x) + \ + jax.vmap(jnp.dot)(grad_g_x, x) + return value.squeeze(0) if squeeze else value + + return value_fn + + def potential_gradient_fn( + self, + params: frozen_dict.FrozenDict[str, jnp.ndarray], + ) -> PotentialGradientFn_t: + """Return a function returning a vector or the gradient of the potential. + + Args: + params: parameters of the module + + Returns: + A function that can be evaluated to obtain the potential's gradient + """ + if self.is_potential: + return jax.vmap(jax.grad(self.potential_value_fn(params))) + return lambda x: self.apply({"params": params}, x) + + def create_train_state( + self, + rng: jax.Array, + optimizer: optax.OptState, + input: Union[int, Tuple[int, ...]], + **kwargs: Any, + ) -> W2NeuralTrainState: + """Create initial training state.""" + params = self.init(rng, jnp.ones(input))["params"] + + return W2NeuralTrainState.create( + apply_fn=self.apply, + params=params, + tx=optimizer, + potential_value_fn=self.potential_value_fn, + potential_gradient_fn=self.potential_gradient_fn, + **kwargs, + ) + + +class W2NeuralDual: + r"""Solver for the Wasserstein-2 Kantorovich dual between Euclidean spaces. + + Learn the Wasserstein-2 optimal transport between two measures + :math:`\alpha` and :math:`\beta` in :math:`n`-dimensional Euclidean space, + denoted source and target, respectively. This is achieved by parameterizing + a Kantorovich potential :math:`f_\theta: \mathbb{R}^n\rightarrow\mathbb{R}` + associated with the :math:`\alpha` measure with an + :class:`~ott.neural.models.ICNN` or :class:`~ott.neural.models.MLP`, where + :math:`\nabla f` transports source to target cells. This potential is learned + by optimizing the dual form associated with the negative inner product cost + + .. math:: + + \text{argsup}_{\theta}\; -\mathbb{E}_{x\sim\alpha}[f_\theta(x)] - + \mathbb{E}_{y\sim\beta}[f^\star_\theta(y)], + + where :math:`f^\star(y) := -\inf_{x\in\mathbb{R}^n} f(x)-\langle x, y\rangle` + is the convex conjugate. + :math:`\nabla f^\star` transports from the target + to source cells and provides the inverse optimal + transport map from :math:`\beta` to :math:`\alpha`. + This solver estimates the conjugate :math:`f^\star` + with a neural approximation :math:`g` that is fine-tuned + with :class:`~ott.neural.solvers.conjugate.FenchelConjugateSolver`, + which is a combination further described in :cite:`amos:23`. + + The :class:`~ott.neural.solvers.neuraldual.BaseW2NeuralDual` potentials for + ``neural_f`` and ``neural_g`` can + + 1. both provide the values of the potentials :math:`f` and :math:`g`, or + 2. one of them can provide the gradient mapping e.g., :math:`\nabla f` + or :math:`\nabla g` where the potential's value can be obtained + via the Fenchel conjugate as discussed in :cite:`amos:23`. + + The potential's value or gradient mapping is specified via + :attr:`~ott.neural.solvers.neuraldual.BaseW2NeuralDual.is_potential`. + + Args: + dim_data: input dimensionality of data required for network init + neural_f: network architecture for potential :math:`f`. + neural_g: network architecture for the conjugate potential + :math:`g\approx f^\star` + optimizer_f: optimizer function for potential :math:`f` + optimizer_g: optimizer function for the conjugate potential :math:`g` + num_train_iters: number of total training iterations + num_inner_iters: number of training iterations of :math:`g` per iteration + of :math:`f` + back_and_forth: alternate between updating the forward and backward + directions. Inspired from :cite:`jacobs:20` + valid_freq: frequency with which model is validated + log_freq: frequency with training and validation are logged + logging: option to return logs + rng: random key used for seeding for network initializations + pos_weights: option to train networks with positive weights or regularizer + beta: regularization parameter when not training with positive weights + conjugate_solver: numerical solver for the Fenchel conjugate. + amortization_loss: amortization loss for the conjugate + :math:`g\approx f^\star`. Options are `'objective'` :cite:`makkuva:20` or + `'regression'` :cite:`amos:23`. + parallel_updates: Update :math:`f` and :math:`g` at the same time + """ + + def __init__( + self, + dim_data: int, + neural_f: Optional[BaseW2NeuralDual] = None, + neural_g: Optional[BaseW2NeuralDual] = None, + optimizer_f: Optional[optax.OptState] = None, + optimizer_g: Optional[optax.OptState] = None, + num_train_iters: int = 20000, + num_inner_iters: int = 1, + back_and_forth: Optional[bool] = None, + valid_freq: int = 1000, + log_freq: int = 1000, + logging: bool = False, + rng: Optional[jax.Array] = None, + pos_weights: bool = True, + beta: float = 1.0, + conjugate_solver: Optional[conjugate.FenchelConjugateSolver + ] = conjugate.DEFAULT_CONJUGATE_SOLVER, + amortization_loss: Literal["objective", "regression"] = "regression", + parallel_updates: bool = True, + ): + self.num_train_iters = num_train_iters + self.num_inner_iters = num_inner_iters + self.back_and_forth = back_and_forth + self.valid_freq = valid_freq + self.log_freq = log_freq + self.logging = logging + self.pos_weights = pos_weights + self.beta = beta + self.parallel_updates = parallel_updates + self.conjugate_solver = conjugate_solver + self.amortization_loss = amortization_loss + + # set default optimizers + if optimizer_f is None: + optimizer_f = optax.adam(learning_rate=0.0001, b1=0.5, b2=0.9, eps=1e-8) + if optimizer_g is None: + optimizer_g = optax.adam(learning_rate=0.0001, b1=0.5, b2=0.9, eps=1e-8) + + # set default neural architectures + if neural_f is None: + neural_f = models.ICNN(dim_data=dim_data, dim_hidden=[64, 64, 64, 64]) + if neural_g is None: + neural_g = models.ICNN(dim_data=dim_data, dim_hidden=[64, 64, 64, 64]) + self.neural_f = neural_f + self.neural_g = neural_g + + # set optimizer and networks + self.setup( + utils.default_prng_key(rng), + neural_f, + neural_g, + dim_data, + optimizer_f, + optimizer_g, + ) + + def setup( + self, + rng: jax.Array, + neural_f: BaseW2NeuralDual, + neural_g: BaseW2NeuralDual, + dim_data: int, + optimizer_f: optax.OptState, + optimizer_g: optax.OptState, + ) -> None: + """Setup all components required to train the network.""" + # split random number generator + rng, rng_f, rng_g = jax.random.split(rng, 3) + + # check setting of network architectures + warn_str = f"Setting of ICNN and the positive weights setting of the " \ + f"`W2NeuralDual` are not consistent. Proceeding with " \ + f"the `W2NeuralDual` setting, with positive weights " \ + f"being {self.pos_weights}." + if isinstance( + neural_f, models.ICNN + ) and neural_f.pos_weights is not self.pos_weights: + warnings.warn(warn_str, stacklevel=2) + neural_f.pos_weights = self.pos_weights + + if isinstance( + neural_g, models.ICNN + ) and neural_g.pos_weights is not self.pos_weights: + warnings.warn(warn_str, stacklevel=2) + neural_g.pos_weights = self.pos_weights + + self.state_f = neural_f.create_train_state( + rng_f, + optimizer_f, + (1, dim_data), # also include the batch dimension + ) + self.state_g = neural_g.create_train_state( + rng_g, + optimizer_g, + (1, dim_data), + ) + + # default to using back_and_forth with the non-convex models + if self.back_and_forth is None: + self.back_and_forth = isinstance(neural_f, models.MLP) + + if self.num_inner_iters == 1 and self.parallel_updates: + self.train_step_parallel = self.get_step_fn( + train=True, to_optimize="both" + ) + self.valid_step_parallel = self.get_step_fn( + train=False, to_optimize="both" + ) + self.train_fn = self.train_neuraldual_parallel + else: + if self.parallel_updates: + warnings.warn( + "parallel_updates set to True but disabling it " + "because num_inner_iters>1", + stacklevel=2 + ) + if self.back_and_forth: + raise NotImplementedError( + "back_and_forth not implemented without parallel updates" + ) + self.train_step_f = self.get_step_fn(train=True, to_optimize="f") + self.valid_step_f = self.get_step_fn(train=False, to_optimize="f") + self.train_step_g = self.get_step_fn(train=True, to_optimize="g") + self.valid_step_g = self.get_step_fn(train=False, to_optimize="g") + self.train_fn = self.train_neuraldual_alternating + + def __call__( # noqa: D102 + self, + trainloader_source: Iterator[jnp.ndarray], + trainloader_target: Iterator[jnp.ndarray], + validloader_source: Iterator[jnp.ndarray], + validloader_target: Iterator[jnp.ndarray], + callback: Optional[Callback_t] = None, + ) -> Union[potentials.DualPotentials, Tuple[potentials.DualPotentials, + Train_t]]: + logs = self.train_fn( + trainloader_source, + trainloader_target, + validloader_source, + validloader_target, + callback=callback, + ) + res = self.to_dual_potentials() + + return (res, logs) if self.logging else res + + def train_neuraldual_parallel( + self, + trainloader_source: Iterator[jnp.ndarray], + trainloader_target: Iterator[jnp.ndarray], + validloader_source: Iterator[jnp.ndarray], + validloader_target: Iterator[jnp.ndarray], + callback: Optional[Callback_t] = None, + ) -> Train_t: + """Training and validation with parallel updates.""" + try: + from tqdm.auto import tqdm + except ImportError: + tqdm = lambda _: _ + # define dict to contain source and target batch + train_batch, valid_batch = {}, {} + + # set logging dictionaries + train_logs = {"loss_f": [], "loss_g": [], "w_dist": [], "directions": []} + valid_logs = {"loss_f": [], "loss_g": [], "w_dist": []} + + for step in tqdm(range(self.num_train_iters)): + update_forward = not self.back_and_forth or step % 2 == 0 + if update_forward: + train_batch["source"] = jnp.asarray(next(trainloader_source)) + train_batch["target"] = jnp.asarray(next(trainloader_target)) + (self.state_f, self.state_g, loss, loss_f, loss_g, + w_dist) = self.train_step_parallel( + self.state_f, + self.state_g, + train_batch, + ) + else: + train_batch["target"] = jnp.asarray(next(trainloader_source)) + train_batch["source"] = jnp.asarray(next(trainloader_target)) + (self.state_g, self.state_f, loss, loss_f, loss_g, + w_dist) = self.train_step_parallel( + self.state_g, + self.state_f, + train_batch, + ) + + if self.logging and step % self.log_freq == 0: + self._update_logs(train_logs, loss_f, loss_g, w_dist) + train_logs["directions"].append( + "forward" if update_forward else "backward" + ) + + if callback is not None: + _ = callback(step, self.to_dual_potentials()) + + if not self.pos_weights: + # Only clip the weights of the f network + self.state_f = self.state_f.replace( + params=self._clip_weights_icnn(self.state_f.params) + ) + + # report the loss on an validation dataset periodically + if step != 0 and step % self.valid_freq == 0: + # get batch + valid_batch["source"] = jnp.asarray(next(validloader_source)) + valid_batch["target"] = jnp.asarray(next(validloader_target)) + + valid_loss_f, valid_loss_g, valid_w_dist = self.valid_step_parallel( + self.state_f, + self.state_g, + valid_batch, + ) + + if self.logging: + self._update_logs( + valid_logs, valid_loss_f, valid_loss_g, valid_w_dist + ) + + return {"train_logs": train_logs, "valid_logs": valid_logs} + + def train_neuraldual_alternating( + self, + trainloader_source: Iterator[jnp.ndarray], + trainloader_target: Iterator[jnp.ndarray], + validloader_source: Iterator[jnp.ndarray], + validloader_target: Iterator[jnp.ndarray], + callback: Optional[Callback_t] = None, + ) -> Train_t: + """Training and validation with alternating updates.""" + try: + from tqdm.auto import tqdm + except ImportError: + tqdm = lambda _: _ + # define dict to contain source and target batch + batch_g, batch_f, valid_batch = {}, {}, {} + + # set logging dictionaries + train_logs = {"loss_f": [], "loss_g": [], "w_dist": []} + valid_logs = {"loss_f": [], "loss_g": [], "w_dist": []} + + for step in tqdm(range(self.num_train_iters)): + # execute training steps + for _ in range(self.num_inner_iters): + # get train batch for potential g + batch_g["source"] = jnp.asarray(next(trainloader_source)) + batch_g["target"] = jnp.asarray(next(trainloader_target)) + + self.state_g, loss_g, _ = self.train_step_g( + self.state_f, self.state_g, batch_g + ) + + # get train batch for potential f + batch_f["source"] = jnp.asarray(next(trainloader_source)) + batch_f["target"] = jnp.asarray(next(trainloader_target)) + + self.state_f, loss_f, w_dist = self.train_step_f( + self.state_f, self.state_g, batch_f + ) + if not self.pos_weights: + # Only clip the weights of the f network + self.state_f = self.state_f.replace( + params=self._clip_weights_icnn(self.state_f.params) + ) + + if callback is not None: + callback(step, self.to_dual_potentials()) + + if self.logging and step % self.log_freq == 0: + self._update_logs(train_logs, loss_f, loss_g, w_dist) + + # report the loss on validation dataset periodically + if step != 0 and step % self.valid_freq == 0: + # get batch + valid_batch["source"] = jnp.asarray(next(validloader_source)) + valid_batch["target"] = jnp.asarray(next(validloader_target)) + + valid_loss_f, _ = self.valid_step_f( + self.state_f, self.state_g, valid_batch + ) + valid_loss_g, valid_w_dist = self.valid_step_g( + self.state_f, self.state_g, valid_batch + ) + + if self.logging: + self._update_logs( + valid_logs, valid_loss_f, valid_loss_g, valid_w_dist + ) + + return {"train_logs": train_logs, "valid_logs": valid_logs} + + def get_step_fn( + self, train: bool, to_optimize: Literal["f", "g", "parallel", "both"] + ): + """Create a parallel training and evaluation function.""" + + def loss_fn(params_f, params_g, f_value, g_value, g_gradient, batch): + """Loss function for both potentials.""" + # get two distributions + source, target = batch["source"], batch["target"] + + init_source_hat = g_gradient(params_g)(target) + + def g_value_partial(y: jnp.ndarray) -> jnp.ndarray: + """Lazy way of evaluating g if f's computation needs it.""" + return g_value(params_g)(y) + + f_value_partial = f_value(params_f, g_value_partial) + if self.conjugate_solver is not None: + finetune_source_hat = lambda y, x_init: self.conjugate_solver.solve( + f_value_partial, y, x_init=x_init + ).grad + finetune_source_hat = jax.vmap(finetune_source_hat) + source_hat_detach = jax.lax.stop_gradient( + finetune_source_hat(target, init_source_hat) + ) + else: + source_hat_detach = init_source_hat + + batch_dot = jax.vmap(jnp.dot) + + f_source = f_value_partial(source) + f_star_target = batch_dot(source_hat_detach, + target) - f_value_partial(source_hat_detach) + dual_source = f_source.mean() + dual_target = f_star_target.mean() + dual_loss = dual_source + dual_target + + if self.amortization_loss == "regression": + amor_loss = ((init_source_hat - source_hat_detach) ** 2).mean() + elif self.amortization_loss == "objective": + f_value_parameters_detached = f_value( + jax.lax.stop_gradient(params_f), g_value_partial + ) + amor_loss = ( + f_value_parameters_detached(init_source_hat) - + batch_dot(init_source_hat, target) + ).mean() + else: + raise ValueError("Amortization loss has been misspecified.") + + if to_optimize == "both": + loss = dual_loss + amor_loss + elif to_optimize == "f": + loss = dual_loss + elif to_optimize == "g": + loss = amor_loss + else: + raise ValueError( + f"Optimization target {to_optimize} has been misspecified." + ) + + if not self.pos_weights: + # Penalize the weights of both networks, even though one + # of them will be exactly clipped. + # Having both here is necessary in case this is being called with + # the potentials reversed with the back_and_forth. + loss += self.beta * self._penalize_weights_icnn(params_f) + \ + self.beta * self._penalize_weights_icnn(params_g) + + # compute Wasserstein-2 distance + C = jnp.mean(jnp.sum(source ** 2, axis=-1)) + \ + jnp.mean(jnp.sum(target ** 2, axis=-1)) + W2_dist = C - 2.0 * (f_source.mean() + f_star_target.mean()) + + return loss, (dual_loss, amor_loss, W2_dist) + + @jax.jit + def step_fn(state_f, state_g, batch): + """Step function of either training or validation.""" + grad_fn = jax.value_and_grad(loss_fn, argnums=[0, 1], has_aux=True) + if train: + # compute loss and gradients + (loss, (loss_f, loss_g, W2_dist)), (grads_f, grads_g) = grad_fn( + state_f.params, + state_g.params, + state_f.potential_value_fn, + state_g.potential_value_fn, + state_g.potential_gradient_fn, + batch, + ) + # update state + if to_optimize == "both": + return ( + state_f.apply_gradients(grads=grads_f), + state_g.apply_gradients(grads=grads_g), loss, loss_f, loss_g, + W2_dist + ) + if to_optimize == "f": + return state_f.apply_gradients(grads=grads_f), loss_f, W2_dist + if to_optimize == "g": + return state_g.apply_gradients(grads=grads_g), loss_g, W2_dist + raise ValueError("Optimization target has been misspecified.") + + # compute loss and gradients + (loss, (loss_f, loss_g, W2_dist)), _ = grad_fn( + state_f.params, + state_g.params, + state_f.potential_value_fn, + state_g.potential_value_fn, + state_g.potential_gradient_fn, + batch, + ) + + # do not update state + if to_optimize == "both": + return loss_f, loss_g, W2_dist + if to_optimize == "f": + return loss_f, W2_dist + if to_optimize == "g": + return loss_g, W2_dist + raise ValueError("Optimization target has been misspecified.") + + return step_fn + + def to_dual_potentials( + self, finetune_g: bool = True + ) -> potentials.DualPotentials: + r"""Return the Kantorovich dual potentials from the trained potentials. + + Args: + finetune_g: Run the conjugate solver to fine-tune the prediction. + + Returns: + A dual potential object + """ + f_value = self.state_f.potential_value_fn(self.state_f.params) + g_value_prediction = self.state_g.potential_value_fn( + self.state_g.params, f_value + ) + + def g_value_finetuned(y: jnp.ndarray) -> jnp.ndarray: + x_hat = jax.grad(g_value_prediction)(y) + grad_g_y = jax.lax.stop_gradient( + self.conjugate_solver.solve(f_value, y, x_init=x_hat).grad + ) + return -f_value(grad_g_y) + jnp.dot(grad_g_y, y) + + return potentials.DualPotentials( + f=f_value, + g=g_value_prediction if not finetune_g or self.conjugate_solver is None + else g_value_finetuned, + cost_fn=costs.SqEuclidean(), + corr=True + ) + + @staticmethod + def _clip_weights_icnn(params): + for k in params: + if k.startswith("w_z"): + params[k]["kernel"] = jnp.clip(params[k]["kernel"], a_min=0) + + return params + + @staticmethod + def _penalize_weights_icnn(params: Dict[str, jnp.ndarray]) -> float: + penalty = 0.0 + for k, param in params.items(): + if k.startswith("w_z"): + penalty += jnp.linalg.norm(jax.nn.relu(-param["kernel"])) + return penalty + + @staticmethod + def _update_logs( + logs: Dict[str, List[Union[float, str]]], + loss_f: jnp.ndarray, + loss_g: jnp.ndarray, + w_dist: jnp.ndarray, + ) -> None: + logs["loss_f"].append(float(loss_f)) + logs["loss_g"].append(float(loss_g)) + logs["w_dist"].append(float(w_dist)) diff --git a/ott/build/lib/ott/problems/__init__.py b/ott/build/lib/ott/problems/__init__.py new file mode 100644 index 0000000..5406247 --- /dev/null +++ b/ott/build/lib/ott/problems/__init__.py @@ -0,0 +1,14 @@ +# 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. +from . import linear, quadratic diff --git a/ott/build/lib/ott/problems/linear/__init__.py b/ott/build/lib/ott/problems/linear/__init__.py new file mode 100644 index 0000000..7d62b13 --- /dev/null +++ b/ott/build/lib/ott/problems/linear/__init__.py @@ -0,0 +1,14 @@ +# 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. +from . import barycenter_problem, linear_problem, potentials diff --git a/ott/build/lib/ott/problems/linear/barycenter_problem.py b/ott/build/lib/ott/problems/linear/barycenter_problem.py new file mode 100644 index 0000000..ca5333a --- /dev/null +++ b/ott/build/lib/ott/problems/linear/barycenter_problem.py @@ -0,0 +1,200 @@ +# 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. +from typing import Any, Dict, Optional, Sequence, Tuple + +import jax +import jax.numpy as jnp + +from ott.geometry import costs, geometry, segment + +__all__ = ["FreeBarycenterProblem", "FixedBarycenterProblem"] + + +@jax.tree_util.register_pytree_node_class +class FreeBarycenterProblem: + """Free Wasserstein barycenter problem :cite:`cuturi:14`. + + Args: + y: Array of shape ``[num_total_points, ndim]`` merging the points of all + measures. Alternatively, already segmented array of shape + ``[num_measures, max_measure_size, ndim]`` can be passed. + See also :func:`~ott.geometry.segment.segment_point_cloud`. + b: Array of shape ``[num_total_points,]`` containing the weights of all + the points within the measures that define the barycenter problem. + Same as ``y``, pre-segmented array of weights of shape + ``[num_measures, max_measure_size]`` can be passed. + If ``y`` is already pre-segmented, this array must be always specified. + weights: Array of shape ``[num_measures,]`` containing the weights of the + measures. + cost_fn: Cost function used. If `None`, + use the :class:`~ott.geometry.costs.SqEuclidean` cost. + epsilon: Epsilon regularization used to solve reg-OT problems. + kwargs: Keyword arguments :func:`~ott.geometry.segment.segment_point_cloud`. + Only used when ``y`` is not already segmented. When passing + ``segment_ids``, 2 arguments must be specified for jitting to work: + + - ``num_segments`` - the total number of measures. + - ``max_measure_size`` - maximum of support sizes of these measures. + """ + + def __init__( + self, + y: jnp.ndarray, + b: Optional[jnp.ndarray] = None, + weights: Optional[jnp.ndarray] = None, + cost_fn: Optional[costs.CostFn] = None, + epsilon: Optional[float] = None, + **kwargs: Any, + ): + self._y = y + if y.ndim == 3 and b is None: + raise ValueError("Specify weights if `y` is already segmented.") + self._b = b + self._weights = weights + self.cost_fn = costs.SqEuclidean() if cost_fn is None else cost_fn + self.epsilon = epsilon + self._kwargs = kwargs + + if self._is_segmented: + # (num_measures, max_measure_size, ndim) + # (num_measures, max_measure_size) + assert self._y.shape[:2] == self._b.shape + else: + # (num_total_points, ndim) + # (num_total_points,) + assert self._b is None or self._y.shape[0] == self._b.shape[0] + + @property + def segmented_y_b(self) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Tuple of arrays containing the segmented measures and weights. + + - Segmented measures of shape ``[num_measures, max_measure_size, ndim]``. + - Segmented weights of shape ``[num_measures, max_measure_size]``. + """ + if self._is_segmented: + y, b = self._y, self._b + else: + y, b = segment.segment_point_cloud( + x=self._y, + a=self._b, + padding_vector=self.cost_fn._padder(self.ndim), + **self._kwargs + ) + return y, b + + @property + def flattened_y(self) -> jnp.ndarray: + """Array of shape ``[num_measures * (N_1 + N_2 + ...), ndim]``.""" + if self._is_segmented: + return self._y.reshape((-1, self._y.shape[-1])) + return self._y + + @property + def flattened_b(self) -> Optional[jnp.ndarray]: + """Array of shape ``[num_measures * (N_1 + N_2 + ...),]``.""" + return None if self._b is None else self._b.ravel() + + @property + def num_measures(self) -> int: + """Number of measures.""" + return self.segmented_y_b[0].shape[0] + + @property + def max_measure_size(self) -> int: + """Maximum number of points across all measures.""" + return self.segmented_y_b[0].shape[1] + + @property + def ndim(self) -> int: + """Number of dimensions of ``y``.""" + return self._y.shape[-1] + + @property + def weights(self) -> jnp.ndarray: + """Barycenter weights of shape ``[num_measures,]`` that sum to 1.""" + if self._weights is None: + return jnp.ones((self.num_measures,)) / self.num_measures + # Check that the number of measures coincides with the weights' size. + assert self._weights.shape[0] == self.num_measures + # By default, we assume that weights sum to 1, and enforce this if needed. + return self._weights / jnp.sum(self._weights) + + @property + def _is_segmented(self) -> bool: + return self._y.ndim == 3 + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return ([self._y, self._b, self._weights], { + "cost_fn": self.cost_fn, + "epsilon": self.epsilon, + **self._kwargs, + }) + + @classmethod + def tree_unflatten( # noqa: D102 + cls, aux_data: Dict[str, Any], children: Sequence[Any] + ) -> "FreeBarycenterProblem": + y, b, weights = children + return cls(y=y, b=b, weights=weights, **aux_data) + + +@jax.tree_util.register_pytree_node_class +class FixedBarycenterProblem: + """Fixed-support Wasserstein barycenter problem. + + Args: + geom: Geometry object. + a: batch of histograms of shape ``[batch, num_a]`` where ``num_a`` matches + the first value of the :attr:`~ott.geometry.Geometry.shape` attribute of + ``geom``. + weights: ``[batch,]`` positive weights summing to :math:`1`. Uniform by + default. + """ + + def __init__( + self, + geom: geometry.Geometry, + a: jnp.ndarray, + weights: Optional[jnp.ndarray] = None, + ): + self.geom = geom + self.a = a + self._weights = weights + + @property + def num_measures(self) -> int: + """Number of measures.""" + return self.a.shape[0] + + @property + def weights(self) -> jnp.ndarray: + """Barycenter weights of shape ``[num_measures,]`` that sum to :math`1`.""" + if self._weights is None: + return jnp.ones((self.num_measures,)) / self.num_measures + + # check that the number of measures coincides with the weights' size + assert self._weights.shape[0] == self.num_measures + # by default, we assume that weights sum to 1, and enforce this if needed + return self._weights / jnp.sum(self._weights) + + def tree_flatten(self): # noqa: D102 + return [self.geom, self.a, self._weights], None + + @classmethod + def tree_unflatten( # noqa: D102 + cls, aux_data: Dict[str, Any], children: Sequence[Any] + ) -> "FixedBarycenterProblem": + del aux_data + geom, a, weights = children + return cls(geom=geom, a=a, weights=weights) diff --git a/ott/build/lib/ott/problems/linear/linear_problem.py b/ott/build/lib/ott/problems/linear/linear_problem.py new file mode 100644 index 0000000..60d8cc8 --- /dev/null +++ b/ott/build/lib/ott/problems/linear/linear_problem.py @@ -0,0 +1,132 @@ +# 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. +from typing import Any, Callable, Dict, Optional, Sequence, Tuple + +import jax +import jax.numpy as jnp + +from ott.geometry import geometry + +__all__ = ["LinearProblem"] + +# TODO(michalk8): move to typing.py when refactoring the types +MarginalFunc = Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray] +TransportAppFunc = Callable[[jnp.ndarray, jnp.ndarray, jnp.ndarray, int], + jnp.ndarray] + + +@jax.tree_util.register_pytree_node_class +class LinearProblem: + r"""Linear OT problem. + + This class describes the main ingredients appearing in a linear OT problem. + Namely, a ``geom`` object (including cost structure/points) describing point + clouds or the support of measures, followed by probability masses ``a`` and + ``b``. Unbalancedness of the problem is also kept track of, through two + coefficients ``tau_a`` and ``tau_b``, which are both kept between 0 and 1 + (1 corresponding to a balanced OT problem). + + Args: + geom: The ground geometry cost of the linear problem. + a: The first marginal. If ``None``, it will be uniform. + b: The second marginal. If ``None``, it will be uniform. + tau_a: If :math:`<1`, defines how much unbalanced the problem is + on the first marginal. + tau_b: If :math:`< 1`, defines how much unbalanced the problem is + on the second marginal. + """ + + def __init__( + self, + geom: geometry.Geometry, + a: Optional[jnp.ndarray] = None, + b: Optional[jnp.ndarray] = None, + tau_a: float = 1.0, + tau_b: float = 1.0, + labels_a: Optional[jnp.ndarray] = None, + labels_b: Optional[jnp.ndarray] = None, + ): + self.geom = geom + self._a = a + self._b = b + self.tau_a = tau_a + self.tau_b = tau_b + self.labels_a = labels_a + self.labels_b = labels_b + + @property + def a(self) -> jnp.ndarray: + """First marginal.""" + num_a = self.geom.shape[0] + return jnp.ones((num_a,)) / num_a if self._a is None else self._a + + @property + def b(self) -> jnp.ndarray: + """Second marginal.""" + num_b = self.geom.shape[1] + return jnp.ones((num_b,)) / num_b if self._b is None else self._b + + @property + def is_balanced(self) -> bool: + """Whether the problem is balanced.""" + return self.tau_a == 1.0 and self.tau_b == 1.0 + + @property + def is_uniform(self) -> bool: + """True if no weights ``a,b`` were passed, and have defaulted to uniform.""" + return self._a is None and self._b is None + + @property + def is_equal_size(self) -> bool: + """True if square shape, i.e. ``n == m``.""" + return self.geom.shape[0] == self.geom.shape[1] + + @property + def epsilon(self) -> float: + """Entropic regularization.""" + return self.geom.epsilon + + def get_transport_functions( + self, lse_mode: bool + ) -> Tuple[MarginalFunc, MarginalFunc, TransportAppFunc]: + """Instantiate useful functions for Sinkhorn depending on lse_mode.""" + geom = self.geom + if lse_mode: + marginal_a = lambda f, g: geom.marginal_from_potentials(f, g, 1) + marginal_b = lambda f, g: geom.marginal_from_potentials(f, g, 0) + app_transport = geom.apply_transport_from_potentials + else: + marginal_a = lambda f, g: geom.marginal_from_scalings( + geom.scaling_from_potential(f), geom.scaling_from_potential(g), 1 + ) + marginal_b = lambda f, g: geom.marginal_from_scalings( + geom.scaling_from_potential(f), geom.scaling_from_potential(g), 0 + ) + app_transport = lambda f, g, z, axis: geom.apply_transport_from_scalings( + geom.scaling_from_potential(f), geom.scaling_from_potential(g), z, + axis + ) + return marginal_a, marginal_b, app_transport + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return ([self.geom, self._a, self._b], { + "tau_a": self.tau_a, + "tau_b": self.tau_b + }) + + @classmethod + def tree_unflatten( # noqa: D102 + cls, aux_data: Dict[str, Any], children: Sequence[Any] + ) -> "LinearProblem": + return cls(*children, **aux_data) diff --git a/ott/build/lib/ott/problems/linear/potentials.py b/ott/build/lib/ott/problems/linear/potentials.py new file mode 100644 index 0000000..6885152 --- /dev/null +++ b/ott/build/lib/ott/problems/linear/potentials.py @@ -0,0 +1,437 @@ +# 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. +from typing import ( + Any, + Callable, + Dict, + Literal, + Optional, + Sequence, + Tuple, +) + +import jax +import jax.numpy as jnp +import jax.scipy as jsp +import jax.tree_util as jtu +import numpy as np + +from ott.geometry import costs +from ott.problems.linear import linear_problem + +try: + import matplotlib as mpl + import matplotlib.pyplot as plt +except ImportError: + mpl = plt = None + +__all__ = ["DualPotentials", "EntropicPotentials"] +Potential_t = Callable[[jnp.ndarray], float] + + +@jtu.register_pytree_node_class +class DualPotentials: + r"""The Kantorovich dual potential functions :math:`f` and :math:`g`. + + :math:`f` and :math:`g` are a pair of functions, candidates for the dual + OT Kantorovich problem, supposedly optimal for a given pair of measures. + + Args: + f: The first dual potential function. + g: The second dual potential function. + cost_fn: The cost function used to solve the OT problem. + corr: Whether the duals solve the problem in distance form, or correlation + form (as used for instance for ICNNs, see, e.g., top right of p.3 in + :cite:`makkuva:20`) + """ + + def __init__( + self, + f: Potential_t, + g: Potential_t, + *, + cost_fn: costs.CostFn, + corr: bool = False + ): + self._f = f + self._g = g + assert ( + not corr or type(cost_fn) == costs.SqEuclidean + ), "Duals in `corr` form can only be used with a squared-Euclidean cost." + self.cost_fn = cost_fn + self._corr = corr + + def transport(self, vec: jnp.ndarray, forward: bool = True) -> jnp.ndarray: + r"""Transport ``vec`` according to Brenier formula :cite:`brenier:91`. + + Uses Theorem 1.17 from :cite:`santambrogio:15` to compute an OT map when + given the Legendre transform of the dual potentials. + + That OT map can be recovered as :math:`x- (\nabla h^*)\circ \nabla f(x)`, + where :math:`h^*` is the Legendre transform of :math:`h`. For instance, + in the case :math:`h(\cdot) = \|\cdot\|^2, \nabla h(\cdot) = 2 \cdot\,`, + one has :math:`h^*(\cdot) = \|.\|^2 / 4`, and therefore + :math:`\nabla h^*(\cdot) = 0.5 \cdot\,`. + + When the dual potentials are solved in correlation form (only in the Sq. + Euclidean distance case), the maps are :math:`\nabla g` for forward, + :math:`\nabla f` for backward. + + Args: + vec: Points to transport, array of shape ``[n, d]``. + forward: Whether to transport the points from source to the target + distribution or vice-versa. + + Returns: + The transported points. + """ + from ott.geometry import costs + + vec = jnp.atleast_2d(vec) + if self._corr and isinstance(self.cost_fn, costs.SqEuclidean): + return self._grad_f(vec) if forward else self._grad_g(vec) + if forward: + return vec - self._grad_h_inv(self._grad_f(vec)) + return vec - self._grad_h_inv(self._grad_g(vec)) + + def distance(self, src: jnp.ndarray, tgt: jnp.ndarray) -> float: + r"""Evaluate Wasserstein distance between samples using dual potentials. + + This uses direct estimation of potentials against measures when dual + functions are provided in usual form. This expression is valid for any + cost function. + + When potentials are given in correlation form, as specified by the flag + ``corr``, the dual potentials solve the dual problem corresponding to the + minimization of the primal OT problem where the ground cost is + :math:`-2\langle x,y\rangle`. To recover the (squared) 2-Wasserstein + distance, terms are re-arranged and contributions from squared norms are + taken into account. + + Args: + src: Samples from the source distribution, array of shape ``[n, d]``. + tgt: Samples from the target distribution, array of shape ``[m, d]``. + + Returns: + Wasserstein distance using specified cost function. + """ + src, tgt = jnp.atleast_2d(src), jnp.atleast_2d(tgt) + f = jax.vmap(self.f) + g = jax.vmap(self.g) + out = jnp.mean(f(src)) + jnp.mean(g(tgt)) + if self._corr: + out = -2.0 * out + jnp.mean(jnp.sum(src ** 2, axis=-1)) + out += jnp.mean(jnp.sum(tgt ** 2, axis=-1)) + return out + + @property + def f(self) -> Potential_t: + """The first dual potential function.""" + return self._f + + @property + def g(self) -> Potential_t: + """The second dual potential function.""" + return self._g + + @property + def _grad_f(self) -> Callable[[jnp.ndarray], jnp.ndarray]: + """Vectorized gradient of the potential function :attr:`f`.""" + return jax.vmap(jax.grad(self.f, argnums=0)) + + @property + def _grad_g(self) -> Callable[[jnp.ndarray], jnp.ndarray]: + """Vectorized gradient of the potential function :attr:`g`.""" + return jax.vmap(jax.grad(self.g, argnums=0)) + + @property + def _grad_h_inv(self) -> Callable[[jnp.ndarray], jnp.ndarray]: + from ott.geometry import costs + + assert isinstance(self.cost_fn, costs.TICost), ( + "Cost must be a `TICost` and " + "provide access to Legendre transform of `h`." + ) + return jax.vmap(jax.grad(self.cost_fn.h_legendre)) + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return [], { + "f": self._f, + "g": self._g, + "cost_fn": self.cost_fn, + "corr": self._corr + } + + @classmethod + def tree_unflatten( # noqa: D102 + cls, aux_data: Dict[str, Any], children: Sequence[Any] + ) -> "DualPotentials": + return cls(*children, **aux_data) + + def plot_ot_map( + self, + source: jnp.ndarray, + target: jnp.ndarray, + samples: Optional[jnp.ndarray] = None, + forward: bool = True, + ax: Optional["plt.Axes"] = None, + scatter_kwargs: Optional[Dict[str, Any]] = None, + legend_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple["plt.Figure", "plt.Axes"]: + """Plot data and learned optimal transport map. + + Args: + source: samples from the source measure + target: samples from the target measure + samples: extra samples to transport, either ``source`` (if ``forward``) or + ``target`` (if not ``forward``) by default. + forward: use the forward map from the potentials if ``True``, + otherwise use the inverse map. + ax: axis to add the plot to + scatter_kwargs: additional kwargs passed into + :meth:`~matplotlib.axes.Axes.scatter` + legend_kwargs: additional kwargs passed into + :meth:`~matplotlib.axes.Axes.legend` + + Returns: + Figure and axes. + """ + if mpl is None: + raise RuntimeError("Please install `matplotlib` first.") + + if scatter_kwargs is None: + scatter_kwargs = {"alpha": 0.5} + if legend_kwargs is None: + legend_kwargs = { + "ncol": 3, + "loc": "upper center", + "bbox_to_anchor": (0.5, -0.05), + "edgecolor": "k" + } + + if ax is None: + fig = plt.figure(facecolor="white") + ax = fig.add_subplot(111) + else: + fig = ax.get_figure() + + # plot the source and target samples + if forward: + label_transport = r"$\nabla f(source)$" + source_color, target_color = "#1A254B", "#A7BED3" + else: + label_transport = r"$\nabla g(target)$" + source_color, target_color = "#A7BED3", "#1A254B" + + ax.scatter( + source[:, 0], + source[:, 1], + color=source_color, + label="source", + **scatter_kwargs, + ) + ax.scatter( + target[:, 0], + target[:, 1], + color=target_color, + label="target", + **scatter_kwargs, + ) + + # plot the transported samples + samples = (source if forward else target) if samples is None else samples + transported_samples = self.transport(samples, forward=forward) + ax.scatter( + transported_samples[:, 0], + transported_samples[:, 1], + color="#F2545B", + label=label_transport, + **scatter_kwargs, + ) + + for i in range(samples.shape[0]): + ax.arrow( + samples[i, 0], + samples[i, 1], + transported_samples[i, 0] - samples[i, 0], + transported_samples[i, 1] - samples[i, 1], + color=[0.5, 0.5, 1], + alpha=0.3, + ) + + ax.legend(**legend_kwargs) + return fig, ax + + def plot_potential( + self, + forward: bool = True, + quantile: float = 0.05, + kantorovich: bool = True, + ax: Optional["mpl.axes.Axes"] = None, + x_bounds: Tuple[float, float] = (-6, 6), + y_bounds: Tuple[float, float] = (-6, 6), + num_grid: int = 50, + contourf_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple["mpl.figure.Figure", "mpl.axes.Axes"]: + r"""Plot the potential. + + Args: + forward: use the forward map from the potentials + if ``True``, otherwise use the inverse map + quantile: quantile to filter the potentials with + kantorovich: whether to plot the Kantorovich potential + ax: axis to add the plot to + x_bounds: x-axis bounds of the plot + :math:`(x_{\text{min}}, x_{\text{max}})` + y_bounds: y-axis bounds of the plot + :math:`(y_{\text{min}}, y_{\text{max}})` + num_grid: number of points to discretize the domain into a grid + along each dimension + contourf_kwargs: additional kwargs passed into + :meth:`~matplotlib.axes.Axes.contourf` + + Returns: + Figure and axes. + """ + if contourf_kwargs is None: + contourf_kwargs = {} + + ax_specified = ax is not None + if not ax_specified: + fig, ax = plt.subplots(figsize=(6, 6), facecolor="white") + else: + fig = ax.get_figure() + + x1 = jnp.linspace(*x_bounds, num=num_grid) + x2 = jnp.linspace(*y_bounds, num=num_grid) + X1, X2 = jnp.meshgrid(x1, x2) + X12flat = jnp.hstack((X1.reshape(-1, 1), X2.reshape(-1, 1))) + Zflat = jax.vmap(self.f if forward else self.g)(X12flat) + if kantorovich: + Zflat = 0.5 * (jnp.linalg.norm(X12flat, axis=-1) ** 2) - Zflat + Zflat = np.asarray(Zflat) + vmin, vmax = np.quantile(Zflat, [quantile, 1.0 - quantile]) + Zflat = Zflat.clip(vmin, vmax) + Z = Zflat.reshape(X1.shape) + + CS = ax.contourf(X1, X2, Z, cmap="Blues", **contourf_kwargs) + ax.set_xlim(*x_bounds) + ax.set_ylim(*y_bounds) + fig.colorbar(CS, ax=ax) + if not ax_specified: + fig.tight_layout() + ax.set_title(r"$f$" if forward else r"$g$") + return fig, ax + + +@jtu.register_pytree_node_class +class EntropicPotentials(DualPotentials): + """Dual potential functions from finite samples :cite:`pooladian:21`. + + Args: + f_xy: The first dual potential vector of shape ``[n,]``. + g_xy: The second dual potential vector of shape ``[m,]``. + prob: Linear problem with :class:`~ott.geometry.pointcloud.PointCloud` + geometry that was used to compute the dual potentials using, e.g., + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. + f_xx: The first dual potential vector of shape ``[n,]`` used for debiasing + :cite:`pooladian:22`. + g_yy: The second dual potential vector of shape ``[m,]`` used for debiasing. + """ + + def __init__( + self, + f_xy: jnp.ndarray, + g_xy: jnp.ndarray, + prob: linear_problem.LinearProblem, + f_xx: Optional[jnp.ndarray] = None, + g_yy: Optional[jnp.ndarray] = None, + ): + # we pass directly the arrays and override the properties + # since only the properties need to be callable + super().__init__(f_xy, g_xy, cost_fn=prob.geom.cost_fn, corr=False) + self._prob = prob + self._f_xx = f_xx + self._g_yy = g_yy + + @property + def f(self) -> Potential_t: # noqa: D102 + return self._potential_fn(kind="f") + + @property + def g(self) -> Potential_t: # noqa: D102 + return self._potential_fn(kind="g") + + def _potential_fn(self, *, kind: Literal["f", "g"]) -> Potential_t: + from ott.geometry import pointcloud + + def callback( + x: jnp.ndarray, + *, + potential: jnp.ndarray, + y: jnp.ndarray, + weights: jnp.ndarray, + epsilon: float, + ) -> float: + x = jnp.atleast_2d(x) + assert x.shape[-1] == y.shape[-1], (x.shape, y.shape) + geom = pointcloud.PointCloud(x, y, cost_fn=self.cost_fn) + cost = geom.cost_matrix + z = (potential - cost) / epsilon + lse = -epsilon * jsp.special.logsumexp(z, b=weights, axis=-1) + return jnp.squeeze(lse) + + assert isinstance( + self._prob.geom, pointcloud.PointCloud + ), f"Expected point cloud geometry, found `{type(self._prob.geom)}`." + x, y = self._prob.geom.x, self._prob.geom.y + a, b = self._prob.a, self._prob.b + + if kind == "f": + # When seeking to evaluate 1st potential function, + # the 2nd set of potential values and support should be used, + # see proof of Prop. 2 in https://arxiv.org/pdf/2109.12004.pdf + potential, arr, weights = self._g, y, b + else: + potential, arr, weights = self._f, x, a + + potential_xy = jax.tree_util.Partial( + callback, + potential=potential, + y=arr, + weights=weights, + epsilon=self.epsilon, + ) + if not self.is_debiased: + return potential_xy + + ep = EntropicPotentials(self._f_xx, self._g_yy, prob=self._prob) + # switch the order because for `kind='f'` we require `f/x/a` in `other` + # which is accessed when `kind='g'` + potential_other = ep._potential_fn(kind="g" if kind == "f" else "f") + + return lambda x: (potential_xy(x) - potential_other(x)) + + @property + def is_debiased(self) -> bool: + """Whether the entropic map is debiased.""" + return self._f_xx is not None and self._g_yy is not None + + @property + def epsilon(self) -> float: + """Entropy regularizer.""" + return self._prob.geom.epsilon + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return [self._f, self._g, self._prob, self._f_xx, self._g_yy], {} diff --git a/ott/build/lib/ott/problems/quadratic/__init__.py b/ott/build/lib/ott/problems/quadratic/__init__.py new file mode 100644 index 0000000..d2bb647 --- /dev/null +++ b/ott/build/lib/ott/problems/quadratic/__init__.py @@ -0,0 +1,14 @@ +# 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. +from . import gw_barycenter, quadratic_costs, quadratic_problem diff --git a/ott/build/lib/ott/problems/quadratic/gw_barycenter.py b/ott/build/lib/ott/problems/quadratic/gw_barycenter.py new file mode 100644 index 0000000..7acf962 --- /dev/null +++ b/ott/build/lib/ott/problems/quadratic/gw_barycenter.py @@ -0,0 +1,325 @@ +# 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 functools +from typing import Any, Dict, Literal, Optional, Sequence, Tuple, Union + +import jax +import jax.numpy as jnp + +from ott.geometry import costs, geometry, pointcloud, segment +from ott.math import utils as mu +from ott.problems.linear import barycenter_problem +from ott.problems.quadratic import quadratic_costs, quadratic_problem + +__all__ = ["GWBarycenterProblem"] + + +# TODO(michalk8): better abstraction (common superclass for Wasserstein bary) +@jax.tree_util.register_pytree_node_class +class GWBarycenterProblem(barycenter_problem.FreeBarycenterProblem): + """(Fused) Gromov-Wasserstein barycenter problem :cite:`peyre:16,vayer:19`. + + Args: + y: Array of shape ``[num_total_points, ndim]`` merging the points of all + measures. Alternatively, already segmented array of shape + ``[num_measures, max_measure_size, ndim]`` can be passed. + See also :func:`~ott.geometry.segment.segment_point_cloud`. + b: Array of shape ``[num_total_points,]`` containing the weights of all + the points within the measures that define the barycenter problem. + Same as ``y``, pre-segmented array of weights of shape + ``[num_measures, max_measure_size]`` can be passed. + If ``y`` is already pre-segmented, this array must be passed. + weights: Array of shape ``[num_measures,]`` containing the weights of the + barycenter problem. + costs: Alternative to ``y``, an array of shape + ``[num_measures, max_measure_size, max_measure_size]`` that defines padded + cost matrices for each measure. Used in the quadratic term. + Only one of ``y`` and ``cost`` can be specified. + y_fused: Array of shape ``[num_total_points, ndim_fused]`` containing + the data of the points of all measures used to define the linear term + in the fused case. Same as ``y``, it can be specified as a pre-segmented + array of shape ``[num_measures, max_measure_size, ndim_fused]``. + gw_loss: Gromov-Wasserstein loss. + fused_penalty: Multiplier of the linear term. Only used when + ``y_fused != None``. + scale_cost: Scaling of cost matrices passed to geometries. + kwargs: Keyword arguments for + :class:`~ott.problems.linear.barycenter_problem.BarycenterProblem`. + """ + + def __init__( + self, + y: Optional[jnp.ndarray] = None, + b: Optional[jnp.ndarray] = None, + weights: Optional[jnp.ndarray] = None, + costs: Optional[jnp.ndarray] = None, + y_fused: Optional[jnp.ndarray] = None, + fused_penalty: float = 1.0, + gw_loss: Literal["sqeucl", "kl"] = "sqeucl", + scale_cost: Union[int, float, Literal["mean", "max_cost"]] = 1.0, + **kwargs: Any, + ): + assert y is None or costs is None, "Cannot specify both `y` and `costs`." + y = y if costs is None else costs + + super().__init__(y=y, b=b, weights=weights, **kwargs) + + self._y_fused = y_fused + self.fused_penalty = fused_penalty + self._loss_name = gw_loss + self.scale_cost = scale_cost + self._y_as_costs = costs is not None + + if self._y_as_costs: + # (num_measures, max_measure_size, max_measure_size) + _, n, m = self._y.shape + assert n == m, "Cost matrices must be square." + if self.is_fused: + seg_y = self._is_segmented + seg_fused = self._y_fused.ndim == 3 + if seg_y and seg_fused: + # (num_measures, max_measure_size, ndim_fused) + # (num_measures, max_measure_size, ndim) + assert self._y_fused.shape[:2] == self._y.shape[:2] + if not seg_y and not seg_fused: + # (num_total_points, ndim_fused), (num_total_points, ndim) + assert self._y_fused.shape[0] == self._y.shape[0] + # TODO(michalk8): in the future, consider checking the other 2 cases + # using `segmented_y` and `segmented_y_fused`? + + def update_barycenter( + self, transports: jnp.ndarray, a: jnp.ndarray + ) -> jnp.ndarray: + """Update the barycenter cost matrix. + + Uses the eq. 14 and 15 of :cite:`peyre:16`. + + Args: + transports: Transport maps of shape + ``[num_measures, bar_size, max_measure_size]``. + a: Barycenter weights of shape ``[bar_size,]``. + + Returns: + Update cost matrix of shape ``[bar_size, bar_size]``. + """ + + @functools.partial(jax.vmap, in_axes=[0, 0, 0, None]) + def project( + y: jnp.ndarray, + b: jnp.ndarray, + transport: jnp.ndarray, + fn: Optional[quadratic_costs.Loss], + ) -> jnp.ndarray: + geom = self._create_y_geometry(y, mask=b > 0.0) + fn, lin = (None, True) if fn is None else (fn.func, fn.is_linear) + + tmp = geom.apply_cost( + transport.T, + axis=0, + fn=fn, + is_linear=lin, + ) + return transport @ tmp + + fn = None if self._loss_name == "sqeucl" else self.gw_loss.h2 + y, b = self.segmented_y_b + weights = self.weights[:, None, None] + + barycenter = jnp.sum(weights * project(y, b, transports, fn), axis=0) + inv_a = jnp.where(a > 0, 1.0 / a, 1.0) + barycenter = (barycenter * inv_a[None, :]) * inv_a[:, None] + + # TODO(michalk8): in future, use `isinstanceof(self.gw_loss, ...)` + # once refactoring has been done + if self._loss_name == "kl": + return jnp.exp(barycenter) + return barycenter + + def update_features(self, transports: jnp.ndarray, + a: jnp.ndarray) -> Optional[jnp.ndarray]: + """Update the barycenter features in the fused case :cite:`vayer:19`. + + Uses :cite:`cuturi:14` eq. 8, and is implemented only + for the :class:`~ott.geometry.costs.SqEuclidean` cost. + + Args: + transports: Transport maps of shape + ``[num_measures, bar_size, max_measure_size]``. + a: Barycenter weights of shape ``[bar_size,]``. + + Returns: + Updated features of shape ``[bar_size, ndim_fused]``. + """ + if not self.is_fused: + raise RuntimeError( + "Updating features is available only in the fused case." + ) + + y_fused = self.segmented_y_fused + weights = self.weights[:, None, None] + inv_a = jnp.where(a > 0, 1.0 / a, 1.0) + transports = transports * inv_a[None, :, None] + + if self._loss_name == "sqeucl": + cost_fn = costs.SqEuclidean() + return jnp.sum( + weights * mu.barycentric_projection(transports, y_fused, cost_fn), + axis=0 + ) + raise NotImplementedError(self._loss_name) + + def _create_bary_geometry( + self, + cost_matrix: jnp.ndarray, + mask: Optional[jnp.ndarray] = None + ) -> geometry.Geometry: + return geometry.Geometry( + cost_matrix=cost_matrix, + src_mask=mask, + tgt_mask=mask, + epsilon=self.epsilon, + scale_cost=self.scale_cost + ) + + def _create_y_geometry( + self, + y: jnp.ndarray, + mask: Optional[jnp.ndarray] = None + ) -> geometry.Geometry: + if self._y_as_costs: + assert y.shape[0] == y.shape[1], y.shape + return geometry.Geometry( + y, + epsilon=self.epsilon, + scale_cost=self.scale_cost, + src_mask=mask, + tgt_mask=mask + ) + return pointcloud.PointCloud( + y, + epsilon=self.epsilon, + scale_cost=self.scale_cost, + cost_fn=self.cost_fn, + src_mask=mask, + tgt_mask=mask + ) + + def _create_fused_geometry( + self, + x: jnp.ndarray, + y: jnp.ndarray, + src_mask: Optional[jnp.ndarray] = None, + tgt_mask: Optional[jnp.ndarray] = None + ) -> pointcloud.PointCloud: + return pointcloud.PointCloud( + x, + y, + cost_fn=self.cost_fn, + epsilon=self.epsilon, + scale_cost=self.scale_cost, + src_mask=src_mask, + tgt_mask=tgt_mask + ) + + def _create_problem( + self, + state: "GWBarycenterState", # noqa: F821 + y: jnp.ndarray, + b: jnp.ndarray, + f: Optional[jnp.ndarray] = None + ) -> quadratic_problem.QuadraticProblem: + # TODO(michalk8): in future, mask in the problem for convenience? + bary_mask = state.a > 0.0 + y_mask = b > 0.0 + + geom_xx = self._create_bary_geometry(state.cost, mask=bary_mask) + geom_yy = self._create_y_geometry(y, mask=y_mask) + if self.is_fused: + assert f is not None + assert state.x.shape[1] == f.shape[1] + geom_xy = self._create_fused_geometry( + state.x, f, src_mask=bary_mask, tgt_mask=y_mask + ) + else: + geom_xy = None + + return quadratic_problem.QuadraticProblem( + geom_xx=geom_xx, + geom_yy=geom_yy, + geom_xy=geom_xy, + a=state.a, + b=b, + fused_penalty=self.fused_penalty, + ) + + @property + def is_fused(self) -> bool: + """Whether the problem is fused.""" + return self._y_fused is not None + + @property + def segmented_y_fused(self) -> Optional[jnp.ndarray]: + """Feature array of shape used in the fused case.""" + if not self.is_fused or self._y_fused.ndim == 3: + return self._y_fused + y_fused, _ = segment.segment_point_cloud( + x=self._y_fused, + padding_vector=self.cost_fn._padder(self.ndim_fused), + **self._kwargs + ) + return y_fused + + @property + def ndim(self) -> Optional[int]: # noqa: D102 + return None if self._y_as_costs else self._y.shape[-1] + + @property + def ndim_fused(self) -> Optional[int]: + """Number of dimensions of the fused term.""" + return self._y_fused.shape[-1] if self.is_fused else None + + @property + def gw_loss(self) -> quadratic_costs.GWLoss: + """Gromov-Wasserstein loss.""" + # TODO(michalk8): custom losses would require inverting some fns; + # `https://jax.readthedocs.io/en/latest/notebooks/ some fns; + # Writing_custom_interpreters_in_Jax.html#your-first-interpreter-invert` + # might be useful + if self._loss_name == "sqeucl": + return quadratic_costs.make_square_loss() + if self._loss_name == "kl": + return quadratic_costs.make_kl_loss() + raise NotImplementedError( + f"Loss `{self._loss_name}` is not yet implemented." + ) + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + (y, b, weights), aux = super().tree_flatten() + if self._y_as_costs: + children = [None, b, weights, y] + else: + children = [y, b, weights, None] + aux["fused_penalty"] = self.fused_penalty + aux["gw_loss"] = self._loss_name + aux["scale_cost"] = self.scale_cost + return children + [self._y_fused], aux + + @classmethod + def tree_unflatten( # noqa: D102 + cls, aux_data: Dict[str, Any], children: Sequence[Any] + ) -> "GWBarycenterProblem": + y, b, weights, costs, y_fused = children + return cls( + y=y, b=b, weights=weights, costs=costs, y_fused=y_fused, **aux_data + ) diff --git a/ott/build/lib/ott/problems/quadratic/quadratic_costs.py b/ott/build/lib/ott/problems/quadratic/quadratic_costs.py new file mode 100644 index 0000000..70f2bf5 --- /dev/null +++ b/ott/build/lib/ott/problems/quadratic/quadratic_costs.py @@ -0,0 +1,78 @@ +# 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. +from typing import Callable, NamedTuple + +import jax.numpy as jnp +import jax.scipy as jsp + +__all__ = ["make_square_loss", "make_kl_loss"] + + +class Loss(NamedTuple): # noqa: D101 + func: Callable[[jnp.ndarray], jnp.ndarray] + is_linear: bool + + +class GWLoss(NamedTuple): + r"""Efficient decomposition of the Gromov-Wasserstein loss function. + + The loss function :math:`L` is assumed to match the form given in eq. 5. of + :cite:`peyre:16`: + + .. math:: + L(x, y) = f_1(x) + f_2(y) - h_1(x) h_2(y) + + Args: + f1: First linear term. + f2: Second linear term. + h1: First quadratic term. + h2: Second quadratic term. + """ + f1: Loss + f2: Loss + h1: Loss + h2: Loss + + +def make_square_loss() -> GWLoss: + """Squared Euclidean loss for Gromov-Wasserstein. + + See Prop. 1 and Remark 1 of :cite:`peyre:16` for more information. + + Returns: + The squared Euclidean loss. + """ + f1 = Loss(lambda x: x ** 2, is_linear=False) + f2 = Loss(lambda y: y ** 2, is_linear=False) + h1 = Loss(lambda x: x, is_linear=True) + h2 = Loss(lambda y: 2.0 * y, is_linear=True) + return GWLoss(f1, f2, h1, h2) + + +def make_kl_loss(clipping_value: float = 1e-8) -> GWLoss: + r"""Kullback-Leibler loss for Gromov-Wasserstein. + + See Prop. 1 and Remark 1 of :cite:`peyre:16` for more information. + + Args: + clipping_value: Value used to avoid :math:`\log(0)`. + + Returns: + The KL loss. + """ + f1 = Loss(lambda x: -jsp.special.entr(x) - x, is_linear=False) + f2 = Loss(lambda y: y, is_linear=True) + h1 = Loss(lambda x: x, is_linear=True) + h2 = Loss(lambda y: jnp.log(jnp.clip(y, clipping_value)), is_linear=False) + return GWLoss(f1, f2, h1, h2) diff --git a/ott/build/lib/ott/problems/quadratic/quadratic_problem.py b/ott/build/lib/ott/problems/quadratic/quadratic_problem.py new file mode 100644 index 0000000..170fb3a --- /dev/null +++ b/ott/build/lib/ott/problems/quadratic/quadratic_problem.py @@ -0,0 +1,588 @@ +# 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. +from typing import TYPE_CHECKING, Literal, Optional, Tuple, Union + +import jax +import jax.numpy as jnp +import jax.scipy as jsp + +from ott import utils +from ott.geometry import epsilon_scheduler, geometry, low_rank, pointcloud +from ott.problems.linear import linear_problem +from ott.problems.quadratic import quadratic_costs +from ott.types import Transport + +if TYPE_CHECKING: + from ott.solvers.linear import sinkhorn_lr + +__all__ = ["QuadraticProblem"] + + +@jax.tree_util.register_pytree_node_class +class QuadraticProblem: + r"""Quadratic OT problem. + + The quadratic loss of a single OT matrix is assumed to + have the form given in :cite:`peyre:16`, eq. 4. + + The two geometries below parameterize matrices :math:`C` and :math:`\bar{C}` + in that equation. The function :math:`L` (of two real values) in that equation + is assumed to match the form given in eq. 5., with our notations: + + .. math:: + L(x, y) = f_1(x) + f_2(y) - h_1(x) h_2(y) + + Args: + geom_xx: Ground geometry of the first space. + geom_yy: Ground geometry of the second space. + geom_xy: Geometry defining the linear penalty term for + fused Gromov-Wasserstein :cite:`vayer:19`. If :obj:`None`, the problem + reduces to a plain Gromov-Wasserstein problem :cite:`peyre:16`. + fused_penalty: Multiplier of the linear term in fused Gromov-Wasserstein, + i.e. ``problem = purely quadratic + fused_penalty * linear problem``. + scale_cost: option to rescale the cost matrices: + + - if :obj:`True`, use the default for each geometry. + - if :obj:`False`, keep the original scaling in geometries. + - if :class:`str`, use a specific method available in + :class:`~ott.geometry.geometry.Geometry` or + :class:`~ott.geometry.pointcloud.PointCloud`. + - if :obj:`None`, do not scale the cost matrices. + + a: The first marginal. If :obj:`None`, it will be uniform. + b: The second marginal. If :obj:`None`, it will be uniform. + loss: Gromov-Wasserstein loss function, see + :class:`~ott.problems.quadratic.quadratic_costs.GWLoss` for more + information. + tau_a: If :math:`< 1.0`, defines how much unbalanced the problem is on + the first marginal. + tau_b: If :math:`< 1.0`, defines how much unbalanced the problem is on + the second marginal. + gw_unbalanced_correction: Whether the unbalanced version of + :cite:`sejourne:21` is used. Otherwise, ``tau_a`` and ``tau_b`` + only affect the inner Sinkhorn loop. + ranks: Ranks of the cost matrices, see + :meth:`~ott.geometry.geometry.Geometry.to_LRCGeometry`. Used when + geometries are *not* :class:`~ott.geometry.pointcloud.PointCloud` with + `'sqeucl'` cost function. If `-1`, the geometries will not be converted + to low-rank. If :class:`tuple`, it specifies the ranks of ``geom_xx``, + ``geom_yy`` and ``geom_xy``, respectively. If :class:`int`, rank is shared + across all geometries. + tolerances: Tolerances used when converting geometries to low-rank. Used + when geometries are not :class:`~ott.geometry.pointcloud.PointCloud` with + `'sqeucl'` cost. If :class:`float`, it is shared across all geometries. + """ + + def __init__( + self, + geom_xx: geometry.Geometry, + geom_yy: geometry.Geometry, + geom_xy: Optional[geometry.Geometry] = None, + fused_penalty: float = 1.0, + scale_cost: Optional[Union[bool, float, str]] = False, + a: Optional[jnp.ndarray] = None, + b: Optional[jnp.ndarray] = None, + loss: Union[Literal["sqeucl", "kl"], quadratic_costs.GWLoss] = "sqeucl", + tau_a: float = 1.0, + tau_b: float = 1.0, + gw_unbalanced_correction: bool = True, + ranks: Union[int, Tuple[int, ...]] = -1, + tolerances: Union[float, Tuple[float, ...]] = 1e-2, + labels_a: Optional[jnp.ndarray] = None, + labels_b: Optional[jnp.ndarray] = None, + n_labels: Optional[jnp.ndarray] = None, + block_diag_mat: Optional[jnp.ndarray] = None, + ): + self._geom_xx = geom_xx.set_scale_cost(scale_cost) + self._geom_yy = geom_yy.set_scale_cost(scale_cost) + self._geom_xy = ( + None if geom_xy is None else geom_xy.set_scale_cost(scale_cost) + ) + self.fused_penalty = fused_penalty + self.scale_cost = scale_cost + self._a = a + self._b = b + self.tau_a = tau_a + self.tau_b = tau_b + self.gw_unbalanced_correction = gw_unbalanced_correction + self.ranks = ranks + self.tolerances = tolerances + + self._loss_name = loss + if self._loss_name == "sqeucl": + self.loss = quadratic_costs.make_square_loss() + elif loss == "kl": + self.loss = quadratic_costs.make_kl_loss() + else: + self.loss = loss + self.labels_a = labels_a + self.labels_b = labels_b + self.n_labels = n_labels + self.block_diag_mat = block_diag_mat + + def marginal_dependent_cost_labeled( + self, + marginal_1: jnp.ndarray, + marginal_2: jnp.ndarray, + labels_1: jnp.ndarray, + labels_2: jnp.ndarray, + ) -> low_rank.LRCGeometry: + r"""Initialize cost term that depends on the marginals of the transport. + + Uses the first term in eq. 6, p. 1 of :cite:`peyre:16`. + + Let :math:`p` be the `[n,]` marginal of the transport matrix for samples + from :attr:`geom_xx` and :math:`q` the `[m,]` marginal of the + transport matrix for samples from :attr:`geom_yy`. + + When ``cost_xx`` (resp. ``cost_yy``) is the cost matrix of :attr:`geom_xx` + (resp. :attr:`geom_yy`), the cost term that depends on these marginals can + be written as: + + .. math:: + + \text{marginal_dep_term} = \text{lin1}(\text{cost_xx}) p \mathbb{1}_{m}^T + + \mathbb{1}_{n}(\text{lin2}(\text{cost_yy}) q)^T + + This helper function instantiates these two low-rank matrices and groups + them into a single low-rank cost geometry object. + + Args: + marginal_1: [n,], first marginal of transport matrix. + marginal_2: [m,], second marginal of transport matrix. + + Returns: + Low-rank geometry of rank 2, storing normalization constants. + """ + geom_xx, geom_yy = self.geom_xx, self.geom_yy + tmp1 = None + tmp2 = None + if self._loss_name == "sqeucl": # quadratic apply, efficient for LR + for l in jnp.unique(labels_1, size = self.n_labels): + tmp1_l = geom_xx.apply_square_cost(marginal_1, axis=1, sub_idx = labels_1 == l) + tmp2_l = geom_yy.apply_square_cost(marginal_2, axis=1, sub_idx = labels_2 == l) + if tmp1 is None: tmp1 = tmp1_l + else: tmp1 += tmp1_l + if tmp2 is None: tmp2 = tmp2_l + else: tmp2 += tmp2_l + else: + f1, f2 = self.linear_loss + for l in jnp.unique(labels_1, size = self.n_labels): + tmp1_l = apply_cost(geom_xx, marginal_1, axis=1, fn=f1, sub_idx = labels_1 == l) + tmp2_l = apply_cost(geom_yy, marginal_2, axis=1, fn=f2, sub_idx = labels_2 == l) + if tmp1 is None: tmp1 = tmp1_l + else: tmp1 += tmp1_l + if tmp2 is None: tmp2 = tmp2_l + else: tmp2 += tmp2_l + x_term = jnp.concatenate((tmp1, jnp.ones_like(tmp1)), axis=1) + y_term = jnp.concatenate((jnp.ones_like(tmp2), tmp2), axis=1) + return low_rank.LRCGeometry(cost_1=x_term, cost_2=y_term) + + def marginal_dependent_cost( + self, + marginal_1: jnp.ndarray, + marginal_2: jnp.ndarray, + ) -> low_rank.LRCGeometry: + r"""Initialize cost term that depends on the marginals of the transport. + + Uses the first term in eq. 6, p. 1 of :cite:`peyre:16`. + + Let :math:`p` be the `[n,]` marginal of the transport matrix for samples + from :attr:`geom_xx` and :math:`q` the `[m,]` marginal of the + transport matrix for samples from :attr:`geom_yy`. + + When ``cost_xx`` (resp. ``cost_yy``) is the cost matrix of :attr:`geom_xx` + (resp. :attr:`geom_yy`), the cost term that depends on these marginals can + be written as: + + .. math:: + + \text{marginal_dep_term} = \text{lin1}(\text{cost_xx}) p \mathbb{1}_{m}^T + + \mathbb{1}_{n}(\text{lin2}(\text{cost_yy}) q)^T + + This helper function instantiates these two low-rank matrices and groups + them into a single low-rank cost geometry object. + + Args: + marginal_1: [n,], first marginal of transport matrix. + marginal_2: [m,], second marginal of transport matrix. + + Returns: + Low-rank geometry of rank 2, storing normalization constants. + """ + geom_xx, geom_yy = self.geom_xx, self.geom_yy + if self._loss_name == "sqeucl": # quadratic apply, efficient for LR + tmp1 = geom_xx.apply_square_cost(marginal_1, axis=1) + tmp2 = geom_yy.apply_square_cost(marginal_2, axis=1) + else: + f1, f2 = self.linear_loss + tmp1 = apply_cost(geom_xx, marginal_1, axis=1, fn=f1) + tmp2 = apply_cost(geom_yy, marginal_2, axis=1, fn=f2) + x_term = jnp.concatenate((tmp1, jnp.ones_like(tmp1)), axis=1) + y_term = jnp.concatenate((jnp.ones_like(tmp2), tmp2), axis=1) + return low_rank.LRCGeometry(cost_1=x_term, cost_2=y_term) + + def cost_unbalanced_correction( + self, + transport_matrix: jnp.ndarray, + marginal_1: jnp.ndarray, + marginal_2: jnp.ndarray, + epsilon: epsilon_scheduler.Epsilon, + ) -> float: + r"""Calculate cost term from the quadratic divergence when unbalanced. + + In the unbalanced setting (``tau_a < 1.0 or tau_b < 1.0``), the + introduction of a quadratic divergence :cite:`sejourne:21` adds a term + to the GW local cost. + + Let :math:`a` [num_a,] be the target weights for samples + from geom_xx and :math:`b` [num_b,] be the target weights + for samples from `geom_yy`. Let :math:`P` [num_a, num_b] be the transport + matrix, :math:`P1` the first marginal and :math:`P^T1` the second marginal. + The term of the cost matrix coming from the quadratic KL in the + unbalanced case can be written as: + + `unbalanced_correction_term` = + :math:`tau_a / (1 - tau_a) * \sum(KL(P1|a))` + :math:`+ tau_b / (1 - tau_b) * \sum(KL(P^T1|b))` + :math:`+ epsilon * \sum(KL(P|ab'))` + + Args: + transport_matrix: jnp.ndarray[num_a, num_b], transport matrix. + marginal_1: jnp.ndarray[num_a,], marginal of the transport matrix + for samples from :attr:`geom_xx`. + marginal_2: jnp.ndarray[num_b,], marginal of the transport matrix + for samples from :attr:`geom_yy`. + epsilon: entropy regularizer. + + Returns: + The cost term. + """ + + def regularizer(tau: float) -> float: + return eps * tau / (1.0 - tau) + + eps = epsilon._target_init + marginal_1loga = jsp.special.xlogy(marginal_1, self.a).sum() + marginal_2logb = jsp.special.xlogy(marginal_2, self.b).sum() + + cost = eps * jsp.special.xlogy(transport_matrix, transport_matrix).sum() + if self.tau_a != 1.0: + cost += regularizer( + self.tau_a + ) * (-jsp.special.entr(marginal_1).sum() - marginal_1loga) + if self.tau_b != 1.0: + cost += regularizer( + self.tau_b + ) * (-jsp.special.entr(marginal_2).sum() - marginal_2logb) + return cost + + # TODO(michalk8): highly coupled to the pre-defined initializer, refactor + def init_transport_mass(self) -> float: + """Initialize the transport mass. + + Returns: + The sum of the elements of the normalized transport matrix. + """ + a = jax.lax.stop_gradient(self.a) + b = jax.lax.stop_gradient(self.b) + return a.sum() * b.sum() + + def update_lr_geom( + self, + lr_sink: "sinkhorn_lr.LRSinkhornOutput", + relative_epsilon: Optional[bool] = None, + ) -> geometry.Geometry: + """Recompute (possibly LRC) linearization using LR Sinkhorn output.""" + marginal_1 = lr_sink.marginal(1) + marginal_2 = lr_sink.marginal(0) + marginal_cost = self.marginal_dependent_cost(marginal_1, marginal_2) + + # Extract factors from LR Sinkhorn output + q, r, inv_sqg = lr_sink.q, lr_sink.r, 1.0 / jnp.sqrt(lr_sink.g) + # Distribute middle marginal evenly across both factors. + q, r = q * inv_sqg[None, :], r * inv_sqg[None, :] + + # Handle LRC Geometry case. + h1, h2 = self.quad_loss + geom_xx, geom_yy, geom_xy = self.geom_xx, self.geom_yy, self.geom_xy + tmp1 = apply_cost(geom_xx, q, axis=1, fn=h1) + tmp2 = apply_cost(geom_yy, r, axis=1, fn=h2) + if self.is_low_rank: + geom = low_rank.LRCGeometry( + cost_1=tmp1, cost_2=-tmp2, relative_epsilon=relative_epsilon + ) + marginal_cost + if self.is_fused: + geom = geom + geom_xy + else: + cost_matrix = marginal_cost.cost_matrix - jnp.dot(tmp1, tmp2.T) + cost_matrix += self.fused_penalty * self._fused_cost_matrix + geom = geometry.Geometry( + cost_matrix=cost_matrix, relative_epsilon=relative_epsilon + ) + return geom # noqa: RET504 + + def update_linearization( + self, + transport: Transport, + epsilon: Optional[Union[epsilon_scheduler.Epsilon, float]] = None, + old_transport_mass: float = 1.0, + relative_epsilon: Optional[bool] = None, + ) -> linear_problem.LinearProblem: + """Update linearization of GW problem by updating cost matrix. + + If the problem is balanced (``tau_a = 1.0 and tau_b = 1.0``), the equation + follows eq. 6, p. 1 of :cite:`peyre:16`. + + If the problem is unbalanced (``tau_a < 1.0 or tau_b < 1.0``), two cases are + possible, as explained in :meth:`init_linearization` above. + Finally, it is also possible to consider a Fused Gromov-Wasserstein problem. + Details about the resulting cost matrix are also given in + :meth:`init_linearization`. + + Args: + transport: Solution of the linearization of the quadratic problem. + epsilon: An epsilon scheduler or a float passed on to the linearization. + old_transport_mass: Sum of the elements of the transport matrix at the + previous iteration. + relative_epsilon: Whether to use relative epsilon in the linearized + geometry. + + Returns: + Updated linear OT problem, a new local linearization of GW problem. + """ + rescale_factor = 1.0 + unbalanced_correction = 0.0 + print("updating linearization") + if not self.is_balanced: + marginal_1 = transport.marginal(axis=1) + transport_mass = jax.lax.stop_gradient(marginal_1.sum()) + rescale_factor = jnp.sqrt(old_transport_mass / transport_mass) + + marginal_1 = transport.marginal(axis=1) * rescale_factor + marginal_2 = transport.marginal(axis=0) * rescale_factor + if self.labels_a is None: + marginal_cost = self.marginal_dependent_cost(marginal_1, marginal_2) + else: + marginal_cost = self.marginal_dependent_cost_labeled(marginal_1, marginal_2, + self.labels_a, self.labels_b) + + transport_matrix = transport.matrix * rescale_factor + if not self.is_balanced: + # Rescales transport for Unbalanced GW according to Sejourne et al. (2021) + transport_mass = jax.lax.stop_gradient(marginal_1.sum()) + epsilon = update_epsilon_unbalanced(epsilon, transport_mass) + unbalanced_correction = self.cost_unbalanced_correction( + transport_matrix, marginal_1, marginal_2, epsilon + ) + + h1, h2 = self.quad_loss + geom_xx, geom_yy = self.geom_xx, self.geom_yy + + tmp = apply_cost(geom_xx, transport_matrix, axis=1, fn=h1) + tmp = apply_cost(geom_yy, tmp.T, axis=1, fn=h2).T + + cost_matrix = marginal_cost.cost_matrix - tmp + unbalanced_correction + cost_matrix += self.fused_penalty * rescale_factor * self._fused_cost_matrix + if self.labels_a is not None: + print("Label considered for Sinkhorn run") + cost_matrix = cost_matrix + (1 / self.block_diag_mat -1) + geom = geometry.Geometry( + cost_matrix=cost_matrix, + epsilon=epsilon, + relative_epsilon=relative_epsilon, + ) + + return linear_problem.LinearProblem( + geom, self.a, self.b, tau_a=self.tau_a, tau_b=self.tau_b, + labels_a = self.labels_a, labels_b = self.labels_b + ) + + def update_lr_linearization( + self, + lr_sink: "sinkhorn_lr.LRSinkhornOutput", + *, + relative_epsilon: Optional[bool] = None, + ) -> linear_problem.LinearProblem: + """Update a Quad problem linearization using a LR Sinkhorn.""" + return linear_problem.LinearProblem( + self.update_lr_geom(lr_sink, relative_epsilon=relative_epsilon), + self.a, + self.b, + tau_a=self.tau_a, + tau_b=self.tau_b + ) + + @property + def _fused_cost_matrix(self) -> Union[float, jnp.ndarray]: + if not self.is_fused: + return 0.0 + geom_xy = self.geom_xy + + if isinstance(geom_xy, pointcloud.PointCloud) and geom_xy.is_online: + return geom_xy._compute_cost_matrix() * geom_xy.inv_scale_cost + return geom_xy.cost_matrix + + @property + def _is_low_rank_convertible(self) -> bool: + + def convertible(geom: geometry.Geometry) -> bool: + return isinstance(geom, low_rank.LRCGeometry) or ( + isinstance(geom, pointcloud.PointCloud) and geom.is_squared_euclidean + ) + + if self.is_low_rank: + return True + + geom_xx, geom_yy, geom_xy = self.geom_xx, self.geom_yy, self.geom_xy + # either explicitly via cost factorization or implicitly (e.g., a PC) + return self.ranks != -1 or ( + convertible(geom_xx) and convertible(geom_yy) and + (geom_xy is None or convertible(geom_xy)) + ) + + def to_low_rank( + self, + rng: Optional[jax.Array] = None, + ) -> "QuadraticProblem": + """Convert geometries to low-rank. + + Args: + rng: Random key for seeding. + + Returns: + Quadratic problem with low-rank geometries. + """ + + def convert( + vals: Union[int, float, Tuple[Union[int, float], ...]] + ) -> Tuple[Union[int, float], ...]: + size = 2 + self.is_fused + if isinstance(vals, (int, float)): + return (vals,) * 3 + assert len(vals) == size, vals + return vals + (None,) * (3 - size) + + if self.is_low_rank: + return self + + rng = utils.default_prng_key(rng) + rng1, rng2, rng3 = jax.random.split(rng, 3) + (geom_xx, geom_yy, geom_xy, *children), aux_data = self.tree_flatten() + (r1, r2, r3), (t1, t2, t3) = convert(self.ranks), convert(self.tolerances) + + geom_xx = geom_xx.to_LRCGeometry(rank=r1, tol=t1, rng=rng1) + geom_yy = geom_yy.to_LRCGeometry(rank=r2, tol=t2, rng=rng2) + if self.is_fused: + if isinstance( + geom_xy, pointcloud.PointCloud + ) and geom_xy.is_squared_euclidean: + geom_xy = geom_xy.to_LRCGeometry(scale=self.fused_penalty) + else: + geom_xy = geom_xy.to_LRCGeometry( + rank=r3, tol=t3, rng=rng3, scale=self.fused_penalty + ) + + return type(self).tree_unflatten( + aux_data, [geom_xx, geom_yy, geom_xy] + children + ) + + @property + def geom_xx(self) -> geometry.Geometry: + """Geometry of the first space.""" + return self._geom_xx + + @property + def geom_yy(self) -> geometry.Geometry: + """Geometry of the second space.""" + return self._geom_yy + + @property + def geom_xy(self) -> Optional[geometry.Geometry]: + """Geometry of the joint space.""" + return self._geom_xy + + @property + def a(self) -> jnp.ndarray: + """First marginal.""" + num_a = self.geom_xx.shape[0] + return jnp.ones((num_a,)) / num_a if self._a is None else self._a + + @property + def b(self) -> jnp.ndarray: + """Second marginal.""" + num_b = self.geom_yy.shape[0] + return jnp.ones((num_b,)) / num_b if self._b is None else self._b + + @property + def is_fused(self) -> bool: + """Whether the problem is fused.""" + return self.geom_xy is not None + + @property + def is_low_rank(self) -> bool: + """Whether all geometries are low-rank.""" + return ( + isinstance(self.geom_xx, low_rank.LRCGeometry) and + isinstance(self.geom_yy, low_rank.LRCGeometry) and + (not self.is_fused or isinstance(self.geom_xy, low_rank.LRCGeometry)) + ) + + @property + def linear_loss(self) -> Tuple[quadratic_costs.Loss, quadratic_costs.Loss]: + """Linear part of the Gromov-Wasserstein loss.""" + return self.loss.f1, self.loss.f2 + + @property + def quad_loss(self) -> Tuple[quadratic_costs.Loss, quadratic_costs.Loss]: + """Quadratic part of the Gromov-Wasserstein loss.""" + return self.loss.h1, self.loss.h2 + + @property + def is_balanced(self) -> bool: + """Whether the problem is balanced.""" + return ((not self.gw_unbalanced_correction) or + (self.tau_a == 1.0 and self.tau_b == 1.0)) + + def tree_flatten(self): # noqa: D102 + return ([self.geom_xx, self.geom_yy, self.geom_xy, self._a, self._b], { + "tau_a": self.tau_a, + "tau_b": self.tau_b, + "loss": self._loss_name, + "fused_penalty": self.fused_penalty, + "scale_cost": self.scale_cost, + "gw_unbalanced_correction": self.gw_unbalanced_correction, + "ranks": self.ranks, + "tolerances": self.tolerances + }) + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + geoms, (a, b) = children[:3], children[3:] + return cls(*geoms, a=a, b=b, **aux_data) + + +def update_epsilon_unbalanced( # noqa: D103 + epsilon: Union[float, epsilon_scheduler.Epsilon], transport_mass: float +) -> epsilon_scheduler.Epsilon: + if not isinstance(epsilon, epsilon_scheduler.Epsilon): + epsilon = epsilon_scheduler.Epsilon(epsilon, scale_epsilon=1.0) + return epsilon.set(scale_epsilon=epsilon._scale_epsilon * transport_mass) + + +def apply_cost( # noqa: D103 + geom: geometry.Geometry, arr: jnp.ndarray, *, axis: int, + fn: quadratic_costs.Loss +) -> jnp.ndarray: + return geom.apply_cost(arr, axis=axis, fn=fn.func, is_linear=fn.is_linear) diff --git a/ott/build/lib/ott/py.typed b/ott/build/lib/ott/py.typed new file mode 100755 index 0000000..e69de29 diff --git a/ott/build/lib/ott/solvers/__init__.py b/ott/build/lib/ott/solvers/__init__.py new file mode 100644 index 0000000..1303312 --- /dev/null +++ b/ott/build/lib/ott/solvers/__init__.py @@ -0,0 +1,14 @@ +# 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. +from . import linear, quadratic, was_solver diff --git a/ott/build/lib/ott/solvers/linear/__init__.py b/ott/build/lib/ott/solvers/linear/__init__.py new file mode 100644 index 0000000..4b7c246 --- /dev/null +++ b/ott/build/lib/ott/solvers/linear/__init__.py @@ -0,0 +1,30 @@ +# 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. +from . import ( + acceleration, + continuous_barycenter, + discrete_barycenter, + implicit_differentiation, + lr_utils, + sinkhorn, + sinkhorn_lr, + univariate, +) +from ._solve import solve + +__all__ = [ + "acceleration", "continuous_barycenter", "discrete_barycenter", + "implicit_differentiation", "lr_utils", "sinkhorn", "sinkhorn_lr", "solve", + "univariate" +] diff --git a/ott/build/lib/ott/solvers/linear/_solve.py b/ott/build/lib/ott/solvers/linear/_solve.py new file mode 100644 index 0000000..2bca6a8 --- /dev/null +++ b/ott/build/lib/ott/solvers/linear/_solve.py @@ -0,0 +1,60 @@ +# 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. +from typing import Any, Optional, Union + +import jax.numpy as jnp + +from ott.geometry import geometry +from ott.problems.linear import linear_problem +from ott.solvers.linear import sinkhorn, sinkhorn_lr + +__all__ = ["solve"] + + +def solve( + geom: geometry.Geometry, + a: Optional[jnp.ndarray] = None, + b: Optional[jnp.ndarray] = None, + tau_a: float = 1.0, + tau_b: float = 1.0, + rank: int = -1, + **kwargs: Any +) -> Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput]: + """Solve linear regularized OT problem using Sinkhorn iterations. + + Args: + geom: The ground geometry of the linear problem. + a: The first marginal. If :obj:`None`, it will be uniform. + b: The second marginal. If :obj:`None`, it will be uniform. + tau_a: If :math:`< 1`, defines how much unbalanced the problem is + on the first marginal. + tau_b: If :math:`< 1`, defines how much unbalanced the problem is + on the second marginal. + rank: + Rank constraint on the coupling to minimize the linear OT problem + :cite:`scetbon:21`. If :math:`-1`, no rank constraint is used. + kwargs: Keyword arguments for + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` or + :class:`~ott.solvers.linear.sinkhorn_lr.LRSinkhorn`, + depending on the ``rank``. + + Returns: + The Sinkhorn output. + """ + prob = linear_problem.LinearProblem(geom, a=a, b=b, tau_a=tau_a, tau_b=tau_b) + if rank > 0: + solver = sinkhorn_lr.LRSinkhorn(rank=rank, **kwargs) + else: + solver = sinkhorn.Sinkhorn(**kwargs) + return solver(prob) diff --git a/ott/build/lib/ott/solvers/linear/acceleration.py b/ott/build/lib/ott/solvers/linear/acceleration.py new file mode 100644 index 0000000..159322d --- /dev/null +++ b/ott/build/lib/ott/solvers/linear/acceleration.py @@ -0,0 +1,172 @@ +# 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. +from typing import TYPE_CHECKING + +import jax +import jax.numpy as jnp + +from ott import utils + +if TYPE_CHECKING: + from ott.problems.linear import linear_problem + from ott.solvers.linear import sinkhorn + +__all__ = ["AndersonAcceleration", "Momentum"] + + +@utils.register_pytree_node +class AndersonAcceleration: + """Implements Anderson acceleration for Sinkhorn.""" + + # TODO(michalk8): use memory=0 as no Anderson acceleration? + memory: int = 2 # Number of iterates considered to form interpolation. + refresh_every: int = 1 # Recompute interpolation periodically. + ridge_identity: float = 1e-2 # Ridge used in the linear system. + + def extrapolation(self, xs: jnp.ndarray, fxs: jnp.ndarray) -> jnp.ndarray: + """Compute Anderson extrapolation from past observations.""" + # Remove -inf values to instantiate quadratic problem. All others + # remain since they might be caused by a valid issue. + fxs_clean = jnp.nan_to_num(fxs, nan=jnp.nan, posinf=jnp.inf, neginf=0.0) + xs_clean = jnp.nan_to_num(xs, nan=jnp.nan, posinf=jnp.inf, neginf=0.0) + residuals = fxs_clean - xs_clean + gram_matrix = jnp.matmul(residuals.T, residuals) + gram_matrix /= jnp.linalg.norm(gram_matrix) + + # Solve linear system to obtain weights + weights = jax.scipy.sparse.linalg.cg( + gram_matrix + self.ridge_identity * jnp.eye(xs.shape[1]), + jnp.ones(xs.shape[1]) + )[0] + weights /= jnp.sum(weights) + + # Recover linear combination and return it with NaN (caused + # by 0 weights leading to -jnp.inf potentials, mixed with weights + # coefficients of different signs), disambiguated to -inf. + combination = jnp.sum(fxs * weights[None, :], axis=1) + return jnp.where(jnp.isfinite(combination), combination, -jnp.inf) + + def update( + self, state: "sinkhorn.SinkhornState", iteration: int, + prob: "linear_problem.LinearProblem", lse_mode: bool + ) -> "sinkhorn.SinkhornState": + """Anderson acceleration update. + + When using Anderson acceleration, first update the dual variable f_u with + previous updates (if iteration count sufficiently large), then record + new iterations in array. + + Anderson acceleration always happens in potentials (not scalings) space, + regardless of the lse_mode setting. If the iteration count is large + enough the update below will output a potential variable. + + Args: + state: Sinkhorn state. + iteration: the current iteration. + prob: linear OT problem. + lse_mode: whether to compute in log-sum-exp or in scalings. + + Returns: + A potential variable. + """ + geom = prob.geom + trigger_update = jnp.logical_and( + iteration > self.memory, iteration % self.refresh_every == 0 + ) + fu = jnp.where( + trigger_update, self.extrapolation(state.old_fus, state.old_mapped_fus), + state.fu + ) + # If the interpolation was triggered, we store it in memory + # Otherwise we add the previous value (converting it to potential form if + # it was initially stored in scaling form). + old_fus = jnp.where( + trigger_update, + jnp.concatenate((state.old_fus[:, 1:], fu[:, None]), axis=1), + jnp.concatenate(( + state.old_fus[:, 1:], + (fu if lse_mode else geom.potential_from_scaling(fu))[:, None] + ), + axis=1) + ) + + # If update was triggered, ensure a scaling is returned, since the result + # from the extrapolation was outputted in potential form. + fu = jnp.where( + trigger_update, fu if lse_mode else geom.scaling_from_potential(fu), fu + ) + return state.set(potentials=(fu, state.gv), old_fus=old_fus) + + def init_maps( + self, pb, state: "sinkhorn.SinkhornState" + ) -> "sinkhorn.SinkhornState": + """Initialize log matrix used in Anderson acceleration with *NaN* values.""" + fus = jnp.ones((pb.geom.shape[0], self.memory)) * jnp.nan + return state.set(old_fus=fus, old_mapped_fus=fus) + + def update_history( + self, state: "sinkhorn.SinkhornState", pb, lse_mode: bool + ) -> "sinkhorn.SinkhornState": + """Update history of mapped dual variables.""" + f = state.fu if lse_mode else pb.geom.potential_from_scaling(state.fu) + mapped = jnp.concatenate((state.old_mapped_fus[:, 1:], f[:, None]), axis=1) + return state.set(old_mapped_fus=mapped) + + +@utils.register_pytree_node +class Momentum: + """Momentum for Sinkhorn updates. + + Can be either constant :cite:`thibault:21` or adaptive :cite:`lehmann:21`. + """ + + start: int = 0 + error_threshold: float = jnp.inf + value: float = 1.0 + inner_iterations: int = 1 + + def weight(self, state: "sinkhorn.SinkhornState", iteration: int) -> float: + """Compute momentum term if needed, using previously seen errors.""" + if self.start == 0: + return self.value + idx = self.start // self.inner_iterations + + return jax.lax.cond( + jnp.logical_and( + iteration >= self.start, state.errors[idx - 1, -1] + < self.error_threshold + ), lambda state: self.lehmann(state), lambda state: self.value, state + ) + + def lehmann(self, state: "sinkhorn.SinkhornState") -> float: + """Momentum formula :cite:`lehmann:21`, eq. 5.""" + idx = self.start // self.inner_iterations + error_ratio = jnp.minimum( + state.errors[idx - 1, -1] / state.errors[idx - 2, -1], 0.99 + ) + power = 1.0 / self.inner_iterations + return 2.0 / (1.0 + jnp.sqrt(1.0 - error_ratio ** power)) + + def __call__( # noqa: D102 + self, + weight: float, + value: jnp.ndarray, + new_value: jnp.ndarray, + lse_mode: bool = True + ) -> jnp.ndarray: + if lse_mode: + value = jnp.where(jnp.isfinite(value), value, 0.0) + return (1.0 - weight) * value + weight * new_value + value = jnp.where(value > 0.0, value, 1.0) + return value ** (1.0 - weight) * new_value ** weight diff --git a/ott/build/lib/ott/solvers/linear/continuous_barycenter.py b/ott/build/lib/ott/solvers/linear/continuous_barycenter.py new file mode 100644 index 0000000..8f60380 --- /dev/null +++ b/ott/build/lib/ott/solvers/linear/continuous_barycenter.py @@ -0,0 +1,233 @@ +# 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 functools +from typing import Any, NamedTuple, Optional, Tuple + +import jax +import jax.numpy as jnp + +from ott import utils +from ott.geometry import pointcloud +from ott.math import fixed_point_loop +from ott.math import utils as mu +from ott.problems.linear import barycenter_problem, linear_problem +from ott.solvers import was_solver + +__all__ = ["FreeBarycenterState", "FreeWassersteinBarycenter"] + + +class FreeBarycenterState(NamedTuple): + """Holds the state of the Wasserstein barycenter solver. + + Args: + costs: Holds the sequence of regularized GW costs seen through the outer + loop of the solver. + linear_convergence: Holds the sequence of bool convergence flags of the + inner Sinkhorn iterations. + errors: Holds sequence of vectors of errors of the Sinkhorn algorithm + at each iteration. + x: barycenter points. + a: barycenter weights. + """ + + costs: Optional[jnp.ndarray] = None + linear_convergence: Optional[jnp.ndarray] = None + errors: Optional[jnp.ndarray] = None + x: Optional[jnp.ndarray] = None + a: Optional[jnp.ndarray] = None + + def set(self, **kwargs: Any) -> "FreeBarycenterState": + """Return a copy of self, possibly with overwrites.""" + return self._replace(**kwargs) + + def update( + self, iteration: int, bar_prob: barycenter_problem.FreeBarycenterProblem, + linear_ot_solver: Any, store_errors: bool + ) -> "FreeBarycenterState": + """Update the state of the solver. + + Args: + iteration: the current iteration of the outer loop. + bar_prob: the barycenter problem. + linear_ot_solver: the linear OT solver to use. + store_errors: whether to store the errors of the inner loop. + + Returns: + The updated state. + """ + seg_y, seg_b = bar_prob.segmented_y_b + + @functools.partial(jax.vmap, in_axes=[None, None, 0, 0]) + def solve_linear_ot( + a: Optional[jnp.ndarray], x: jnp.ndarray, b: jnp.ndarray, y: jnp.ndarray + ): + out = linear_ot_solver( + linear_problem.LinearProblem( + pointcloud.PointCloud( + x, + y, + src_mask=a > 0.0, + tgt_mask=b > 0.0, + cost_fn=bar_prob.cost_fn, + epsilon=bar_prob.epsilon + ), a, b + ) + ) + return ( + out.reg_ot_cost, out.converged, out.matrix, + out.errors if store_errors else None + ) + + reg_ot_costs, convergeds, matrices, errors = solve_linear_ot( + self.a, self.x, seg_b, seg_y + ) + + cost = jnp.sum(reg_ot_costs * bar_prob.weights) + updated_costs = self.costs.at[iteration].set(cost) + converged = jnp.all(convergeds) + linear_convergence = self.linear_convergence.at[iteration].set(converged) + + if store_errors and self.errors is not None: + errors = self.errors.at[iteration, :, :].set(errors) + else: + errors = None + + # Approximation of barycenter as barycenter of barycenters per measure. + + barycenters_per_measure = mu.barycentric_projection( + matrices, seg_y, bar_prob.cost_fn + ) + + x_new = jax.vmap( + lambda w, y: bar_prob.cost_fn.barycenter(w, y)[0], in_axes=[None, 1] + )(bar_prob.weights, barycenters_per_measure) + + return self.set( + costs=updated_costs, + linear_convergence=linear_convergence, + errors=errors, + x=x_new + ) + + +@jax.tree_util.register_pytree_node_class +class FreeWassersteinBarycenter(was_solver.WassersteinSolver): + """Continuous Wasserstein barycenter solver :cite:`cuturi:14`.""" + + def __call__( # noqa: D102 + self, + bar_prob: barycenter_problem.FreeBarycenterProblem, + bar_size: int = 100, + x_init: Optional[jnp.ndarray] = None, + rng: Optional[jax.Array] = None, + ) -> FreeBarycenterState: + # TODO(michalk8): no reason for iterations to be outside this class + rng = utils.default_prng_key(rng) + return iterations(self, bar_size, bar_prob, x_init, rng) + + def init_state( + self, + bar_prob: barycenter_problem.FreeBarycenterProblem, + bar_size: int, + x_init: Optional[jnp.ndarray] = None, + rng: Optional[jax.Array] = None, + ) -> FreeBarycenterState: + """Initialize the state of the Wasserstein barycenter iterations. + + Args: + bar_prob: The barycenter problem. + bar_size: Size of the barycenter. + x_init: Initial barycenter estimate of shape ``[bar_size, ndim]``. + If `None`, ``bar_size`` points will be sampled from the input + measures according to their weights + :attr:`~ott.problems.linear.barycenter_problem.FreeBarycenterProblem.flattened_y`. + rng: Random key for seeding. + + Returns: + The initial barycenter state. + """ + if x_init is not None: + assert bar_size == x_init.shape[0] + x = x_init + else: + # sample randomly points in the support of the y measures + rng = utils.default_prng_key(rng) + indices_subset = jax.random.choice( + rng, + a=bar_prob.flattened_y.shape[0], + shape=(bar_size,), + replace=False, + p=bar_prob.flattened_b + ) + x = bar_prob.flattened_y[indices_subset, :] + + # TODO(cuturi) expand to non-uniform weights for barycenter. + a = jnp.ones((bar_size,)) / bar_size + num_iter = self.max_iterations + if self.store_inner_errors: + errors = -jnp.ones(( + num_iter, bar_prob.num_measures, + self.linear_ot_solver.outer_iterations + )) + else: + errors = None + return FreeBarycenterState( + -jnp.ones((num_iter,)), -jnp.ones((num_iter,)), errors, x, a + ) + + def output_from_state( # noqa: D102 + self, state: FreeBarycenterState + ) -> FreeBarycenterState: + # TODO(michalk8): create an output variable to match rest of the framework + return state + + +def iterations( + solver: FreeWassersteinBarycenter, bar_size: int, + bar_prob: barycenter_problem.FreeBarycenterProblem, x_init: jnp.ndarray, + rng: jax.Array +) -> FreeBarycenterState: + """Jittable Wasserstein barycenter outer loop.""" + + def cond_fn( + iteration: int, + constants: Tuple[FreeWassersteinBarycenter, + barycenter_problem.FreeBarycenterProblem], + state: FreeBarycenterState + ) -> bool: + solver, _ = constants + return solver._continue(state, iteration) + + def body_fn( + iteration, constants: Tuple[FreeWassersteinBarycenter, + barycenter_problem.FreeBarycenterProblem], + state: FreeBarycenterState, compute_error: bool + ) -> FreeBarycenterState: + del compute_error # Always assumed True + solver, bar_prob = constants + return state.update( + iteration, bar_prob, solver.linear_ot_solver, solver.store_inner_errors + ) + + state = fixed_point_loop.fixpoint_iter( + cond_fn=cond_fn, + body_fn=body_fn, + min_iterations=solver.min_iterations, + max_iterations=solver.max_iterations, + inner_iterations=1, + constants=(solver, bar_prob), + state=solver.init_state(bar_prob, bar_size, x_init, rng) + ) + + return solver.output_from_state(state) diff --git a/ott/build/lib/ott/solvers/linear/discrete_barycenter.py b/ott/build/lib/ott/solvers/linear/discrete_barycenter.py new file mode 100644 index 0000000..dcfdc14 --- /dev/null +++ b/ott/build/lib/ott/solvers/linear/discrete_barycenter.py @@ -0,0 +1,251 @@ +# 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 functools +from typing import NamedTuple, Optional, Sequence + +import jax +import jax.numpy as jnp + +from ott.geometry import geometry +from ott.math import fixed_point_loop +from ott.problems.linear import barycenter_problem +from ott.solvers.linear import sinkhorn + +__all__ = ["SinkhornBarycenterOutput", "FixedBarycenter"] + + +class SinkhornBarycenterOutput(NamedTuple): # noqa: D101 + f: jnp.ndarray + g: jnp.ndarray + histogram: jnp.ndarray + errors: jnp.ndarray + + +@jax.tree_util.register_pytree_node_class +class FixedBarycenter: + """A Wasserstein barycenter solver for histograms on a common geometry. + + This solver uses a variant of the + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` algorithm proposed in + :cite:`janati:20a` to compute the barycenter of various measures supported on + the same (common to all) geometry. The geometry is assumed to be either + symmetric, or to describe costs between a set of points and another. In that + case all reference measures have support on the first measure, whereas the + barycenter is supported on the second. + + Args: + threshold: convergence threshold. The algorithm stops when the marginal + violations of all transport plans computed for that barycenter go below + that threshold. + norm_error: norm used to compute marginal deviation. + inner_iterations: number of iterations run before recomputing errors. + min_iterations: number of iterations run without checking whether + termination criterion is true. + max_iterations: maximal number of iterations. + lse_mode: sets computations in kernel (``False``) or log-sum-exp mode. + debiased: uses debiasing correction to avoid blur due to entropic + regularization. + """ + + def __init__( + self, + threshold: float = 1e-2, + norm_error: int = 1, + inner_iterations: float = 10, + min_iterations: int = 0, + max_iterations: int = 2000, + lse_mode: bool = True, + debiased: bool = False + ): + self.threshold = threshold + self.norm_error = norm_error + self.inner_iterations = inner_iterations + self.min_iterations = min_iterations + self.max_iterations = max_iterations + self.lse_mode = lse_mode + self.debiased = debiased + + def __call__( + self, + fixed_bp: barycenter_problem.FixedBarycenterProblem, + dual_initialization: Optional[jnp.ndarray] = None, + ) -> SinkhornBarycenterOutput: + """Solve barycenter problem, possibly using clever initialization. + + Args: + fixed_bp: Fixed barycenter problem. + dual_initialization: Initial value for the g_v potential/scalings, + one for each of the histograms described in ``fixed_bp``. If ``None``, + use initialization from :cite:`cuturi:15`, eq. 3.6. + + Returns: + The barycenter. + """ + geom = fixed_bp.geom + a = fixed_bp.a + num_a, num_b = geom.shape + + weights = fixed_bp.weights + + if dual_initialization is None: + # initialization strategy from :cite:`cuturi:15`, (3.6). + dual_initialization = geom.apply_cost(a.T, axis=0).T + dual_initialization -= jnp.average( + dual_initialization, weights=weights, axis=0 + )[jnp.newaxis, :] + + if self.debiased and not geom.is_symmetric: + raise ValueError("Geometry must be symmetric to use debiased option.") + norm_error = (self.norm_error,) + return _discrete_barycenter( + geom, a, weights, dual_initialization, self.threshold, norm_error, + self.inner_iterations, self.min_iterations, self.max_iterations, + self.lse_mode, self.debiased, num_a, num_b + ) + + def tree_flatten(self): # noqa: D102 + aux = vars(self).copy() + aux.pop("threshold") + return [ + self.threshold, + ], aux + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + return cls(**aux_data, threshold=children[0]) + + +@functools.partial(jax.jit, static_argnums=(5, 6, 7, 8, 9, 10, 11, 12)) +def _discrete_barycenter( + geom: geometry.Geometry, a: jnp.ndarray, weights: jnp.ndarray, + dual_initialization: jnp.ndarray, threshold: float, + norm_error: Sequence[int], inner_iterations: int, min_iterations: int, + max_iterations: int, lse_mode: bool, debiased: bool, num_a: int, num_b: int +) -> SinkhornBarycenterOutput: + """Jit'able function to compute discrete barycenters.""" + if lse_mode: + f_u = jnp.zeros_like(a) + g_v = dual_initialization + else: + f_u = jnp.ones_like(a) + g_v = geom.scaling_from_potential(dual_initialization) + # d below is as described in https://arxiv.org/abs/2006.02575. Note that + # d should be considered to be equal to eps log(d) with those notations + # if running in log-sum-exp mode. + d = jnp.zeros((num_b,)) if lse_mode else jnp.ones((num_b,)) + + if lse_mode: + parallel_update = jax.vmap( + lambda f, g, marginal, iter: geom. + update_potential(f, g, jnp.log(marginal), axis=1), + in_axes=[0, 0, 0, None] + ) + parallel_apply = jax.vmap( + lambda f_, g_, eps_: geom. + apply_lse_kernel(f_, g_, eps_, vec=None, axis=0)[0], + in_axes=[0, 0, None] + ) + else: + parallel_update = jax.vmap( + lambda f, g, marginal, iter: geom.update_scaling(g, marginal, axis=1), + in_axes=[0, 0, 0, None] + ) + parallel_apply = jax.vmap( + lambda f_, g_, eps_: geom.apply_kernel(f_, eps_, axis=0), + in_axes=[0, 0, None] + ) + + errors_fn = jax.vmap( + functools.partial( + sinkhorn.marginal_error, + geom=geom, + axis=1, + norm_error=norm_error, + lse_mode=lse_mode + ), + in_axes=[0, 0, 0] + ) + errors = -jnp.ones((max_iterations // inner_iterations + 1, len(norm_error))) + + const = (geom, a, weights) + + def cond_fn(iteration, const, state): # pylint: disable=unused-argument + errors = state[0] + return jnp.logical_or( + iteration == 0, errors[iteration // inner_iterations - 1, 0] > threshold + ) + + def body_fn(iteration, const, state, compute_error): + geom, a, weights = const + errors, d, f_u, g_v = state + + eps = geom._epsilon.at(iteration) # pylint: disable=protected-access + f_u = parallel_update(f_u, g_v, a, iteration) + # kernel_f_u stands for K times potential u if running in scaling mode, + # eps log K exp f / eps in lse mode. + kernel_f_u = parallel_apply(f_u, g_v, eps) + # b below is the running estimate for the barycenter if running in scaling + # mode, eps log b if running in lse mode. + if lse_mode: + b = jnp.average(kernel_f_u, weights=weights, axis=0) + else: + b = jnp.prod(kernel_f_u ** weights[:, jnp.newaxis], axis=0) + + if debiased: + if lse_mode: + b += d + d = 0.5 * ( + d + geom.update_potential( + jnp.zeros((num_a,)), d, b / eps, iteration=iteration, axis=0 + ) + ) + else: + b *= d + d = jnp.sqrt(d * geom.update_scaling(d, b, iteration=iteration, axis=0)) + if lse_mode: + g_v = b[jnp.newaxis, :] - kernel_f_u + else: + g_v = b[jnp.newaxis, :] / kernel_f_u + + # re-compute error if compute_error is True, else set to inf. + err = jnp.where( + jnp.logical_and(compute_error, iteration >= min_iterations), + jnp.mean(errors_fn(f_u, g_v, a)), jnp.inf + ) + + errors = errors.at[iteration // inner_iterations, :].set(err) + return errors, d, f_u, g_v + + state = (errors, d, f_u, g_v) + + state = fixed_point_loop.fixpoint_iter_backprop( + cond_fn, body_fn, min_iterations, max_iterations, inner_iterations, const, + state + ) + + errors, d, f_u, g_v = state + kernel_f_u = parallel_apply(f_u, g_v, geom.epsilon) + if lse_mode: + b = jnp.average(kernel_f_u, weights=weights, axis=0) + else: + b = jnp.prod(kernel_f_u ** weights[:, jnp.newaxis], axis=0) + + if debiased: + if lse_mode: + b += d + else: + b *= d + if lse_mode: + b = jnp.exp(b / geom.epsilon) + return SinkhornBarycenterOutput(f_u, g_v, b, errors) diff --git a/ott/build/lib/ott/solvers/linear/implicit_differentiation.py b/ott/build/lib/ott/solvers/linear/implicit_differentiation.py new file mode 100644 index 0000000..fbf98ce --- /dev/null +++ b/ott/build/lib/ott/solvers/linear/implicit_differentiation.py @@ -0,0 +1,324 @@ +# 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 dataclasses +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple + +import jax +import jax.numpy as jnp + +from ott import utils +from ott.math import unbalanced_functions as uf + +if TYPE_CHECKING: + from ott.problems.linear import linear_problem + +LinOp_t = Callable[[jnp.ndarray], jnp.ndarray] +Solver_t = Callable[[LinOp_t, jnp.ndarray, Optional[LinOp_t], bool], + jnp.ndarray] + +__all__ = ["ImplicitDiff", "solve_jax_cg"] + + +@utils.register_pytree_node +class ImplicitDiff: + """Implicit differentiation of Sinkhorn algorithm. + + Args: + solver: Callable to compute the solution to a linear problem. The callable + expects a linear function, a vector, optionally another linear function + that implements the transpose of that function, and a boolean flag to + specify symmetry. This solver is by default one of :class:`lineax.CG` or + :class:`lineax.NormalCG` solvers, if the package can be imported, as + described in :func:`~ott.solvers.linear.lineax_implicit.solve_lineax`. + The :mod:`jax` alternative is described in + :func:`~ott.solvers.linear.implicit_differentiation.solve_jax_cg`. + Note that `lineax` solvers handle better poorly conditioned problems, + which arise typically when differentiating the solutions of balanced OT + problems (when ``tau_a==tau_b==1.0``). Relying on + :func:`~ott.solvers.linear.implicit_differentiation.solve_jax_cg` + for such cases might require hand-tuning ridge parameters, + in particular ``ridge_kernel`` and ``ridge_identity`` as described in its + doc. These parameters can be passed using ``solver_kwargs`` below. + solver_kwargs: keyword arguments passed on to the solver. + symmetric: flag used to figure out whether the linear system solved in the + implicit function theorem is symmetric or not. This happens when + ``tau_a==tau_b``, and when ``a == b``, or the precondition_fun + is the identity. The flag is False by default, and is also tested against + ``tau_a==tau_b``. It needs to be set manually by the user in the more + favorable case where the system is guaranteed to be symmetric. + precondition_fun: Function used to precondition, on both sides, the linear + system derived from first-order conditions of the regularized OT problem. + That linear system typically involves an equality between marginals (or + simple transform of these marginals when the problem is unbalanced) and + another function of the potentials. When that function is specified, that + function is applied on both sides of the equality, before being further + differentiated to provide the Jacobians needed for implicit function + theorem differentiation. + """ + + solver: Optional[Solver_t] = None + solver_kwargs: Optional[Dict[str, Any]] = None + symmetric: bool = False + precondition_fun: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None + + def solve( + self, + gr: Tuple[jnp.ndarray, jnp.ndarray], + ot_prob: "linear_problem.LinearProblem", + f: jnp.ndarray, + g: jnp.ndarray, + lse_mode: bool, + ) -> jnp.ndarray: + r"""Apply minus inverse of [hessian ``reg_ot_cost`` w.r.t. ``f``, ``g``]. + + This function is used to carry out implicit differentiation of ``sinkhorn`` + outputs, notably optimal potentials ``f`` and ``g``. That differentiation + requires solving a linear system, using (and inverting) the Jacobian of + (preconditioned) first-order conditions w.r.t. the reg-OT problem. + + Given a ``precondition_fun``, written here for short as :math:`h`, + the first order conditions for the dual energy + + .. math:: + + E(K, \epsilon, a, b, f, g) :=- + - \langle\exp^{f/\epsilon}, K \exp^{g/\epsilon}> + + form the basis of the Sinkhorn algorithm. To differentiate optimal solutions + to that problem, we exploit the fact that :math:`h(\nabla E = 0)` and + differentiate that identity to recover variations (Jacobians) of optimal + solutions :math:`f^\star, g^\star$` as a function of changes in the inputs. + The Jacobian of :math:`h(\nabla_{f,g} E = 0)` is a linear operator which, if + it were to be instantiated as a matrix, would be of size + :math:`(n+m) \times (n+m)`. When :math:`h` is the identity, that matrix is + the Hessian of :math:`E`, is symmetric and negative-definite + (:math:`E` is concave) and is structured as :math:`[A, B; B^T, D]`. More + generally, for other functions :math:`h`, the Jacobian of these + preconditioned + first order conditions is no longer symmetric (except if ``a==b``), and + has now a structure as :math:`[A, B; C, D]`. That system can + be still inverted more generic solvers. By default, :math:`h = \epsilon + \log`, as proposed in :cite:`cuturi:20a`. + + In both cases :math:`A` and :math:`D` are diagonal matrices, equal to the + row and + column marginals respectively, multiplied by the derivatives of :math:`h` + evaluated at those marginals, corrected (if handling the unbalanced case) + by the second derivative of the part of the objective that ties potentials + to the marginals (terms in ``phi_star``). When :math:`h` is the identity, + :math:`B` and :math:`B^T` are equal respectively to the OT matrix and its + transpose, i.e. :math:`n \times m` and :math:`m \times n` matrices. + When :math:`h` is not the identity, :math:`B` (resp. :math:`C`) is equal + to the OT matrix (resp. its transpose), rescaled on the left by the + application elementwise of :math:`h'` to the row (respectively column) + marginal sum of the transport. + + Note that we take great care in not instantiating these transport + matrices, to rely instead on calls to the ``app_transport`` method from the + ``Geometry`` object ``geom`` (which will either use potentials or scalings, + depending on ``lse_mode``) + + The Jacobian's diagonal + off-diagonal blocks structure allows to exploit + Schur complements. Depending on the sizes involved, it is better to + instantiate the Schur complement of the first or of the second diagonal + block. + + These linear systems are solved using the user-defined ``solver``, using + by default :mod:`lineax` solvers when available, or falling back on + :mod:`jax` when not. + + Args: + gr: 2-tuple, (vector of size ``n``, vector of size ``m``). + ot_prob: the instantiation of the regularized transport problem. + f: potential, w.r.t marginal a. + g: potential, w.r.t marginal b. + lse_mode: bool, log-sum-exp mode if True, kernel else. + + Returns: + A tuple of two vectors, of the same size as ``gr``. + """ + solver = _get_solver() if self.solver is None else self.solver + solver_kwargs = {} if self.solver_kwargs is None else self.solver_kwargs + geom = ot_prob.geom + marginal_a, marginal_b, app_transport = ( + ot_prob.get_transport_functions(lse_mode) + ) + if self.precondition_fun is None: + precond_fun = lambda x: geom.epsilon * jnp.log(x) + symmetric = False + else: + precond_fun = self.precondition_fun + symmetric = self.symmetric + + derivative = jax.vmap(jax.grad(precond_fun)) + + n, m = geom.shape + # pylint: disable=g-long-lambda + vjp_fg = lambda z: app_transport( + f, g, z * derivative(marginal_b(f, g)), axis=1 + ) / geom.epsilon + vjp_gf = lambda z: app_transport( + f, g, z * derivative(marginal_a(f, g)), axis=0 + ) / geom.epsilon + + if not symmetric: + vjp_fgt = lambda z: app_transport( + f, g, z, axis=0 + ) * derivative(marginal_b(f, g)) / geom.epsilon + vjp_gft = lambda z: app_transport( + f, g, z, axis=1 + ) * derivative(marginal_a(f, g)) / geom.epsilon + + diag_hess_a = ( + marginal_a(f, g) * derivative(marginal_a(f, g)) / geom.epsilon + + uf.diag_jacobian_of_marginal_fit( + ot_prob.a, f, ot_prob.tau_a, geom.epsilon, derivative + ) + ) + diag_hess_b = ( + marginal_b(f, g) * derivative(marginal_b(f, g)) / geom.epsilon + + uf.diag_jacobian_of_marginal_fit( + ot_prob.b, g, ot_prob.tau_b, geom.epsilon, derivative + ) + ) + n, m = geom.shape + # TODO(cuturi) consider materializing linear operator schur if size allows. + # Forks on using Schur complement of either A or D, depending on size. + if n > m: # if n is bigger, run m x m linear system. + inv_vjp_ff = lambda z: z / diag_hess_a + vjp_gg = lambda z: z * diag_hess_b + schur = lambda z: vjp_gg(z) - vjp_gf(inv_vjp_ff(vjp_fg(z))) + if not symmetric: + schur_t = lambda z: vjp_gg(z) - vjp_fgt(inv_vjp_ff(vjp_gft(z))) + else: + schur_t = None + res = gr[1] - vjp_gf(inv_vjp_ff(gr[0])) + sch = solver(schur, res, schur_t, symmetric, **solver_kwargs) + vjp_gr_f = inv_vjp_ff(gr[0] - vjp_fg(sch)) + vjp_gr_g = sch + else: + vjp_ff = lambda z: z * diag_hess_a + inv_vjp_gg = lambda z: z / diag_hess_b + schur = lambda z: vjp_ff(z) - vjp_fg(inv_vjp_gg(vjp_gf(z))) + + if not symmetric: + schur_t = lambda z: vjp_ff(z) - vjp_gft(inv_vjp_gg(vjp_fgt(z))) + else: + schur_t = None + res = gr[0] - vjp_fg(inv_vjp_gg(gr[1])) + sch = solver(schur, res, schur_t, symmetric, **solver_kwargs) + vjp_gr_g = inv_vjp_gg(gr[1] - vjp_gf(sch)) + vjp_gr_f = sch + + return jnp.concatenate((-vjp_gr_f, -vjp_gr_g)) + + def first_order_conditions( + self, prob, f: jnp.ndarray, g: jnp.ndarray, lse_mode: bool + ): + r"""Compute vector of first order conditions for the reg-OT problem. + + The output of this vector should be close to zero at optimality. + Upon completion of the Sinkhorn forward pass, its norm (using the norm + parameter defined using ``norm_error``) should be below the threshold + parameter. + + This error will be itself assumed to be close to zero when using implicit + differentiation. + + Args: + prob: definition of the linear optimal transport problem. + f: jnp.ndarray, first potential + g: jnp.ndarray, second potential + lse_mode: bool + + Returns: + a jnp.ndarray of size (size of ``n + m``) quantifying deviation to + optimality for variables ``f`` and ``g``. + """ + geom = prob.geom + marginal_a, marginal_b, _ = prob.get_transport_functions(lse_mode) + grad_a = uf.grad_of_marginal_fit(prob.a, f, prob.tau_a, geom.epsilon) + grad_b = uf.grad_of_marginal_fit(prob.b, g, prob.tau_b, geom.epsilon) + if self.precondition_fun is None: + precond_fun = lambda x: geom.epsilon * jnp.log(x) + else: + precond_fun = self.precondition_fun + + result_a = jnp.where( + prob.a > 0, + precond_fun(marginal_a(f, g)) - precond_fun(grad_a), 0.0 + ) + result_b = jnp.where( + prob.b > 0, + precond_fun(marginal_b(f, g)) - precond_fun(grad_b), 0.0 + ) + return jnp.concatenate((result_a, result_b)) + + def gradient( + self, prob: "linear_problem.LinearProblem", f: jnp.ndarray, + g: jnp.ndarray, lse_mode: bool, gr: Tuple[jnp.ndarray, jnp.ndarray] + ) -> "linear_problem.LinearProblem": + """Apply VJP to recover gradient in reverse mode differentiation.""" + # Applies first part of vjp to gr: inverse part of implicit function theorem + vjp_gr = self.solve(gr, prob, f, g, lse_mode) + + # Instantiates vjp of first order conditions of the objective, as a + # function of geom, a and b parameters (against which we differentiate) + foc_prob = lambda prob: self.first_order_conditions(prob, f, g, lse_mode) + + # Carries pullback onto original inputs, here geom, a and b. + _, pull_prob = jax.vjp(foc_prob, prob) + return pull_prob(vjp_gr) + + def replace(self, **kwargs: Any) -> "ImplicitDiff": # noqa: D102 + return dataclasses.replace(self, **kwargs) + + +def solve_jax_cg( + lin: LinOp_t, + b: jnp.ndarray, + lin_t: Optional[LinOp_t] = None, + symmetric: bool = False, + ridge_identity: float = 0.0, + ridge_kernel: float = 0.0, + **kwargs: Any +) -> jnp.ndarray: + """Wrapper around JAX native linear solvers. + + Args: + lin: Linear operator + b: vector. Returned `x` is such that `lin(x)=b` + lin_t: Linear operator, corresponding to transpose of `lin`. + symmetric: whether `lin` is symmetric. + ridge_kernel: promotes zero-sum solutions. Only use if `tau_a = tau_b = 1.0` + ridge_identity: handles rank deficient transport matrices (this happens + typically when rows/cols in cost/kernel matrices are collinear, or, + equivalently when two points from either measure are close). + kwargs: arguments passed to :func:`~jax.scipy.sparse.linalg.cg` + """ + op = lin if symmetric else lambda x: lin_t(lin(x)) + if ridge_kernel > 0.0 or ridge_identity > 0.0: + lin_reg = lambda x: op(x) + ridge_kernel * jnp.sum(x) + ridge_identity * x + else: + lin_reg = op + return jax.scipy.sparse.linalg.cg(lin_reg, b, **kwargs)[0] + + +def _get_solver() -> Solver_t: + """Get lineax solver when possible, default to jax.scipy else.""" + try: + from ott.solvers.linear import lineax_implicit + return lineax_implicit.solve_lineax + except ImportError: + return solve_jax_cg diff --git a/ott/build/lib/ott/solvers/linear/lineax_implicit.py b/ott/build/lib/ott/solvers/linear/lineax_implicit.py new file mode 100644 index 0000000..79b9e7c --- /dev/null +++ b/ott/build/lib/ott/solvers/linear/lineax_implicit.py @@ -0,0 +1,98 @@ +# 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. +from typing import Any, Callable, Optional, TypeVar + +import equinox as eqx +import jax +import jax.numpy as jnp +import jax.tree_util as jtu +import lineax as lx +from jaxtyping import Array, Float, PyTree + +_T = TypeVar("_T") +_FlatPyTree = tuple[list[_T], jtu.PyTreeDef] + +__all__ = ["CustomTransposeLinearOperator", "solve_lineax"] + + +class CustomTransposeLinearOperator(lx.FunctionLinearOperator): + """Implement a linear operator that can specify its transpose directly.""" + fn: Callable[[PyTree[Float[Array, "..."]]], PyTree[Float[Array, "..."]]] + fn_t: Callable[[PyTree[Float[Array, "..."]]], PyTree[Float[Array, "..."]]] + input_structure: _FlatPyTree[jax.ShapeDtypeStruct] = eqx.static_field() + input_structure_t: _FlatPyTree[jax.ShapeDtypeStruct] = eqx.static_field() + tags: frozenset[object] + + def __init__(self, fn, fn_t, input_structure, input_structure_t, tags=()): + super().__init__(fn, input_structure, tags) + self.fn_t = eqx.filter_closure_convert(fn_t, input_structure_t) + self.input_structure_t = input_structure_t + + def transpose(self): + """Provide custom transposition operator from function.""" + return lx.FunctionLinearOperator(self.fn_t, self.input_structure_t) + + +def solve_lineax( + lin: Callable, + b: jnp.ndarray, + lin_t: Optional[Callable] = None, + symmetric: bool = False, + nonsym_solver: Optional[lx.AbstractLinearSolver] = None, + ridge_identity: float = 0.0, + ridge_kernel: float = 0.0, + **kwargs: Any +) -> jnp.ndarray: + """Wrapper around lineax solvers. + + Args: + lin: Linear operator + b: vector. Returned `x` is such that `lin(x)=b` + lin_t: Linear operator, corresponding to transpose of `lin`. + symmetric: whether `lin` is symmetric. + nonsym_solver: solver used when handling non-symmetric cases. Note that + :class:`~lineax.CG` is used by default in the symmetric case. + ridge_kernel: promotes zero-sum solutions. Only use if `tau_a = tau_b = 1.0` + ridge_identity: handles rank deficient transport matrices (this happens + typically when rows/cols in cost/kernel matrices are collinear, or, + equivalently when two points from either measure are close). + kwargs: arguments passed to :class:`~lineax.AbstractLinearSolver` linear + solver. + """ + input_structure = jax.eval_shape(lambda: b) + kwargs.setdefault("rtol", 1e-6) + kwargs.setdefault("atol", 1e-6) + + if ridge_kernel > 0.0 or ridge_identity > 0.0: + lin_reg = lambda x: lin(x) + ridge_kernel * jnp.sum(x) + ridge_identity * x + lin_t_reg = lambda x: lin_t(x) + ridge_kernel * jnp.sum( + x + ) + ridge_identity * x + else: + lin_reg, lin_t_reg = lin, lin_t + + if symmetric: + solver = lx.CG(**kwargs) + fn_operator = lx.FunctionLinearOperator( + lin_reg, input_structure, tags=lx.positive_semidefinite_tag + ) + return lx.linear_solve(fn_operator, b, solver).value + # In the non-symmetric case, use NormalCG by default, but consider + # user defined choice of alternative lx solver. + solver_type = lx.NormalCG if nonsym_solver is None else nonsym_solver + solver = solver_type(**kwargs) + fn_operator = CustomTransposeLinearOperator( + lin_reg, lin_t_reg, input_structure, input_structure + ) + return lx.linear_solve(fn_operator, b, solver).value diff --git a/ott/build/lib/ott/solvers/linear/lr_utils.py b/ott/build/lib/ott/solvers/linear/lr_utils.py new file mode 100644 index 0000000..8ade265 --- /dev/null +++ b/ott/build/lib/ott/solvers/linear/lr_utils.py @@ -0,0 +1,371 @@ +# 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. +from typing import NamedTuple, Optional, Tuple + +import jax +import jax.numpy as jnp +import jax.scipy as jsp + +from ott.math import fixed_point_loop +from ott.problems.linear import linear_problem + +__all__ = ["unbalanced_dykstra_lse", "unbalanced_dykstra_kernel"] + + +class State(NamedTuple): # noqa: D101 + v1: jnp.ndarray + v2: jnp.ndarray + u1: jnp.ndarray + u2: jnp.ndarray + g: jnp.ndarray + err: float + + +class Constants(NamedTuple): # noqa: D101 + a: jnp.ndarray + b: jnp.ndarray + rho_a: float + rho_b: float + supp_a: Optional[jnp.ndarray] = None + supp_b: Optional[jnp.ndarray] = None + + +def unbalanced_dykstra_lse( + c_q: jnp.ndarray, + c_r: jnp.ndarray, + c_g: jnp.ndarray, + gamma: float, + ot_prob: linear_problem.LinearProblem, + translation_invariant: bool = True, + tolerance: float = 1e-3, + min_iter: int = 0, + inner_iter: int = 10, + max_iter: int = 10000 +) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Dykstra's algorithm for the unbalanced + :class:`~ott.solvers.linear.sinkhorn_lr.LRSinkhorn` in LSE mode. + + Args: + c_q: Cost associated with :math:`Q`. + c_r: Cost associated with :math:`R`. + c_g: Cost associated with :math:`g`. + gamma: The (inverse of) the gradient step. + ot_prob: Unbalanced OT problem. + translation_invariant: Whether to use the translation invariant objective, + see :cite:`scetbon:23`, alg. 3. + tolerance: Convergence threshold. + min_iter: Minimum number of iterations. + inner_iter: Compute error every ``inner_iter``. + max_iter: Maximum number of iterations. + + Returns: + The :math:`Q`, :math:`R` and :math:`g` factors. + """ # noqa: D205 + + def _softm( + v: jnp.ndarray, + c: jnp.ndarray, + axis: int, + ) -> jnp.ndarray: + v = jnp.expand_dims(v, axis=1 - axis) + return jsp.special.logsumexp(v + c, axis=axis) + + def _error( + gamma: float, + new_state: State, + old_state: State, + ) -> float: + u1_err = jnp.linalg.norm(new_state.u1 - old_state.u1, ord=jnp.inf) + u2_err = jnp.linalg.norm(new_state.u2 - old_state.u2, ord=jnp.inf) + v1_err = jnp.linalg.norm(new_state.v1 - old_state.v1, ord=jnp.inf) + v2_err = jnp.linalg.norm(new_state.v2 - old_state.v2, ord=jnp.inf) + return (1.0 / gamma) * jnp.max(jnp.array([u1_err, u2_err, v1_err, v2_err])) + + def cond_fn( + iteration: int, + const: Constants, + state: State, + ) -> bool: + del iteration, const + return tolerance < state.err + + def body_fn( + iteration: int, const: Constants, state: State, compute_error: bool + ) -> State: + log_a, log_b = jnp.log(const.a), jnp.log(const.b) + rho_a, rho_b = const.rho_a, const.rho_b + + c_a = _get_ratio(const.rho_a, gamma) + c_b = _get_ratio(const.rho_b, gamma) + + if translation_invariant: + lam_a, lam_b = compute_lambdas(const, state, gamma, g=c_g, lse_mode=True) + + u1 = c_a * (log_a - _softm(state.v1, c_q, axis=1)) + u1 = u1 - lam_a / ((1.0 / gamma) + rho_a) + u2 = c_b * (log_b - _softm(state.v2, c_r, axis=1)) + u2 = u2 - lam_b / ((1.0 / gamma) + rho_b) + + state_lam = State( + v1=state.v1, v2=state.v2, u1=u1, u2=u2, g=state.g, err=state.err + ) + lam_a, lam_b = compute_lambdas( + const, state_lam, gamma, g=c_g, lse_mode=True + ) + + v1_trans = _softm(u1, c_q, axis=0) + v2_trans = _softm(u2, c_r, axis=0) + + g_trans = gamma * (lam_a + lam_b) + c_g + else: + u1 = c_a * (log_a - _softm(state.v1, c_q, axis=1)) + u2 = c_b * (log_b - _softm(state.v2, c_r, axis=1)) + + v1_trans = _softm(u1, c_q, axis=0) + v2_trans = _softm(u2, c_r, axis=0) + g_trans = c_g + + g = (1.0 / 3.0) * (g_trans + v1_trans + v2_trans) + v1 = g - v1_trans + v2 = g - v2_trans + + new_state = State(v1=v1, v2=v2, u1=u1, u2=u2, g=g, err=jnp.inf) + err = jax.lax.cond( + jnp.logical_and(compute_error, iteration >= min_iter), + _error, + lambda *_: state.err, + gamma, + new_state, + state, + ) + return State(v1=v1, v2=v2, u1=u1, u2=u2, g=g, err=err) + + n, m, r = c_q.shape[0], c_r.shape[0], c_g.shape[0] + constants = Constants( + a=ot_prob.a, + b=ot_prob.b, + rho_a=_rho(ot_prob.tau_a), + rho_b=_rho(ot_prob.tau_b), + supp_a=ot_prob.a > 0, + supp_b=ot_prob.b > 0, + ) + init_state = State( + v1=jnp.zeros(r), + v2=jnp.zeros(r), + u1=jnp.zeros(n), + u2=jnp.zeros(m), + g=c_g, + err=jnp.inf, + ) + + state: State = fixed_point_loop.fixpoint_iter_backprop( + cond_fn, body_fn, min_iter, max_iter, inner_iter, constants, init_state + ) + + q = jnp.exp(state.u1[:, None] + c_q + state.v1[None, :]) + r = jnp.exp(state.u2[:, None] + c_r + state.v2[None, :]) + g = jnp.exp(state.g) + + return q, r, g + + +def unbalanced_dykstra_kernel( + k_q: jnp.ndarray, + k_r: jnp.ndarray, + k_g: jnp.ndarray, + gamma: float, + ot_prob: linear_problem.LinearProblem, + translation_invariant: bool = True, + tolerance: float = 1e-3, + min_iter: int = 0, + inner_iter: int = 10, + max_iter: int = 10000 +) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Dykstra's algorithm for the unbalanced + :class:`~ott.solvers.linear.sinkhorn_lr.LRSinkhorn` in kernel mode. + + Args: + k_q: Kernel associated with :math:`Q`. + k_r: Kernel associated with :math:`R`. + k_g: Kernel associated with :math:`g`. + gamma: The (inverse of) the gradient step. + ot_prob: Unbalanced OT problem. + translation_invariant: Whether to use the translation invariant objective, + see :cite:`scetbon:23`, alg. 3. + tolerance: Convergence threshold. + min_iter: Minimum number of iterations. + inner_iter: Compute error every ``inner_iter``. + max_iter: Maximum number of iterations. + + Returns: + The :math:`Q`, :math:`R` and :math:`g` factors. + """ # noqa: D205 + + def _error( + gamma: float, + new_state: State, + old_state: State, + ) -> float: + u1_err = jnp.linalg.norm( + jnp.log(new_state.u1) - jnp.log(old_state.u1), ord=jnp.inf + ) + u2_err = jnp.linalg.norm( + jnp.log(new_state.u2) - jnp.log(old_state.u2), ord=jnp.inf + ) + v1_err = jnp.linalg.norm( + jnp.log(new_state.v1) - jnp.log(old_state.v1), ord=jnp.inf + ) + v2_err = jnp.linalg.norm( + jnp.log(new_state.v2) - jnp.log(old_state.v2), ord=jnp.inf + ) + return (1.0 / gamma) * jnp.max(jnp.array([u1_err, u2_err, v1_err, v2_err])) + + def cond_fn( + iteration: int, + const: Constants, + state: State, + ) -> bool: + del iteration, const + return tolerance < state.err + + def body_fn( + iteration: int, const: Constants, state: State, compute_error: bool + ) -> State: + c_a = _get_ratio(const.rho_a, gamma) + c_b = _get_ratio(const.rho_b, gamma) + + if translation_invariant: + lam_a, lam_b = compute_lambdas(const, state, gamma, g=k_g, lse_mode=False) + + u1 = jnp.where(const.supp_a, (const.a / (k_q @ state.v1)) ** c_a, 0.0) + u1 = u1 * jnp.exp(-lam_a / ((1.0 / gamma) + const.rho_a)) + u2 = jnp.where(const.supp_b, (const.b / (k_r @ state.v2)) ** c_b, 0.0) + u2 = u2 * jnp.exp(-lam_b / ((1.0 / gamma) + const.rho_b)) + + state_lam = State( + v1=state.v1, v2=state.v2, u1=u1, u2=u2, g=state.g, err=state.err + ) + lam_a, lam_b = compute_lambdas( + const, state_lam, gamma, g=k_g, lse_mode=False + ) + + v1_trans = k_q.T @ u1 + v2_trans = k_r.T @ u2 + + k_trans = jnp.exp(gamma * (lam_a + lam_b)) * k_g + g = (k_trans * v1_trans * v2_trans) ** (1.0 / 3.0) + else: + u1 = jnp.where(const.supp_a, (const.a / (k_q @ state.v1)) ** c_a, 0.0) + u2 = jnp.where(const.supp_b, (const.b / (k_r @ state.v2)) ** c_b, 0.0) + + v1_trans = k_q.T @ u1 + v2_trans = k_r.T @ u2 + + g = (k_g * v1_trans * v2_trans) ** (1.0 / 3.0) + + v1 = g / v1_trans + v2 = g / v2_trans + + new_state = State(v1=v1, v2=v2, u1=u1, u2=u2, g=g, err=jnp.inf) + err = jax.lax.cond( + jnp.logical_and(compute_error, iteration >= min_iter), + _error, + lambda *_: state.err, + gamma, + new_state, + state, + ) + return State(v1=v1, v2=v2, u1=u1, u2=u2, g=g, err=err) + + n, m, r = k_q.shape[0], k_r.shape[0], k_g.shape[0] + constants = Constants( + a=ot_prob.a, + b=ot_prob.b, + rho_a=_rho(ot_prob.tau_a), + rho_b=_rho(ot_prob.tau_b), + supp_a=ot_prob.a > 0.0, + supp_b=ot_prob.b > 0.0, + ) + init_state = State( + v1=jnp.ones(r), + v2=jnp.ones(r), + u1=jnp.ones(n), + u2=jnp.ones(m), + g=k_g, + err=jnp.inf + ) + + state: State = fixed_point_loop.fixpoint_iter_backprop( + cond_fn, body_fn, min_iter, max_iter, inner_iter, constants, init_state + ) + + q = state.u1[:, None] * k_q * state.v1[None, :] + r = state.u2[:, None] * k_r * state.v2[None, :] + + return q, r, state.g + + +def compute_lambdas( + const: Constants, state: State, gamma: float, g: jnp.ndarray, *, + lse_mode: bool +) -> Tuple[float, float]: + """TODO.""" + gamma_inv = 1.0 / gamma + rho_a = const.rho_a + rho_b = const.rho_b + + if lse_mode: + num_1 = jsp.special.logsumexp((-gamma_inv / rho_a) * state.u1, b=const.a) + num_2 = jsp.special.logsumexp((-gamma_inv / rho_b) * state.u2, b=const.b) + den = jsp.special.logsumexp(g - (state.v1 + state.v2)) + const_1 = num_1 - den + const_2 = num_2 - den + + ratio_1 = _get_ratio(rho_a, gamma) + ratio_2 = _get_ratio(rho_b, gamma) + harmonic = 1.0 / (1.0 - (ratio_1 * ratio_2)) + lam_1 = harmonic * gamma_inv * ratio_1 * (const_1 - ratio_2 * const_2) + lam_2 = harmonic * gamma_inv * ratio_2 * (const_2 - ratio_1 * const_1) + return lam_1, lam_2 + + num_1 = jnp.sum( + jnp.where( + const.supp_a, ((state.u1 ** (-gamma_inv / rho_a)) * const.a), 0.0 + ) + ) + num_2 = jnp.sum( + jnp.where( + const.supp_b, ((state.u2 ** (-gamma_inv / rho_b)) * const.b), 0.0 + ) + ) + den = jnp.sum(g / (state.v1 * state.v2)) + const_1 = jnp.log(num_1 / den) + const_2 = jnp.log(num_2 / den) + + ratio_1 = _get_ratio(rho_a, gamma) + ratio_2 = _get_ratio(rho_b, gamma) + harmonic = 1.0 / (1.0 - (ratio_1 * ratio_2)) + lam_1 = harmonic * gamma_inv * ratio_1 * (const_1 - ratio_2 * const_2) + lam_2 = harmonic * gamma_inv * ratio_2 * (const_2 - ratio_1 * const_1) + return lam_1, lam_2 + + +def _rho(tau: float) -> float: + tau = jnp.asarray(tau) # avoid division by 0 in Python, get NaN instead + return tau / (1.0 - tau) + + +def _get_ratio(rho: float, gamma: float) -> float: + gamma_inv = 1.0 / gamma + return jnp.where(jnp.isfinite(rho), rho / (rho + gamma_inv), 1.0) diff --git a/ott/build/lib/ott/solvers/linear/sinkhorn.py b/ott/build/lib/ott/solvers/linear/sinkhorn.py new file mode 100644 index 0000000..beb3b61 --- /dev/null +++ b/ott/build/lib/ott/solvers/linear/sinkhorn.py @@ -0,0 +1,1230 @@ +# 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. +from typing import ( + Any, + Callable, + Literal, + Mapping, + NamedTuple, + Optional, + Sequence, + Tuple, + Union, +) + +import jax +import jax.experimental +import jax.numpy as jnp +import jax.scipy as jsp +import numpy as np + +from ott import utils +from ott.geometry import geometry +from ott.initializers.linear import initializers as init_lib +from ott.math import fixed_point_loop +from ott.math import unbalanced_functions as uf +from ott.math import utils as mu +from ott.problems.linear import linear_problem, potentials +from ott.solvers.linear import acceleration +from ott.solvers.linear import implicit_differentiation as implicit_lib + +__all__ = ["Sinkhorn", "SinkhornOutput"] + +ProgressCallbackFn_t = Callable[ + [Tuple[np.ndarray, np.ndarray, np.ndarray, "SinkhornState"]], None] + + +class SinkhornState(NamedTuple): + """Holds the state variables used to solve OT with Sinkhorn.""" + + potentials: Tuple[jnp.ndarray, ...] + errors: Optional[jnp.ndarray] = None + old_fus: Optional[jnp.ndarray] = None + old_mapped_fus: Optional[jnp.ndarray] = None + + def set(self, **kwargs: Any) -> "SinkhornState": + """Return a copy of self, with potential overwrites.""" + return self._replace(**kwargs) + + def solution_error( + self, + ot_prob: linear_problem.LinearProblem, + norm_error: Sequence[int], + *, + lse_mode: bool, + parallel_dual_updates: bool, + recenter: bool, + ) -> jnp.ndarray: + """State dependent function to return error.""" + fu, gv = self.fu, self.gv + if recenter and lse_mode: + fu, gv = self.recenter(fu, gv, ot_prob=ot_prob) + + return solution_error( + fu, + gv, + ot_prob, + norm_error=norm_error, + lse_mode=lse_mode, + parallel_dual_updates=parallel_dual_updates + ) + + def compute_kl_reg_cost( # noqa: D102 + self, ot_prob: linear_problem.LinearProblem, lse_mode: bool + ) -> float: + return compute_kl_reg_cost(self.fu, self.gv, ot_prob, lse_mode) + + def recenter( + self, + f: jnp.ndarray, + g: jnp.ndarray, + ot_prob: linear_problem.LinearProblem, + ) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Re-center dual potentials. + + If the ``ot_prob`` is balanced, the ``f`` potential is zero-centered. + Otherwise, use prop. 2 of :cite:`sejourne:22` re-center the potentials iff + ``tau_a < 1`` and ``tau_b < 1``. + + Args: + f: The first dual potential. + g: The second dual potential. + ot_prob: Linear OT problem. + + Returns: + The centered potentials. + """ + if ot_prob.is_balanced: + # center the potentials for numerical stability + is_finite = jnp.isfinite(f) + shift = jnp.sum(jnp.where(is_finite, f, 0.0)) / jnp.sum(is_finite) + return f - shift, g + shift + + if ot_prob.tau_a == 1.0 or ot_prob.tau_b == 1.0: + # re-centering wasn't done during the lse-step, ignore + return f, g + + rho_a = uf.rho(ot_prob.epsilon, ot_prob.tau_a) + rho_b = uf.rho(ot_prob.epsilon, ot_prob.tau_b) + tau = rho_a * rho_b / (rho_a + rho_b) + + shift = tau * ( + mu.logsumexp(-f / rho_a, b=ot_prob.a) - + mu.logsumexp(-g / rho_b, b=ot_prob.b) + ) + return f + shift, g - shift + + @property + def fu(self) -> jnp.ndarray: + """The first dual potential or scaling.""" + return self.potentials[0] + + @property + def gv(self) -> jnp.ndarray: + """The second dual potential or scaling.""" + return self.potentials[1] + + +def solution_error( + f_u: jnp.ndarray, + g_v: jnp.ndarray, + ot_prob: linear_problem.LinearProblem, + *, + norm_error: Sequence[int], + lse_mode: bool, + parallel_dual_updates: bool, +) -> jnp.ndarray: + """Given two potential/scaling solutions, computes deviation to optimality. + + When the ``ot_prob`` problem is balanced and the usual Sinkhorn updates are + used, this is simply deviation of the coupling's marginal to ``ot_prob.b``. + This is the case because the second (and last) update of the Sinkhorn + algorithm equalizes the row marginal of the coupling to ``ot_prob.a``. To + simplify the logic, this is parameterized by checking whether + `parallel_dual_updates = False`. + + When that flag is `True`, or when the problem is unbalanced, + additional quantities to qualify optimality must be taken into account. + + Args: + f_u: jnp.ndarray, potential or scaling + g_v: jnp.ndarray, potential or scaling + ot_prob: linear OT problem + norm_error: int, p-norm used to compute error. + lse_mode: True if log-sum-exp operations, False if kernel vector products. + parallel_dual_updates: Whether potentials/scalings were computed in + parallel. + + Returns: + a positive number quantifying how far from optimality current solution is. + """ + if ot_prob.is_balanced and not parallel_dual_updates: + return marginal_error( + f_u, g_v, ot_prob.b, ot_prob.geom, 0, norm_error, lse_mode + ) + + # In the unbalanced case, we compute the norm of the gradient. + # the gradient is equal to the marginal of the current plan minus + # the gradient of < z, rho_z(exp^(-h/rho_z) -1> where z is either a or b + # and h is either f or g. Note this is equal to z if rho_z → inf, which + # is the case when tau_z → 1.0 + if lse_mode: + grad_a = uf.grad_of_marginal_fit( + ot_prob.a, f_u, ot_prob.tau_a, ot_prob.epsilon + ) + grad_b = uf.grad_of_marginal_fit( + ot_prob.b, g_v, ot_prob.tau_b, ot_prob.epsilon + ) + else: + u = ot_prob.geom.potential_from_scaling(f_u) + v = ot_prob.geom.potential_from_scaling(g_v) + grad_a = uf.grad_of_marginal_fit( + ot_prob.a, u, ot_prob.tau_a, ot_prob.epsilon + ) + grad_b = uf.grad_of_marginal_fit( + ot_prob.b, v, ot_prob.tau_b, ot_prob.epsilon + ) + err = marginal_error(f_u, g_v, grad_a, ot_prob.geom, 1, norm_error, lse_mode) + err += marginal_error(f_u, g_v, grad_b, ot_prob.geom, 0, norm_error, lse_mode) + return err + + +def marginal_error( + f_u: jnp.ndarray, + g_v: jnp.ndarray, + target: jnp.ndarray, + geom: geometry.Geometry, + axis: int = 0, + norm_error: Sequence[int] = (1,), + lse_mode: bool = True +) -> jnp.asarray: + """Output how far Sinkhorn solution is w.r.t target. + + Args: + f_u: a vector of potentials or scalings for the first marginal. + g_v: a vector of potentials or scalings for the second marginal. + target: target marginal. + geom: Geometry object. + axis: axis (0 or 1) along which to compute marginal. + norm_error: (tuple of int) p's to compute p-norm between marginal/target + lse_mode: whether operating on scalings or potentials + + Returns: + Array of floats, quantifying difference between target / marginal. + """ + if lse_mode: + marginal = geom.marginal_from_potentials(f_u, g_v, axis=axis) + else: + marginal = geom.marginal_from_scalings(f_u, g_v, axis=axis) + norm_error = jnp.asarray(norm_error) + return jnp.sum( + jnp.abs(marginal - target) ** norm_error[:, jnp.newaxis], axis=1 + ) ** (1.0 / norm_error) + + +def compute_kl_reg_cost( + f: jnp.ndarray, g: jnp.ndarray, ot_prob: linear_problem.LinearProblem, + lse_mode: bool +) -> float: + r"""Compute objective of Sinkhorn for OT problem given dual solutions. + + The objective is evaluated for dual solution ``f`` and ``g``, using + information contained in ``ot_prob``. The objective is the regularized + optimal transport cost (i.e. the cost itself plus entropic and unbalanced + terms). Situations where marginals ``a`` or ``b`` in ot_prob have zero + coordinates are reflected in minus infinity entries in their corresponding + dual potentials. To avoid NaN that may result when multiplying 0's by infinity + values, ``jnp.where`` is used to cancel these contributions. + + Args: + f: jnp.ndarray, potential + g: jnp.ndarray, potential + ot_prob: linear optimal transport problem. + lse_mode: bool, whether to compute total mass in lse or kernel mode. + + Returns: + The regularized transport cost. + """ + supp_a = ot_prob.a > 0 + supp_b = ot_prob.b > 0 + fa = ot_prob.geom.potential_from_scaling(ot_prob.a) + if ot_prob.tau_a == 1.0: + div_a = jnp.sum(jnp.where(supp_a, ot_prob.a * (f - fa), 0.0)) + else: + rho_a = uf.rho(ot_prob.epsilon, ot_prob.tau_a) + div_a = -jnp.sum( + jnp.where(supp_a, ot_prob.a * uf.phi_star(-(f - fa), rho_a), 0.0) + ) + + gb = ot_prob.geom.potential_from_scaling(ot_prob.b) + if ot_prob.tau_b == 1.0: + div_b = jnp.sum(jnp.where(supp_b, ot_prob.b * (g - gb), 0.0)) + else: + rho_b = uf.rho(ot_prob.epsilon, ot_prob.tau_b) + div_b = -jnp.sum( + jnp.where(supp_b, ot_prob.b * uf.phi_star(-(g - gb), rho_b), 0.0) + ) + + # Using https://arxiv.org/pdf/1910.12958.pdf (24) + if lse_mode: + total_sum = jnp.sum(ot_prob.geom.marginal_from_potentials(f, g)) + else: + u = ot_prob.geom.scaling_from_potential(f) + v = ot_prob.geom.scaling_from_potential(g) + total_sum = jnp.sum(ot_prob.geom.marginal_from_scalings(u, v)) + return div_a + div_b + ot_prob.epsilon * ( + jnp.sum(ot_prob.a) * jnp.sum(ot_prob.b) - total_sum + ) + + +class SinkhornOutput(NamedTuple): + """Holds the output of a Sinkhorn solver applied to a problem. + + Objects of this class contain both solutions and problem definition of a + regularized OT problem, along several methods that can be used to access its + content, to, for instance, materialize an OT matrix or apply it to a vector + (without having to materialize it when not needed). + + Args: + f: dual variables vector of size ``ot.prob.shape[0]`` returned by Sinkhorn + g: dual variables vector of size ``ot.prob.shape[1]`` returned by Sinkhorn + errors: vector or errors, along iterations. This vector is of size + ``max_iterations // inner_iterations`` where those were the parameters + passed on to the :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. + For each entry indexed at ``i``, ``errors[i]`` can be either a real + non-negative value (meaning the algorithm recorded that error at the + ``i * inner_iterations`` iteration), a ``jnp.inf`` value (meaning the + algorithm computed that iteration but did not compute its error, because, + for instance, ``i < min_iterations // inner_iterations``), or a ``-1``, + meaning that execution was terminated before that iteration, because the + criterion was found to be smaller than ``threshold``. + reg_ot_cost: the regularized optimal transport cost. By default this is + the linear contribution + KL term. See + :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.ent_reg_cost`, + :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.primal_cost` and + :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.dual_cost` for other + objective values. + ot_prob: stores the definition of the OT problem, including geometry, + marginals, unbalanced regularizers, etc. + threshold: convergence threshold used to control the termination of the + algorithm. + converged: whether the output corresponds to a solution whose error is + below the convergence threshold. + inner_iterations: number of iterations that were run between two + computations of errors. + """ + + potentials: Tuple[jnp.ndarray, ...] + errors: Optional[jnp.ndarray] = None + reg_ot_cost: Optional[float] = None + ot_prob: Optional[linear_problem.LinearProblem] = None + threshold: Optional[jnp.ndarray] = None + converged: Optional[bool] = None + inner_iterations: Optional[int] = None + + def set(self, **kwargs: Any) -> "SinkhornOutput": + """Return a copy of self, with potential overwrites.""" + return self._replace(**kwargs) + + def set_cost( # noqa: D102 + self, ot_prob: linear_problem.LinearProblem, lse_mode: bool, + use_danskin: bool + ) -> "SinkhornOutput": + f = jax.lax.stop_gradient(self.f) if use_danskin else self.f + g = jax.lax.stop_gradient(self.g) if use_danskin else self.g + return self.set(reg_ot_cost=compute_kl_reg_cost(f, g, ot_prob, lse_mode)) + + @property + def dual_cost(self) -> jnp.ndarray: + """Return dual transport cost, without considering regularizer.""" + a, b = self.ot_prob.a, self.ot_prob.b + dual_cost = jnp.sum(jnp.where(a > 0.0, a * self.f, 0)) + dual_cost += jnp.sum(jnp.where(b > 0.0, b * self.g, 0)) + return dual_cost + + @property + def primal_cost(self) -> float: + """Return transport cost of current transport solution at geometry.""" + return self.transport_cost_at_geom(other_geom=self.geom) + + @property + def ent_reg_cost(self) -> float: + r"""Entropy regularized cost. + + This outputs + + .. math:: + + \langle P^{\star},C\rangle - \varepsilon H(P^{\star}) + + \rho_a\text{KL}(P^{\star} 1|a) + \rho_b\text{KL}(1^T P^{\star}|b), + + where :math:`P^{\star}, a, b` is the coupling returned by the + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` and the two marginal weight + vectors; :math:`\rho_a=\varepsilon \tau_a / (1-\tau_a)` and + :m ath:`\rho_b=\varepsilon \tau_b / (1-\tau_b)` are obtained when the problem + is unbalanced from parameters ``tau_a`` and ``tau_b``. Note that the last + two terms vanish in the balanced case, when ``tau_a==tau_b==1``. + """ + ent_a = jnp.sum(jsp.special.entr(self.ot_prob.a)) + ent_b = jnp.sum(jsp.special.entr(self.ot_prob.b)) + return self.reg_ot_cost - self.geom.epsilon * (ent_a + ent_b) + + @property + def kl_reg_cost(self) -> float: + r"""KL regularized OT transport cost. + + This outputs + + .. math:: + + \langle P^{\star}, C \rangle + \varepsilon KL(P^{\star},ab^T) + + \rho_a\text{KL}(P^{\star} 1|a) + \rho_b\text{KL}(1^T P^{\star}|b), + + where :math:`P^{\star}, a, b` are the coupling returned by the + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` algorithm and the two + marginal weight vectors, respectively, and + :math:`\rho_a=\varepsilon \tau_a / (1-\tau_a)` and + :math:`\rho_b=\varepsilon \tau_b / (1-\tau_b)` are obtained when the problem + is unbalanced from parameters ``tau_a`` and ``tau_b``. Note that the last + two terms vanish in the balanced case, when ``tau_a==tau_b==1``. This + quantity coincides with :attr:`reg_ot_cost`, which is computed using + dual variables. + """ + return self.reg_ot_cost + + def transport_cost_at_geom( + self, other_geom: geometry.Geometry + ) -> jnp.ndarray: + r"""Return bare transport cost of current solution at any geometry. + + In order to compute cost, we check first if the geometry can be converted + to a low-rank cost geometry in order to speed up computations, without + having to materialize the full cost matrix. If this is not possible, + we resort to instantiating both transport matrix and cost matrix. + + Args: + other_geom: geometry whose cost matrix is used to evaluate the transport + cost. + + Returns: + the transportation cost at :math:`C`, i.e. :math:`\langle P, C \rangle`. + """ + # TODO(cuturi): handle online mode for non Euclidean pointcloud geometries. + # TODO(michalk8): handle SqEucl point cloud is not converted to LRCGeom + if other_geom.can_LRC: + geom = other_geom.to_LRCGeometry() + return jnp.sum(self.apply(geom.cost_1.T) * geom.cost_2.T) + return jnp.sum(self.matrix * other_geom.cost_matrix) + + @property + def geom(self) -> geometry.Geometry: # noqa: D102 + return self.ot_prob.geom + + @property + def a(self) -> jnp.ndarray: # noqa: D102 + return self.ot_prob.a + + @property + def b(self) -> jnp.ndarray: # noqa: D102 + return self.ot_prob.b + + @property + def n_iters(self) -> int: # noqa: D102 + """Returns the total number of iterations that were needed to terminate.""" + return jnp.sum(self.errors != -1) * self.inner_iterations + + @property + def scalings(self) -> Tuple[jnp.ndarray, jnp.ndarray]: # noqa: D102 + u = self.ot_prob.geom.scaling_from_potential(self.f) + v = self.ot_prob.geom.scaling_from_potential(self.g) + return u, v + + @property + def matrix(self) -> jnp.ndarray: + """Transport matrix if it can be instantiated.""" + try: + return self.ot_prob.geom.transport_from_potentials(self.f, self.g) + except ValueError: + return self.ot_prob.geom.transport_from_scalings(*self.scalings) + + @property + def transport_mass(self) -> float: + """Sum of transport matrix.""" + return self.marginal(0).sum() + + def apply( + self, + inputs: jnp.ndarray, + axis: int = 0, + lse_mode: bool = True + ) -> jnp.ndarray: + """Apply the transport to a ndarray; axis=1 for its transpose.""" + geom = self.ot_prob.geom + if lse_mode: + return geom.apply_transport_from_potentials( + self.f, self.g, inputs, axis=axis + ) + u = geom.scaling_from_potential(self.f) + v = geom.scaling_from_potential(self.g) + return geom.apply_transport_from_scalings(u, v, inputs, axis=axis) + + def marginal(self, axis: int) -> jnp.ndarray: # noqa: D102 + return self.ot_prob.geom.marginal_from_potentials(self.f, self.g, axis=axis) + + def cost_at_geom(self, other_geom: geometry.Geometry) -> float: + """Return reg-OT cost for matrix, evaluated at other cost matrix.""" + return ( + jnp.sum(self.matrix * other_geom.cost_matrix) - + self.geom.epsilon * jnp.sum(jax.scipy.special.entr(self.matrix)) + ) + + def to_dual_potentials(self) -> potentials.EntropicPotentials: + """Return the entropic map estimator.""" + return potentials.EntropicPotentials(self.f, self.g, self.ot_prob) + + @property + def f(self) -> jnp.ndarray: + """The first dual potential.""" + return self.potentials[0] + + @property + def g(self) -> jnp.ndarray: + """The second dual potential.""" + return self.potentials[1] + + +@jax.tree_util.register_pytree_node_class +class Sinkhorn: + r"""Sinkhorn solver. + + The Sinkhorn algorithm is a fixed point iteration that solves a regularized + optimal transport (reg-OT) problem between two measures. + The optimization variables are a pair of vectors (called potentials, or + scalings when parameterized as exponential of the former). Calling this + function returns therefore a pair of optimal vectors. In addition to these, + it also returns the objective value achieved by these optimal vectors; + a vector of size ``max_iterations/inner_iterations`` that records the vector + of values recorded to monitor convergence, throughout the execution of the + algorithm (padded with `-1` if convergence happens before), as well as a + boolean to signify whether the algorithm has converged within the number of + iterations specified by the user. + + The reg-OT problem is specified by two measures, of respective sizes ``n`` and + ``m``. From the viewpoint of the ``sinkhorn`` function, these two measures are + only seen through a triplet (``geom``, ``a``, ``b``), where ``geom`` is a + ``Geometry`` object, and ``a`` and ``b`` are weight vectors of respective + sizes ``n`` and ``m``. Starting from two initial values for those potentials + or scalings (both can be defined by the user by passing value in + ``init_dual_a`` or ``init_dual_b``), the Sinkhorn algorithm will use + elementary operations that are carried out by the ``geom`` object. + + Math: + Given a geometry ``geom``, which provides a cost matrix :math:`C` with its + regularization parameter :math:`\varepsilon`, (or a kernel matrix :math:`K`) + the reg-OT problem consists in finding two vectors `f`, `g` of size ``n``, + ``m`` that maximize the following criterion. + + .. math:: + + \arg\max_{f, g}{- \langle a, \phi_a^{*}(-f) \rangle - \langle b, + \phi_b^{*}(-g) \rangle - \varepsilon \langle e^{f/\varepsilon}, + e^{-C/\varepsilon} e^{-g/\varepsilon}} \rangle + + where :math:`\phi_a(z) = \rho_a z(\log z - 1)` is a scaled entropy, and + :math:`\phi_a^{*}(z) = \rho_a e^{z/\varepsilon}`, its Legendre transform. + + That problem can also be written, instead, using positive scaling vectors + `u`, `v` of size ``n``, ``m``, handled with the kernel + :math:`K := e^{-C/\varepsilon}`, + + .. math:: + + \arg\max_{u, v >0} - \langle a,\phi_a^{*}(-\varepsilon\log u) \rangle + + \langle b, \phi_b^{*}(-\varepsilon\log v) \rangle - \langle u, K v \rangle + + Both of these problems corresponds, in their *primal* formulation, to + solving the unbalanced optimal transport problem with a variable matrix + :math:`P` of size ``n`` x ``m``: + + .. math:: + + \arg\min_{P>0} \langle P,C \rangle -\varepsilon \text{KL}(P | ab^T) + + \rho_a \text{KL}(P\mathbf{1}_m | a) + \rho_b \text{KL}(P^T \mathbf{1}_n + | b) + + where :math:`KL` is the generalized Kullback-Leibler divergence. + + The very same primal problem can also be written using a kernel :math:`K` + instead of a cost :math:`C` as well: + + .. math:: + + \arg\min_{P} \varepsilon \text{KL}(P|K) + + \rho_a \text{KL}(P\mathbf{1}_m | a) + + \rho_b \text{KL}(P^T \mathbf{1}_n | b) + + The *original* OT problem taught in linear programming courses is recovered + by using the formulation above relying on the cost :math:`C`, and letting + :math:`\varepsilon \rightarrow 0`, and :math:`\rho_a, \rho_b \rightarrow + \infty`. + In that case the entropy disappears, whereas the :math:`KL` regularization + above become constraints on the marginals of :math:`P`: This results in a + standard min cost flow problem. This problem is not handled for now in this + toolbox, which focuses exclusively on the case :math:`\varepsilon > 0`. + + The *balanced* regularized OT problem is recovered for finite + :math:`\varepsilon > 0` but letting :math:`\rho_a, \rho_b \rightarrow + \infty`. This problem can be shown to be equivalent to a matrix scaling + problem, which can be solved using the Sinkhorn fixed-point algorithm. + To handle the case :math:`\rho_a, \rho_b \rightarrow \infty`, the + ``sinkhorn`` function uses parameters ``tau_a`` and ``tau_b`` equal + respectively to :math:`\rho_a /(\varepsilon + \rho_a)` and + :math:`\rho_b / (\varepsilon + \rho_b)` instead. Setting either of these + parameters to 1 corresponds to setting the corresponding + :math:`\rho_a, \rho_b` to :math:`\infty`. + + The Sinkhorn algorithm solves the reg-OT problem by seeking optimal + :math:`f`, :math:`g` potentials (or alternatively their parameterization + as positive scaling vectors :math:`u`, :math:`v`), rather than solving the + primal problem in :math:`P`. This is mostly for efficiency (potentials and + scalings have a ``n + m`` memory footprint, rather than ``n m`` required + to store `P`). This is also because both problems are, in fact, equivalent, + since the optimal transport :math:`P^{\star}` can be recovered from + optimal potentials :math:`f^{\star}`, :math:`g^{\star}` or scaling + :math:`u^{\star}`, :math:`v^{\star}`, using the geometry's cost or kernel + matrix respectively: + + .. math:: + + P^{\star} = \exp\left(\frac{f^{\star}\mathbf{1}_m^T + \mathbf{1}_n g^{*T}- + C}{\varepsilon}\right) \text{ or } P^{\star} = \text{diag}(u^{\star}) K + \text{diag}(v^{\star}) + + By default, the Sinkhorn algorithm solves this dual problem in :math:`f, g` + or :math:`u, v` using block coordinate ascent, i.e. devising an update for + each :math:`f` and :math:`g` (resp. :math:`u` and :math:`v`) that cancels + their respective gradients, one at a time. These two iterations are repeated + ``inner_iterations`` times, after which the norm of these gradients will be + evaluated and compared with the ``threshold`` value. The iterations are then + repeated as long as that error exceeds ``threshold``. + + Note on Sinkhorn updates: + The boolean flag ``lse_mode`` sets whether the algorithm is run in either: + + - log-sum-exp mode (``lse_mode=True``), in which case it is directly + defined in terms of updates to `f` and `g`, using log-sum-exp + computations. This requires access to the cost matrix :math:`C`, as it is + stored, or possibly computed on the fly by ``geom``. + + - kernel mode (``lse_mode=False``), in which case it will require access + to a matrix vector multiplication operator :math:`z \rightarrow K z`, + where :math:`K` is either instantiated from :math:`C` as + :math:`\exp(-C/\varepsilon)`, or provided directly. In that case, rather + than optimizing on :math:`f` and :math:`g`, it is more convenient to + optimize on their so called scaling formulations, + :math:`u := \exp(f / \varepsilon)` and :math:`v := \exp(g / \varepsilon)`. + While faster (applying matrices is faster than applying ``lse`` repeatedly + over lines), this mode is also less stable numerically, notably for + smaller :math:`\varepsilon`. + + In the source code, the variables ``f_u`` or ``g_v`` can be either regarded + as potentials (real) or scalings (positive) vectors, depending on the choice + of ``lse_mode`` by the user. Once optimization is carried out, we only + return dual variables in potential form, i.e. ``f`` and ``g``. + + In addition to standard Sinkhorn updates, the user can also use heavy-ball + type updates using a ``momentum`` parameter in ]0,2[. We also implement a + strategy that tries to set that parameter adaptively at + ``chg_momentum_from`` iterations, as a function of progress in the error, + as proposed in the literature. + + Another upgrade to the standard Sinkhorn updates provided to the users lies + in using Anderson acceleration. This can be parameterized by setting the + otherwise null ``anderson`` to a positive integer. When selected,the + algorithm will recompute, every ``refresh_anderson_frequency`` (set by + default to 1) an extrapolation of the most recently computed ``anderson`` + iterates. When using that option, notice that differentiation (if required) + can only be carried out using implicit differentiation, and that all + momentum related parameters are ignored. + + The ``parallel_dual_updates`` flag is set to ``False`` by default. In that + setting, ``g_v`` is first updated using the latest values for ``f_u`` and + ``g_v``, before proceeding to update ``f_u`` using that new value for + ``g_v``. When the flag is set to ``True``, both ``f_u`` and ``g_v`` are + updated simultaneously. Note that setting that choice to ``True`` requires + using some form of averaging (e.g. ``momentum=0.5``). Without this, and on + its own ``parallel_dual_updates`` won't work. + + Differentiation: + The optimal solutions ``f`` and ``g`` and the optimal objective + (``reg_ot_cost``) outputted by the Sinkhorn algorithm can be differentiated + w.r.t. relevant inputs ``geom``, ``a`` and ``b``. In the default setting, + implicit differentiation of the optimality conditions (``implicit_diff`` + not equal to ``None``), this has two consequences, treating ``f`` and ``g`` + differently from ``reg_ot_cost``. + + - The termination criterion used to stop Sinkhorn (cancellation of + gradient of objective w.r.t. ``f_u`` and ``g_v``) is used to differentiate + ``f`` and ``g``, given a change in the inputs. These changes are computed + by solving a linear system. The arguments starting with + ``implicit_solver_*`` allow to define the linear solver that is used, and + to control for two types or regularization (we have observed that, + depending on the architecture, linear solves may require higher ridge + parameters to remain stable). The optimality conditions in Sinkhorn can be + analyzed as satisfying a ``z=z'`` condition, which are then + differentiated. It might be beneficial (e.g., as in :cite:`cuturi:20a`) + to use a preconditioning function ``precondition_fun`` to differentiate + instead ``h(z) = h(z')``. + + - The objective ``reg_ot_cost`` returned by Sinkhorn uses the so-called + envelope (or Danskin's) theorem. In that case, because it is assumed that + the gradients of the dual variables ``f_u`` and ``g_v`` w.r.t. dual + objective are zero (reflecting the fact that they are optimal), small + variations in ``f_u`` and ``g_v`` due to changes in inputs (such as + ``geom``, ``a`` and ``b``) are considered negligible. As a result, + ``stop_gradient`` is applied on dual variables ``f_u`` and ``g_v`` when + evaluating the ``reg_ot_cost`` objective. Note that this approach is + `invalid` when computing higher order derivatives. In that case the + ``use_danskin`` flag must be set to ``False``. + + An alternative yet more costly way to differentiate the outputs of the + Sinkhorn iterations is to use unrolling, i.e. reverse mode differentiation + of the Sinkhorn loop. This is possible because Sinkhorn iterations are + wrapped in a custom fixed point iteration loop, defined in + ``fixed_point_loop``, rather than a standard while loop. This is to ensure + the end result of this fixed point loop can also be differentiated, if + needed, using standard JAX operations. To ensure differentiability, + the ``fixed_point_loop.fixpoint_iter_backprop`` loop does checkpointing of + state variables (here ``f_u`` and ``g_v``) every ``inner_iterations``, and + backpropagates automatically, block by block, through blocks of + ``inner_iterations`` at a time. + + Note: + * The Sinkhorn algorithm may not converge within the maximum number of + iterations for possibly several reasons: + + 1. the regularizer (defined as ``epsilon`` in the geometry ``geom`` + object) is too small. Consider either switching to ``lse_mode=True`` + (at the price of a slower execution), increasing ``epsilon``, or, + alternatively, if you are unable or unwilling to increase ``epsilon``, + either increase ``max_iterations`` or ``threshold``. + 2. the probability weights ``a`` and ``b`` do not have the same total + mass, while using a balanced (``tau_a=tau_b=1.0``) setup. + Consider either normalizing ``a`` and ``b``, or set either ``tau_a`` + and/or ``tau_b<1.0``. + 3. OOMs issues may arise when storing either cost or kernel matrices that + are too large in ``geom``. In the case where, the ``geom`` geometry is + a ``PointCloud``, some of these issues might be solved by setting the + ``online`` flag to ``True``. This will trigger a re-computation on the + fly of the cost/kernel matrix. + + * The weight vectors ``a`` and ``b`` can be passed on with coordinates that + have zero weight. This is then handled by relying on simple arithmetic for + ``inf`` values that will likely arise (due to :math:`\log 0` when + ``lse_mode`` is ``True``, or divisions by zero when ``lse_mode`` is + ``False``). Whenever that arithmetic is likely to produce ``NaN`` values + (due to ``-inf * 0``, or ``-inf - -inf``) in the forward pass, we use + ``jnp.where`` conditional statements to carry ``inf`` rather than ``NaN`` + values. In the reverse mode differentiation, the inputs corresponding to + these 0 weights (a location `x`, or a row in the corresponding cost/kernel + matrix), and the weight itself will have ``NaN`` gradient values. This is + reflects that these gradients are undefined, since these points were not + considered in the optimization and have therefore no impact on the output. + + Args: + lse_mode: ``True`` for log-sum-exp computations, ``False`` for kernel + multiplication. + threshold: tolerance used to stop the Sinkhorn iterations. This is + typically the deviation between a target marginal and the marginal of the + current primal solution when either or both tau_a and tau_b are 1.0 + (balanced or semi-balanced problem), or the relative change between two + successive solutions in the unbalanced case. + norm_error: power used to define p-norm of error for marginal/target. + inner_iterations: the Sinkhorn error is not recomputed at each + iteration but every ``inner_iterations`` instead. + min_iterations: the minimum number of Sinkhorn iterations carried + out before the error is computed and monitored. + max_iterations: the maximum number of Sinkhorn iterations. If + ``max_iterations`` is equal to ``min_iterations``, Sinkhorn iterations are + run by default using a :func:`jax.lax.scan` loop rather than a custom, + unroll-able :func:`jax.lax.while_loop` that monitors convergence. + In that case the error is not monitored and the ``converged`` + flag will return ``False`` as a consequence. + momentum: Momentum instance. + anderson: AndersonAcceleration instance. + implicit_diff: instance used to solve implicit differentiation. Unrolls + iterations if None. + parallel_dual_updates: updates potentials or scalings in parallel if True, + sequentially (in Gauss-Seidel fashion) if False. + recenter_potentials: Whether to re-center the dual potentials. + If the problem is balanced, the ``f`` potential is zero-centered for + numerical stability. Otherwise, use the approach of :cite:`sejourne:22` + to achieve faster convergence. Only used when ``lse_mode = True`` and + ``tau_a < 1`` and ``tau_b < 1``. + use_danskin: when ``True``, it is assumed the entropy regularized cost + is evaluated using optimal potentials that are frozen, i.e. whose + gradients have been stopped. This is useful when carrying out first order + differentiation, and is only valid (as with ``implicit_differentiation``) + when the algorithm has converged with a low tolerance. + initializer: how to compute the initial potentials/scalings. This refers to + a few possible classes implemented following the template in + :class:`~ott.initializers.linear.SinkhornInitializer`. + progress_fn: callback function which gets called during the Sinkhorn + iterations, so the user can display the error at each iteration, + e.g., using a progress bar. See :func:`~ott.utils.default_progress_fn` + for a basic implementation. + kwargs_init: keyword arguments when creating the initializer. + """ + + def __init__( + self, + lse_mode: bool = True, + threshold: float = 1e-3, + norm_error: int = 1, + inner_iterations: int = 10, + min_iterations: int = 0, + max_iterations: int = 2000, + momentum: Optional[acceleration.Momentum] = None, + anderson: Optional[acceleration.AndersonAcceleration] = None, + parallel_dual_updates: bool = False, + recenter_potentials: bool = False, + use_danskin: Optional[bool] = None, + implicit_diff: Optional[implicit_lib.ImplicitDiff + ] = implicit_lib.ImplicitDiff(), # noqa: B008 + initializer: Union[Literal["default", "gaussian", "sorting", "subsample"], + init_lib.SinkhornInitializer] = "default", + progress_fn: Optional[ProgressCallbackFn_t] = None, + kwargs_init: Optional[Mapping[str, Any]] = None, + ): + self.lse_mode = lse_mode + self.threshold = threshold + self.inner_iterations = inner_iterations + self.min_iterations = min_iterations + self.max_iterations = max_iterations + self._norm_error = norm_error + self.anderson = anderson + self.implicit_diff = implicit_diff + + if momentum is not None: + self.momentum = acceleration.Momentum( + momentum.start, momentum.error_threshold, momentum.value, + self.inner_iterations + ) + else: + # Use no momentum if using Anderson or unrolling. + if self.anderson is not None or self.implicit_diff is None: + self.momentum = acceleration.Momentum( + inner_iterations=self.inner_iterations + ) + else: + # no momentum + self.momentum = acceleration.Momentum() + + self.parallel_dual_updates = parallel_dual_updates + self.recenter_potentials = recenter_potentials + self.initializer = initializer + self.progress_fn = progress_fn + self.kwargs_init = {} if kwargs_init is None else kwargs_init + + # Force implicit_differentiation to True when using Anderson acceleration, + # Reset all momentum parameters to default (i.e. no momentum) + if anderson: + self.implicit_diff = ( + implicit_lib.ImplicitDiff() + if self.implicit_diff is None else self.implicit_diff + ) + self.momentum = acceleration.Momentum( + inner_iterations=self.inner_iterations + ) + + # By default, use Danskin theorem to differentiate + # the objective when using implicit_lib. + self.use_danskin = ((self.implicit_diff is not None) + if use_danskin is None else use_danskin) + + def __call__( + self, + ot_prob: linear_problem.LinearProblem, + init: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray]] = (None, None), + rng: Optional[jax.Array] = None, + ) -> SinkhornOutput: + """Run Sinkhorn algorithm. + + Args: + ot_prob: Linear OT problem. + init: Initial dual potentials/scalings f_u and g_v, respectively. + Any `None` values will be initialized using the initializer. + rng: Random number generator key for stochastic initialization. + + Returns: + The Sinkhorn output. + """ + rng = utils.default_prng_key(rng) + initializer = self.create_initializer() + init_dual_a, init_dual_b = initializer( + ot_prob, *init, lse_mode=self.lse_mode, rng=rng + ) + return run(ot_prob, self, (init_dual_a, init_dual_b)) + + def lse_step( + self, ot_prob: linear_problem.LinearProblem, state: SinkhornState, + iteration: int + ) -> SinkhornState: + """Sinkhorn LSE update.""" + + def k(tau_i: float, tau_j: float) -> float: + num = -tau_j * (tau_a - 1) * (tau_b - 1) * (tau_i - 1) + denom = (tau_j - 1) * (tau_a * (tau_b - 1) + tau_b * (tau_a - 1)) + return num / denom + + def xi(tau_i: float, tau_j: float) -> float: + k_ij = k(tau_i, tau_j) + return k_ij / (1.0 - k_ij) + + def smin( + potential: jnp.ndarray, marginal: jnp.ndarray, tau: float + ) -> float: + rho = uf.rho(ot_prob.epsilon, tau) + return -rho * mu.logsumexp(-potential / rho, b=marginal) + + # only for an unbalanced problems with `tau_{a,b} < 1` + recenter = ( + self.recenter_potentials and ot_prob.tau_a < 1.0 and ot_prob.tau_b < 1.0 + ) + w = self.momentum.weight(state, iteration) + tau_a, tau_b = ot_prob.tau_a, ot_prob.tau_b + old_fu, old_gv = state.fu, state.gv + + if recenter: + k11, k22 = k(tau_a, tau_a), k(tau_b, tau_b) + xi12, xi21 = xi(tau_a, tau_b), xi(tau_b, tau_a) + + # update g potential + new_gv = tau_b * ot_prob.geom.update_potential( + old_fu, old_gv, jnp.log(ot_prob.b), iteration, axis=0 + ) + if recenter: + new_gv -= k22 * smin(old_fu, ot_prob.a, tau_a) + new_gv += xi21 * smin(new_gv, ot_prob.b, tau_b) + gv = self.momentum(w, old_gv, new_gv, self.lse_mode) + + if not self.parallel_dual_updates: + old_gv = gv + + # update f potential + new_fu = tau_a * ot_prob.geom.update_potential( + old_fu, old_gv, jnp.log(ot_prob.a), iteration, axis=1 + ) + if recenter: + new_fu -= k11 * smin(old_gv, ot_prob.b, tau_b) + new_fu += xi12 * smin(new_fu, ot_prob.a, tau_a) + fu = self.momentum(w, old_fu, new_fu, self.lse_mode) + + return state.set(potentials=(fu, gv)) + + def kernel_step( + self, ot_prob: linear_problem.LinearProblem, state: SinkhornState, + iteration: int + ) -> SinkhornState: + """Sinkhorn multiplicative update.""" + w = self.momentum.weight(state, iteration) + old_gv = state.gv + new_gv = ot_prob.geom.update_scaling( + state.fu, ot_prob.b, iteration, axis=0 + ) ** ot_prob.tau_b + gv = self.momentum(w, state.gv, new_gv, self.lse_mode) + new_fu = ot_prob.geom.update_scaling( + old_gv if self.parallel_dual_updates else gv, + ot_prob.a, + iteration, + axis=1 + ) ** ot_prob.tau_a + fu = self.momentum(w, state.fu, new_fu, self.lse_mode) + return state.set(potentials=(fu, gv)) + + def one_iteration( + self, ot_prob: linear_problem.LinearProblem, state: SinkhornState, + iteration: int, compute_error: bool + ) -> SinkhornState: + """Carries out one Sinkhorn iteration. + + Depending on lse_mode, these iterations can be either in: + + - log-space for numerical stability. + - scaling space, using standard kernel-vector multiply operations. + + Args: + ot_prob: the transport problem definition + state: SinkhornState named tuple. + iteration: the current iteration of the Sinkhorn loop. + compute_error: flag to indicate this iteration computes/stores an error + + Returns: + The updated state. + """ + # When running updates in parallel (Gauss-Seidel mode), old_g_v will be + # used to update f_u, rather than the latest g_v computed in this loop. + # Unused otherwise. + if self.anderson: + state = self.anderson.update(state, iteration, ot_prob, self.lse_mode) + + if self.lse_mode: # In lse_mode, run additive updates. + print("lse step") + state = self.lse_step(ot_prob, state, iteration) + else: + state = self.kernel_step(ot_prob, state, iteration) + + if self.anderson: + state = self.anderson.update_history(state, ot_prob, self.lse_mode) + + # re-computes error if compute_error is True, else set it to inf. + err = jax.lax.cond( + jnp.logical_or( + iteration == self.max_iterations - 1, + jnp.logical_and(compute_error, iteration >= self.min_iterations) + ), + lambda state, prob: state.solution_error( + prob, + self.norm_error, + lse_mode=self.lse_mode, + parallel_dual_updates=self.parallel_dual_updates, + recenter=self.recenter_potentials + )[0], + lambda *_: jnp.inf, + state, + ot_prob, + ) + errors = state.errors.at[iteration // self.inner_iterations, :].set(err) + state = state.set(errors=errors) + + if self.progress_fn is not None: + jax.experimental.io_callback( + self.progress_fn, None, + (iteration, self.inner_iterations, self.max_iterations, state) + ) + return state + + def _converged(self, state: SinkhornState, iteration: int) -> bool: + err = state.errors[iteration // self.inner_iterations - 1, 0] + return jnp.logical_and(iteration > 0, err < self.threshold) + + def _diverged(self, state: SinkhornState, iteration: int) -> bool: + err = state.errors[iteration // self.inner_iterations - 1, 0] + return jnp.logical_not(jnp.isfinite(err)) + + def _continue(self, state: SinkhornState, iteration: int) -> bool: + """Continue while not(converged) and not(diverged).""" + return jnp.logical_and( + jnp.logical_not(self._diverged(state, iteration)), + jnp.logical_not(self._converged(state, iteration)) + ) + + @property + def outer_iterations(self) -> int: + """Upper bound on number of times inner_iterations are carried out. + + This integer can be used to set constant array sizes to track the algorithm + progress, notably errors. + """ + return np.ceil(self.max_iterations / self.inner_iterations).astype(int) + + def init_state( + self, ot_prob: linear_problem.LinearProblem, init: Tuple[jnp.ndarray, + jnp.ndarray] + ) -> SinkhornState: + """Return the initial state of the loop.""" + errors = -jnp.ones((self.outer_iterations, len(self.norm_error))) + state = SinkhornState(init, errors=errors) + return self.anderson.init_maps(ot_prob, state) if self.anderson else state + + def output_from_state( + self, ot_prob: linear_problem.LinearProblem, state: SinkhornState + ) -> SinkhornOutput: + """Create an output from a loop state. + + Note: + When differentiating the regularized OT cost, and assuming Sinkhorn has + run to convergence, Danskin's (or the envelope) + `theorem `_ + :cite:`danskin:67,bertsekas:71` + states that the resulting OT cost as a function of the inputs + (``geometry``, ``a``, ``b``) behaves locally as if the dual optimal + potentials were frozen and did not vary with those inputs. + + Notice this is only valid, as when using ``implicit_differentiation`` + mode, if the Sinkhorn algorithm outputs potentials that are near optimal. + namely when the threshold value is set to a small tolerance. + + The flag ``use_danskin`` controls whether that assumption is made. By + default, that flag is set to the value of ``implicit_differentiation`` if + not specified. If you wish to compute derivatives of order 2 and above, + set ``use_danskin`` to ``False``. + + Args: + ot_prob: the transport problem. + state: a SinkhornState. + + Returns: + A SinkhornOutput. + """ + geom = ot_prob.geom + + f = state.fu if self.lse_mode else geom.potential_from_scaling(state.fu) + g = state.gv if self.lse_mode else geom.potential_from_scaling(state.gv) + if self.recenter_potentials: + f, g = state.recenter(f, g, ot_prob=ot_prob) + + # By convention, the algorithm is said to have converged if the algorithm + # has not nan'ed during iterations (notice some errors might be infinite, + # this convention is used when the error is not recomputed), and if the + # last recorded error is lower than the threshold. Note that this will be + # the case if either the algorithm terminated earlier (in which case the + # last state.errors[-1] = -1 by convention) or if the algorithm carried out + # the maximal number of iterations and its last recorded error (at -1 + # position) is lower than the threshold. + + converged = jnp.logical_and( + jnp.logical_not(jnp.any(jnp.isnan(state.errors))), state.errors[-1] + < self.threshold + )[0] + + return SinkhornOutput((f, g), + errors=state.errors[:, 0], + threshold=jnp.array(self.threshold), + converged=converged, + inner_iterations=self.inner_iterations) + + @property + def norm_error(self) -> Tuple[int, ...]: + """Powers used to compute the p-norm between marginal/target.""" + # To change momentum adaptively, one needs errors in ||.||_1 norm. + # In that case, we add this exponent to the list of errors to compute, + # notably if that was not the error requested by the user. + if self.momentum and self.momentum.start > 0 and self._norm_error != 1: + return self._norm_error, 1 + return self._norm_error, + + # TODO(michalk8): in the future, enforce this (+ in GW) via abstract method + def create_initializer(self) -> init_lib.SinkhornInitializer: # noqa: D102 + if isinstance(self.initializer, init_lib.SinkhornInitializer): + return self.initializer + if self.initializer == "default": + return init_lib.DefaultInitializer() + if self.initializer == "gaussian": + return init_lib.GaussianInitializer() + if self.initializer == "sorting": + return init_lib.SortingInitializer(**self.kwargs_init) + if self.initializer == "subsample": + return init_lib.SubsampleInitializer(**self.kwargs_init) + raise NotImplementedError( + f"Initializer `{self.initializer}` is not yet implemented." + ) + + def tree_flatten(self): # noqa: D102 + aux = vars(self).copy() + aux["norm_error"] = aux.pop("_norm_error") + aux.pop("threshold") + return [self.threshold], aux + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + return cls(**aux_data, threshold=children[0]) + + +def run( + ot_prob: linear_problem.LinearProblem, solver: Sinkhorn, + init: Tuple[jnp.ndarray, ...] +) -> SinkhornOutput: + """Run loop of the solver, outputting a state upgraded to an output.""" + #iter_fun = _iterations_implicit if solver.implicit_diff else iterations + iter_fun = iterations + out = iter_fun(ot_prob, solver, init) + # Be careful here, the geom and the cost are injected at the end, where it + # does not interfere with the implicit differentiation. + out = out.set_cost(ot_prob, solver.lse_mode, solver.use_danskin) + return out.set(ot_prob=ot_prob) + + +def iterations( + ot_prob: linear_problem.LinearProblem, solver: Sinkhorn, + init: Tuple[jnp.ndarray, ...] +) -> SinkhornOutput: + """Jittable Sinkhorn loop. args contain initialization variables.""" + + def cond_fn( + iteration: int, const: Tuple[linear_problem.LinearProblem, Sinkhorn], + state: SinkhornState + ) -> bool: + _, solver = const + return solver._continue(state, iteration) + + def body_fn( + iteration: int, const: Tuple[linear_problem.LinearProblem, Sinkhorn], + state: SinkhornState, compute_error: bool + ) -> SinkhornState: + ot_prob, solver = const + return solver.one_iteration(ot_prob, state, iteration, compute_error) + + # Run the Sinkhorn loop. Choose either a standard fixpoint_iter loop if + # differentiation is implicit, otherwise switch to the backprop friendly + # version of that loop if unrolling to differentiate. + if solver.implicit_diff: + fix_point = fixed_point_loop.fixpoint_iter + else: + fix_point = fixed_point_loop.fixpoint_iter_backprop + + const = ot_prob, solver + state = solver.init_state(ot_prob, init) + state = fix_point( + cond_fn, body_fn, solver.min_iterations, solver.max_iterations, + solver.inner_iterations, const, state + ) + return solver.output_from_state(ot_prob, state) + + +def _iterations_taped( + ot_prob: linear_problem.LinearProblem, solver: Sinkhorn, + init: Tuple[jnp.ndarray, ...] +) -> Tuple[SinkhornOutput, Tuple[jnp.ndarray, jnp.ndarray, + linear_problem.LinearProblem, Sinkhorn]]: + """Run forward pass of the Sinkhorn algorithm storing side information.""" + state = iterations(ot_prob, solver, init) + return state, (state.f, state.g, ot_prob, solver) + + +def _iterations_implicit_bwd(res, gr: SinkhornOutput): + """Run Sinkhorn in backward mode, using implicit differentiation. + + Args: + res: residual data sent from fwd pass, used for computations below. In this + case consists in the output itself, as well as inputs against which we + wish to differentiate. + gr: gradients w.r.t outputs of fwd pass, here w.r.t size f, g, errors. Note + that differentiability w.r.t. errors is not handled, and only f, g is + considered. + + Returns: + a tuple of gradients: PyTree for geom, one jnp.ndarray for each of a and b. + """ + f, g, ot_prob, solver = res + out = solver.implicit_diff.gradient( + ot_prob, f, g, solver.lse_mode, gr.potentials + ) + return *out, None, None + + +# sets threshold, norm_errors, geom, a and b to be differentiable, as those are +# non-static. Only differentiability w.r.t. geom, a and b will be used. +_iterations_implicit = jax.custom_vjp(iterations) +_iterations_implicit.defvjp(_iterations_taped, _iterations_implicit_bwd) diff --git a/ott/build/lib/ott/solvers/linear/sinkhorn_lr.py b/ott/build/lib/ott/solvers/linear/sinkhorn_lr.py new file mode 100644 index 0000000..da949da --- /dev/null +++ b/ott/build/lib/ott/solvers/linear/sinkhorn_lr.py @@ -0,0 +1,822 @@ +# 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. +from typing import ( + Any, + Callable, + Literal, + Mapping, + NamedTuple, + Optional, + Tuple, + Union, +) + +import jax +import jax.experimental +import jax.numpy as jnp +import jax.scipy as jsp +import numpy as np + +from ott.geometry import geometry +from ott.initializers.linear import initializers_lr +from ott.math import fixed_point_loop +from ott.math import utils as mu +from ott.problems.linear import linear_problem +from ott.solvers.linear import lr_utils, sinkhorn + +__all__ = ["LRSinkhorn", "LRSinkhornOutput"] + +ProgressCallbackFn_t = Callable[ + [Tuple[np.ndarray, np.ndarray, np.ndarray, "LRSinkhornState"]], None] + + +class LRSinkhornState(NamedTuple): + """State of the Low Rank Sinkhorn algorithm.""" + q: jnp.ndarray + r: jnp.ndarray + g: jnp.ndarray + gamma: float + costs: jnp.ndarray + errors: jnp.ndarray + crossed_threshold: bool + + def compute_error( # noqa: D102 + self, previous_state: "LRSinkhornState" + ) -> float: + err_q = mu.gen_js(self.q, previous_state.q, c=1.0) + err_r = mu.gen_js(self.r, previous_state.r, c=1.0) + err_g = mu.gen_js(self.g, previous_state.g, c=1.0) + + return ((1.0 / self.gamma) ** 2) * (err_q + err_r + err_g) + + def reg_ot_cost( # noqa: D102 + self, + ot_prob: linear_problem.LinearProblem, + *, + epsilon: float, + use_danskin: bool = False + ) -> float: + """For LR Sinkhorn, this defaults to the primal cost of LR solution.""" + return compute_reg_ot_cost( + self.q, + self.r, + self.g, + ot_prob, + epsilon=epsilon, + use_danskin=use_danskin + ) + + def solution_error( # noqa: D102 + self, ot_prob: linear_problem.LinearProblem, norm_error: Tuple[int, ...] + ) -> jnp.ndarray: + return solution_error(self.q, self.r, ot_prob, norm_error) + + def set(self, **kwargs: Any) -> "LRSinkhornState": + """Return a copy of self, with potential overwrites.""" + return self._replace(**kwargs) + + +def compute_reg_ot_cost( + q: jnp.ndarray, + r: jnp.ndarray, + g: jnp.ndarray, + ot_prob: linear_problem.LinearProblem, + epsilon: float, + use_danskin: bool = False +) -> float: + """Compute the regularized OT cost, here the primal cost of the LR solution. + + Args: + q: first factor of solution + r: second factor of solution + g: weights of solution + ot_prob: linear problem + epsilon: Entropic regularization. + use_danskin: if True, use Danskin's theorem :cite:`danskin:67,bertsekas:71` + to avoid computing the gradient of the cost function. + + Returns: + regularized OT cost, the (primal) transport cost of the low-rank solution. + """ + + def ent(x: jnp.ndarray) -> float: + # generalized entropy + return jnp.sum(jsp.special.entr(x) + x) + + tau_a, tau_b = ot_prob.tau_a, ot_prob.tau_b + + q = jax.lax.stop_gradient(q) if use_danskin else q + r = jax.lax.stop_gradient(r) if use_danskin else r + g = jax.lax.stop_gradient(g) if use_danskin else g + + cost = jnp.sum(ot_prob.geom.apply_cost(r, axis=1) * q * (1.0 / g)[None, :]) + cost -= epsilon * (ent(q) + ent(r) + ent(g)) + if tau_a != 1.0: + cost += tau_a / (1.0 - tau_a) * mu.gen_kl(jnp.sum(q, axis=1), ot_prob.a) + if tau_b != 1.0: + cost += tau_b / (1.0 - tau_b) * mu.gen_kl(jnp.sum(r, axis=1), ot_prob.b) + + return cost + + +def solution_error( + q: jnp.ndarray, r: jnp.ndarray, ot_prob: linear_problem.LinearProblem, + norm_error: Tuple[int, ...] +) -> jnp.ndarray: + """Compute solution error. + + Since only balanced case is available for LR, this is marginal deviation. + + Args: + q: first factor of solution. + r: second factor of solution. + ot_prob: linear problem. + norm_error: int, p-norm used to compute error. + + Returns: + one or possibly many numbers quantifying deviation to true marginals. + """ + norm_error = jnp.array(norm_error) + # Update the error + err = jnp.sum( + jnp.abs(jnp.sum(q, axis=1) - ot_prob.a) ** norm_error[:, None], axis=1 + ) ** (1.0 / norm_error) + err += jnp.sum( + jnp.abs(jnp.sum(r, axis=1) - ot_prob.b) ** norm_error[:, None], axis=1 + ) ** (1.0 / norm_error) + err += jnp.sum( + jnp.abs(jnp.sum(q, axis=0) - jnp.sum(r, axis=0)) ** norm_error[:, None], + axis=1 + ) ** (1.0 / norm_error) + + return err + + +class LRSinkhornOutput(NamedTuple): + """Transport interface for a low-rank Sinkhorn solution.""" + + q: jnp.ndarray + r: jnp.ndarray + g: jnp.ndarray + costs: jnp.ndarray + # TODO(michalk8): must be called `errors`, because of `store_inner_errors` + # in future, enforce via class hierarchy + errors: jnp.ndarray + ot_prob: linear_problem.LinearProblem + epsilon: float + inner_iterations: int + # TODO(michalk8): Optional is an artifact of the current impl., refactor + reg_ot_cost: Optional[float] = None + + def set(self, **kwargs: Any) -> "LRSinkhornOutput": + """Return a copy of self, with potential overwrites.""" + return self._replace(**kwargs) + + def set_cost( # noqa: D102 + self, + ot_prob: linear_problem.LinearProblem, + lse_mode: bool, + use_danskin: bool = False + ) -> "LRSinkhornOutput": + del lse_mode + return self.set(reg_ot_cost=self.compute_reg_ot_cost(ot_prob, use_danskin)) + + def compute_reg_ot_cost( # noqa: D102 + self, + ot_prob: linear_problem.LinearProblem, + use_danskin: bool = False, + ) -> float: + return compute_reg_ot_cost( + self.q, + self.r, + self.g, + ot_prob, + epsilon=self.epsilon, + use_danskin=use_danskin + ) + + @property + def geom(self) -> geometry.Geometry: # noqa: D102 + return self.ot_prob.geom + + @property + def a(self) -> jnp.ndarray: # noqa: D102 + return self.ot_prob.a + + @property + def b(self) -> jnp.ndarray: # noqa: D102 + return self.ot_prob.b + + @property + def n_iters(self) -> int: # noqa: D102 + return jnp.sum(self.errors != -1) * self.inner_iterations + + @property + def converged(self) -> bool: # noqa: D102 + return jnp.logical_and( + jnp.any(self.costs == -1), jnp.all(jnp.isfinite(self.costs)) + ) + + @property + def matrix(self) -> jnp.ndarray: + """Transport matrix if it can be instantiated.""" + return (self.q * self._inv_g) @ self.r.T + + def apply(self, inputs: jnp.ndarray, axis: int = 0) -> jnp.ndarray: + """Apply the transport to a array; axis=1 for its transpose.""" + q, r = (self.q, self.r) if axis == 1 else (self.r, self.q) + # for `axis=0`: (batch, m), (m, r), (r,), (r, n) + return ((inputs @ r) * self._inv_g) @ q.T + + def marginal(self, axis: int) -> jnp.ndarray: # noqa: D102 + length = self.q.shape[0] if axis == 0 else self.r.shape[0] + return self.apply(jnp.ones(length,), axis=axis) + + def cost_at_geom(self, other_geom: geometry.Geometry) -> float: + """Return OT cost for current solution, evaluated at any cost matrix.""" + return jnp.sum(self.q * other_geom.apply_cost(self.r, axis=1) * self._inv_g) + + def transport_cost_at_geom(self, other_geom: geometry.Geometry) -> float: + """Return (by recomputing it) bare transport cost of current solution.""" + return self.cost_at_geom(other_geom) + + @property + def primal_cost(self) -> float: + """Return (by recomputing it) transport cost of current solution.""" + return self.transport_cost_at_geom(other_geom=self.geom) + + @property + def transport_mass(self) -> float: + """Sum of transport matrix.""" + return self.marginal(0).sum() + + @property + def _inv_g(self) -> jnp.ndarray: + return 1.0 / self.g + + +@jax.tree_util.register_pytree_node_class +class LRSinkhorn(sinkhorn.Sinkhorn): + r"""Low-Rank Sinkhorn solver for linear reg-OT problems. + + The algorithm is described in :cite:`scetbon:21` and the implementation + contained here is adapted from `LOT `_. + + The algorithm minimizes a non-convex problem. It therefore requires special + care to initialization and convergence. Convergence is evaluated on successive + evaluations of the objective. + + Args: + rank: Rank constraint on the coupling to minimize the linear OT problem + gamma: The (inverse of) gradient step size used by mirror descent. + gamma_rescale: Whether to rescale :math:`\gamma` every iteration as + described in :cite:`scetbon:22b`. + epsilon: Entropic regularization added on top of low-rank problem. + initializer: How to initialize the :math:`Q`, :math:`R` and :math:`g` + factors. + lse_mode: Whether to run computations in LSE or kernel mode. + inner_iterations: Number of inner iterations used by the algorithm before + re-evaluating progress. + use_danskin: Use Danskin theorem to evaluate gradient of objective w.r.t. + input parameters. Only `True` handled at this moment. + progress_fn: callback function which gets called during the Sinkhorn + iterations, so the user can display the error at each iteration, + e.g., using a progress bar. See :func:`~ott.utils.default_progress_fn` + for a basic implementation. + kwargs_dys: Keyword arguments passed to :meth:`dykstra_update_lse`, + :meth:`dykstra_update_kernel` or one of the functions defined in + :mod:`ott.solvers.linear`, depending on whether the problem + is balanced and on the ``lse_mode``. + kwargs_init: Keyword arguments for + :class:`~ott.initializers.linear.initializers_lr.LRInitializer`. + kwargs: Keyword arguments for + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. + """ + + def __init__( + self, + rank: int, + gamma: float = 10.0, + gamma_rescale: bool = True, + epsilon: float = 0.0, + initializer: Union[Literal["random", "rank2", "k-means", + "generalized-k-means"], + initializers_lr.LRInitializer] = "random", + lse_mode: bool = True, + inner_iterations: int = 10, + use_danskin: bool = True, + kwargs_dys: Optional[Mapping[str, Any]] = None, + kwargs_init: Optional[Mapping[str, Any]] = None, + progress_fn: Optional[ProgressCallbackFn_t] = None, + **kwargs: Any, + ): + kwargs["implicit_diff"] = None # not yet implemented + super().__init__( + lse_mode=lse_mode, + inner_iterations=inner_iterations, + use_danskin=use_danskin, + **kwargs + ) + self.rank = rank + self.gamma = gamma + self.gamma_rescale = gamma_rescale + self.epsilon = epsilon + self.initializer = initializer + self.progress_fn = progress_fn + # can be `None` + self.kwargs_dys = {} if kwargs_dys is None else kwargs_dys + self.kwargs_init = {} if kwargs_init is None else kwargs_init + + def __call__( + self, + ot_prob: linear_problem.LinearProblem, + init: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray], + Optional[jnp.ndarray]] = (None, None, None), + rng: Optional[jax.Array] = None, + **kwargs: Any, + ) -> LRSinkhornOutput: + """Run low-rank Sinkhorn. + + Args: + ot_prob: Linear OT problem. + init: Initial values for the low-rank factors: + + - :attr:`~ott.solvers.linear.sinkhorn_lr.LRSinkhornOutput.q`. + - :attr:`~ott.solvers.linear.sinkhorn_lr.LRSinkhornOutput.r`. + - :attr:`~ott.solvers.linear.sinkhorn_lr.LRSinkhornOutput.g`. + + Any `None` values will be initialized using the initializer. + rng: Random key for seeding. + kwargs: Additional arguments when calling the initializer. + + Returns: + The low-rank Sinkhorn output. + """ + initializer = self.create_initializer(ot_prob) + init = initializer(ot_prob, *init, rng=rng, **kwargs) + return run(ot_prob, self, init) + + def _get_costs( + self, + ot_prob: linear_problem.LinearProblem, + state: LRSinkhornState, + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, float]: + log_q, log_r, log_g = ( + mu.safe_log(state.q), mu.safe_log(state.r), mu.safe_log(state.g) + ) + + inv_g = 1.0 / state.g[None, :] + tmp = ot_prob.geom.apply_cost(state.r, axis=1) + + grad_q = tmp * inv_g + grad_r = ot_prob.geom.apply_cost(state.q) * inv_g + grad_g = -jnp.sum(state.q * tmp, axis=0) / (state.g ** 2) + + grad_q += self.epsilon * log_q + grad_r += self.epsilon * log_r + grad_g += self.epsilon * log_g + + if self.gamma_rescale: + norm_q = jnp.max(jnp.abs(grad_q)) ** 2 + norm_r = jnp.max(jnp.abs(grad_r)) ** 2 + norm_g = jnp.max(jnp.abs(grad_g)) ** 2 + gamma = self.gamma / jnp.max(jnp.array([norm_q, norm_r, norm_g])) + else: + gamma = self.gamma + + eps_factor = 1.0 / (self.epsilon * gamma + 1.0) + gamma *= eps_factor + + c_q = -gamma * grad_q + eps_factor * log_q + c_r = -gamma * grad_r + eps_factor * log_r + c_g = -gamma * grad_g + eps_factor * log_g + + return c_q, c_r, c_g, gamma + + # TODO(michalk8): move to `lr_utils` when refactoring this + def dykstra_update_lse( + self, + c_q: jnp.ndarray, + c_r: jnp.ndarray, + h: jnp.ndarray, + gamma: float, + ot_prob: linear_problem.LinearProblem, + min_entry_value: float = 1e-6, + tolerance: float = 1e-3, + min_iter: int = 0, + inner_iter: int = 10, + max_iter: int = 10000 + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Run Dykstra's algorithm.""" + # shortcuts for problem's definition. + r = self.rank + n, m = ot_prob.geom.shape + loga, logb = jnp.log(ot_prob.a), jnp.log(ot_prob.b) + + h_old = h + g1_old, g2_old = jnp.zeros(r), jnp.zeros(r) + f1, f2 = jnp.zeros(n), jnp.zeros(m) + + w_gi, w_gp = jnp.zeros(r), jnp.zeros(r) + w_q, w_r = jnp.zeros(r), jnp.zeros(r) + err = jnp.inf + state_inner = f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err + constants = c_q, c_r, loga, logb + + def cond_fn( + iteration: int, constants: Tuple[jnp.ndarray, ...], + state_inner: Tuple[jnp.ndarray, ...] + ) -> bool: + del iteration, constants + *_, err = state_inner + return err > tolerance + + def _softm( + f: jnp.ndarray, g: jnp.ndarray, c: jnp.ndarray, axis: int + ) -> jnp.ndarray: + return jsp.special.logsumexp( + gamma * (f[:, None] + g[None, :] - c), axis=axis + ) + + def body_fn( + iteration: int, constants: Tuple[jnp.ndarray, ...], + state_inner: Tuple[jnp.ndarray, ...], compute_error: bool + ) -> Tuple[jnp.ndarray, ...]: + # TODO(michalk8): in the future, use `NamedTuple` + f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err = state_inner + c_q, c_r, loga, logb = constants + + # First Projection + f1 = jnp.where( + jnp.isfinite(loga), + (loga - _softm(f1, g1_old, c_q, axis=1)) / gamma + f1, loga + ) + f2 = jnp.where( + jnp.isfinite(logb), + (logb - _softm(f2, g2_old, c_r, axis=1)) / gamma + f2, logb + ) + + h = h_old + w_gi + h = jnp.maximum(jnp.log(min_entry_value) / gamma, h) + w_gi += h_old - h + h_old = h + + # Update couplings + g_q = _softm(f1, g1_old, c_q, axis=0) + g_r = _softm(f2, g2_old, c_r, axis=0) + + # Second Projection + h = (1.0 / 3.0) * (h_old + w_gp + w_q + w_r) + h += g_q / (3.0 * gamma) + h += g_r / (3.0 * gamma) + g1 = h + g1_old - g_q / gamma + g2 = h + g2_old - g_r / gamma + + w_q = w_q + g1_old - g1 + w_r = w_r + g2_old - g2 + w_gp = h_old + w_gp - h + + q, r, _ = recompute_couplings(f1, g1, c_q, f2, g2, c_r, h, gamma) + + g1_old = g1 + g2_old = g2 + h_old = h + + err = jax.lax.cond( + jnp.logical_and(compute_error, iteration >= min_iter), + lambda: solution_error(q, r, ot_prob, self.norm_error)[0], lambda: err + ) + + return f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err + + def recompute_couplings( + f1: jnp.ndarray, + g1: jnp.ndarray, + c_q: jnp.ndarray, + f2: jnp.ndarray, + g2: jnp.ndarray, + c_r: jnp.ndarray, + h: jnp.ndarray, + gamma: float, + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + q = jnp.exp(gamma * (f1[:, None] + g1[None, :] - c_q)) + r = jnp.exp(gamma * (f2[:, None] + g2[None, :] - c_r)) + g = jnp.exp(gamma * h) + return q, r, g + + state_inner = fixed_point_loop.fixpoint_iter_backprop( + cond_fn, body_fn, min_iter, max_iter, inner_iter, constants, state_inner + ) + + f1, f2, g1_old, g2_old, h_old, _, _, _, _, _ = state_inner + return recompute_couplings(f1, g1_old, c_q, f2, g2_old, c_r, h_old, gamma) + + def dykstra_update_kernel( + self, + k_q: jnp.ndarray, + k_r: jnp.ndarray, + k_g: jnp.ndarray, + gamma: float, + ot_prob: linear_problem.LinearProblem, + min_entry_value: float = 1e-6, + tolerance: float = 1e-3, + min_iter: int = 0, + inner_iter: int = 10, + max_iter: int = 10000 + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Run Dykstra's algorithm.""" + # shortcuts for problem's definition. + rank = self.rank + n, m = ot_prob.geom.shape + a, b = ot_prob.a, ot_prob.b + supp_a, supp_b = a > 0, b > 0 + + g_old = k_g + v1_old, v2_old = jnp.ones(rank), jnp.ones(rank) + u1, u2 = jnp.ones(n), jnp.ones(m) + + q_gi, q_gp = jnp.ones(rank), jnp.ones(rank) + q_q, q_r = jnp.ones(rank), jnp.ones(rank) + err = jnp.inf + state_inner = u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err + constants = k_q, k_r, k_g, a, b + + def cond_fn( + iteration: int, constants: Tuple[jnp.ndarray, ...], + state_inner: Tuple[jnp.ndarray, ...] + ) -> bool: + del iteration, constants + *_, err = state_inner + return err > tolerance + + def body_fn( + iteration: int, constants: Tuple[jnp.ndarray, ...], + state_inner: Tuple[jnp.ndarray, ...], compute_error: bool + ) -> Tuple[jnp.ndarray, ...]: + # TODO(michalk8): in the future, use `NamedTuple` + u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err = state_inner + k_q, k_r, k_g, a, b = constants + + # First Projection + u1 = jnp.where(supp_a, a / jnp.dot(k_q, v1_old), 0.0) + u2 = jnp.where(supp_b, b / jnp.dot(k_r, v2_old), 0.0) + g = jnp.maximum(min_entry_value, g_old * q_gi) + q_gi = (g_old * q_gi) / g + g_old = g + + # Second Projection + v1_trans = jnp.dot(k_q.T, u1) + v2_trans = jnp.dot(k_r.T, u2) + g = (g_old * q_gp * v1_old * q_q * v1_trans * v2_old * q_r * + v2_trans) ** (1 / 3) + v1 = g / v1_trans + v2 = g / v2_trans + q_gp = (g_old * q_gp) / g + q_q = (q_q * v1_old) / v1 + q_r = (q_r * v2_old) / v2 + v1_old = v1 + v2_old = v2 + g_old = g + + # Compute Couplings + q, r, _ = recompute_couplings(u1, v1, k_q, u2, v2, k_r, g) + + err = jax.lax.cond( + jnp.logical_and(compute_error, iteration >= min_iter), + lambda: solution_error(q, r, ot_prob, self.norm_error)[0], lambda: err + ) + + return u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err + + def recompute_couplings( + u1: jnp.ndarray, + v1: jnp.ndarray, + k_q: jnp.ndarray, + u2: jnp.ndarray, + v2: jnp.ndarray, + k_r: jnp.ndarray, + g: jnp.ndarray, + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + q = u1.reshape((-1, 1)) * k_q * v1.reshape((1, -1)) + r = u2.reshape((-1, 1)) * k_r * v2.reshape((1, -1)) + return q, r, g + + state_inner = fixed_point_loop.fixpoint_iter_backprop( + cond_fn, body_fn, min_iter, max_iter, inner_iter, constants, state_inner + ) + + u1, u2, v1_old, v2_old, g_old, _, _, _, _, _ = state_inner + return recompute_couplings(u1, v1_old, k_q, u2, v2_old, k_r, g_old) + + def lse_step( + self, ot_prob: linear_problem.LinearProblem, state: LRSinkhornState, + iteration: int + ) -> LRSinkhornState: + """LR Sinkhorn LSE update.""" + c_q, c_r, c_g, gamma = self._get_costs(ot_prob, state) + + if ot_prob.is_balanced: + c_q, c_r, h = c_q / -gamma, c_r / -gamma, c_g / gamma + q, r, g = self.dykstra_update_lse( + c_q, c_r, h, gamma, ot_prob, **self.kwargs_dys + ) + else: + q, r, g = lr_utils.unbalanced_dykstra_lse( + c_q, c_r, c_g, gamma, ot_prob, **self.kwargs_dys + ) + return state.set(q=q, g=g, r=r, gamma=gamma) + + def kernel_step( + self, ot_prob: linear_problem.LinearProblem, state: LRSinkhornState, + iteration: int + ) -> LRSinkhornState: + """LR Sinkhorn Kernel update.""" + c_q, c_r, c_g, gamma = self._get_costs(ot_prob, state) + c_q, c_r, c_g = jnp.exp(c_q), jnp.exp(c_r), jnp.exp(c_g) + + if ot_prob.is_balanced: + q, r, g = self.dykstra_update_kernel( + c_q, c_r, c_g, gamma, ot_prob, **self.kwargs_dys + ) + else: + q, r, g = lr_utils.unbalanced_dykstra_kernel( + c_q, c_r, c_g, gamma, ot_prob, **self.kwargs_dys + ) + return state.set(q=q, g=g, r=r, gamma=gamma) + + def one_iteration( + self, ot_prob: linear_problem.LinearProblem, state: LRSinkhornState, + iteration: int, compute_error: bool + ) -> LRSinkhornState: + """Carries out one low-rank Sinkhorn iteration. + + Depending on lse_mode, these iterations can be either in: + + - log-space for numerical stability. + - scaling space, using standard kernel-vector multiply operations. + + Args: + ot_prob: the transport problem definition + state: LRSinkhornState named tuple. + iteration: the current iteration of the Sinkhorn outer loop. + compute_error: flag to indicate this iteration computes/stores an error + + Returns: + The updated state. + """ + previous_state = state + it = iteration // self.inner_iterations + if self.lse_mode: # In lse_mode, run additive updates. + state = self.lse_step(ot_prob, state, iteration) + else: + state = self.kernel_step(ot_prob, state, iteration) + + # re-computes error if compute_error is True, else set it to inf. + cost = jax.lax.cond( + jnp.logical_and(compute_error, iteration >= self.min_iterations), + lambda: state.reg_ot_cost(ot_prob, epsilon=self.epsilon), + lambda: jnp.inf + ) + error = state.compute_error(previous_state) + crossed_threshold = jnp.logical_or( + state.crossed_threshold, + jnp.logical_and( + state.errors[it - 1] >= self.threshold, error < self.threshold + ) + ) + + state = state.set( + costs=state.costs.at[it].set(cost), + errors=state.errors.at[it].set(error), + crossed_threshold=crossed_threshold, + ) + + if self.progress_fn is not None: + jax.experimental.io_callback( + self.progress_fn, None, + (iteration, self.inner_iterations, self.max_iterations, state) + ) + + return state + + @property + def norm_error(self) -> Tuple[int]: # noqa: D102 + return self._norm_error, + + def create_initializer( + self, prob: linear_problem.LinearProblem + ) -> initializers_lr.LRInitializer: + """Create a low-rank Sinkhorn initializer. + + Args: + prob: Linear OT problem used to determine the initializer. + + Returns: + Low-rank initializer. + """ + if isinstance(self.initializer, initializers_lr.LRInitializer): + assert self.initializer.rank == self.rank, \ + f"Expected initializer's rank to be `{self.rank}`," \ + f"found `{self.initializer.rank}`." + return self.initializer + return initializers_lr.LRInitializer.from_solver( + self, kind=self.initializer, **self.kwargs_init + ) + + def init_state( + self, ot_prob: linear_problem.LinearProblem, + init: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] + ) -> LRSinkhornState: + """Return the initial state of the loop.""" + q, r, g = init + return LRSinkhornState( + q=q, + r=r, + g=g, + gamma=self.gamma, + costs=-jnp.ones(self.outer_iterations), + errors=-jnp.ones(self.outer_iterations), + crossed_threshold=False, + ) + + def output_from_state( + self, ot_prob: linear_problem.LinearProblem, state: LRSinkhornState + ) -> LRSinkhornOutput: + """Create an output from a loop state. + + Args: + ot_prob: the transport problem. + state: a LRSinkhornState. + + Returns: + A LRSinkhornOutput. + """ + return LRSinkhornOutput( + q=state.q, + r=state.r, + g=state.g, + ot_prob=ot_prob, + costs=state.costs, + errors=state.errors, + epsilon=self.epsilon, + inner_iterations=self.inner_iterations, + ) + + def _converged(self, state: LRSinkhornState, iteration: int) -> bool: + + def conv_crossed(prev_err: float, curr_err: float) -> bool: + return jnp.logical_and( + prev_err < self.threshold, curr_err < self.threshold + ) + + def conv_not_crossed(prev_err: float, curr_err: float) -> bool: + return jnp.logical_and(curr_err < prev_err, curr_err < self.threshold) + + # for convergence error, we consider 2 possibilities: + # 1. we either crossed the convergence threshold; in this case we require + # that the previous error was also below the threshold + # 2. we haven't crossed the threshold; in this case, we can be below or + # above the threshold: + # if we're above, we wait until we reach the convergence threshold and + # then, the above condition applies + # if we're below and we improved w.r.t. the previous iteration, + # we have converged; otherwise we continue, since we may be stuck + # in a local minimum (e.g., during the initial iterations) + + it = iteration // self.inner_iterations + return jax.lax.cond( + state.crossed_threshold, conv_crossed, conv_not_crossed, + state.errors[it - 2], state.errors[it - 1] + ) + + def _diverged(self, state: LRSinkhornState, iteration: int) -> bool: + it = iteration // self.inner_iterations + return jnp.logical_and( + jnp.logical_not(jnp.isfinite(state.errors[it - 1])), + jnp.logical_not(jnp.isfinite(state.costs[it - 1])) + ) + + +def run( + ot_prob: linear_problem.LinearProblem, + solver: LRSinkhorn, + init: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray], + Optional[jnp.ndarray]], +) -> LRSinkhornOutput: + """Run loop of the solver, outputting a state upgraded to an output.""" + out = sinkhorn.iterations(ot_prob, solver, init) + out = out.set_cost( + ot_prob, lse_mode=solver.lse_mode, use_danskin=solver.use_danskin + ) + return out.set(ot_prob=ot_prob) diff --git a/ott/build/lib/ott/solvers/linear/univariate.py b/ott/build/lib/ott/solvers/linear/univariate.py new file mode 100644 index 0000000..1bb7b13 --- /dev/null +++ b/ott/build/lib/ott/solvers/linear/univariate.py @@ -0,0 +1,346 @@ +# 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. +from typing import NamedTuple, Optional, Tuple, Union + +import jax +import jax.numpy as jnp + +from ott import utils +from ott.geometry import costs, pointcloud +from ott.math import utils as mu +from ott.problems.linear import linear_problem + +__all__ = [ + "UnivariateOutput", "UnivariateSolver", "uniform_distance", + "quantile_distance" +] + +Distance_t = Tuple[float, Optional[jnp.ndarray], Optional[jnp.ndarray]] + + +class UnivariateOutput(NamedTuple): # noqa: D101 + """Output of the :class:`~ott.solvers.linear.UnivariateSolver`. + + Objects of this class contain both solutions and problem definition of a + univariate OT problem. + + Args: + prob: OT problem between 2 weighted ``[n, d]`` and ``[m, d]`` point clouds. + ot_costs: ``[d,]`` optimal transport cost values, computed independently + along each of the ``d`` slices. + paired_indices: ``None`` if no transport was computed / recorded (e.g. when + using quantiles or subsampling approximations). Otherwise, output a tensor + of shape ``[d, 2, m+n]``, of ``m+n`` pairs of indices, for which the + optimal transport assigns mass, on each slice of the ``d`` slices + described in the dataset. Namely, for each index ``0<=k jnp.ndarray: + """Outputs a ``[d, n, m]`` tensor of all ``[n, m]`` transport matrices. + + This tensor will be extremely sparse, since it will have at most ``d(n+m)`` + non-zero values, out of ``dnm`` total entries. + """ + assert self.paired_indices is not None, \ + "[d, n, m] tensor of transports cannot be computed, likely because an" \ + " approximate method was used (using either subsampling or quantiles)." + + n, m = self.prob.geom.shape + if self.prob.is_equal_size and self.prob.is_uniform: + transport_matrices_from_indices = jax.vmap( + lambda idx, idy: jnp.eye(n)[idx, :][:, idy].T, in_axes=[0, 0] + ) + return transport_matrices_from_indices( + self.paired_indices[:, 0, :], self.paired_indices[:, 1, :] + ) + + # raveled indexing of entries. + indices = self.paired_indices[:, 0] * m + self.paired_indices[:, 1] + # segment sum is needed to collect several contributions + return jax.vmap( + lambda idx, mass: jax.ops.segment_sum( + mass, idx, indices_are_sorted=True, num_segments=n * m + ).reshape(n, m), + in_axes=[0, 0] + )(indices, self.mass_paired_indices) + + @property + def mean_transport_matrix(self) -> jnp.ndarray: + """Return the mean transport matrix, averaged over slices.""" + return jnp.mean(self.transport_matrices, axis=0) + + +@jax.tree_util.register_pytree_node_class +class UnivariateSolver: + r"""Univariate solver to compute 1D OT distance over slices of data. + + Computes 1-Dimensional optimal transport distance between two $d$-dimensional + point clouds. The total distance is the sum of univariate Wasserstein + distances on the $d$ slices of data: given two weighted point-clouds, stored + as ``[n, d]`` and ``[m, d]`` in a + :class:`~ott.problems.linear.linear_problem.LinearProblem` object, with + respective weights ``a`` and ``b``, the solver + computes ``d`` OT distances between each of these ``[n, 1]`` and ``[m, 1]`` + slices. The distance is computed using the analytical formula by default, + which involves sorting each of the slices independently. The optimal transport + matrices are also outputted when possible (described in sparse form, i.e. + pairs of indices and mass transferred between those indices). + + When weights ``a`` and ``b`` are uniform, and ``n=m``, the computation only + involves comparing sorted entries per slice, and ``d`` assignments are given. + + The user may also supply a ``num_subsamples`` parameter to extract as many + points from the original point cloud, sampled with probability masses ``a`` + and ``b``. This then simply applied the method above to the subsamples, to + output ``d`` costs, but assignments are not provided. + + When the problem is not uniform or not of equal size, the method defaults to + an inversion of the CDF, and outputs both costs and transport matrix in sparse + form. + + When a ``quantiles`` argument is passed, either specifying explicit quantiles + or a grid of quantiles, the distance is evaluated by comparing the quantiles + of the two point clouds on each slice. The OT costs are returned but + assignments are not provided. + + Args: + num_subsamples: Option to reduce the size of inputs by doing random + subsampling, taken into account marginal probabilities. + quantiles: When a vector or a number of quantiles is passed, the distance + is computed by evaluating the cost function on the sectional (one for each + dimension) quantiles of the two point cloud distributions described in the + problem. + """ + + def __init__( + self, + num_subsamples: Optional[int] = None, + quantiles: Optional[Union[int, jnp.ndarray]] = None, + ): + self._quantiles = quantiles + self.num_subsamples = num_subsamples + + @property + def quantiles(self) -> Optional[jnp.ndarray]: + """Quantiles' values used to evaluate OT cost.""" + if self._quantiles is None: + return None + if isinstance(self._quantiles, int): + return jnp.linspace(0.0, 1.0, self._quantiles) + return self._quantiles + + @property + def num_quantiles(self) -> int: + """Number of quantiles used to evaluate OT cost.""" + return 0 if self.quantiles is None else self.quantiles.shape[0] + + def __call__( + self, + prob: linear_problem.LinearProblem, + return_transport: bool = True, + rng: Optional[jax.Array] = None, + ) -> UnivariateOutput: + """Computes Univariate Distance between the ``d`` dimensional slices. + + Args: + prob: Problem with a :attr:`~ott.problems.linear.LinearProblem.geom` + attribute, the two point clouds ``x`` and ``y`` + (of respective sizes ``[n, d]`` and ``[m, d]``) and a ground + `TI cost ` between two scalars. + The ``[n,]`` and ``[m,]`` size probability weights vectors are stored + in attributes `:attr:`~ott.problems.linear.LinearProblem.a` and + :attr:`~ott.problems.linear.LinearProblem.b`. + return_transport: Whether to also return pairs of matched indices used to + compute optimal transport matrices. + rng: Used for random downsampling, if specified in the solver. + + Returns: + An output object, that computes ``d`` OT costs, in addition to, possibly, + paired lists of indices and their corresponding masses, on each of the + ``d`` dimensional slices of the input. + """ + geom = prob.geom + assert isinstance(geom, pointcloud.PointCloud), \ + "Geometry object in problem must be a PointCloud." + assert isinstance(geom.cost_fn, costs.TICost), \ + "Geometry's cost must be translation invariant." + + rng = utils.default_prng_key(rng) + + if self.num_subsamples: + x, y = self._subsample(prob, rng) + is_uniform_same_size = True + else: + # check if problem has the property uniform / same number of points + x, y = geom.x, geom.y + is_uniform_same_size = prob.is_uniform and prob.is_equal_size + + if self.quantiles is not None: + assert prob.is_uniform, \ + "The 'quantiles' method can only be used with uniform marginals." + out = _quant_dist(x, y, geom.cost_fn, self.quantiles, self.num_quantiles) + elif is_uniform_same_size: + return_transport = return_transport and not self.num_subsamples + out = uniform_distance(x, y, geom.cost_fn, return_transport) + else: + fn = jax.vmap(quantile_distance, in_axes=[1, 1, None, None, None, None]) + out = fn(x, y, geom.cost_fn, prob.a, prob.b, return_transport) + + return UnivariateOutput(prob, *out) + + def _subsample(self, prob: linear_problem.LinearProblem, + rng: jax.Array) -> Tuple[jnp.ndarray, jnp.ndarray]: + n, m = prob.geom.shape + x, y = prob.geom.x, prob.geom.y + + if prob.is_uniform: + x = x[jnp.linspace(0, n, num=self.num_subsamples).astype(int), :] + y = y[jnp.linspace(0, m, num=self.num_subsamples).astype(int), :] + return x, y + + rng1, rng2 = jax.random.split(rng, 2) + x = jax.random.choice(rng1, x, (self.num_subsamples,), p=prob.a, axis=0) + y = jax.random.choice(rng2, y, (self.num_subsamples,), p=prob.b, axis=0) + return x, y + + def tree_flatten(self): # noqa: D102 + return None, (self.num_subsamples, self._quantiles) + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + del children + return cls(*aux_data) + + +def uniform_distance( + x: jnp.ndarray, + y: jnp.ndarray, + cost_fn: costs.TICost, + return_transport: bool = True +) -> Distance_t: + """Distance between two equal-size families of uniformly weighted values x/y. + + Args: + x: Vector ``[n,]`` of real values. + y: Vector ``[n,]`` of real values. + cost_fn: Translation invariant cost function, i.e. ``c(x, y) = h(x - y)``. + return_transport: whether to return mapped pairs. + + Returns: + optimal transport cost, a list of ``n+m`` paired indices, and their + corresponding transport mass. Note that said mass can be null in some + entries, but sums to 1.0 + """ + n = x.shape[0] + i_x, i_y = jnp.argsort(x, axis=0), jnp.argsort(y, axis=0) + x = jnp.take_along_axis(x, i_x, axis=0) + y = jnp.take_along_axis(y, i_y, axis=0) + ot_costs = jax.vmap(cost_fn.h, in_axes=[0])(x.T - y.T) / n + + if return_transport: + paired_indices = jnp.stack([i_x, i_y]).transpose([2, 0, 1]) + mass_paired_indices = jnp.ones((n,)) / n + return ot_costs, paired_indices, mass_paired_indices + + return ot_costs, None, None + + +def quantile_distance( + x: jnp.ndarray, + y: jnp.ndarray, + cost_fn: costs.TICost, + a: jnp.ndarray, + b: jnp.ndarray, + return_transport: bool = True, +) -> Distance_t: + """Computes distance between quantile functions of distributions (a,x)/(b,y). + + Args: + x: Vector ``[n,]`` of real values. + y: Vector ``[m,]`` of real values. + cost_fn: Translation invariant cost function, i.e. ``c(x, y) = h(x - y)``. + a: Vector ``[n,]`` of non-negative weights summing to 1. + b: Vector ``[m,]`` of non-negative weights summing to 1. + return_transport: whether to return mapped pairs. + + Returns: + optimal transport cost, a list of ``n + m`` paired indices, and their + corresponding transport mass. Note that said mass can be null in some + entries, but sums to 1.0 + + Notes: + Inspired by :func:`~scipy.stats.wasserstein_distance`, + but can be used with other costs, not just :math:`c(x, y) = |x - y|`. + """ + x, i_x = mu.sort_and_argsort(x, argsort=True) + y, i_y = mu.sort_and_argsort(y, argsort=True) + + all_values = jnp.concatenate([x, y]) + all_values_sorted, all_values_sorter = mu.sort_and_argsort( + all_values, argsort=True + ) + + x_pdf = jnp.concatenate([a[i_x], jnp.zeros_like(b)])[all_values_sorter] + y_pdf = jnp.concatenate([jnp.zeros_like(a), b[i_y]])[all_values_sorter] + + x_cdf = jnp.cumsum(x_pdf) + y_cdf = jnp.cumsum(y_pdf) + + x_y_cdfs = jnp.concatenate([x_cdf, y_cdf]) + quantile_levels, _ = mu.sort_and_argsort(x_y_cdfs, argsort=False) + + i_x_cdf_inv = jnp.searchsorted(x_cdf, quantile_levels) + x_cdf_inv = all_values_sorted[i_x_cdf_inv] + i_y_cdf_inv = jnp.searchsorted(y_cdf, quantile_levels) + y_cdf_inv = all_values_sorted[i_y_cdf_inv] + + diff_q = jnp.diff(quantile_levels) + cost = jnp.sum( + jax.vmap(cost_fn.h)(y_cdf_inv[1:, None] - x_cdf_inv[1:, None]) * diff_q + ) + if not return_transport: + return cost, None, None + + n = x.shape[0] + + i_in_sorted_x_of_quantile = all_values_sorter[i_x_cdf_inv] % n + i_in_sorted_y_of_quantile = all_values_sorter[i_y_cdf_inv] - n + + orig_i = i_x[i_in_sorted_x_of_quantile][1:] + orig_j = i_y[i_in_sorted_y_of_quantile][1:] + + return cost, jnp.stack([orig_i, orig_j]), diff_q + + +def _quant_dist( + x: jnp.ndarray, y: jnp.ndarray, cost_fn: costs.TICost, q: jnp.ndarray, + n_q: int +) -> Tuple[jnp.ndarray, None, None]: + x_q = jnp.quantile(x, q, axis=0) + y_q = jnp.quantile(y, q, axis=0) + ot_costs = jax.vmap(cost_fn.pairwise, in_axes=[1, 1])(x_q, y_q) + + return ot_costs / n_q, None, None diff --git a/ott/build/lib/ott/solvers/quadratic/__init__.py b/ott/build/lib/ott/solvers/quadratic/__init__.py new file mode 100644 index 0000000..560ac3d --- /dev/null +++ b/ott/build/lib/ott/solvers/quadratic/__init__.py @@ -0,0 +1,20 @@ +# 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. +from . import ( + gromov_wasserstein, + gromov_wasserstein_lr, + gw_barycenter, + lower_bound, +) +from ._solve import solve diff --git a/ott/build/lib/ott/solvers/quadratic/_solve.py b/ott/build/lib/ott/solvers/quadratic/_solve.py new file mode 100644 index 0000000..9cdefec --- /dev/null +++ b/ott/build/lib/ott/solvers/quadratic/_solve.py @@ -0,0 +1,91 @@ +# 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. +from typing import Any, Literal, Optional, Union + +import jax.numpy as jnp + +from ott.geometry import geometry +from ott.problems.quadratic import quadratic_costs, quadratic_problem +from ott.solvers.quadratic import gromov_wasserstein as gw +from ott.solvers.quadratic import gromov_wasserstein_lr as lrgw + +__all__ = ["solve"] + + +def solve( + geom_xx: geometry.Geometry, + geom_yy: geometry.Geometry, + geom_xy: Optional[geometry.Geometry] = None, + fused_penalty: float = 1.0, + a: Optional[jnp.ndarray] = None, + b: Optional[jnp.ndarray] = None, + tau_a: float = 1.0, + tau_b: float = 1.0, + loss: Union[Literal["sqeucl", "kl"], quadratic_costs.GWLoss] = "sqeucl", + gw_unbalanced_correction: bool = True, + rank: int = -1, + **kwargs: Any, +) -> Union[gw.GWOutput, lrgw.LRGWOutput]: + """Solve quadratic regularized OT problem using a Gromov-Wasserstein solver. + + Args: + geom_xx: Ground geometry of the first space. + geom_yy: Ground geometry of the second space. + geom_xy: Geometry defining the linear penalty term for + fused Gromov-Wasserstein :cite:`vayer:19`. If :obj:`None`, the problem + reduces to a plain Gromov-Wasserstein problem :cite:`peyre:16`. + fused_penalty: Multiplier of the linear term in fused Gromov-Wasserstein, + i.e. ``problem = purely quadratic + fused_penalty * linear problem``. + a: The first marginal. If :obj:`None`, it will be uniform. + b: The second marginal. If :obj:`None`, it will be uniform. + tau_a: If :math:`< 1`, defines how much unbalanced the problem is + on the first marginal. + tau_b: If :math:`< 1`, defines how much unbalanced the problem is + on the second marginal. + loss: Gromov-Wasserstein loss function, see + :class:`~ott.problems.quadratic.quadratic_costs.GWLoss` for more + information. If ``rank > 0``, ``'sqeucl'`` is always used. + gw_unbalanced_correction: Whether the unbalanced version of + :cite:`sejourne:21` is used. Otherwise, ``tau_a`` and ``tau_b`` + only affect the resolution of the linearization of the GW problem + in the inner loop. Only used when ``rank = -1``. + rank: Rank constraint on the coupling to minimize the quadratic OT problem + :cite:`scetbon:22`. If :math:`-1`, no rank constraint is used. + kwargs: Keyword arguments for + :class:`~ott.solvers.quadratic.gromov_wasserstein.GromovWasserstein` or + :class:`~ott.solvers.quadratic.gromov_wasserstein_lr.LRGromovWasserstein`, + depending on the ``rank`` + + Returns: + The Gromov-Wasserstein output. + """ + prob = quadratic_problem.QuadraticProblem( + geom_xx=geom_xx, + geom_yy=geom_yy, + geom_xy=geom_xy, + fused_penalty=fused_penalty, + a=a, + b=b, + tau_a=tau_a, + tau_b=tau_b, + loss=loss, + gw_unbalanced_correction=gw_unbalanced_correction + ) + + if rank > 0: + solver = lrgw.LRGromovWasserstein(rank=rank, **kwargs) + else: + solver = gw.GromovWasserstein(rank=rank, **kwargs) + + return solver(prob) diff --git a/ott/build/lib/ott/solvers/quadratic/gromov_wasserstein.py b/ott/build/lib/ott/solvers/quadratic/gromov_wasserstein.py new file mode 100644 index 0000000..7272bd1 --- /dev/null +++ b/ott/build/lib/ott/solvers/quadratic/gromov_wasserstein.py @@ -0,0 +1,461 @@ +# 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. +from typing import ( + Any, + Callable, + Dict, + Literal, + Mapping, + NamedTuple, + Optional, + Sequence, + Tuple, + Union, +) + +import jax +import jax.experimental +import jax.numpy as jnp +import numpy as np + +from ott import utils +from ott.geometry import geometry +from ott.initializers.quadratic import initializers as quad_initializers +from ott.math import fixed_point_loop +from ott.problems.linear import linear_problem +from ott.problems.quadratic import quadratic_problem +from ott.solvers import was_solver +from ott.solvers.linear import sinkhorn, sinkhorn_lr + +__all__ = ["GromovWasserstein", "GWOutput"] + +LinearOutput = Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput] + +ProgressCallbackFn_t = Callable[ + [Tuple[np.ndarray, np.ndarray, np.ndarray, "GWState"]], None +] + + +class GWOutput(NamedTuple): + """Holds the output of the Gromov-Wasserstein solver. + + Args: + costs: Holds the sequence of regularized GW costs seen through the outer + loop of the solver. + linear_convergence: Holds the sequence of bool convergence flags of the + inner Sinkhorn iterations. + converged: Convergence flag for the outer GW iterations. + errors: Holds sequence of vectors of errors of the Sinkhorn algorithm + at each iteration. + linear_state: State used to solve and store solutions to the local + linearization of GW. + geom: The geometry underlying the local linearization. + old_transport_mass: Holds total mass of transport at previous iteration. + """ + + costs: Optional[jnp.ndarray] = None + linear_convergence: Optional[jnp.ndarray] = None + converged: bool = False + errors: Optional[jnp.ndarray] = None + linear_state: Optional[LinearOutput] = None + geom: Optional[geometry.Geometry] = None + # Intermediate values. + old_transport_mass: float = 1.0 + + def set(self, **kwargs: Any) -> "GWOutput": + """Return a copy of self, possibly with overwrites.""" + return self._replace(**kwargs) + + @property + def matrix(self) -> jnp.ndarray: + """Transport matrix.""" + return self._rescale_factor * self.linear_state.matrix + + def apply(self, inputs: jnp.ndarray, axis: int = 0) -> jnp.ndarray: + """Apply the transport to an array; axis=1 for its transpose.""" + return self._rescale_factor * self.linear_state.apply(inputs, axis=axis) + + @property + def reg_gw_cost(self) -> float: + """Regularized optimal transport cost of the linearization.""" + return self.linear_state.reg_ot_cost + + @property + def _rescale_factor(self) -> float: + return jnp.sqrt( + self.old_transport_mass / self.linear_state.transport_mass + ) + + @property + def primal_cost(self) -> float: + """Return transport cost of current linear OT solution at geometry.""" + return self.linear_state.transport_cost_at_geom(other_geom=self.geom) + + @property + def n_iters(self) -> int: # noqa: D102 + if self.errors is None: + return -1 + return jnp.sum(self.errors[:, 0] != -1) + + +class GWState(NamedTuple): + """State of the Gromov-Wasserstein solver. + + Attributes: + costs: Holds the sequence of regularized GW costs seen through the outer + loop of the solver. + linear_convergence: Holds the sequence of bool convergence flags of the + inner Sinkhorn iterations. + linear_state: State used to solve and store solutions to the local + linearization of GW. + linear_pb: Local linearization of the quadratic GW problem. + old_transport_mass: Intermediary value of the mass of the transport matrix. + rngs: Random keys passed to low-rank initializers at every GW iteration + when not using warm start. + errors: Holds sequence of vectors of errors of the Sinkhorn algorithm + at each iteration. + """ + + costs: jnp.ndarray + linear_convergence: jnp.ndarray + linear_state: LinearOutput + linear_pb: linear_problem.LinearProblem + old_transport_mass: float + rngs: Optional[jax.Array] = None + errors: Optional[jnp.ndarray] = None + + def set(self, **kwargs: Any) -> "GWState": + """Return a copy of self, possibly with overwrites.""" + return self._replace(**kwargs) + + def update( # noqa: D102 + self, + iteration: int, + linear_sol: LinearOutput, + linear_pb: linear_problem.LinearProblem, + store_errors: bool, + old_transport_mass: float, + ) -> "GWState": + costs = self.costs.at[iteration].set(linear_sol.reg_ot_cost) + errors = None + if store_errors and self.errors is not None: + errors = self.errors.at[iteration, :].set(linear_sol.errors) + linear_convergence = self.linear_convergence.at[iteration].set( + linear_sol.converged + ) + + return self.set( + linear_state=linear_sol, + linear_pb=linear_pb, + costs=costs, + linear_convergence=linear_convergence, + errors=errors, + old_transport_mass=old_transport_mass, + ) + + +@jax.tree_util.register_pytree_node_class +class GromovWasserstein(was_solver.WassersteinSolver): + """Gromov-Wasserstein solver :cite:`peyre:16`. + + .. seealso:: + Low-rank Gromov-Wasserstein :cite:`scetbon:23` is implemented in + :class:`~ott.solvers.quadratic.gromov_wasserstein_lr.LRGromovWasserstein`. + + Args: + args: Positional arguments for + :class:`~ott.solvers.was_solver.WassersteinSolver`. + warm_start: Whether to initialize Sinkhorn calls using values + from the previous iteration. If :obj:`None`, warm starts are not used for + standard Sinkhorn. + relative_epsilon: Whether to use relative epsilon in the linearized + geometry. + quad_initializer: Quadratic initializer. If the solver is entropic, + :class:`~ott.initializers.quadratic.initializers.QuadraticInitializer` + is always used. + progress_fn: callback function which gets called during the + Gromov-Wasserstein iterations, so the user can display the error at each + iteration, e.g., using a progress bar. + See :func:`~ott.utils.default_progress_fn` for a basic implementation. + kwargs_init: Keyword arguments when creating the initializer. + kwargs: Keyword arguments for + :class:`~ott.solvers.was_solver.WassersteinSolver`. + """ + + def __init__( + self, + *args: Any, + warm_start: Optional[bool] = None, + relative_epsilon: Optional[bool] = None, + quad_initializer: Optional[ + Union[ + Literal["random", "rank2", "k-means", "generalized-k-means"], + quad_initializers.BaseQuadraticInitializer, + ] + ] = None, + progress_fn: Optional[ProgressCallbackFn_t] = None, + kwargs_init: Optional[Mapping[str, Any]] = None, + kwargs_sinkhorn: Any = {}, + **kwargs: Any, + ): + super().__init__(kwargs=kwargs_sinkhorn, *args, **kwargs) + assert not self.is_low_rank, ( + "For low-rank GW, use " + "`ott.solvers.quadratic.gromov_wasserstein_lr.LRGromovWasserstein`." + ) + self._warm_start = warm_start + self.relative_epsilon = relative_epsilon + self.quad_initializer = quad_initializer + self.progress_fn = progress_fn + self.kwargs_init = {} if kwargs_init is None else kwargs_init + + def __call__( + self, + prob: quadratic_problem.QuadraticProblem, + init: Optional[linear_problem.LinearProblem] = None, + rng: Optional[jax.Array] = None, + **kwargs: Any, + ) -> GWOutput: + """Run the Gromov-Wasserstein solver. + + Args: + prob: Quadratic OT problem. + init: Initial linearization of the quadratic problem. If `None`, it will + be computed using the initializer. + rng: Random number key. + kwargs: Keyword arguments used when calling the initializer. + + Returns: + The Gromov-Wasserstein output. + """ + rng = utils.default_prng_key(rng) + rng1, rng2 = jax.random.split(rng, 2) + + # if prob._is_low_rank_convertible: + # prob = prob.to_low_rank() + + if init is None: + initializer = self.create_initializer(prob) + init = initializer( + prob, + epsilon=self.epsilon, + rng=rng1, + relative_epsilon=self.relative_epsilon, + **kwargs, + ) + print("GW called") + out = iterations(self, prob, init, rng2) + # TODO(lpapaxanthoos): remove stop_gradient when using backprop + if self.is_low_rank: + linearization = prob.update_lr_linearization( + jax.lax.stop_gradient(out.linear_state), + relative_epsilon=self.relative_epsilon, + ) + else: + linearization = prob.update_linearization( + jax.lax.stop_gradient(out.linear_state), + epsilon=self.epsilon, + old_transport_mass=jax.lax.stop_gradient( + out.old_transport_mass + ), + relative_epsilon=self.relative_epsilon, + ) + + linear_state = out.linear_state.set_cost(linearization, True, True) + iteration = jnp.sum(out.costs != -1) + converged = jnp.logical_and( + iteration < self.max_iterations, jnp.all(out.linear_convergence) + ) + return out.set( + linear_state=linear_state, + geom=linearization.geom, + converged=converged, + ) + + def init_state( + self, + prob: quadratic_problem.QuadraticProblem, + init: linear_problem.LinearProblem, + rng: jax.Array, + ) -> GWState: + """Initialize the state of the Gromov-Wasserstein iterations. + + Args: + prob: Quadratic OT problem. + init: Initial linearization of the quadratic problem. + rng: Random key for low-rank initializers. Only used when + :attr:`warm_start` is `False`. + + Returns: + The initial Gromov-Wasserstein state. + """ + linear_state = self.linear_ot_solver(init) + num_iter = self.max_iterations + transport_mass = prob.init_transport_mass() + if self.store_inner_errors: + errors = -jnp.ones( + (num_iter, self.linear_ot_solver.outer_iterations) + ) + else: + errors = None + + return GWState( + costs=-jnp.ones((num_iter,)), + linear_convergence=-jnp.ones((num_iter,)), + linear_state=linear_state, + linear_pb=init, + old_transport_mass=transport_mass, + rngs=jax.random.split(rng, num_iter), + errors=errors, + ) + + def output_from_state( + self, + state: GWState, + ) -> GWOutput: + """Create an output from a loop state. + + Arguments: + state: A GWState. + + Returns: + A GWOutput. + """ + return GWOutput( + costs=state.costs, + linear_convergence=state.linear_convergence, + errors=state.errors, + linear_state=state.linear_state, + geom=state.linear_pb.geom, + old_transport_mass=state.old_transport_mass, + ) + + def create_initializer( + self, prob: quadratic_problem.QuadraticProblem + ) -> quad_initializers.BaseQuadraticInitializer: + """Create quadratic, possibly low-rank initializer. + + Args: + prob: Quadratic OT problem used to determine the initializer. + + Returns: + The initializer. + """ + if isinstance( + self.quad_initializer, quad_initializers.BaseQuadraticInitializer + ): + return self.quad_initializer + # no other options implemented, use the default + return quad_initializers.QuadraticInitializer(**self.kwargs_init) + + @property + def warm_start(self) -> bool: + """Whether to initialize Sinkhorn using previous solutions.""" + return ( + self.is_low_rank if self._warm_start is None else self._warm_start + ) + + def tree_flatten( + self, + ) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + children, aux_data = super().tree_flatten() + aux_data["warm_start"] = self._warm_start + aux_data["progress_fn"] = self.progress_fn + aux_data["relative_epsilon"] = self.relative_epsilon + aux_data["quad_initializer"] = self.quad_initializer + aux_data["kwargs_init"] = self.kwargs_init + return children, aux_data + + +def iterations( + solver: GromovWasserstein, + prob: quadratic_problem.QuadraticProblem, + init: linear_problem.LinearProblem, + rng: jax.Array, +) -> GWOutput: + """Jittable Gromov-Wasserstein outer loop.""" + + def cond_fn( + iteration: int, solver: GromovWasserstein, state: GWState + ) -> bool: + return solver._continue(state, iteration) + + def body_fn( + iteration: int, + solver: GromovWasserstein, + state: GWState, + compute_error: bool, + ) -> GWState: + del compute_error # always assumed true for the outer loop of GW + + lin_state = state.linear_state + if solver.is_low_rank: + rng = state.rngs[iteration] + init = ( + (lin_state.q, lin_state.r, lin_state.g) + if solver.warm_start + else (None, None, None) + ) + linear_pb = prob.update_lr_linearization( + state.linear_state, relative_epsilon=solver.relative_epsilon + ) + out = solver.linear_ot_solver(linear_pb, init=init, rng=rng) + else: + init = ( + (lin_state.f, lin_state.g) + if solver.warm_start + else (None, None) + ) + linear_pb = prob.update_linearization( + lin_state, + solver.epsilon, + state.old_transport_mass, + relative_epsilon=solver.relative_epsilon, + ) + prob.linear_problem = linear_pb + out = solver.linear_ot_solver(linear_pb, init=init) + + old_transport_mass = jax.lax.stop_gradient( + state.linear_state.transport_mass + ) + new_state = state.update( + iteration, + out, + linear_pb, + solver.store_inner_errors, + old_transport_mass, + ) + + # Inner iterations is currently fixed to 1. + inner_iterations = 1 + if solver.progress_fn is not None: + jax.experimental.io_callback( + solver.progress_fn, + None, + (iteration, inner_iterations, solver.max_iterations, state), + ) + + return new_state + + state = fixed_point_loop.fixpoint_iter( + cond_fn=cond_fn, + body_fn=body_fn, + min_iterations=solver.min_iterations, + max_iterations=solver.max_iterations, + inner_iterations=1, + constants=solver, + state=solver.init_state(prob, init, rng=rng), + ) + + return solver.output_from_state(state) diff --git a/ott/build/lib/ott/solvers/quadratic/gromov_wasserstein_lr.py b/ott/build/lib/ott/solvers/quadratic/gromov_wasserstein_lr.py new file mode 100644 index 0000000..cb12911 --- /dev/null +++ b/ott/build/lib/ott/solvers/quadratic/gromov_wasserstein_lr.py @@ -0,0 +1,897 @@ +# 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. +"""A Jax implementation of the unbalanced low-rank GW algorithm.""" +from typing import ( + Any, + Callable, + Literal, + Mapping, + NamedTuple, + Optional, + Tuple, + Union, +) + +import jax +import jax.experimental +import jax.numpy as jnp +import jax.scipy as jsp +import numpy as np + +from ott import utils +from ott.geometry import geometry, low_rank +from ott.initializers.linear import initializers_lr +from ott.math import fixed_point_loop +from ott.math import unbalanced_functions as uf +from ott.math import utils as mu +from ott.problems.quadratic import quadratic_problem +from ott.solvers.linear import lr_utils, sinkhorn + +__all__ = ["LRGromovWasserstein", "LRGWOutput"] + +ProgressCallbackFn_t = Callable[ + [Tuple[np.ndarray, np.ndarray, np.ndarray, "LRGWState"]], None] + + +class LRGWState(NamedTuple): + """State of the low-rank GW algorithm.""" + q: jnp.ndarray + r: jnp.ndarray + g: jnp.ndarray + gamma: float + costs: jnp.ndarray + errors: jnp.ndarray + crossed_threshold: bool + + def compute_error( # noqa: D102 + self, previous_state: "LRGWState" + ) -> float: + err_q = mu.gen_js(self.q, previous_state.q, c=1.0) + err_r = mu.gen_js(self.r, previous_state.r, c=1.0) + err_g = mu.gen_js(self.g, previous_state.g, c=1.0) + + return ((1.0 / self.gamma) ** 2) * (err_q + err_r + err_g) + + def reg_gw_cost( # noqa: D102 + self, + ot_prob: quadratic_problem.QuadraticProblem, + *, + epsilon: float, + use_danskin: bool = False + ) -> float: + return compute_reg_gw_cost( + self.q, + self.r, + self.g, + ot_prob, + epsilon=epsilon, + use_danskin=use_danskin + ) + + def set(self, **kwargs: Any) -> "LRGWState": + """Return a copy of self, with potential overwrites.""" + return self._replace(**kwargs) + + +def compute_reg_gw_cost( + q: jnp.ndarray, + r: jnp.ndarray, + g: jnp.ndarray, + ot_prob: quadratic_problem.QuadraticProblem, + epsilon: float, + use_danskin: bool = False +) -> float: + """Compute the regularized OT cost, here the primal cost of the LR solution. + + Args: + q: first factor of solution + r: second factor of solution + g: weights of solution + ot_prob: linear problem + epsilon: Entropic regularization. + use_danskin: if True, use Danskin's theorem :cite:`danskin:67,bertsekas:71` + to avoid computing the gradient of the cost function. + + Returns: + regularized OT cost, the (primal) transport cost of the low-rank solution. + """ + + def ent(x: jnp.ndarray) -> float: + # generalized entropy + return jnp.sum(jsp.special.entr(x) + x) + + q = jax.lax.stop_gradient(q) if use_danskin else q + r = jax.lax.stop_gradient(r) if use_danskin else r + g = jax.lax.stop_gradient(g) if use_danskin else g + + out = LRGWOutput( + q=q, + r=r, + g=g, + ot_prob=ot_prob, + costs=None, + errors=None, + epsilon=None, + inner_iterations=None, + ) + + cost = out.primal_cost - epsilon * (ent(q) + ent(r) + ent(g)) + if ot_prob.tau_a != 1.0: + rho_a = uf.rho(1.0, ot_prob.tau_a) + cost += rho_a * mu.gen_kl(jnp.sum(q, axis=1), ot_prob.a) + if ot_prob.tau_b != 1.0: + rho_b = uf.rho(1.0, ot_prob.tau_b) + cost += rho_b * mu.gen_kl(jnp.sum(r, axis=1), ot_prob.b) + + return cost + + +class LRGWOutput(NamedTuple): + """Transport interface for a low-rank GW solution.""" + q: jnp.ndarray + r: jnp.ndarray + g: jnp.ndarray + costs: jnp.ndarray + # TODO(michalk8): must be called `errors`, because of `store_inner_errors` + # in future, enforce via class hierarchy + errors: jnp.ndarray + ot_prob: quadratic_problem.QuadraticProblem + epsilon: float + inner_iterations: int + reg_gw_cost: Optional[float] = None + + def set(self, **kwargs: Any) -> "LRGWOutput": + """Return a copy of self, with potential overwrites.""" + return self._replace(**kwargs) + + def set_cost( # noqa: D102 + self, + ot_prob: quadratic_problem.QuadraticProblem, + lse_mode: bool, + use_danskin: bool = False + ) -> "LRGWOutput": + del lse_mode + return self.set(reg_gw_cost=self.compute_reg_gw_cost(ot_prob, use_danskin)) + + def compute_reg_gw_cost( # noqa: D102 + self, + ot_prob: quadratic_problem.QuadraticProblem, + use_danskin: bool = False, + ) -> float: + return compute_reg_gw_cost( + self.q, + self.r, + self.g, + ot_prob, + epsilon=self.epsilon, + use_danskin=use_danskin + ) + + @property + def geom(self) -> geometry.Geometry: # noqa: D102 + """Linearized geometry.""" + return _linearized_geometry(self.ot_prob, q=self.q, r=self.r, g=self.g) + + @property + def a(self) -> jnp.ndarray: # noqa: D102 + return self.ot_prob.a + + @property + def b(self) -> jnp.ndarray: # noqa: D102 + return self.ot_prob.b + + @property + def n_iters(self) -> int: # noqa: D102 + return jnp.sum(self.errors != -1) * self.inner_iterations + + @property + def converged(self) -> bool: # noqa: D102 + return jnp.logical_and( + jnp.any(self.costs == -1), jnp.all(jnp.isfinite(self.costs)) + ) + + @property + def matrix(self) -> jnp.ndarray: + """Transport matrix if it can be instantiated.""" + return (self.q * self._inv_g) @ self.r.T + + def apply(self, inputs: jnp.ndarray, axis: int = 0) -> jnp.ndarray: + """Apply the transport to a array; axis=1 for its transpose.""" + q, r = (self.q, self.r) if axis == 1 else (self.r, self.q) + # for `axis=0`: (batch, m), (m, r), (r,), (r, n) + return ((inputs @ r) * self._inv_g) @ q.T + + def marginal(self, axis: int) -> jnp.ndarray: # noqa: D102 + length = self.q.shape[0] if axis == 0 else self.r.shape[0] + return self.apply(jnp.ones(length,), axis=axis) + + def cost_at_geom(self, other_geom: geometry.Geometry) -> float: + """Return OT cost for current solution, evaluated at any cost matrix.""" + return jnp.sum(self.q * other_geom.apply_cost(self.r, axis=1) * self._inv_g) + + def transport_cost_at_geom(self, other_geom: geometry.Geometry) -> float: + """Return (by recomputing it) bare transport cost of current solution.""" + return self.cost_at_geom(other_geom) + + @property + def primal_cost(self) -> float: + """Return (by recomputing it) transport cost of current solution.""" + geom_xx, geom_yy = self.ot_prob.geom_xx, self.ot_prob.geom_yy + marginal_a = self.ot_prob.a if self.ot_prob.tau_a == 1.0 else self.q.sum(1) + marginal_b = self.ot_prob.b if self.ot_prob.tau_b == 1.0 else self.r.sum(1) + + quad_cost = 0.5 * self.transport_cost_at_geom(other_geom=self.geom) + quad_cost += jnp.vdot(geom_xx.apply_square_cost(marginal_a), marginal_a) + quad_cost += jnp.vdot(geom_yy.apply_square_cost(marginal_b), marginal_b) + + if not self.ot_prob.is_fused: + return quad_cost + + alpha = self.ot_prob.fused_penalty / (self.ot_prob.fused_penalty + 1.0) + norm_g = jnp.linalg.norm(self.g, ord=1) + + lin_cost = self.cost_at_geom(self.ot_prob.geom_xy) + return alpha * norm_g * lin_cost + (1.0 - alpha) * quad_cost + + @property + def transport_mass(self) -> float: + """Sum of transport matrix.""" + return self.marginal(0).sum() + + @property + def _inv_g(self) -> jnp.ndarray: + return 1.0 / self.g + + +@jax.tree_util.register_pytree_node_class +class LRGromovWasserstein(sinkhorn.Sinkhorn): + r"""Low-rank Gromov-Wasserstein solver :cite:`scetbon:23`. + + The algorithm minimizes a non-convex problem. It therefore requires special + care to initialization and convergence. Convergence is evaluated on successive + evaluations of the objective. + + .. warning:: + This solver only for the **unbalanced** case. Balanced case is implemented + in :class:`~ott.solvers.quadratic.gromov_wasserstein.GromovWasserstein` + and will be unified here in the future release. + + Args: + rank: Rank constraint on the coupling to minimize the linear OT problem + gamma: The (inverse of) gradient step size used by mirror descent. + gamma_rescale: Whether to rescale :math:`\gamma` every iteration as + described in :cite:`scetbon:22b`. + epsilon: Entropic regularization added on top of low-rank problem. + initializer: How to initialize the :math:`Q`, :math:`R` and :math:`g` + factors. + lse_mode: Whether to run computations in LSE or kernel mode. + inner_iterations: Number of inner iterations used by the algorithm before + re-evaluating progress. + use_danskin: Use Danskin theorem to evaluate gradient of objective w.r.t. + input parameters. Only `True` handled at this moment. + implicit_diff: Whether to use implicit differentiation. Currently, only + ``implicit_diff = False`` is implemented. + progress_fn: callback function which gets called during the GW + iterations, so the user can display the error at each iteration, + e.g., using a progress bar. See :func:`~ott.utils.default_progress_fn` + for a basic implementation. + kwargs_dys: Keyword arguments passed to :meth:`dykstra_update_lse`, + :meth:`dykstra_update_kernel` or one of the functions defined in + :mod:`ott.solvers.linear`, depending on the ``lse_mode``. + kwargs_init: Keyword arguments for + :class:`~ott.initializers.linear.initializers_lr.LRInitializer`. + kwargs: Keyword arguments for + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. + """ + + def __init__( + self, + rank: int, + gamma: float = 10.0, + gamma_rescale: bool = True, + epsilon: float = 0.0, + initializer: Union[Literal["random", "rank2", "k-means", + "generalized-k-means"], + initializers_lr.LRInitializer] = "random", + lse_mode: bool = True, + inner_iterations: int = 10, + use_danskin: bool = True, + implicit_diff: bool = False, + kwargs_dys: Optional[Mapping[str, Any]] = None, + kwargs_init: Optional[Mapping[str, Any]] = None, + progress_fn: Optional[ProgressCallbackFn_t] = None, + **kwargs: Any, + ): + assert not implicit_diff, "Implicit diff. not yet implemented." + super().__init__( + lse_mode=lse_mode, + inner_iterations=inner_iterations, + use_danskin=use_danskin, + implicit_diff=implicit_diff, + **kwargs + ) + self.rank = rank + self.gamma = gamma + self.gamma_rescale = gamma_rescale + self.epsilon = epsilon + self.initializer = initializer + self.progress_fn = progress_fn + # can be `None` + self.kwargs_dys = {} if kwargs_dys is None else kwargs_dys + self.kwargs_init = {} if kwargs_init is None else kwargs_init + + def __call__( + self, + ot_prob: quadratic_problem.QuadraticProblem, + init: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray], + Optional[jnp.ndarray]] = (None, None, None), + rng: Optional[jax.Array] = None, + **kwargs: Any, + ) -> LRGWOutput: + """Run low-rank Gromov-Wasserstein solver. + + Args: + ot_prob: Linear OT problem. + init: Initial values for the low-rank factors: + + - :attr:`~ott.solvers.linear.sinkhorn_lr.LRGWOutput.q`. + - :attr:`~ott.solvers.linear.sinkhorn_lr.LRGWOutput.r`. + - :attr:`~ott.solvers.linear.sinkhorn_lr.LRGWOutput.g`. + + Any `None` values will be initialized using the initializer. + rng: Random key for seeding. + kwargs: Additional arguments when calling the initializer. + + Returns: + The low-rank GW output. + """ + rng = utils.default_prng_key(rng) + rng_lrc, rng_init = jax.random.split(rng) + + if ot_prob._is_low_rank_convertible: + ot_prob = ot_prob.to_low_rank(rng=rng_lrc) + + initializer = self.create_initializer(ot_prob) + init = initializer(ot_prob, *init, rng=rng_init, **kwargs) + return run(ot_prob, self, init) + + def _get_costs( + self, + ot_prob: quadratic_problem.QuadraticProblem, + state: LRGWState, + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, float]: + q, r, g = state.q, state.r, state.g + log_q, log_r, log_g = mu.safe_log(q), mu.safe_log(r), mu.safe_log(g) + inv_g = 1.0 / g[None, :] + lin_geom = _linearized_geometry(ot_prob, q=q, r=r, g=g) + + tmp = lin_geom.apply_cost(r, axis=1) + grad_q = tmp * inv_g + if ot_prob.tau_a != 1.0: # unbalanced grad + grad_q += 2.0 * ot_prob.geom_xx.apply_square_cost(q.sum(1), axis=1) + + grad_r = lin_geom.apply_cost(q, axis=0) * inv_g + if ot_prob.tau_b != 1.0: # unbalanced grad + grad_r += 2.0 * ot_prob.geom_yy.apply_square_cost(r.sum(1), axis=1) + + omega_quad = jnp.sum(q * tmp, axis=0) + grad_g = -omega_quad / (g ** 2) + + if ot_prob.is_fused: + alpha = ot_prob.fused_penalty / (ot_prob.fused_penalty + 1.0) + norm_g = jnp.linalg.norm(g, ord=1) + + tmp = ot_prob.geom_xy.apply_cost(r, axis=1) + lin_grad_q = tmp * inv_g * norm_g + lin_grad_r = ot_prob.geom_xy.apply_cost(q) * inv_g * norm_g + + omega_lin = jnp.sum(q * tmp, axis=0) + lin_grad_g = -omega_lin / (g ** 2) * norm_g + jnp.sum(q * tmp * inv_g) + + grad_q = alpha * lin_grad_q + (1.0 - alpha) * grad_q + grad_r = alpha * lin_grad_r + (1.0 - alpha) * grad_r + grad_g = alpha * lin_grad_g + (1.0 - alpha) * grad_g + + grad_q += self.epsilon * log_q + grad_r += self.epsilon * log_r + grad_g += self.epsilon * log_g + + if self.gamma_rescale: + norm_q = jnp.max(jnp.abs(grad_q)) ** 2 + norm_r = jnp.max(jnp.abs(grad_r)) ** 2 + norm_g = jnp.max(jnp.abs(grad_g)) ** 2 + gamma = self.gamma / jnp.max(jnp.array([norm_q, norm_r, norm_g])) + else: + gamma = self.gamma + + eps_factor = 1.0 / (self.epsilon * gamma + 1.0) + gamma *= eps_factor + + c_q = -gamma * grad_q + eps_factor * log_q + c_r = -gamma * grad_r + eps_factor * log_r + c_g = -gamma * grad_g + eps_factor * log_g + + return c_q, c_r, c_g, gamma + + # TODO(michalk8): move to `lr_utils` when refactoring this the future + def dykstra_update_lse( + self, + c_q: jnp.ndarray, + c_r: jnp.ndarray, + h: jnp.ndarray, + gamma: float, + ot_prob: quadratic_problem.QuadraticProblem, + min_entry_value: float = 1e-6, + tolerance: float = 1e-3, + min_iter: int = 0, + inner_iter: int = 10, + max_iter: int = 10000 + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Run Dykstra's algorithm.""" + # shortcuts for problem's definition. + r = self.rank + n, m = ot_prob.geom_xx.shape[0], ot_prob.geom_yy.shape[0] + loga, logb = jnp.log(ot_prob.a), jnp.log(ot_prob.b) + + h_old = h + g1_old, g2_old = jnp.zeros(r), jnp.zeros(r) + f1, f2 = jnp.zeros(n), jnp.zeros(m) + + w_gi, w_gp = jnp.zeros(r), jnp.zeros(r) + w_q, w_r = jnp.zeros(r), jnp.zeros(r) + err = jnp.inf + state_inner = f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err + constants = c_q, c_r, loga, logb + + def cond_fn( + iteration: int, constants: Tuple[jnp.ndarray, ...], + state_inner: Tuple[jnp.ndarray, ...] + ) -> bool: + del iteration, constants + *_, err = state_inner + return err > tolerance + + def _softm( + f: jnp.ndarray, g: jnp.ndarray, c: jnp.ndarray, axis: int + ) -> jnp.ndarray: + return jsp.special.logsumexp( + gamma * (f[:, None] + g[None, :] - c), axis=axis + ) + + def body_fn( + iteration: int, constants: Tuple[jnp.ndarray, ...], + state_inner: Tuple[jnp.ndarray, ...], compute_error: bool + ) -> Tuple[jnp.ndarray, ...]: + # TODO(michalk8): in the future, use `NamedTuple` + f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err = state_inner + c_q, c_r, loga, logb = constants + + # First Projection + f1 = jnp.where( + jnp.isfinite(loga), + (loga - _softm(f1, g1_old, c_q, axis=1)) / gamma + f1, loga + ) + f2 = jnp.where( + jnp.isfinite(logb), + (logb - _softm(f2, g2_old, c_r, axis=1)) / gamma + f2, logb + ) + + h = h_old + w_gi + h = jnp.maximum(jnp.log(min_entry_value) / gamma, h) + w_gi += h_old - h + h_old = h + + # Update couplings + g_q = _softm(f1, g1_old, c_q, axis=0) + g_r = _softm(f2, g2_old, c_r, axis=0) + + # Second Projection + h = (1.0 / 3.0) * (h_old + w_gp + w_q + w_r) + h += g_q / (3.0 * gamma) + h += g_r / (3.0 * gamma) + g1 = h + g1_old - g_q / gamma + g2 = h + g2_old - g_r / gamma + + w_q = w_q + g1_old - g1 + w_r = w_r + g2_old - g2 + w_gp = h_old + w_gp - h + + q, r, _ = recompute_couplings(f1, g1, c_q, f2, g2, c_r, h, gamma) + + g1_old = g1 + g2_old = g2 + h_old = h + + err = jax.lax.cond( + jnp.logical_and(compute_error, iteration >= min_iter), + lambda: dykstra_solution_error(q, r, ot_prob, self.norm_error)[0], + lambda: err + ) + + return f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err + + def recompute_couplings( + f1: jnp.ndarray, + g1: jnp.ndarray, + c_q: jnp.ndarray, + f2: jnp.ndarray, + g2: jnp.ndarray, + c_r: jnp.ndarray, + h: jnp.ndarray, + gamma: float, + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + q = jnp.exp(gamma * (f1[:, None] + g1[None, :] - c_q)) + r = jnp.exp(gamma * (f2[:, None] + g2[None, :] - c_r)) + g = jnp.exp(gamma * h) + return q, r, g + + state_inner = fixed_point_loop.fixpoint_iter_backprop( + cond_fn, body_fn, min_iter, max_iter, inner_iter, constants, state_inner + ) + + f1, f2, g1_old, g2_old, h_old, _, _, _, _, _ = state_inner + return recompute_couplings(f1, g1_old, c_q, f2, g2_old, c_r, h_old, gamma) + + def dykstra_update_kernel( + self, + k_q: jnp.ndarray, + k_r: jnp.ndarray, + k_g: jnp.ndarray, + gamma: float, + ot_prob: quadratic_problem.QuadraticProblem, + min_entry_value: float = 1e-6, + tolerance: float = 1e-3, + min_iter: int = 0, + inner_iter: int = 10, + max_iter: int = 10000 + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Run Dykstra's algorithm.""" + # shortcuts for problem's definition. + del gamma + rank = self.rank + n, m = ot_prob.geom_xx.shape[0], ot_prob.geom_yy.shape[0] + a, b = ot_prob.a, ot_prob.b + supp_a, supp_b = a > 0, b > 0 + + g_old = k_g + v1_old, v2_old = jnp.ones(rank), jnp.ones(rank) + u1, u2 = jnp.ones(n), jnp.ones(m) + + q_gi, q_gp = jnp.ones(rank), jnp.ones(rank) + q_q, q_r = jnp.ones(rank), jnp.ones(rank) + err = jnp.inf + state_inner = u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err + constants = k_q, k_r, k_g, a, b + + def cond_fn( + iteration: int, constants: Tuple[jnp.ndarray, ...], + state_inner: Tuple[jnp.ndarray, ...] + ) -> bool: + del iteration, constants + *_, err = state_inner + return err > tolerance + + def body_fn( + iteration: int, constants: Tuple[jnp.ndarray, ...], + state_inner: Tuple[jnp.ndarray, ...], compute_error: bool + ) -> Tuple[jnp.ndarray, ...]: + # TODO(michalk8): in the future, use `NamedTuple` + u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err = state_inner + k_q, k_r, k_g, a, b = constants + + # First Projection + u1 = jnp.where(supp_a, a / jnp.dot(k_q, v1_old), 0.0) + u2 = jnp.where(supp_b, b / jnp.dot(k_r, v2_old), 0.0) + g = jnp.maximum(min_entry_value, g_old * q_gi) + q_gi = (g_old * q_gi) / g + g_old = g + + # Second Projection + v1_trans = jnp.dot(k_q.T, u1) + v2_trans = jnp.dot(k_r.T, u2) + g = (g_old * q_gp * v1_old * q_q * v1_trans * v2_old * q_r * + v2_trans) ** (1 / 3) + v1 = g / v1_trans + v2 = g / v2_trans + q_gp = (g_old * q_gp) / g + q_q = (q_q * v1_old) / v1 + q_r = (q_r * v2_old) / v2 + v1_old = v1 + v2_old = v2 + g_old = g + + # Compute Couplings + q, r, _ = recompute_couplings(u1, v1, k_q, u2, v2, k_r, g) + + err = jax.lax.cond( + jnp.logical_and(compute_error, iteration >= min_iter), + lambda: dykstra_solution_error(q, r, ot_prob, self.norm_error)[0], + lambda: err + ) + + return u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err + + def recompute_couplings( + u1: jnp.ndarray, + v1: jnp.ndarray, + k_q: jnp.ndarray, + u2: jnp.ndarray, + v2: jnp.ndarray, + k_r: jnp.ndarray, + g: jnp.ndarray, + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + q = u1.reshape((-1, 1)) * k_q * v1.reshape((1, -1)) + r = u2.reshape((-1, 1)) * k_r * v2.reshape((1, -1)) + return q, r, g + + state_inner = fixed_point_loop.fixpoint_iter_backprop( + cond_fn, body_fn, min_iter, max_iter, inner_iter, constants, state_inner + ) + + u1, u2, v1_old, v2_old, g_old, _, _, _, _, _ = state_inner + return recompute_couplings(u1, v1_old, k_q, u2, v2_old, k_r, g_old) + + def lse_step( + self, ot_prob: quadratic_problem.QuadraticProblem, state: LRGWState, + iteration: int + ) -> LRGWState: + """Low-rank GW LSE update.""" + c_q, c_r, c_g, gamma = self._get_costs(ot_prob, state) + + if ot_prob.is_balanced: + c_q, c_r, h = c_q / -gamma, c_r / -gamma, c_g / gamma + q, r, g = self.dykstra_update_lse( + c_q, c_r, h, gamma, ot_prob, **self.kwargs_dys + ) + else: + q, r, g = lr_utils.unbalanced_dykstra_lse( + c_q, c_r, c_g, gamma, ot_prob, **self.kwargs_dys + ) + return state.set(q=q, g=g, r=r, gamma=gamma) #, (c_q, c_r, c_g) + + def kernel_step( + self, ot_prob: quadratic_problem.QuadraticProblem, state: LRGWState, + iteration: int + ) -> LRGWState: + """Low-rank GW kernel update.""" + c_q, c_r, c_g, gamma = self._get_costs(ot_prob, state) + c_q, c_r, c_g = jnp.exp(c_q), jnp.exp(c_r), jnp.exp(c_g) + + if ot_prob.is_balanced: + q, r, g = self.dykstra_update_kernel( + c_q, c_r, c_g, gamma, ot_prob, **self.kwargs_dys + ) + else: + q, r, g = lr_utils.unbalanced_dykstra_kernel( + c_q, c_r, c_g, gamma, ot_prob, **self.kwargs_dys + ) + return state.set(q=q, g=g, r=r, gamma=gamma) + + def one_iteration( + self, ot_prob: quadratic_problem.QuadraticProblem, state: LRGWState, + iteration: int, compute_error: bool + ) -> LRGWState: + """Carries out one low-rank GW iteration. + + Depending on lse_mode, these iterations can be either in: + + - log-space for numerical stability. + - scaling space, using standard kernel-vector multiply operations. + + Args: + ot_prob: the transport problem definition + state: the current state. + iteration: the current iteration of the GW outer loop. + compute_error: flag to indicate this iteration computes/stores an error + + Returns: + The updated state. + """ + previous_state = state + + it = iteration // self.inner_iterations + if self.lse_mode: # In lse_mode, run additive updates. + state = self.lse_step(ot_prob, state, iteration) + else: + state = self.kernel_step(ot_prob, state, iteration) + + # re-computes error if compute_error is True, else set it to inf. + cost = jax.lax.cond( + jnp.logical_and(compute_error, iteration >= self.min_iterations), + lambda: state.reg_gw_cost(ot_prob, epsilon=self.epsilon), + lambda: jnp.inf + ) + error = state.compute_error(previous_state) + crossed_threshold = jnp.logical_or( + state.crossed_threshold, + jnp.logical_and( + state.errors[it - 1] >= self.threshold, error < self.threshold + ) + ) + + state = state.set( + costs=state.costs.at[it].set(cost), + errors=state.errors.at[it].set(error), + crossed_threshold=crossed_threshold, + ) + + if self.progress_fn is not None: + jax.experimental.io_callback( + self.progress_fn, None, + (iteration, self.inner_iterations, self.max_iterations, state) + ) + + return state + + @property + def norm_error(self) -> Tuple[int]: # noqa: D102 + return self._norm_error, + + def create_initializer( + self, + prob: quadratic_problem.QuadraticProblem, + ) -> initializers_lr.LRInitializer: + """Create a low-rank GW initializer. + + Args: + prob: Quadratic OT problem used to determine the initializer. + + Returns: + Low-rank initializer. + """ + if isinstance(self.initializer, initializers_lr.LRInitializer): + assert self.initializer.rank == self.rank, \ + f"Expected initializer's rank to be `{self.rank}`," \ + f"found `{self.initializer.rank}`." + return self.initializer + + return initializers_lr.LRInitializer.from_solver( + self, kind=self.initializer, **self.kwargs_init + ) + + def init_state( + self, ot_prob: quadratic_problem.QuadraticProblem, + init: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] + ) -> LRGWState: + """Return the initial state of the loop.""" + q, r, g = init + return LRGWState( + q=q, + r=r, + g=g, + gamma=self.gamma, + costs=-jnp.ones(self.outer_iterations), + errors=-jnp.ones(self.outer_iterations), + crossed_threshold=False, + ) + + def output_from_state( + self, ot_prob: quadratic_problem.QuadraticProblem, state: LRGWState + ) -> LRGWOutput: + """Create an output from a loop state. + + Args: + ot_prob: the transport problem. + state: GW state. + + Returns: + A LRGWOutput. + """ + return LRGWOutput( + q=state.q, + r=state.r, + g=state.g, + ot_prob=ot_prob, + costs=state.costs, + errors=state.errors, + epsilon=self.epsilon, + inner_iterations=self.inner_iterations, + ) + + def _converged(self, state: LRGWState, iteration: int) -> bool: + + def conv_crossed(prev_err: float, curr_err: float) -> bool: + return jnp.logical_and( + prev_err < self.threshold, curr_err < self.threshold + ) + + def conv_not_crossed(prev_err: float, curr_err: float) -> bool: + return jnp.logical_and(curr_err < prev_err, curr_err < self.threshold) + + # for convergence error, we consider 2 possibilities: + # 1. we either crossed the convergence threshold; in this case we require + # that the previous error was also below the threshold + # 2. we haven't crossed the threshold; in this case, we can be below or + # above the threshold: + # if we're above, we wait until we reach the convergence threshold and + # then, the above condition applies + # if we're below and we improved w.r.t. the previous iteration, + # we have converged; otherwise we continue, since we may be stuck + # in a local minimum (e.g., during the initial iterations) + + it = iteration // self.inner_iterations + return jax.lax.cond( + state.crossed_threshold, conv_crossed, conv_not_crossed, + state.errors[it - 2], state.errors[it - 1] + ) + + def _diverged(self, state: LRGWState, iteration: int) -> bool: + it = iteration // self.inner_iterations + return jnp.logical_and( + jnp.logical_not(jnp.isfinite(state.errors[it - 1])), + jnp.logical_not(jnp.isfinite(state.costs[it - 1])) + ) + + +def run( + ot_prob: quadratic_problem.QuadraticProblem, + solver: LRGromovWasserstein, + init: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray], + Optional[jnp.ndarray]], +) -> LRGWOutput: + """Run loop of the solver, outputting a state upgraded to an output.""" + out = sinkhorn.iterations(ot_prob, solver, init) + out = out.set_cost( + ot_prob, lse_mode=solver.lse_mode, use_danskin=solver.use_danskin + ) + return out.set(ot_prob=ot_prob) + + +def dykstra_solution_error( + q: jnp.ndarray, r: jnp.ndarray, ot_prob: quadratic_problem.QuadraticProblem, + norm_error: Tuple[int, ...] +) -> jnp.ndarray: + """Compute solution error. + + Since only balanced case is available for LR, this is marginal deviation. + + Args: + q: first factor of solution. + r: second factor of solution. + ot_prob: linear problem. + norm_error: int, p-norm used to compute error. + + Returns: + one or possibly many numbers quantifying deviation to true marginals. + """ + norm_error = jnp.array(norm_error) + # Update the error + err = jnp.sum( + jnp.abs(jnp.sum(q, axis=1) - ot_prob.a) ** norm_error[:, None], axis=1 + ) ** (1.0 / norm_error) + err += jnp.sum( + jnp.abs(jnp.sum(r, axis=1) - ot_prob.b) ** norm_error[:, None], axis=1 + ) ** (1.0 / norm_error) + err += jnp.sum( + jnp.abs(jnp.sum(q, axis=0) - jnp.sum(r, axis=0)) ** norm_error[:, None], + axis=1 + ) ** (1.0 / norm_error) + + return err + + +def _linearized_geometry( + prob: quadratic_problem.QuadraticProblem, + *, + q: jnp.ndarray, + r: jnp.ndarray, + g: jnp.ndarray, +) -> low_rank.LRCGeometry: + inv_sqrt_g = 1.0 / jnp.sqrt(g[None, :]) + + # TODO(michalk8): below is for squared loss, handle KL loss in the future; + # will need to be updated in many other places as well + tmp1 = -4.0 * prob.geom_xx.apply_cost(q, axis=1) * inv_sqrt_g + tmp2 = prob.geom_yy.apply_cost(r, axis=1) * inv_sqrt_g + return low_rank.LRCGeometry(tmp1, tmp2) diff --git a/ott/build/lib/ott/solvers/quadratic/gw_barycenter.py b/ott/build/lib/ott/solvers/quadratic/gw_barycenter.py new file mode 100644 index 0000000..f0d350b --- /dev/null +++ b/ott/build/lib/ott/solvers/quadratic/gw_barycenter.py @@ -0,0 +1,339 @@ +# 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. +from functools import partial +from typing import Any, Dict, NamedTuple, Optional, Sequence, Tuple, Union + +import jax +import jax.numpy as jnp + +from ott import utils +from ott.geometry import pointcloud +from ott.math import fixed_point_loop +from ott.problems.linear import linear_problem +from ott.problems.quadratic import gw_barycenter +from ott.solvers import was_solver +from ott.solvers.quadratic import gromov_wasserstein + +__all__ = ["GWBarycenterState", "GromovWassersteinBarycenter"] + + +class GWBarycenterState(NamedTuple): + """State of the GW barycenter problem. + + Args: + cost: Barycenter cost matrix of shape ``[bar_size, bar_size]``. + x: Barycenter features of shape ``[bar_size, ndim_fused]``. + Only used in the fused case. + a: Weights of the barycenter of shape ``[bar_size,]``. + errors: Array of shape + ``[max_iter, num_measures, quad_max_iter, lin_outer_iter]`` containing + the GW errors at each iteration. + costs: Array of shape ``[max_iter,]`` containing the cost at each iteration. + costs_bary: Array of shape ``[max_iter, num_measures]`` containing the + cost between the individual measures and the barycenter at each iteration. + gw_convergence: Array of shape ``[max_iter,]`` containing the convergence + of all GW problems at each iteration. + """ + cost: Optional[jnp.ndarray] = None + x: Optional[jnp.ndarray] = None + a: Optional[jnp.ndarray] = None + errors: Optional[jnp.ndarray] = None + costs: Optional[jnp.ndarray] = None + costs_bary: Optional[jnp.ndarray] = None + gw_convergence: Optional[jnp.ndarray] = None + + def set(self, **kwargs: Any) -> "GWBarycenterState": + """Return a copy of self, possibly with overwrites.""" + return self._replace(**kwargs) + + @property + def n_iters(self) -> int: + """Number of iterations.""" + if self.gw_convergence is None: + return -1 + return jnp.sum(self.gw_convergence != -1) + + +@jax.tree_util.register_pytree_node_class +class GromovWassersteinBarycenter(was_solver.WassersteinSolver): + """Gromov-Wasserstein barycenter solver. + + Args: + epsilon: Entropy regularizer. + min_iterations: Minimum number of iterations. + max_iterations: Maximum number of outermost iterations. + threshold: Convergence threshold. + store_inner_errors: Whether to store the errors of the GW solver, as well + as its linear solver, at each iteration for each measure. + quad_solver: The GW solver. + kwargs: Keyword argument for + :class:`~ott.solvers.quadratic.gromov_wasserstein.GromovWasserstein`. + Only used when ``quad_solver = None``. + """ + + def __init__( + self, + epsilon: Optional[float] = None, + min_iterations: int = 5, + max_iterations: int = 50, + threshold: float = 1e-3, + store_inner_errors: bool = False, + quad_solver: Optional[gromov_wasserstein.GromovWasserstein] = None, + # TODO(michalk8): maintain the API compatibility with `was_solver` + # but makes passing kwargs with the same name to `quad_solver` impossible + # will be fixed when refactoring the solvers + # note that `was_solver` also suffers from this + **kwargs: Any, + ): + super().__init__( + epsilon=epsilon, + min_iterations=min_iterations, + max_iterations=max_iterations, + threshold=threshold, + store_inner_errors=store_inner_errors, + ) + if quad_solver is None: + kwargs["epsilon"] = epsilon + # TODO(michalk8): store only GW errors? + kwargs["store_inner_errors"] = store_inner_errors + self._quad_solver = gromov_wasserstein.GromovWasserstein(**kwargs) + else: + self._quad_solver = quad_solver + + def __call__( + self, problem: gw_barycenter.GWBarycenterProblem, bar_size: int, + **kwargs: Any + ) -> GWBarycenterState: + """Solver the (fused) GW barycenter problem. + + Args: + problem: The GW barycenter problem. + bar_size: Size of the barycenter. + kwargs: Keyword arguments for :meth:`init_state`. + + Returns: + The solution. + """ + state = self.init_state(problem, bar_size, **kwargs) + state = iterations(self, problem, state) + return self.output_from_state(state) + + def init_state( + self, + problem: gw_barycenter.GWBarycenterProblem, + bar_size: int, + bar_init: Optional[Union[jnp.ndarray, Tuple[jnp.ndarray, + jnp.ndarray]]] = None, + a: Optional[jnp.ndarray] = None, + rng: Optional[jax.Array] = None, + ) -> GWBarycenterState: + """Initialize the (fused) Gromov-Wasserstein barycenter state. + + Args: + problem: The barycenter problem. + bar_size: Size of the barycenter. + bar_init: Initial barycenter value. Can be one of the following: + + - ``None`` - randomly initialize the barycenter. + - :class:`jax.numpy.ndarray` - barycenter cost matrix of shape + ``[bar_size, bar_size]``. + Only used in the non-fused case. + - :class:`tuple` of :class:`jax.numpy.ndarray` - the first array + corresponds to a cost matrix of shape ``[bar_size, bar_size]``, + the second array is a ``[bar_size, ndim_fused]`` feature matrix used + in the fused case. + + a: An array of shape ``[bar_size,]`` containing the barycenter weights. + rng: Random key for seeding used when ``bar_init = None``. + + Returns: + The initial barycenter state. + """ + if a is None: + a = jnp.ones((bar_size,)) / bar_size + else: + assert a.shape == (bar_size,) + + if bar_init is None: + rng = utils.default_prng_key(rng) + _, b = problem.segmented_y_b + rngs = jax.random.split(rng, problem.num_measures) + linear_solver = self._quad_solver.linear_ot_solver + + transports = init_transports(linear_solver, rngs, a, b, problem.epsilon) + x = problem.update_features(transports, a) if problem.is_fused else None + cost = problem.update_barycenter(transports, a) + else: + cost, x = bar_init if isinstance(bar_init, tuple) else (bar_init, None) + assert cost.shape == (bar_size, bar_size) + if problem.is_fused: + assert x is not None, "Barycenter features are not initialized." + assert x.shape == (bar_size, problem.ndim_fused) + + num_iter = self.max_iterations + if self.store_inner_errors: + # TODO(michalk8): in the future, think about how to do this in general + errors = -jnp.ones(( + num_iter, problem.num_measures, self._quad_solver.max_iterations, + self._quad_solver.linear_ot_solver.outer_iterations + )) + else: + errors = None + + costs = -jnp.ones((num_iter,)) + costs_bary = -jnp.ones((num_iter, problem.num_measures)) + gw_convergence = -jnp.ones((num_iter,)) + return GWBarycenterState( + cost=cost, + x=x, + a=a, + errors=errors, + costs=costs, + costs_bary=costs_bary, + gw_convergence=gw_convergence + ) + + def update_state( + self, + state: GWBarycenterState, + iteration: int, + problem: gw_barycenter.GWBarycenterProblem, + store_errors: bool = True, + ) -> Tuple[float, bool, jnp.ndarray, Optional[jnp.ndarray]]: + """Solve the (fused) Gromov-Wasserstein barycenter problem.""" + + def solve_gw( + state: GWBarycenterState, b: jnp.ndarray, y: jnp.ndarray, + f: Optional[jnp.ndarray] + ) -> Tuple[float, bool, jnp.ndarray, Optional[jnp.ndarray]]: + quad_problem = problem._create_problem(state, y=y, b=b, f=f) + out = self._quad_solver(quad_problem) + return ( + out.reg_gw_cost, out.converged, out.matrix, + out.errors if store_errors else None + ) + + in_axes = [None, 0, 0] + in_axes += [0] if problem.is_fused else [None] + solve_fn = jax.vmap(solve_gw, in_axes=in_axes) + + y, b = problem.segmented_y_b + y_f = problem.segmented_y_fused + costs, convergeds, transports, errors = solve_fn(state, b, y, y_f) + + cost = jnp.sum(costs * problem.weights) + costs_bary = state.costs_bary.at[iteration].set(costs) + costs = state.costs.at[iteration].set(cost) + + converged = jnp.all(convergeds) + gw_convergence = state.gw_convergence.at[iteration].set(converged) + + if self.store_inner_errors: + errors = state.errors.at[iteration, ...].set(errors) + else: + errors = None + + x = problem.update_features( + transports, state.a + ) if problem.is_fused else state.x + cost = problem.update_barycenter(transports, state.a) + return state.set( + cost=cost, + x=x, + costs=costs, + costs_bary=costs_bary, + errors=errors, + gw_convergence=gw_convergence + ) + + def output_from_state(self, state: GWBarycenterState) -> GWBarycenterState: + """No-op.""" + # TODO(michalk8): just for consistency with continuous barycenter + # will be refactored in the future to create an output + return state + + def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + children, aux = super().tree_flatten() + return children + [self._quad_solver], aux + + @classmethod + def tree_unflatten( # noqa: D102 + cls, aux_data: Dict[str, Any], children: Sequence[Any] + ) -> "GromovWassersteinBarycenter": + epsilon, _, threshold, quad_solver = children + return cls( + epsilon=epsilon, + threshold=threshold, + quad_solver=quad_solver, + **aux_data, + ) + + +@partial(jax.vmap, in_axes=[None, 0, None, 0, None]) +def init_transports( + solver, rng: jax.Array, a: jnp.ndarray, b: jnp.ndarray, + epsilon: Optional[float] +) -> jnp.ndarray: + """Initialize random 2D point cloud and solve the linear OT problem. + + Args: + solver: Linear OT solver. + rng: Random key for seeding. + a: Source marginals (e.g., for barycenter) of shape ``[bar_size,]``. + b: Target marginals of shape ``[max_measure_size,]``. + epsilon: Entropy regularization. + + Returns: + Transport map of shape ``[bar_size, max_measure_size]``. + """ + rng1, rng2 = jax.random.split(rng, 2) + x = jax.random.normal(rng1, shape=(len(a), 2)) + y = jax.random.normal(rng2, shape=(len(b), 2)) + geom = pointcloud.PointCloud( + x, y, epsilon=epsilon, src_mask=a > 0, tgt_mask=b > 0 + ) + problem = linear_problem.LinearProblem(geom, a=a, b=b) + return solver(problem).matrix + + +def iterations( # noqa: D103 + solver: GromovWassersteinBarycenter, + problem: gw_barycenter.GWBarycenterProblem, init_state: GWBarycenterState +) -> GWBarycenterState: + + def cond_fn( + iteration: int, constants: GromovWassersteinBarycenter, + state: GWBarycenterState + ) -> bool: + solver, _ = constants + return solver._continue(state, iteration) + + def body_fn( + iteration, constants: Tuple[GromovWassersteinBarycenter, + gw_barycenter.GWBarycenterProblem], + state: GWBarycenterState, compute_error: bool + ) -> GWBarycenterState: + del compute_error # always assumed true + solver, problem = constants + return solver.update_state(state, iteration, problem) + + return fixed_point_loop.fixpoint_iter( + cond_fn=cond_fn, + body_fn=body_fn, + min_iterations=solver.min_iterations, + max_iterations=solver.max_iterations, + inner_iterations=1, + constants=(solver, problem), + state=init_state, + ) diff --git a/ott/build/lib/ott/solvers/quadratic/lower_bound.py b/ott/build/lib/ott/solvers/quadratic/lower_bound.py new file mode 100644 index 0000000..f0868ad --- /dev/null +++ b/ott/build/lib/ott/solvers/quadratic/lower_bound.py @@ -0,0 +1,96 @@ +# 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. +from typing import TYPE_CHECKING, Any, Optional + +import jax +import jax.tree_util as jtu + +from ott.geometry import pointcloud +from ott.problems.quadratic import quadratic_problem +from ott.solvers import linear +from ott.solvers.linear import sinkhorn + +if TYPE_CHECKING: + from ott.geometry import distrib_costs + +__all__ = ["LowerBoundSolver"] + + +@jtu.register_pytree_node_class +class LowerBoundSolver: + """Lower bound OT solver. + + Computes the third lower bound distance from :cite:`memoli:11`, def. 6.3. + + Args: + epsilon: Entropy regularization for the resulting linear problem. + distrib_cost: Univariate Wasserstein cost, used to compare two point clouds + in different spaces, where each point is seen as its distribution of costs + to other points in its point cloud. + """ + + def __init__( + self, + epsilon: Optional[float] = None, + distrib_cost: Optional["distrib_costs.UnivariateWasserstein"] = None, + ): + from ott.geometry import distrib_costs + + self.epsilon = epsilon + self.distrib_cost = ( + distrib_costs.UnivariateWasserstein() + if distrib_cost is None else distrib_cost + ) + + def __call__( + self, + prob: quadratic_problem.QuadraticProblem, + epsilon: Optional[float] = None, + rng: Optional[jax.Array] = None, + **kwargs: Any + ) -> sinkhorn.SinkhornOutput: + """Compute a lower-bound for the GW problem using a simple linearization. + + This solver handles a quadratic problem by computing a proxy ``[n, m]`` + cost-matrix, injecting it into a linear OT solver to output a first an OT + matrix that can be used either to linearize/initialize the resolution + ot the GW problem, or more simply as a simple GW solution. + + Args: + prob: Quadratic OT problem. + epsilon: Entropic regularization passed on to solve the linearization of + the quadratic problem using 1D costs. + rng: Random key, possibly used when computing 1D costs when using + subsampling. + kwargs: Keyword arguments for :func:`~ott.solvers.linear.solve`. + + Returns: + A linear OT output, an approximation of the OT coupling obtained using + the lower bound provided by :cite:`memoli:11`. + """ + dists_xx = prob.geom_xx.cost_matrix + dists_yy = prob.geom_yy.cost_matrix + + geom_xy = pointcloud.PointCloud( + dists_xx, dists_yy, cost_fn=self.distrib_cost, epsilon=self.epsilon + ) + return linear.solve(geom_xy, **kwargs) + + def tree_flatten(self): # noqa: D102 + return (self.epsilon, self.distrib_cost), None + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + del aux_data + return cls(*children) diff --git a/ott/build/lib/ott/solvers/was_solver.py b/ott/build/lib/ott/solvers/was_solver.py new file mode 100644 index 0000000..6b25718 --- /dev/null +++ b/ott/build/lib/ott/solvers/was_solver.py @@ -0,0 +1,131 @@ +# 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. +from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple, Union + +import jax +import jax.numpy as jnp + +if TYPE_CHECKING: + from ott.solvers.linear import continuous_barycenter, sinkhorn, sinkhorn_lr + +__all__ = ["WassersteinSolver"] + +State = Union[ + "sinkhorn.SinkhornState", + "sinkhorn_lr.LRSinkhornState", + "continuous_barycenter.FreeBarycenterState", +] + + +# TODO(michalk8): refactor to have generic nested solver API +@jax.tree_util.register_pytree_node_class +class WassersteinSolver: + """A generic solver for problems that use a linear problem in inner loop.""" + + def __init__( + self, + epsilon: Optional[float] = None, + rank: int = -1, + linear_ot_solver: Optional[ + Union["sinkhorn.Sinkhorn", "sinkhorn_lr.LRSinkhorn"] + ] = None, + min_iterations: int = 5, + max_iterations: int = 50, + threshold: float = 1e-3, + store_inner_errors: bool = False, + kwargs: Any = {}, + ): + from ott.solvers.linear import sinkhorn, sinkhorn_lr + + default_epsilon = 1.0 + # Set epsilon value to default if needed, but keep track of whether None was + # passed to handle the case where a linear_ot_solver is passed or not. + self.epsilon = epsilon if epsilon is not None else default_epsilon + self.rank = rank + self.linear_ot_solver = linear_ot_solver + if self.linear_ot_solver is None: + # Detect if user requests low-rank solver. In that case the + # default_epsilon makes little sense, since it was designed for GW. + if self.is_low_rank: + if epsilon is None: + # Use default entropic regularization in LRSinkhorn if None was passed + self.linear_ot_solver = sinkhorn_lr.LRSinkhorn( + rank=self.rank, **kwargs + ) + else: + # If epsilon is passed, use it to replace the default LRSinkhorn value + self.linear_ot_solver = sinkhorn_lr.LRSinkhorn( + rank=self.rank, epsilon=self.epsilon, **kwargs + ) + else: + # When using Entropic GW, epsilon is not handled inside Sinkhorn, + # but rather added back to the Geometry object re-instantiated + # when linearizing the problem. Therefore, no need to pass it to solver. + self.linear_ot_solver = sinkhorn.Sinkhorn(**kwargs) + + self.min_iterations = min_iterations + self.max_iterations = max_iterations + self.threshold = threshold + self.store_inner_errors = store_inner_errors + self._kwargs = kwargs + + @property + def is_low_rank(self) -> bool: + """Whether the solver is low-rank.""" + return self.rank > 0 + + def tree_flatten( + self, + ) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + return ( + [self.epsilon, self.linear_ot_solver, self.threshold], + { + "min_iterations": self.min_iterations, + "max_iterations": self.max_iterations, + "rank": self.rank, + "store_inner_errors": self.store_inner_errors, + **self._kwargs, + }, + ) + + @classmethod + def tree_unflatten( # noqa: D102 + cls, aux_data: Dict[str, Any], children: Sequence[Any] + ) -> "WassersteinSolver": + epsilon, linear_ot_solver, threshold = children + return cls( + epsilon=epsilon, + linear_ot_solver=linear_ot_solver, + threshold=threshold, + **aux_data, + ) + + def _converged(self, state: State, iteration: int) -> bool: + costs, i, tol = state.costs, iteration, self.threshold + return jnp.logical_and( + i >= 2, jnp.isclose(costs[i - 2], costs[i - 1], rtol=tol) + ) + + def _diverged(self, state: State, iteration: int) -> bool: + return jnp.logical_not(jnp.isfinite(state.costs[iteration - 1])) + + def _continue(self, state: State, iteration: int) -> bool: + """Continue while not(converged) and not(diverged).""" + return jnp.logical_or( + iteration <= 2, + jnp.logical_and( + jnp.logical_not(self._diverged(state, iteration)), + jnp.logical_not(self._converged(state, iteration)), + ), + ) diff --git a/ott/build/lib/ott/tools/__init__.py b/ott/build/lib/ott/tools/__init__.py new file mode 100644 index 0000000..75dbfae --- /dev/null +++ b/ott/build/lib/ott/tools/__init__.py @@ -0,0 +1,21 @@ +# 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. +from . import ( + gaussian_mixture, + k_means, + plot, + segment_sinkhorn, + sinkhorn_divergence, + soft_sort, +) diff --git a/ott/build/lib/ott/tools/gaussian_mixture/__init__.py b/ott/build/lib/ott/tools/gaussian_mixture/__init__.py new file mode 100644 index 0000000..3195d77 --- /dev/null +++ b/ott/build/lib/ott/tools/gaussian_mixture/__init__.py @@ -0,0 +1,14 @@ +# 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. +from . import fit_gmm_pair, gaussian, gaussian_mixture, gaussian_mixture_pair diff --git a/ott/build/lib/ott/tools/gaussian_mixture/fit_gmm.py b/ott/build/lib/ott/tools/gaussian_mixture/fit_gmm.py new file mode 100644 index 0000000..fe869ab --- /dev/null +++ b/ott/build/lib/ott/tools/gaussian_mixture/fit_gmm.py @@ -0,0 +1,303 @@ +# 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. +r"""Fit a Gaussian mixture model. + +Sample usage: + +# initialize GMM with K-means++ +gmm_init = fit_gmm.initialize( + rng=rng, + points=my_points, + point_weights=None, + n_components=COMPONENTS) + +# refine GMM parameters using EM +gmm = fit_gmm.fit_model_em( + gmm=gmm_init, + points=my_points, + point_weights=None, + steps=10, + verbose=True) + + +We fit the model using EM. Below we'll use notation following +https://en.wikipedia.org/wiki/Expectation%E2%80%93maximization_algorithm + +Our data X is generated by a Gaussian mixture with unknown parameters \Theta. +We denote the (unobserved) component that gave rise to each point as Z. + +In EM we start with an initial estimate of $\Theta$, $\Theta^{(0)}$, and we +then iteratively update it via + +$$\Theta^{(t+1)} = \argmax_{\Theta} Q(\Theta|\Theta^{(t)})$$ + +where + +$$ +Q(\Theta|\Theta(t)) = E_{Z|X,\Theta^{(t)}} \left[ \log L(\Theta; X, Z) \right] +$$ +""" + +from typing import Optional + +import jax +import jax.numpy as jnp + +from ott.tools.gaussian_mixture import gaussian_mixture + +__all__ = ["initialize", "fit_model_em"] + +# EM algorithm for parameter estimation + + +def get_assignment_probs( + gmm: gaussian_mixture.GaussianMixture, points: jnp.ndarray +) -> jnp.ndarray: + r"""Get component assignment probabilities used in the E step of EM. + + Here we compute the component assignment probabilities p(Z|X, \Theta^{(t)}) + that we need to compute the expectation used for Q(\Theta|\Theta^{(t)}). + + Args: + gmm: GMM model + points: set of samples being fitted, shape (n, n_dimensions) + + Returns: + An array of assignment probabilities with shape (n, n_components) + """ + return jnp.exp(gmm.get_log_component_posterior(points)) + + +def get_q( + gmm: gaussian_mixture.GaussianMixture, + assignment_probs: jnp.ndarray, + points: jnp.ndarray, + point_weights: Optional[jnp.ndarray] = None, +) -> float: + r"""Get Q(\Theta|\Theta^{(t)}). + + Args: + gmm: GaussianMixture with parameters \Theta + assignment_probs: p(Z|X, \Theta^{(t)}) as computed by get_assignment_probs + points: observations X + point_weights: optional set of weights for the samples. If None, use + a weight of 1/n where n is the number of points. + + Returns: + Q(\Theta|\Theta^{(t)}) + """ + # log P(X, Z| \Theta) = log P(X|Z, \Theta) + log P(Z|\Theta) + loglik = (gmm.conditional_log_prob(points) + gmm.log_component_weights()) + if point_weights is None: + point_weights = jnp.ones(points.shape[0]) + return ( + jnp.sum(point_weights * jnp.sum(assignment_probs * loglik, axis=-1)) / + jnp.sum(point_weights) + ) + + +def log_prob_loss( + gmm: gaussian_mixture.GaussianMixture, + points: jnp.ndarray, + point_weights: Optional[jnp.ndarray] = None, +) -> float: + """Loss function: weighted mean of (-log prob of observations). + + Args: + gmm: GMM model + points: set of samples being fitted + point_weights: optional set of weights for the samples. If None, use + a weight of 1/n where n is the number of points. + + Returns: + The GMM loss for the points. + """ + if point_weights is None: + return -jnp.mean(gmm.log_prob(points)) + return -jnp.sum(point_weights * gmm.log_prob(points)) / jnp.sum(point_weights) + + +def fit_model_em( + gmm: gaussian_mixture.GaussianMixture, + points: jnp.ndarray, + point_weights: Optional[jnp.ndarray], + steps: int, + jit: bool = True, + verbose: bool = False, +) -> gaussian_mixture.GaussianMixture: + """Fit a GMM using the EM algorithm. + + Args: + gmm: initial GMM model + points: set of samples to fit, shape (n, n_dimensions) + point_weights: optional set of weights for points, shape (n,). If None, + uses equal weights for all points. + steps: number of steps of EM to perform + jit: if True, compile functions + verbose: if True, print the loss at each step + + Returns: + A GMM with updated parameters. + """ + if point_weights is None: + point_weights = jnp.ones(points.shape[:-1]) + loss_fn = log_prob_loss + get_q_fn = get_q + e_step_fn = get_assignment_probs + m_step_fn = gaussian_mixture.GaussianMixture.from_points_and_assignment_probs + if jit: + loss_fn = jax.jit(loss_fn) + get_q_fn = jax.jit(get_q_fn) + e_step_fn = jax.jit(e_step_fn) + m_step_fn = jax.jit(m_step_fn) + + for i in range(steps): + assignment_probs = e_step_fn(gmm, points) + gmm_new = m_step_fn(points, point_weights, assignment_probs) + if gmm_new.has_nans(): + raise ValueError("NaNs in fit.") + if verbose: + loss = loss_fn(gmm_new, points, point_weights) + q = get_q_fn( + gmm=gmm_new, + assignment_probs=assignment_probs, + points=points, + point_weights=point_weights + ) + print(f"{i} q={q} -log prob={loss}") # noqa: T201 + gmm = gmm_new + return gmm + + +# KMeans++ for initialization +# See https://en.wikipedia.org/wiki/K-means%2B%2B for details + + +def _get_dist_sq(points: jnp.ndarray, loc: jnp.ndarray) -> jnp.ndarray: + """Get the squared distance from each point to each loc.""" + + def _dist_sq_one_loc(points: jnp.ndarray, loc: jnp.ndarray) -> jnp.ndarray: + return jnp.sum((points - loc[None]) ** 2, axis=-1) + + dist_sq_fn = jax.vmap(_dist_sq_one_loc, in_axes=(None, 0), out_axes=1) + return dist_sq_fn(points, loc) + + +def _get_locs( + rng: jax.Array, points: jnp.ndarray, n_components: int +) -> jnp.ndarray: + """Get the initial component means. + + Args: + rng: jax.random key + points: (n, n_dimensions) array of observations + n_components: desired number of components + + Returns: + (n_components, n_dimensions) array of means. + """ + points = points.copy() + n_points = points.shape[0] + weights = jnp.ones(n_points) / n_points + rng, subrng = jax.random.split(rng) + index = jax.random.choice(subrng, a=points.shape[0], p=weights) + loc = points[index] + points = jnp.concatenate([points[:index], points[index + 1:]], axis=0) + + locs = loc[None] + for _ in range(n_components - 1): + dist_sq = _get_dist_sq(points, locs) + min_dist_sq = jnp.min(dist_sq, axis=-1) + weights = min_dist_sq / jnp.sum(min_dist_sq) + rng, subrng = jax.random.split(rng) + index = jax.random.choice(subrng, a=points.shape[0], p=weights) + loc = points[index] + points = jnp.concatenate([points[:index], points[index + 1:]], axis=0) + locs = jnp.concatenate([locs, loc[None]], axis=0) + return locs + + +def from_kmeans_plusplus( + rng: jax.Array, + points: jnp.ndarray, + point_weights: Optional[jnp.ndarray], + n_components: int, +) -> gaussian_mixture.GaussianMixture: + """Initialize a GMM via a single pass of K-means++. + + Args: + rng: jax.random key + points: (n, n_dimensions) array of observations + point_weights: (n,) array of weights for points + n_components: desired number of components + + Returns: + An initial Gaussian mixture model. + + Raises: + ValueError if any fitted parameters are non-finite. + """ + rng, subrng = jax.random.split(rng) + locs = _get_locs(rng=subrng, points=points, n_components=n_components) + dist_sq = _get_dist_sq(points, locs) + assignment_prob = (dist_sq == jnp.min(dist_sq, + axis=-1)[:, None]).astype(points.dtype) + del dist_sq + + if point_weights is None: + point_weights = jnp.ones_like(points[..., 0]) + return gaussian_mixture.GaussianMixture.from_points_and_assignment_probs( + points=points, + point_weights=point_weights, + assignment_probs=assignment_prob + ) + + +def initialize( + rng: jax.Array, + points: jnp.ndarray, + point_weights: Optional[jnp.ndarray], + n_components: int, + n_attempts: int = 50, + verbose: bool = False +) -> gaussian_mixture.GaussianMixture: + """Initialize a GMM via K-means++ with retries on failure. + + Args: + rng: jax.random key + points: (n, n_dimensions) array of observations + point_weights: (n,) array of weights for points + n_components: desired number of components + n_attempts: number of attempts to initialize before failing + verbose: if True, print status information + + Returns: + An initial Gaussian mixture model. + + Raises: + ValueError if initialization was unsuccessful after n_attempts attempts. + """ + for attempt in range(n_attempts): + rng, subrng = jax.random.split(rng) + try: + return from_kmeans_plusplus( + rng=subrng, + points=points, + point_weights=point_weights, + n_components=n_components + ) + except ValueError: + if verbose: + print(f"Failed to initialize, attempt {attempt}.") # noqa: T201 + raise ValueError("Failed to initialize.") diff --git a/ott/build/lib/ott/tools/gaussian_mixture/fit_gmm_pair.py b/ott/build/lib/ott/tools/gaussian_mixture/fit_gmm_pair.py new file mode 100644 index 0000000..7ecde26 --- /dev/null +++ b/ott/build/lib/ott/tools/gaussian_mixture/fit_gmm_pair.py @@ -0,0 +1,393 @@ +# 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. +r"""Fit 2 GMMs to 2 point clouds using likelihood and (approx) W2 distance. + +Suppose we have two large point clouds and want to estimate a coupling and a +W2 distance between them. :cite:`delon:20` propose fitting a GMM to each +point cloud while simultaneously minimizing a Wasserstein-like distance +called MW2 between the fitted GMMs. MW2 is an upper bound on W2, +the Wasserstein distance between the GMMs. Here we implement +their algorithm as well as a generalization that allows for reweightings using +generalized, penalized expectation-maximization +(see section 6.2 of :cite:`delon:20`). + +As in `fit_gmm.py`, we assume that the observations $X_0$ and $X_1$ from +batches 0 and 1 are generated by GMMs with parameters $\Theta_0$ and $\Theta_1$, +respectively. We will use $\Theta$ to denote the combined parameters +for the two GMMs. We denote the (unobserved) components that gave rise to the +observations $X_i$ as $Z_i$. + +Our goal is to maximize a weighted sum of the likelihood of the observations $X$ +under the fitted GMMs and a measure of distance, $MW_2$, between the fitted +GMMs. The problem would be a straightforward maximization exercise if we knew +the components $Z$ that generated each observation $X$. Because the $Z$ are +unobserved, however, we use EM: + +We start with an initial estimate of $\Theta$, $\Theta^{(t)}$. + +* The E-step: We use the current $\Theta^{(t)}$ to estimate the likelihood of + all possible cluster attributions for each observation $X$. + +* The M-step: We form the function $Q(\Theta|\Theta^{(t)})$, + the log likelihood of our observations averaged over all possible + assignments. We then obtain an updated parameter estimate, $\Theta^{(t+1)}$, + by numerically maximizing the sum of $Q$ and our GMM distance penalty. + +It can be shown that if we maximize the penalized $Q$ above, this procedure will +increase or leave unchanged the penalized log likelihood for $\Theta$. We +iterate over these two steps until convergence. Note that the resulting +estimate for $\Theta$ may only be a *local* maximum of the penalized +likelihood function. + + +Sample usage: + +# (Note that we usually initialize a pair to a single GMM that we fit to a +# pooled set, then the two GMMs separate as we optimize the pair.) +pair_init = gaussian_mixture_pair.GaussianMixturePair( + gmm0=gmm0, + gmm1=gmm1, + epsilon=1.e-2, + tau=1.) +fit_model_em_fn = fit_gmm_pair.get_fit_model_em_fn( + weight_transport=0.1, + weight_splitting=1., + epsilon=pair_init.epsilon, + jit=True) +pair, loss = fit_model_em_fn( + pair=pair_init, + points0=samples_gmm0, + points1=samples_gmm1, + point_weights0=None, + point_weights1=None, + em_steps=30, + m_steps=20, + verbose=True) +""" +# TODO(geoffd): look into refactoring so we jit higher level functions + +import functools +import math +from typing import Callable, NamedTuple, Optional, Tuple + +import jax +import jax.numpy as jnp + +from ott.tools.gaussian_mixture import ( + fit_gmm, + gaussian_mixture, + gaussian_mixture_pair, +) + +__all__ = ["get_fit_model_em_fn"] + +LOG2 = math.log(2) + + +class Observations(NamedTuple): + """Weighted observations and their E-step assignment probabilities.""" + + points: jnp.ndarray + point_weights: jnp.ndarray + assignment_probs: jnp.ndarray + + +# Model fit + + +def get_q( + gmm: gaussian_mixture.GaussianMixture, obs: Observations +) -> jnp.ndarray: + r"""Get Q(\Theta|\Theta^{(t)}). + + Here Q is the log likelihood for our observations based on the current + parameter estimates for \Theta and averaged over the current component + assignment probabilities. See the overview of EM above for more details. + + Args: + gmm: GMM model parameterized by Theta + obs: weighted observations with component assignments computed in the E step + for \Theta^{(t)} + + Returns: + Q(\Theta|\Theta^{(t)}) + """ + # Q = E_Z log p(X, Z| Theta) + # = \sum_Z P(Z|X, Theta^(t)) [log p(X, Z | Theta)] + # Here P(Z|X, theta^(t)) is the set of assignment probabilities + # we computed in the E step. + # log p(X, Z| theta) is given by + log_p_x_z = ( + gmm.conditional_log_prob(obs.points) + # p(X | Z, theta) + gmm.log_component_weights() + ) # p(Z | theta) + return ( + jnp.sum( + obs.point_weights * + jnp.sum(log_p_x_z * obs.assignment_probs, axis=-1), + axis=0 + ) / jnp.sum(obs.point_weights, axis=0) + ) + + +# Objective function + + +@functools.lru_cache +def get_objective_fn(weight_transport: float): + """Get the total loss function with static parameters in a closure. + + Args: + weight_transport: weight for the transport penalty + + Returns: + A function that returns the objective for a GaussianMixturePair. + """ + + def _objective_fn( + pair: gaussian_mixture_pair.GaussianMixturePair, + obs0: Observations, + obs1: Observations, + ) -> jnp.ndarray: + """Compute the objective function for a pair of GMMs. + + Args: + pair: pair of GMMs + coupling for which to evaluate the objective + obs0: first set of observations + obs1: second set of observations + + Returns: + The objective to be minimized in the M-step. + """ + q0 = get_q(gmm=pair.gmm0, obs=obs0) + q1 = get_q(gmm=pair.gmm1, obs=obs1) + cost_matrix = pair.get_cost_matrix() + sinkhorn_output = pair.get_sinkhorn(cost_matrix=cost_matrix) + transport_penalty = sinkhorn_output.reg_ot_cost + return q0 + q1 - weight_transport * transport_penalty + + return _objective_fn + + +def print_losses( + iteration: int, weight_transport: float, + pair: gaussian_mixture_pair.GaussianMixturePair, obs0: Observations, + obs1: Observations +): + """Print the loss components for diagnostic purposes.""" + q0 = get_q(gmm=pair.gmm0, obs=obs0) + q1 = get_q(gmm=pair.gmm1, obs=obs1) + cost_matrix = pair.get_cost_matrix() + sinkhorn_output = pair.get_sinkhorn(cost_matrix=cost_matrix) + transport_penalty = sinkhorn_output.reg_ot_cost + objective = q0 + q1 - weight_transport * transport_penalty + + print( # noqa: T201 + f"{iteration:3d} {q0:.3f} {q1:.3f} " + f"transport:{transport_penalty:.3f} " + f"objective:{objective:.3f}" + ) + + +# The E-step for a single GMM + + +def do_e_step( # noqa: D103 + e_step_fn: Callable[[gaussian_mixture.GaussianMixture, jnp.ndarray], + jnp.ndarray], + gmm: gaussian_mixture.GaussianMixture, + points: jnp.ndarray, + point_weights: jnp.ndarray, +) -> Observations: + assignment_probs = e_step_fn(gmm, points) + return Observations( + points=points, + point_weights=point_weights, + assignment_probs=assignment_probs + ) + + +# The M-step + + +def get_m_step_fn(learning_rate: float, objective_fn, jit: bool): + """Get a function that performs the M-step of the EM algorithm. + + We precompile and precompute a few quantities that we put into a closure. + + Args: + learning_rate: learning rate to use for the Adam optimizer + objective_fn: the objective function to maximize + jit: if True, precompile key methods + + Returns: + A function that performs the M-step of EM. + """ + import optax + + def _m_step_fn( + pair: gaussian_mixture_pair.GaussianMixturePair, + obs0: Observations, + obs1: Observations, + steps: int, + ) -> gaussian_mixture_pair.GaussianMixturePair: + """Perform the M-step on a pair of Gaussian mixtures. + + Args: + pair: GMM parameters to optimize + obs0: first set of observations + obs1: second set of observations + steps: number of optimization steps to use when maximizing the objective + + Returns: + A GaussianMixturePair with updated parameters. + """ + state = opt_init((pair,)) + + for _ in range(steps): + grad_objective = grad_objective_fn(pair, obs0, obs1) + updates, state = opt_update(grad_objective, state, (pair,)) + (pair,) = optax.apply_updates((pair,), updates) + for j, gmm in enumerate((pair.gmm0, pair.gmm1)): + if gmm.has_nans(): + raise ValueError(f"NaN in gmm{j}") + return pair + + grad_objective_fn = jax.grad(objective_fn, argnums=(0,)) + if jit: + grad_objective_fn = jax.jit(grad_objective_fn) + + opt_init, opt_update = optax.chain( + # Set the parameters of Adam. Note the learning_rate is not here. + optax.scale_by_adam(b1=0.9, b2=0.999, eps=1e-8), + optax.scale(learning_rate) + ) + + return _m_step_fn + + +def get_fit_model_em_fn( + weight_transport: float, + learning_rate: float = 0.001, + jit: bool = True, +): + """Get a function that performs penalized EM. + + We precompile and precompute a few quantities that we put into a closure. + + Args: + weight_transport: weight for the transportation loss in the total loss + learning_rate: learning rate to use for the Adam optimizer + jit: if True, precompile key methods + + Returns: + A function that performs generalized, penalized EM. + """ + objective_fn = get_objective_fn(weight_transport=weight_transport) + e_step_fn = fit_gmm.get_assignment_probs + + if jit: + objective_fn = jax.jit(objective_fn) + e_step_fn = jax.jit(e_step_fn) + + m_step_fn = get_m_step_fn( + learning_rate=learning_rate, objective_fn=objective_fn, jit=jit + ) + + def _fit_model_em( + pair: gaussian_mixture_pair.GaussianMixturePair, + points0: jnp.ndarray, + points1: jnp.ndarray, + point_weights0: Optional[jnp.ndarray], + point_weights1: Optional[jnp.ndarray], + em_steps: int, + m_steps: int = 50, + verbose: bool = False, + ) -> Tuple[gaussian_mixture_pair.GaussianMixturePair, float]: + """Optimize a GaussianMixturePair using penalized EM. + + Args: + pair: GaussianMixturePair to optimize + points0: observations associated with pair.gmm0 + points1: observations associated with pair.gmm1 + point_weights0: weights for points0 + point_weights1: weights for points1 + em_steps: number of EM steps to perform + m_steps: number of gradient descent steps to perform in the M-step + verbose: if True, print status messages + + Returns: + An updated GaussianMixturePair and the final loss. + """ + if point_weights0 is None: + point_weights0 = jnp.ones(points0.shape[0]) + if point_weights1 is None: + point_weights1 = jnp.ones(points1.shape[0]) + + if pair.lock_gmm1: + obs1 = do_e_step( + e_step_fn=e_step_fn, + gmm=pair.gmm1, + points=points1, + point_weights=point_weights1 + ) + + for i in range(em_steps): + # E-step + obs0 = do_e_step( + e_step_fn=e_step_fn, + gmm=pair.gmm0, + points=points0, + point_weights=point_weights0 + ) + if not pair.lock_gmm1: + obs1 = do_e_step( + e_step_fn=e_step_fn, + gmm=pair.gmm1, + points=points1, + point_weights=point_weights1 + ) + + # print current losses + if verbose: + print_losses( + iteration=i, + weight_transport=weight_transport, + pair=pair, + obs0=obs0, + obs1=obs1 + ) + + # the M-step + pair = m_step_fn(pair=pair, obs0=obs0, obs1=obs1, steps=m_steps) + + # final E-step before computing the loss + obs0 = do_e_step( + e_step_fn=e_step_fn, + gmm=pair.gmm0, + points=points0, + point_weights=point_weights0 + ) + if not pair.lock_gmm1: + obs1 = do_e_step( + e_step_fn=e_step_fn, + gmm=pair.gmm1, + points=points1, + point_weights=point_weights1 + ) + + loss = objective_fn(pair=pair, obs0=obs0, obs1=obs1) + return pair, loss + + return _fit_model_em diff --git a/ott/build/lib/ott/tools/gaussian_mixture/gaussian.py b/ott/build/lib/ott/tools/gaussian_mixture/gaussian.py new file mode 100644 index 0000000..722f465 --- /dev/null +++ b/ott/build/lib/ott/tools/gaussian_mixture/gaussian.py @@ -0,0 +1,217 @@ +# 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 Optional, Union + +import jax +import jax.numpy as jnp + +from ott.tools.gaussian_mixture import scale_tril + +__all__ = ["Gaussian"] + +LOG2PI = math.log(2.0 * math.pi) + + +@jax.tree_util.register_pytree_node_class +class Gaussian: + """Normal distribution.""" + + def __init__(self, loc: jnp.ndarray, scale: scale_tril.ScaleTriL): + self._loc = loc + self._scale = scale + + @classmethod + def from_samples( + cls, + points: jnp.ndarray, + weights: Optional[jnp.ndarray] = None + ) -> "Gaussian": + """Construct a Gaussian from weighted samples. + + Unbiased, weighted covariance formula from `GSL + `_. + + Args: + points: [n x d] array of samples + weights: [n] array of weights + + Returns: + Gaussian. + """ + n = points.shape[0] + if weights is None: + weights = jnp.ones(n) / n + + mean = weights.dot(points) + centered_x = (points - mean) + scaled_centered_x = centered_x * weights.reshape(-1, 1) + cov = scaled_centered_x.T.dot(centered_x) / (1 - weights.dot(weights)) + return cls.from_mean_and_cov(mean=mean, cov=cov) + + @classmethod + def from_random( + cls, + rng: jax.Array, + n_dimensions: int, + stdev_mean: float = 0.1, + stdev_cov: float = 0.1, + ridge: Union[float, jnp.ndarray] = 0, + ) -> "Gaussian": + """Construct a random Gaussian. + + Args: + rng: jax.random key + n_dimensions: desired covariance dimensions + stdev_mean: standard deviation of location and log eigenvalues + (means for both are 0) + stdev_cov: standard deviated of the covariance + ridge: Offset for means. + + Returns: + A random Gaussian. + """ + rng, subrng0, subrng1 = jax.random.split(rng, num=3) + loc = jax.random.normal(subrng0, shape=(n_dimensions,)) * stdev_mean + ridge + scale = scale_tril.ScaleTriL.from_random( + subrng1, n_dimensions=n_dimensions, stdev=stdev_cov + ) + return cls(loc=loc, scale=scale) + + @classmethod + def from_mean_and_cov(cls, mean: jnp.ndarray, cov: jnp.ndarray) -> "Gaussian": + """Construct a Gaussian from a mean and covariance.""" + scale = scale_tril.ScaleTriL.from_covariance(cov) + return cls(loc=mean, scale=scale) + + @property + def loc(self) -> jnp.ndarray: + """Mean of the Gaussian.""" + return self._loc + + @property + def scale(self) -> scale_tril.ScaleTriL: + """Scale of the Gaussian.""" + return self._scale + + @property + def n_dimensions(self) -> int: + """Dimensionality of the Gaussian.""" + return self.loc.shape[-1] + + def covariance(self) -> jnp.ndarray: + """Covariance of the Gaussian.""" + return self.scale.covariance() + + def to_z(self, x: jnp.ndarray) -> jnp.ndarray: + r"""Transform :math:`x` to :math:`z = \frac{x - loc}{scale}`.""" + return self.scale.centered_to_z(x_centered=x - self.loc) + + def from_z(self, z: jnp.ndarray) -> jnp.ndarray: + r"""Transform :math:`z` to :math:`x = loc + scale \cdot z`.""" + return self.scale.z_to_centered(z=z) + self.loc + + def log_prob( + self, + x: jnp.ndarray, # (?, d) + ) -> jnp.ndarray: # (?, d) + """Log probability for a Gaussian with a diagonal covariance.""" + d = x.shape[-1] + z = self.to_z(x) + log_det = self.scale.log_det_covariance() + return ( + -0.5 * (d * LOG2PI + log_det[None] + jnp.sum(z ** 2, axis=-1)) + ) # (?, k) + + def sample(self, rng: jax.Array, size: int) -> jnp.ndarray: + """Generate samples from the distribution.""" + std_samples_t = jax.random.normal(rng, shape=(self.n_dimensions, size)) + return self.loc[None] + ( + jnp.swapaxes( + jnp.matmul(self.scale.cholesky(), std_samples_t), + axis1=-2, + axis2=-1 + ) + ) + + def w2_dist(self, other: "Gaussian") -> jnp.ndarray: + r"""Wasserstein distance :math:`W_2^2` to another Gaussian. + + .. math:: + + W_2^2 = ||\mu_0-\mu_1||^2 + + \text{trace} ( (\Lambda_0^\frac{1}{2} - \Lambda_1^\frac{1}{2})^2 ) + + Args: + other: other Gaussian + + Returns: + The :math:`W_2^2` distance between self and other + """ + delta_mean = jnp.sum((self.loc - other.loc) ** 2, axis=-1) + delta_sigma = self.scale.w2_dist(other.scale) + return delta_mean + delta_sigma + + def f_potential(self, dest: "Gaussian", points: jnp.ndarray) -> jnp.ndarray: + """Optimal potential for W2 distance between Gaussians. Evaluated on points. + + Args: + dest: Gaussian object + points: samples + + Returns: + Dual potential, f + """ + scale_matrix = self.scale.gaussian_map(dest_scale=dest.scale) + centered_x = points - self.loc + scaled_x = (scale_matrix @ centered_x.T) + + @jax.vmap + def batch_inner_product(x, y): + return x.dot(y) + + return ( + 0.5 * batch_inner_product(points, points) - + 0.5 * batch_inner_product(centered_x, scaled_x.T) - + points.dot(dest.loc) + ) + + def transport(self, dest: "Gaussian", points: jnp.ndarray) -> jnp.ndarray: + """Transport points according to map between two Gaussian measures. + + Args: + dest: Gaussian object + points: samples + + Returns: + Transported samples + """ + return self.scale.transport( + dest_scale=dest.scale, points=points - self.loc[None] + ) + dest.loc[None] + + def tree_flatten(self): # noqa: D102 + children = (self.loc, self.scale) + aux_data = {} + return children, aux_data + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + return cls(*children, **aux_data) + + def __hash__(self): + return jax.tree_util.tree_flatten(self).__hash__() + + def __eq__(self, other): + return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture.py b/ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture.py new file mode 100644 index 0000000..27a5689 --- /dev/null +++ b/ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture.py @@ -0,0 +1,329 @@ +# 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. +from typing import List, Tuple, Union + +import jax +import jax.numpy as jnp + +from ott.tools.gaussian_mixture import ( + gaussian, + linalg, + probabilities, + scale_tril, +) + +__all__ = ["GaussianMixture"] + + +def get_summary_stats_from_points_and_assignment_probs( + points: jnp.ndarray, point_weights: jnp.ndarray, + assignment_probs: jnp.ndarray +) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Get component summary stats from points and component probabilities. + + Args: + points: array of points, shape (n, n_dim) + point_weights: array of weights for the points, shape (n,) + assignment_probs: array of component assignment probabilities for the + points, shape (n, n_components) + + Returns: + Tuple containing for each component, + * the sample mean for each component, shape (n_components, n_dim) + * the sample covariance for each component, + shape (n_components, n_dim, n_dim) + * the weight for each component, + shape (n_components,) + """ + + def component_from_points(points, point_weights, assignment_probs): + component_weight = ( + jnp.sum(point_weights * assignment_probs) / jnp.sum(point_weights) + ) + component_mean, component_cov = linalg.get_mean_and_cov( + points=points, weights=point_weights * assignment_probs + ) + return component_mean, component_cov, component_weight + + components_from_points_fn = jax.vmap( + component_from_points, in_axes=(None, None, 1), out_axes=0 + ) + + return components_from_points_fn(points, point_weights, assignment_probs) + + +@jax.tree_util.register_pytree_node_class +class GaussianMixture: + """Gaussian Mixture model.""" + + def __init__( + self, loc: jnp.ndarray, scale_params: jnp.ndarray, + component_weight_ob: probabilities.Probabilities + ): + self._loc = loc + self._scale_params = scale_params + self._component_weight_ob = component_weight_ob + + @classmethod + def from_random( + cls, + rng: jax.Array, + n_components: int, + n_dimensions: int, + stdev_mean: float = 0.1, + stdev_cov: float = 0.1, + stdev_weights: float = 0.1, + ridge: Union[float, jnp.array] = 0, + ) -> "GaussianMixture": + """Construct a random GMM.""" + loc = [] + scale_params = [] + for _ in range(n_components): + rng, subrng = jax.random.split(rng) + component = gaussian.Gaussian.from_random( + subrng, + n_dimensions=n_dimensions, + stdev_mean=stdev_mean, + stdev_cov=stdev_cov, + ridge=ridge, + ) + loc.append(component.loc) + scale_params.append(component.scale.params) + loc = jnp.stack(loc, axis=0) + scale_params = jnp.stack(scale_params, axis=0) + weight_ob = probabilities.Probabilities.from_random( + subrng, + n_dimensions=n_components, + stdev=stdev_weights, + ) + return cls( + loc=loc, scale_params=scale_params, component_weight_ob=weight_ob + ) + + @classmethod + def from_mean_cov_component_weights( + cls, mean: jnp.ndarray, cov: jnp.ndarray, component_weights: jnp.ndarray + ): + """Construct a GMM from means, covariances, and component weights.""" + scale_params = [] + for i in range(cov.shape[0]): + scale_params.append(scale_tril.ScaleTriL.from_covariance(cov[i]).params) + scale_params = jnp.stack(scale_params, axis=0) + weight_ob = probabilities.Probabilities.from_probs(component_weights) + return cls( + loc=mean, scale_params=scale_params, component_weight_ob=weight_ob + ) + + @classmethod + def from_points_and_assignment_probs( + cls, + points: jnp.ndarray, + point_weights: jnp.ndarray, + assignment_probs: jnp.ndarray, + ) -> "GaussianMixture": + """Estimate a GMM from points and a set of component probabilities.""" + mean, cov, wts = get_summary_stats_from_points_and_assignment_probs( + points=points, + point_weights=point_weights, + assignment_probs=assignment_probs + ) + return cls.from_mean_cov_component_weights( + mean=mean, cov=cov, component_weights=wts + ) + + @property + def dtype(self): + """Dtype of the GMM parameters.""" + return self.loc.dtype + + @property + def n_dimensions(self): + """Number of dimensions of the GMM parameters.""" + return self._loc.shape[-1] + + @property + def n_components(self): + """Number of components of the GMM parameters.""" + return self._loc.shape[-2] + + @property + def loc(self) -> jnp.ndarray: + """Location parameters of the GMM.""" + return self._loc + + @property + def scale_params(self) -> jnp.ndarray: + """Scale parameters of the GMM.""" + return self._scale_params + + @property + def cholesky(self) -> jnp.ndarray: + """Cholesky decomposition of the GMM covariance matrices.""" + size = self.n_dimensions + + def _get_cholesky(scale_params): + return scale_tril.ScaleTriL(params=scale_params, size=size).cholesky() + + return jax.vmap(_get_cholesky, in_axes=0, out_axes=0)(self.scale_params) + + @property + def covariance(self) -> jnp.ndarray: + """Covariance matrices of the GMM.""" + size = self.n_dimensions + + def _get_covariance(scale_params): + return scale_tril.ScaleTriL(params=scale_params, size=size).covariance() + + return jax.vmap(_get_covariance, in_axes=0, out_axes=0)(self.scale_params) + + @property + def component_weight_ob(self) -> probabilities.Probabilities: + """Component weight object.""" + return self._component_weight_ob + + @property + def component_weights(self) -> jnp.ndarray: + """Component weights probabilities.""" + return self._component_weight_ob.probs() + + def log_component_weights(self) -> jnp.ndarray: + """Log component weights probabilities.""" + return self._component_weight_ob.log_probs() + + def _get_normal( + self, loc: jnp.ndarray, scale_params: jnp.ndarray + ) -> gaussian.Gaussian: + size = loc.shape[-1] + return gaussian.Gaussian( + loc=loc, scale=scale_tril.ScaleTriL(params=scale_params, size=size) + ) + + def get_component(self, index: int) -> gaussian.Gaussian: + """Specified GMM component.""" + return self._get_normal( + loc=self.loc[index], scale_params=self.scale_params[index] + ) + + def components(self) -> List[gaussian.Gaussian]: + """List of all GMM components.""" + return [self.get_component(i) for i in range(self.n_components)] + + def sample(self, rng: jax.Array, size: int) -> jnp.ndarray: + """Generate samples from the distribution.""" + subrng0, subrng1 = jax.random.split(rng) + component = self.component_weight_ob.sample(rng=subrng0, size=size) + std_samples = jax.random.normal(subrng1, shape=(size, self.n_dimensions)) + + def _transform_single_component(k, scale, loc): + + def _transform_single_value(single_component, single_x): + return jax.lax.cond( + single_component == k, + lambda x: jnp.matmul(scale, x[:, None])[:, 0] + loc, jnp.zeros_like, + single_x + ) + + return jax.vmap(_transform_single_value)(component, std_samples) + + return jnp.sum( + jax.vmap(_transform_single_component) + (jnp.arange(self.n_components), self.cholesky, self.loc), + axis=0 + ) + + def conditional_log_prob(self, x: jnp.ndarray) -> jnp.ndarray: + """Compute the component-conditional log probability of x. + + Args: + x: (n, n_dimensions) array of points + + Returns: + (n, n_components) array of the log probability of x conditioned on it + having come from each component. + """ + + def _log_prob_single_component( + loc: jnp.ndarray, scale_params: jnp.ndarray, x: jnp.ndarray + ): + norm = self._get_normal(loc=loc, scale_params=scale_params) + return norm.log_prob(x) + + conditional_log_prob_fn = jax.vmap( + _log_prob_single_component, in_axes=(0, 0, None), out_axes=1 + ) + return conditional_log_prob_fn(self._loc, self._scale_params, x) + + def log_prob(self, x: jnp.ndarray) -> jnp.ndarray: + """Compute the log probability of the observations x. + + Args: + x: (n, n_dimensions) array of points + + Returns: + (n,) array of log probabilities. + """ + # p(x) = \sum_i p(x|c_i) p(c_i) + log_prob_conditional = self.conditional_log_prob(x) + log_component_weight = self.log_component_weights() + return jax.scipy.special.logsumexp( + log_prob_conditional + log_component_weight[None, :], axis=-1 + ) + + def get_log_component_posterior(self, x: jnp.ndarray) -> jnp.ndarray: + """Compute the posterior probability that x came from each component. + + Args: + x: (n, n_dimensions) array of points + + Returns: + (n, n_components) array of poster component log probabilities. + """ + # p(x | c_i) = p(x, c_i) / p(c_i) => p(x, c_i) = p(x | c_i) p(c_i) + # p(c_i | x) = p(x, c_i) / p(x) + # = p(x | c_i) p(c_i) / sum_j(p(x | c_j)p(c_j)) + log_prob_conditional = self.conditional_log_prob(x) + log_component_weight = self.log_component_weights() + log_prob_unnorm = log_prob_conditional + log_component_weight[None, :] + return log_prob_unnorm - jax.scipy.special.logsumexp( + log_prob_unnorm, axis=-1, keepdims=True + ) + + def has_nans(self) -> bool: # noqa: D102 + for leaf in jax.tree_util.tree_leaves(self): + if jnp.any(~jnp.isfinite(leaf)): + return True + return False + + def tree_flatten(self): # noqa: D102 + children = (self.loc, self.scale_params, self.component_weight_ob) + aux_data = {} + return children, aux_data + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + return cls(*children, **aux_data) + + def __repr__(self): + class_name = type(self).__name__ + children, aux = self.tree_flatten() + return "{}({})".format( + class_name, ", ".join([repr(c) for c in children] + + [f"{k}: {repr(v)}" for k, v in aux.items()]) + ) + + def __hash__(self): + return jax.tree_util.tree_flatten(self).__hash__() + + def __eq__(self, other): + return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture_pair.py b/ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture_pair.py new file mode 100644 index 0000000..62e04c0 --- /dev/null +++ b/ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture_pair.py @@ -0,0 +1,222 @@ +# 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. +from typing import Any + +import jax +import jax.numpy as jnp + +from ott.geometry import costs, geometry, pointcloud +from ott.problems.linear import linear_problem +from ott.solvers.linear import sinkhorn +from ott.tools.gaussian_mixture import gaussian_mixture + +__all__ = ["GaussianMixturePair"] + + +@jax.tree_util.register_pytree_node_class +class GaussianMixturePair: + """Coupled pair of Gaussian mixture models. + + Includes methods used in estimating an optimal pairing between GMM components + using the Wasserstein-like method described in :cite:`delon:20`, + as well as generalization that allows for the reweighting of components. + + :cite:`delon:20` propose fitting a pair of GMMs to a pair + of point clouds in such a way that the sum of the log likelihood of the + points minus a weighted penalty involving a Wasserstein-like distance between + the GMMs. Their proposed algorithm involves using EM in which a balanced + Sinkhorn algorithm is used to estimate a coupling between the GMMs at each + step of EM. + + Our generalization of this algorithm allows for a mismatch between the + marginals of the coupling and the GMM component weights. This mismatch can be + interpreted as components being reweighted rather than being transported. + We penalize reweighting with a generalized KL-divergence penalty, and we give + the option to use the unbalanced Sinkhorn algorithm rather than the balanced + to compute the divergence between GMMs. + """ + + def __init__( + self, + gmm0: gaussian_mixture.GaussianMixture, + gmm1: gaussian_mixture.GaussianMixture, + epsilon: float = 1e-2, + tau: float = 1.0, + lock_gmm1: bool = False, + ): + """Constructor. + + When fitting a pair of coupled GMMs with *no* reweighting of components + using the algorithm in :cite:`delon:20`, set tau = 1. The coupling between + components will be determined via the balanced Sinkhorn algorithm. + + When fitting a pair of coupled GMMs in which reweighting of components is + allowed, set tau to a value in (0, 1). The resulting coupling will penalize + the generalized KL divergence between the coupling's marginals and the GMM + component weights with a weight of rho = epsilon tau / (1 - tau). + + Args: + gmm0: first GMM in the pair + gmm1: second GMM in the pair + epsilon: regularization weight to use for the Sinkhorn algorithm + tau: encodes the weight, rho, to use for the generalized KL divergence + between the coupling's marginals and GMM component weights as + rho = epsilon tau / (1 - tau) + lock_gmm1: indicates whether the parameters of gmm1 should be modified + during optimization + """ # noqa: D401 + self._gmm0 = gmm0 + self._gmm1 = gmm1 + self._epsilon = epsilon + self._tau = tau + self._lock_gmm1 = lock_gmm1 + + @property + def dtype(self): # noqa: D102 + return self.gmm0.dtype + + @property + def gmm0(self): # noqa: D102 + return self._gmm0 + + @property + def gmm1(self): # noqa: D102 + return self._gmm1 + + @property + def epsilon(self): # noqa: D102 + return self._epsilon + + @property + def tau(self): # noqa: D102 + return self._tau + + @property + def rho(self): # noqa: D102 + return self.epsilon * self.tau / (1.0 - self.tau) + + @property + def lock_gmm1(self): # noqa: D102 + return self._lock_gmm1 + + def get_bures_geometry(self) -> pointcloud.PointCloud: + """Get a Bures Geometry for the two GMMs.""" + mean0 = self.gmm0.loc + dimension = mean0.shape[-1] + cov0 = self.gmm0.covariance + cov0 = cov0.reshape(cov0.shape[:-2] + (dimension * dimension,)) + x = jnp.concatenate([mean0, cov0], axis=-1) + mean1 = self.gmm1.loc + cov1 = self.gmm1.covariance + cov1 = cov1.reshape(cov1.shape[:-2] + (dimension * dimension,)) + y = jnp.concatenate([mean1, cov1], axis=-1) + return pointcloud.PointCloud( + x=x, + y=y, + cost_fn=costs.Bures(dimension=dimension), + epsilon=self.epsilon + ) + + def get_cost_matrix(self) -> jnp.ndarray: + """Get matrix of :math:`W_2^2` costs between all pairs of components.""" + return self.get_bures_geometry().cost_matrix + + def get_sinkhorn( + self, cost_matrix: jnp.ndarray, **kwargs: Any + ) -> sinkhorn.SinkhornOutput: + """Get the output of Sinkhorn's method for a given cost matrix.""" + # We use a Geometry here rather than the PointCloud created in + # get_bures_geometry to avoid recomputing the cost matrix, since + # the cost matrix is quite expensive + geom = geometry.Geometry(cost_matrix=cost_matrix, epsilon=self.epsilon) + prob = linear_problem.LinearProblem( + geom, + a=self.gmm0.component_weights, + b=self.gmm1.component_weights, + tau_a=self.tau, + tau_b=self.tau + ) + return sinkhorn.Sinkhorn(**kwargs)(prob) + + def get_normalized_sinkhorn_coupling( + self, + sinkhorn_output: sinkhorn.SinkhornOutput, + ) -> jnp.ndarray: + """Get the normalized coupling matrix for the specified Sinkhorn output. + + Args: + sinkhorn_output: Sinkhorn algorithm output as returned by + :meth:`get_sinkhorn`. + + Returns: + A coupling matrix that tells how much of the mass of each component of + :attr:`gmm0` is mapped to each component of :attr:`gmm1`. + """ + return sinkhorn_output.matrix / jnp.sum(sinkhorn_output.matrix) + + def tree_flatten(self): + """Method used by jax.tree_util to flatten a GaussianMixturePair. + + We control the subset of parameters that we will optimize in fit_gmm_pair + by selectively placing them in either children (the parameters to optimize) + or aux_data (the parameters to leave alone). + + Returns: + A tuple of child pytrees and a dict of auxiliary data. + """ # noqa: D401 + children = [self.gmm0] + aux_data = { + "epsilon": self.epsilon, + "tau": self.tau, + "lock_gmm1": self.lock_gmm1 + } + if self.lock_gmm1: + aux_data["gmm1"] = self.gmm1 + else: + children.append(self.gmm1) + return tuple(children), aux_data + + @classmethod + def tree_unflatten(cls, aux_data, children): + """Method used by jax.tree_util to unflatten a GaussianMixturePair. + + tree_flatten controls which parameters get optimized by placing them in + either children or aux_data; here we invert the process. + + Args: + aux_data: auxiliary data that is passed to the constructor as kwargs + children: child pytrees passed to the constructor as args + + Returns: + A GaussianMixturePair. + """ # noqa: D401 + children = list(children) + if "gmm1" in aux_data: + gmm1 = aux_data.pop("gmm1") + children.insert(1, gmm1) + return cls(*children, **aux_data) + + def __repr__(self): + class_name = type(self).__name__ + children, aux = self.tree_flatten() + return "{}({})".format( + class_name, ", ".join([repr(c) for c in children] + + [f"{k}: {repr(v)}" for k, v in aux.items()]) + ) + + def __hash__(self): + return jax.tree_util.tree_flatten(self).__hash__() + + def __eq__(self, other): + return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/ott/build/lib/ott/tools/gaussian_mixture/linalg.py b/ott/build/lib/ott/tools/gaussian_mixture/linalg.py new file mode 100644 index 0000000..dfc7d4a --- /dev/null +++ b/ott/build/lib/ott/tools/gaussian_mixture/linalg.py @@ -0,0 +1,138 @@ +# 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. +from typing import Callable, Iterable, List, Tuple + +import jax +import jax.numpy as jnp + + +def get_mean_and_var( + points: jnp.ndarray, # (n, d) + weights: jnp.ndarray, # (n,) +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Get the mean and variance of a weighted set of points.""" + weights_sum = jnp.sum(weights, axis=-1) # (1,) + mean = ( + # matmul((1, n), (n, d)) -> (1, d) + jnp.matmul(weights, points) / weights_sum + ) + # center points + centered = points - mean[None, :] # (n, d) - (1, d) + var = ( + # matmul((1, n), (n, d)) -> (1, d) + jnp.matmul(weights, centered ** 2) / weights_sum + ) + return mean, var + + +def get_mean_and_cov( + points: jnp.ndarray, # (n, d) + weights: jnp.ndarray, # (n,) +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Get the mean and covariance of a weighted set of points.""" + weights_sum = jnp.sum(weights, axis=-1, keepdims=True) # (1,) + mean = ( + # matmul((1, n), (n, d)) -> (1, d) + jnp.matmul(weights, points) / weights_sum + ) + # center points + centered = points - mean[None, :] # (n, d) - (1, d) + cov = ( + jnp.matmul( + # (1, n) (d, n) + weights[None, :] * jnp.swapaxes(centered, axis1=-2, axis2=-1), + # (n, d) + centered + ) / weights_sum + ) + return mean, cov + + +def flat_to_tril(x: jnp.ndarray, size: int) -> jnp.ndarray: + """Map flat values to lower triangular matrices. + + Args: + x: flat values + size: size of lower triangular matrices. x should have shape + (..., size(size+1)/2), and the final matrices should have shape + (..., size, size). + + Returns: + Lower triangular matrices. + """ + m = jnp.zeros(x.shape[:-1] + (size, size)) + tril = jnp.tril_indices(size) + return m.at[..., tril[0], tril[1]].set(x) + + +def tril_to_flat(m: jnp.ndarray) -> jnp.ndarray: + """Flatten lower triangular matrices. + + Args: + m: lower triangular matrices of shape (..., size, size) + + Returns: + A vector of shape (..., size (size+1) // 2) + """ + size = m.shape[-1] + tril = jnp.tril_indices(size) + return m[..., tril[0], tril[1]] + + +def apply_to_diag( + m: jnp.ndarray, fn: Callable[[jnp.ndarray], jnp.ndarray] +) -> jnp.ndarray: + """Apply a function to the diagonal of a matrix.""" + size = m.shape[-1] + diag = jnp.diagonal(m, axis1=-2, axis2=-1) + ind = jnp.arange(size) + return m.at[..., ind, ind].set(fn(diag)) + + +def matrix_powers( + m: jnp.ndarray, + powers: Iterable[float], +) -> List[jnp.ndarray]: + """Raise a real, symmetric matrix to multiple powers.""" + eigs, q = jnp.linalg.eigh(m) + qt = jnp.swapaxes(q, axis1=-2, axis2=-1) + ret = [] + for power in powers: + ret.append(jnp.matmul(jnp.expand_dims(eigs ** power, -2) * q, qt)) + return ret + + +def invmatvectril( + m: jnp.ndarray, x: jnp.ndarray, lower: bool = True +) -> jnp.ndarray: + """Multiply x by the inverse of a triangular matrix. + + Args: + m: triangular matrix, shape (d, d) + x: array of points, shape (n, d) + lower: if True, m is lower triangular; otherwise m is upper triangular + + Returns: + m^{-1} x + """ + return jnp.transpose( + jax.scipy.linalg.solve_triangular(m, jnp.transpose(x), lower=lower) + ) + + +def get_random_orthogonal(rng: jax.Array, dim: int) -> jnp.ndarray: + """Get a random orthogonal matrix with the specified dimension.""" + m = jax.random.normal(rng, shape=[dim, dim]) + q, _ = jnp.linalg.qr(m) + return q diff --git a/ott/build/lib/ott/tools/gaussian_mixture/probabilities.py b/ott/build/lib/ott/tools/gaussian_mixture/probabilities.py new file mode 100644 index 0000000..c0ae2b2 --- /dev/null +++ b/ott/build/lib/ott/tools/gaussian_mixture/probabilities.py @@ -0,0 +1,100 @@ +# 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. +from typing import Optional + +import jax +import jax.numpy as jnp + +__all__ = ["Probabilities"] + + +@jax.tree_util.register_pytree_node_class +class Probabilities: + """Parameterized array of probabilities of length n. + + The internal representation is a length n-1 unconstrained array. We convert + to a length n simplex by appending a 0 and taking a softmax. + """ + + _params: jnp.ndarray + + def __init__(self, params): + self._params = params + + @classmethod + def from_random( + cls, + rng: jax.Array, + n_dimensions: int, + stdev: Optional[float] = 0.1, + ) -> "Probabilities": + """Construct a random Probabilities.""" + return cls(params=jax.random.normal(rng, shape=(n_dimensions - 1,)) * stdev) + + @classmethod + def from_probs(cls, probs: jnp.ndarray) -> "Probabilities": + """Construct Probabilities from a vector of probabilities.""" + log_probs = jnp.log(probs) + log_probs_normalized, norm = log_probs[:-1], log_probs[-1] + log_probs_normalized -= norm + return cls(params=log_probs_normalized) + + @property + def params(self): # noqa: D102 + return self._params + + @property + def dtype(self): # noqa: D102 + return self._params.dtype + + def unnormalized_log_probs(self) -> jnp.ndarray: + """Get the unnormalized log probabilities.""" + return jnp.concatenate([self._params, jnp.zeros((1,))], axis=-1) + + def log_probs(self) -> jnp.ndarray: + """Get the log probabilities.""" + return jax.nn.log_softmax(self.unnormalized_log_probs()) + + def probs(self) -> jnp.ndarray: + """Get the probabilities.""" + return jax.nn.softmax(self.unnormalized_log_probs()) + + def sample(self, rng: jax.Array, size: int) -> jnp.ndarray: + """Sample from the distribution.""" + return jax.random.categorical( + rng, logits=self.unnormalized_log_probs(), shape=(size,) + ) + + def tree_flatten(self): # noqa: D102 + children = (self.params,) + aux_data = {} + return children, aux_data + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + return cls(*children, **aux_data) + + def __repr__(self): + class_name = type(self).__name__ + children, aux = self.tree_flatten() + return "{}({})".format( + class_name, ", ".join([repr(c) for c in children] + + [f"{k}: {repr(v)}" for k, v in aux.items()]) + ) + + def __hash__(self): + return jax.tree_util.tree_flatten(self).__hash__() + + def __eq__(self, other): + return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/ott/build/lib/ott/tools/gaussian_mixture/scale_tril.py b/ott/build/lib/ott/tools/gaussian_mixture/scale_tril.py new file mode 100644 index 0000000..b26dded --- /dev/null +++ b/ott/build/lib/ott/tools/gaussian_mixture/scale_tril.py @@ -0,0 +1,214 @@ +# 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. +from typing import Optional, Tuple + +import jax +import jax.numpy as jnp + +from ott.geometry import costs +from ott.math import matrix_square_root +from ott.tools.gaussian_mixture import linalg + +__all__ = ["ScaleTriL"] + + +@jax.tree_util.register_pytree_node_class +class ScaleTriL: + """Pytree for a lower triangular Cholesky-factored covariance matrix.""" + + def __init__(self, params: jnp.ndarray, size: int): + self._params = params + self._size = size + + @classmethod + def from_points_and_weights( + cls, + points: jnp.ndarray, + weights: jnp.ndarray, + ) -> Tuple[jnp.ndarray, "ScaleTriL"]: + """Get a mean and a ScaleTriL from a set of points and weights.""" + mean, cov = linalg.get_mean_and_cov(points=points, weights=weights) + return mean, cls.from_covariance(cov) + + @classmethod + def from_random( + cls, + rng: jax.Array, + n_dimensions: int, + stdev: Optional[float] = 0.1, + ) -> "ScaleTriL": + """Construct a random ScaleTriL. + + Args: + rng: pseudo-random number generator key + n_dimensions: number of dimensions + stdev: desired standard deviation (around 0) for the log eigenvalues + + Returns: + A ScaleTriL. + """ + # generate a random orthogonal matrix + rng, subrng = jax.random.split(rng) + q = linalg.get_random_orthogonal(subrng, dim=n_dimensions) + + # generate random eigenvalues + eigs = stdev * jnp.exp(jax.random.normal(rng, shape=(n_dimensions,))) + + # random positive definite matrix + sigma = q * jnp.expand_dims(eigs, -2) @ q.T + + # cholesky factorization + chol = jnp.linalg.cholesky(sigma) + # flatten + m = linalg.apply_to_diag(chol, jnp.log) + flat = linalg.tril_to_flat(m) + return cls(params=flat, size=n_dimensions) + + @classmethod + def from_cholesky(cls, cholesky: jnp.ndarray) -> "ScaleTriL": + """Construct ScaleTriL from a Cholesky factor of a covariance matrix.""" + m = linalg.apply_to_diag(cholesky, jnp.log) + flat = linalg.tril_to_flat(m) + return cls(params=flat, size=cholesky.shape[-1]) + + @classmethod + def from_covariance( + cls, + covariance: jnp.ndarray, + ) -> "ScaleTriL": + """Construct ScaleTriL from a covariance matrix.""" + cholesky = jnp.linalg.cholesky(covariance) + return cls.from_cholesky(cholesky) + + @property + def params(self) -> jnp.ndarray: + """Internal representation.""" + return self._params + + @property + def size(self) -> int: + """Size of the covariance matrix.""" + return self._size + + @property + def dtype(self): + """Data type of the covariance matrix.""" + return self._params.dtype + + def cholesky(self) -> jnp.ndarray: + """Get a lower triangular Cholesky factor for the covariance matrix.""" + m = linalg.flat_to_tril(self._params, size=self._size) + return linalg.apply_to_diag(m, jnp.exp) + + def covariance(self) -> jnp.ndarray: + """Get the covariance matrix.""" + cholesky = self.cholesky() + return cholesky @ cholesky.T + + def covariance_sqrt(self) -> jnp.ndarray: + """Get the square root of the covariance matrix.""" + return linalg.matrix_powers(self.covariance(), (0.5,))[0] + + def log_det_covariance(self) -> jnp.ndarray: + """Get the log of the determinant of the covariance matrix.""" + diag = jnp.diagonal(self.cholesky(), axis1=-2, axis2=-1) + return 2.0 * jnp.sum(jnp.log(diag), axis=-1) + + def centered_to_z(self, x_centered: jnp.ndarray) -> jnp.ndarray: + """Map centered points to standardized centered points (i.e. cov(z) = I).""" + return linalg.invmatvectril(m=self.cholesky(), x=x_centered, lower=True) + + def z_to_centered(self, z: jnp.ndarray) -> jnp.ndarray: + """Scale standardized points to points with the specified covariance.""" + return (self.cholesky() @ z.T).T + + def w2_dist(self, other: "ScaleTriL") -> jnp.ndarray: + r"""Wasserstein distance W_2^2 to another Gaussian with same mean. + + Args: + other: Scale for the other Gaussian + + Returns: + The W_2^2 distance + """ + dimension = self.size + + def _flatten_cov(cov: jnp.ndarray) -> jnp.ndarray: + cov = cov.reshape(cov.shape[:-2] + (dimension * dimension,)) + return jnp.concatenate([jnp.zeros(dimension), cov], axis=-1) + + x0 = _flatten_cov(self.covariance()) + x1 = _flatten_cov(other.covariance()) + cost_fn = costs.Bures(dimension=dimension) + return (cost_fn.norm(x0) + cost_fn.norm(x1) + cost_fn.pairwise(x0, x1))[ + ..., + ] + + def gaussian_map(self, dest_scale: "ScaleTriL") -> jnp.ndarray: + """Scaling matrix used in transport between 0-mean Gaussians. + + Sigma_mu^{-1/2} @ + [Sigma_mu ^{1/2} Sigma_nu Sigma_mu ^{1/2}]^{1/2} + @ Sigma_mu ^{-1/2} + + Args: + dest_scale: destination Scale + + Returns: + Gaussian scaling matrix, same dimension as self.covaraince + """ + sqrt0, sqrt0_inv = linalg.matrix_powers(self.covariance(), (0.5, -0.5)) + sigma1 = dest_scale.covariance() + m = matrix_square_root.sqrtm_only( + jnp.matmul(sqrt0, jnp.matmul(sigma1, sqrt0)) + ) + return jnp.matmul(sqrt0_inv, jnp.matmul(m, sqrt0_inv)) + + def transport( + self, dest_scale: "ScaleTriL", points: jnp.ndarray + ) -> jnp.ndarray: + """Apply Monge map, computed between two 0-mean Gaussians, to points. + + Args: + dest_scale: destination Scale + points: points to transport + + Returns: + Points transported to a Gaussian with the new scale. + """ + m = self.gaussian_map(dest_scale) + return (m @ points.T).T + + def tree_flatten(self): # noqa: D102 + children = (self.params,) + aux_data = {"size": self.size} + return children, aux_data + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + return cls(*children, **aux_data) + + def __repr__(self): + class_name = type(self).__name__ + children, aux = self.tree_flatten() + return "{}({})".format( + class_name, ", ".join([repr(c) for c in children] + + [f"{k}: {repr(v)}" for k, v in aux.items()]) + ) + + def __hash__(self): + return jax.tree_util.tree_flatten(self).__hash__() + + def __eq__(self, other): + return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/ott/build/lib/ott/tools/k_means.py b/ott/build/lib/ott/tools/k_means.py new file mode 100644 index 0000000..2a12f6e --- /dev/null +++ b/ott/build/lib/ott/tools/k_means.py @@ -0,0 +1,412 @@ +# 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 functools +import math +from typing import Callable, Literal, NamedTuple, Optional, Tuple, Union + +import jax +import jax.numpy as jnp + +from ott import utils +from ott.geometry import costs, pointcloud +from ott.math import fixed_point_loop + +__all__ = ["k_means", "KMeansOutput"] + +Init_t = Union[Literal["k-means++", "random"], + Callable[[pointcloud.PointCloud, int, jnp.ndarray], jnp.ndarray]] + + +class KPPState(NamedTuple): # noqa: D101 + rng: jax.Array + centroids: jnp.ndarray + centroid_dists: jnp.ndarray + + +class KMeansState(NamedTuple): # noqa: D101 + centroids: jnp.ndarray + prev_assignment: jnp.ndarray + assignment: jnp.ndarray + errors: jnp.ndarray + center_shift: float + + +class KMeansConst(NamedTuple): # noqa: D101 + geom: pointcloud.PointCloud + x_weights: jnp.ndarray + + @property + def x(self) -> jnp.ndarray: + """Array of shape ``[n, ndim]`` containing the unweighted point cloud.""" + return self.geom.x + + @property + def weighted_x(self): + """Array of shape ``[n, ndim]`` containing the weighted point cloud.""" + return self.x_weights[:, :-1] + + @property + def weights(self) -> jnp.ndarray: + """Array of shape ``[n, 1]`` containing weights for each point.""" + return self.x_weights[:, -1:] + + +class KMeansOutput(NamedTuple): + """Output of the :func:`~ott.tools.k_means.k_means` algorithm. + + Args: + centroids: Array of shape ``[k, ndim]`` containing the centroids. + assignment: Array of shape ``[n,]`` containing the labels. + converged: Whether the algorithm has converged. + iteration: The number of iterations run. + error: (Weighted) sum of squared distances from each point to its closest + center. + inner_errors: Array of shape ``[max_iterations,]`` containing the ``error`` + at every iteration. + """ + centroids: jnp.ndarray + assignment: jnp.ndarray + converged: bool + iteration: int + error: float + inner_errors: Optional[jnp.ndarray] + + @classmethod + def _from_state( + cls, + state: KMeansState, + *, + tol: float, + store_inner_errors: bool = False + ) -> "KMeansOutput": + errs = state.errors + mask = errs == -1 + error = jnp.nanmin(jnp.where(mask, jnp.nan, errs)) + + assignment_same = jnp.all(state.prev_assignment == state.assignment) + tol_satisfied = jnp.logical_or(jnp.any(mask), (errs[-2] - errs[-1]) <= tol) + converged = jnp.logical_or(assignment_same, tol_satisfied) + + return cls( + centroids=state.centroids, + assignment=state.assignment.astype(int), + converged=converged, + iteration=jnp.sum(~mask), + error=error, + inner_errors=errs if store_inner_errors else None, + ) + + +def _random_init( + geom: pointcloud.PointCloud, k: int, rng: jax.Array +) -> jnp.ndarray: + ixs = jnp.arange(geom.shape[0]) + ixs = jax.random.choice(rng, ixs, shape=(k,), replace=False) + return geom.subset(ixs, None).x + + +def _k_means_plus_plus( + geom: pointcloud.PointCloud, + k: int, + rng: jax.Array, + n_local_trials: Optional[int] = None, +) -> jnp.ndarray: + + def init_fn(geom: pointcloud.PointCloud, rng: jax.Array) -> KPPState: + rng, next_rng = jax.random.split(rng, 2) + ix = jax.random.choice(rng, jnp.arange(geom.shape[0]), shape=()) + centroids = jnp.full((k, geom.cost_rank), jnp.inf).at[0].set(geom.x[ix]) + dists = geom.subset([ix], None).cost_matrix[0] + return KPPState(rng=next_rng, centroids=centroids, centroid_dists=dists) + + def body_fn( + iteration: int, const: Tuple[pointcloud.PointCloud, jnp.ndarray], + state: KPPState, compute_error: bool + ) -> KPPState: + del compute_error + rng, next_rng = jax.random.split(state.rng, 2) + geom, ixs = const + + # no need to normalize when `replace=True` + probs = state.centroid_dists + ixs = jax.random.choice( + rng, ixs, shape=(n_local_trials,), p=probs, replace=True + ) + geom = geom.subset(ixs, None) + + candidate_dists = jnp.minimum(geom.cost_matrix, state.centroid_dists) + best_ix = jnp.argmin(candidate_dists.sum(1)) + + centroids = state.centroids.at[iteration + 1].set(geom.x[best_ix]) + centroid_dists = candidate_dists[best_ix] + + return KPPState( + rng=next_rng, centroids=centroids, centroid_dists=centroid_dists + ) + + if n_local_trials is None: + n_local_trials = 2 + int(math.log(k)) + assert n_local_trials > 0, n_local_trials + + state = init_fn(geom, rng) + constants = (geom, jnp.arange(geom.shape[0])) + state = fixed_point_loop.fixpoint_iter( + lambda *_, **__: True, + body_fn, + min_iterations=k - 1, + max_iterations=k - 1, + inner_iterations=1, + constants=constants, + state=state + ) + + return state.centroids + + +@functools.partial(jax.vmap, in_axes=[None, 0, 0, 0], out_axes=0) +def _reallocate_centroids( + const: KMeansConst, + ix: jnp.ndarray, + centroid: jnp.ndarray, + weight: jnp.ndarray, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + is_empty = weight <= 0.0 + new_centroid = (1 - is_empty) * centroid + is_empty * const.x[ix] # (ndim,) + centroid_to_remove = is_empty * const.weighted_x[ix] # (ndim,) + weight_to_remove = is_empty * const.weights[ix] # (1,) + return new_centroid, jnp.concatenate([centroid_to_remove, weight_to_remove]) + + +def _update_assignment( + const: KMeansConst, + centroids: jnp.ndarray, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + (x, _, *args), aux_data = const.geom.tree_flatten() + cost_matrix = type( + const.geom + ).tree_unflatten(aux_data, [x, centroids] + args).cost_matrix + + assignment = jnp.argmin(cost_matrix, axis=1) + dist_to_centers = cost_matrix[jnp.arange(len(assignment)), assignment] + return assignment, dist_to_centers + + +def _update_centroids( + const: KMeansConst, k: int, assignment: jnp.ndarray, + dist_to_centers: jnp.ndarray +) -> jnp.ndarray: + # TODO(michalk8): + # cannot put `k` into `const`, see https://github.com/ott-jax/ott/issues/129 + x_weights = jax.ops.segment_sum(const.x_weights, assignment, num_segments=k) + centroids, ws = x_weights[:, :-1], x_weights[:, -1:] + + far_ixs = jnp.argsort(dist_to_centers)[-k:][::-1] + centroids, to_remove = _reallocate_centroids(const, far_ixs, centroids, ws) + to_remove = jax.ops.segment_sum( + to_remove, assignment[far_ixs], num_segments=k + ) + centroids -= to_remove[:, :-1] + ws -= to_remove[:, -1:] + + return centroids * jnp.where(ws > 0.0, 1.0 / ws, 1.0) + + +@functools.partial(jax.vmap, in_axes=[0] + [None] * 9) +def _k_means( + rng: jax.Array, + geom: pointcloud.PointCloud, + k: int, + weights: Optional[jnp.ndarray] = None, + init: Init_t = "k-means++", + n_local_trials: Optional[int] = None, + tol: float = 1e-4, + min_iterations: int = 0, + max_iterations: int = 300, + store_inner_errors: bool = False, +) -> KMeansOutput: + + def init_fn(init: Init_t) -> KMeansState: + if init == "k-means++": + init = functools.partial( + _k_means_plus_plus, n_local_trials=n_local_trials + ) + elif init == "random": + init = _random_init + if not callable(init): + raise TypeError( + f"Expected `init` to be 'k-means++', 'random' " + f"or a callable, found `{init_fn!r}`." + ) + + centroids = init(geom, k, rng) + if centroids.shape != (k, geom.cost_rank): + raise ValueError( + f"Expected initial centroids to have shape " + f"`{k, geom.cost_rank}`, found `{centroids.shape}`." + ) + n = geom.shape[0] + # TODO(michalk8): find a better solution for the below error + # not using floats for the assignment when `fixpoint_iter_backprop` is used: + + # .../jax/_src/dtypes.py:370: + # .0 = + # > CUB = set.intersection(*(UB[n] for n in N)) + # E jax._src.traceback_util.UnfilteredStackTrace: + # KeyError: dtype([('float0', 'V')]) + # E The stack trace below excludes JAX-internal frames. + # E The preceding is the original exception that occurred, unmodified. + prev_assignment = jnp.full((n,), -2.0) + assignment = jnp.full((n,), -1.0) + errors = jnp.full((max_iterations,), -1.0) + + return KMeansState( + centroids=centroids, + prev_assignment=prev_assignment, + assignment=assignment, + center_shift=jnp.inf, + errors=errors, + ) + + def cond_fn(iteration: int, const: KMeansConst, state: KMeansState) -> bool: + del iteration, const + assignment_not_same = jnp.any(state.prev_assignment != state.assignment) + tol_not_satisfied = state.center_shift > tol + return jnp.logical_and(assignment_not_same, tol_not_satisfied) + + def body_fn( + iteration: int, const: KMeansConst, state: KMeansState, + compute_error: bool + ) -> KMeansState: + del compute_error + + assignment, dist_to_centers = _update_assignment(const, state.centroids) + centroids = _update_centroids(const, k, assignment, dist_to_centers) + err = jnp.sum(const.weights[:, 0] * dist_to_centers) + center_shift = jnp.linalg.norm(state.centroids - centroids, ord="fro") ** 2 + + return KMeansState( + centroids=centroids, + prev_assignment=state.assignment, + assignment=assignment.astype(float), + center_shift=center_shift, + errors=state.errors.at[iteration].set(err) + ) + + def finalize_fn(const: KMeansConst, state: KMeansState) -> KMeansState: + last_iter = jnp.sum(state.errors != -1) - 1 + + assignment, dist_to_centers = _update_assignment(const, state.centroids) + err = jnp.sum(const.weights[:, 0] * dist_to_centers) + + return state._replace( + assignment=assignment.astype(float), + errors=state.errors.at[last_iter].set(err) + ) + + force_scan = min_iterations == max_iterations + fixpoint_fn = ( # prefer auto-diff if possible + fixed_point_loop.fixpoint_iter if force_scan else + fixed_point_loop.fixpoint_iter_backprop + ) + x_weights = jnp.hstack([weights[:, None] * geom.x, weights[:, None]]) + const = KMeansConst(geom, x_weights) + + state = fixpoint_fn( + cond_fn, + body_fn, + min_iterations=min_iterations, + max_iterations=max_iterations, + inner_iterations=1, + constants=const, + state=init_fn(init) + ) + state = jax.lax.cond( + jnp.all(state.prev_assignment == state.assignment), (lambda _, s: s), + finalize_fn, const, state + ) + + return KMeansOutput._from_state( + state, tol=tol, store_inner_errors=store_inner_errors + ) + + +def k_means( + geom: Union[jnp.ndarray, pointcloud.PointCloud], + k: int, + weights: Optional[jnp.ndarray] = None, + init: Init_t = "k-means++", + n_init: int = 10, + n_local_trials: Optional[int] = None, + tol: float = 1e-4, + min_iterations: int = 0, + max_iterations: int = 300, + store_inner_errors: bool = False, + rng: Optional[jax.Array] = None, +) -> KMeansOutput: + r"""K-means clustering using Lloyd's algorithm :cite:`lloyd:82`. + + Args: + geom: Point cloud of shape ``[n, ndim]`` to cluster. If passed as an array, + :class:`~ott.geometry.costs.SqEuclidean` cost is assumed. + k: The number of clusters. + weights: The weights of input points. These weights are considered when + computing the centroids and inertia. If ``None``, use uniform weights. + init: Initialization method. Can be one of the following: + + - **'k-means++'** - select initial centroids that are + :math:`\mathcal{O}(\log k)`-optimal :cite:`arthur:07`. + - **'random'** - randomly select ``k`` points from the ``geom``. + - :func:`callable` - a function which takes the point cloud, the number of + clusters and a random key and returns the centroids as an array of shape + ``[k, ndim]``. + + n_init: Number of times k-means will run with different initial seeds. + n_local_trials: Number of local trials when ``init = 'k-means++'``. + If ``None``, :math:`2 + \lfloor log(k) \rfloor` is used. + tol: Relative tolerance with respect to the Frobenius norm of the centroids' + shift between two consecutive iterations. + min_iterations: Minimum number of iterations. + max_iterations: Maximum number of iterations. + store_inner_errors: Whether to store the errors (inertia) at each iteration. + rng: Random key for seeding the initializations. + + Returns: + The k-means clustering. + """ + assert geom.shape[ + 0] >= k, f"Cannot cluster `{geom.shape[0]}` points into `{k}` clusters." + if isinstance(geom, jnp.ndarray): + geom = pointcloud.PointCloud(geom) + if isinstance(geom.cost_fn, costs.Cosine): + geom = geom._cosine_to_sqeucl() + assert geom.is_squared_euclidean + rng = utils.default_prng_key(rng) + + if geom.is_online: + # to allow materializing the cost matrix + children, aux_data = geom.tree_flatten() + aux_data["batch_size"] = None + geom = type(geom).tree_unflatten(aux_data, children) + + if weights is None: + weights = jnp.ones(geom.shape[0]) + assert weights.shape == (geom.shape[0],) + + rngs = jax.random.split(rng, n_init) + out = _k_means( + rngs, geom, k, weights, init, n_local_trials, tol, min_iterations, + max_iterations, store_inner_errors + ) + best_ix = jnp.argmin(out.error) + return jax.tree_util.tree_map(lambda arr: arr[best_ix], out) diff --git a/ott/build/lib/ott/tools/plot.py b/ott/build/lib/ott/tools/plot.py new file mode 100644 index 0000000..bd1f42e --- /dev/null +++ b/ott/build/lib/ott/tools/plot.py @@ -0,0 +1,247 @@ +# 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. +from typing import List, Optional, Sequence, Tuple, Union + +import jax.numpy as jnp +import numpy as np +import scipy + +from ott.geometry import pointcloud +from ott.solvers.linear import sinkhorn, sinkhorn_lr +from ott.solvers.quadratic import gromov_wasserstein + +try: + import matplotlib.pyplot as plt + from matplotlib import animation +except ImportError: + plt = animation = None + +# TODO(michalk8): make sure all outputs conform to a unified transport interface +Transport = Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput, + gromov_wasserstein.GWOutput] + + +def bidimensional(x: jnp.ndarray, + y: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Apply PCA to reduce to bi-dimensional data.""" + if x.shape[1] < 3: + return x, y + + u, s, _ = scipy.sparse.linalg.svds( + np.array(jnp.concatenate([x, y], axis=0)), k=2 + ) + proj = u * s + k = x.shape[0] + return proj[:k], proj[k:] + + +class Plot: + """Plot an optimal transport map between two point clouds. + + This object can either plot or update a plot, to create animations as a + :class:`~matplotlib.animation.FuncAnimation`, which can in turned be saved to + disk at will. There are two design principles here: + + #. we do not rely on saving to/loading from disk to create animations + #. we try as much as possible to disentangle the transport problem from + its visualization. + + We use 2D scatter plots by default, relying on PCA visualization for d>3 data. + This step requires a conversion to a numpy array, in order to compute leading + singular values. This tool is therefore not designed having performance in + mind. + + Args: + fig: Specify figure object. Created by default + ax: Specify axes objects. Created by default + threshold: value below which links in transportation matrix won't be + plotted. This value should be negative when using animations. + scale: scale used for marker plots. + show_lines: whether to show OT lines, as described in ``ot.matrix`` argument + cmap: color map used to plot line colors. + scale_alpha_by_coupling: use or not the coupling's value as proxy for alpha + alpha: default alpha value for lines. + title: title of the plot. + """ + + def __init__( + self, + fig: Optional["plt.Figure"] = None, + ax: Optional["plt.Axes"] = None, + threshold: float = -1.0, + scale: int = 200, + show_lines: bool = True, + cmap: str = "cool", + scale_alpha_by_coupling: bool = False, + alpha: float = 0.7, + title: Optional[str] = None + ): + if plt is None: + raise RuntimeError("Please install `matplotlib` first.") + + if ax is None and fig is None: + fig, ax = plt.subplots() + elif fig is None: + fig = plt.gcf() + elif ax is None: + ax = plt.gca() + self.fig = fig + self.ax = ax + self._show_lines = show_lines + self._lines = [] + self._points_x = None + self._points_y = None + self._threshold = threshold + self._scale = scale + self._cmap = cmap + self._scale_alpha_by_coupling = scale_alpha_by_coupling + self._alpha = alpha + self._title = title + + def _scatter(self, ot: Transport): + """Compute the position and scales of the points on a 2D plot.""" + if not isinstance(ot.geom, pointcloud.PointCloud): + raise ValueError("So far we only plot PointCloud geometry.") + + x, y = ot.geom.x, ot.geom.y + a, b = ot.a, ot.b + x, y = bidimensional(x, y) + scales_x = a * self._scale * a.shape[0] + scales_y = b * self._scale * b.shape[0] + return x, y, scales_x, scales_y + + def _mapping(self, x: jnp.ndarray, y: jnp.ndarray, matrix: jnp.ndarray): + """Compute the lines representing the mapping between the 2 point clouds.""" + # Only plot the lines with a cost above the threshold. + u, v = jnp.where(matrix > self._threshold) + c = matrix[jnp.where(matrix > self._threshold)] + xy = jnp.concatenate([x[u], y[v]], axis=-1) + + # Check if we want to adjust transparency. + scale_alpha_by_coupling = self._scale_alpha_by_coupling + + # We can only adjust transparency if max(c) != min(c). + if scale_alpha_by_coupling: + min_matrix, max_matrix = jnp.min(c), jnp.max(c) + scale_alpha_by_coupling = max_matrix != min_matrix + + result = [] + for i in range(xy.shape[0]): + strength = jnp.max(jnp.array(matrix.shape)) * c[i] + if scale_alpha_by_coupling: + normalized_strength = (c[i] - min_matrix) / (max_matrix - min_matrix) + alpha = self._alpha * float(normalized_strength) + else: + alpha = self._alpha + + # Matplotlib's transparency is sensitive to numerical errors. + alpha = np.clip(alpha, 0.0, 1.0) + + start, end = xy[i, [0, 2]], xy[i, [1, 3]] + result.append((start, end, strength, alpha)) + + return result + + def __call__(self, ot: Transport) -> List["plt.Artist"]: + """Plot couplings in 2-D, using PCA if data is higher dimensional.""" + x, y, sx, sy = self._scatter(ot) + self._points_x = self.ax.scatter( + *x.T, s=sx, edgecolors="k", marker="o", label="x" + ) + self._points_y = self.ax.scatter( + *y.T, s=sy, edgecolors="k", marker="X", label="y" + ) + self.ax.legend(fontsize=15) + if not self._show_lines: + return [] + + lines = self._mapping(x, y, ot.matrix) + cmap = plt.get_cmap(self._cmap) + self._lines = [] + for start, end, strength, alpha in lines: + line, = self.ax.plot( + start, + end, + linewidth=0.5 + 4 * strength, + color=cmap(strength), + zorder=0, + alpha=alpha + ) + self._lines.append(line) + if self._title is not None: + self.ax.set_title(self._title) + return [self._points_x, self._points_y] + self._lines + + def update(self, + ot: Transport, + title: Optional[str] = None) -> List["plt.Artist"]: + """Update a plot with a transport instance.""" + x, y, _, _ = self._scatter(ot) + self._points_x.set_offsets(x) + self._points_y.set_offsets(y) + if not self._show_lines: + return [] + + new_lines = self._mapping(x, y, ot.matrix) + cmap = plt.get_cmap(self._cmap) + for line, new_line in zip(self._lines, new_lines): + start, end, strength, alpha = new_line + + line.set_data(start, end) + line.set_linewidth(0.5 + 4 * strength) + line.set_color(cmap(strength)) + line.set_alpha(alpha) + + # Maybe add new lines to the plot. + num_lines = len(self._lines) + num_to_plot = len(new_lines) if self._show_lines else 0 + for i in range(num_lines, num_to_plot): + start, end, strength, alpha = new_lines[i] + + line, = self.ax.plot( + start, + end, + linewidth=0.5 + 4 * strength, + color=cmap(strength), + zorder=0, + alpha=alpha + ) + self._lines.append(line) + + self._lines = self._lines[:num_to_plot] # Maybe remove some + if title is not None: + self.ax.set_title(title) + return [self._points_x, self._points_y] + self._lines + + def animate( + self, + transports: Sequence[Transport], + titles: Optional[Sequence[str]] = None, + frame_rate: float = 10.0 + ) -> "animation.FuncAnimation": + """Make an animation from several transports.""" + _ = self(transports[0]) + if titles is None: + titles = [None for _ in np.arange(0, len(transports))] + assert len(titles) == len(transports), ( + f"titles/transports lengths differ `{len(titles)}`/`{len(transports)}`." + ) + return animation.FuncAnimation( + self.fig, + lambda i: self.update(transports[i], titles[i]), + np.arange(0, len(transports)), + init_func=lambda: self.update(transports[0], titles[0]), + interval=1000 / frame_rate, + blit=True + ) diff --git a/ott/build/lib/ott/tools/segment_sinkhorn.py b/ott/build/lib/ott/tools/segment_sinkhorn.py new file mode 100644 index 0000000..3e6b29d --- /dev/null +++ b/ott/build/lib/ott/tools/segment_sinkhorn.py @@ -0,0 +1,143 @@ +# 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. +from types import MappingProxyType +from typing import Any, Mapping, Optional, Tuple + +import jax.numpy as jnp + +from ott.geometry import costs, pointcloud, segment +from ott.problems.linear import linear_problem +from ott.solvers.linear import sinkhorn + + +def segment_sinkhorn( + x: jnp.ndarray, + y: jnp.ndarray, + num_segments: Optional[int] = None, + max_measure_size: Optional[int] = None, + cost_fn: Optional[costs.CostFn] = None, + segment_ids_x: Optional[jnp.ndarray] = None, + segment_ids_y: Optional[jnp.ndarray] = None, + indices_are_sorted: bool = False, + num_per_segment_x: Optional[Tuple[int, ...]] = None, + num_per_segment_y: Optional[Tuple[int, ...]] = None, + weights_x: Optional[jnp.ndarray] = None, + weights_y: Optional[jnp.ndarray] = None, + sinkhorn_kwargs: Mapping[str, Any] = MappingProxyType({}), + **kwargs: Any +) -> jnp.ndarray: + """Compute regularized OT cost between subsets of vectors in `x` and `y`. + + Helper function designed to compute Sinkhorn regularized OT cost between + several point clouds of varying size, in parallel, using padding. + In practice, The inputs `x` and `y` (and their weight vectors `weights_x` and + `weights_y`) are assumed to be large weighted point clouds, that describe + points taken from multiple measures. To extract several subsets of points, we + provide two interfaces. The first interface assumes that a vector of id's is + passed, describing for each point of `x` (resp. `y`) to which measure the + point belongs to. The second interface assumes that `x` and `y` were simply + formed by concatenating several measures contiguously, and that only indices + that segment these groups are needed to recover them. + + For both interfaces, both `x` and `y` should contain the same total number of + segments. Each segment will be padded as necessary, all segments rearranged as + a tensor, and :func:`jax.vmap` used to evaluate Sinkhorn divergences in + parallel. + + Args: + x: Array of input points, of shape `[num_x, feature]`. + Multiple segments are held in this single array. + y: Array of target points, of shape `[num_y, feature]`. + num_segments: Number of segments contained in `x` and `y`. + Providing this is required for JIT compilation to work, + see also :func:`~ott.geometry.segment.segment_point_cloud`. + max_measure_size: Total size of measures after padding. Should ideally be + set to an upper bound on points clouds processed with the segment + interface. Providing this is required for JIT compilation to work. + cost_fn: Cost function, defaults to + :class:`~ott.geometry.costs.SqEuclidean`. + segment_ids_x: **1st interface** The segment ID for which each row of `x` + belongs. This is a similar interface to `jax.ops.segment_sum`. + segment_ids_y: **1st interface** The segment ID for which each row of `y` + belongs. + indices_are_sorted: **1st interface** Whether `segment_ids_x` and + `segment_ids_y` are sorted. + num_per_segment_x: **2nd interface** Number of points in each segment in + `x`. For example, [100, 20, 30] would imply that `x` is segmented into + three arrays of length `[100]`, `[20]`, and `[30]` respectively. + num_per_segment_y: **2nd interface** Number of points in each segment in + `y`. + weights_x: Weights of each input points, arranged in the same segmented + order as `x`. + weights_y: Weights of each input points, arranged in the same segmented + order as `y`. + sinkhorn_kwargs: Optionally a dict containing the keywords arguments for + calls for the :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver, + called three times to evaluate for each segment the Sinkhorn regularized + OT cost between `x`/`y`, `x`/`x`, and `y`/`y` (except when `static_b` is + `True`, in which case `y`/`y` is not evaluated). + kwargs: keywords arguments passed to form + :class:`~ott.geometry.pointcloud.PointCloud` geometry objects from the + subsets of points and masses selected in `x` and `y`, possibly a + :class:`~ott.geometry.costs.CostFn` or an entropy regularizer. + + Returns: + An array of Sinkhorn regularized OT costs for each segment. + """ + # instantiate padding vector + dim = x.shape[1] + if cost_fn is None: + # default padder + padding_vector = costs.CostFn._padder(dim=dim) + else: + padding_vector = cost_fn._padder(dim=dim) + + def eval_fn( + padded_x: jnp.ndarray, + padded_y: jnp.ndarray, + padded_weight_x: jnp.ndarray, + padded_weight_y: jnp.ndarray, + ) -> float: + mask_x = padded_weight_x > 0.0 + mask_y = padded_weight_y > 0.0 + + geom = pointcloud.PointCloud( + padded_x, + padded_y, + cost_fn=cost_fn, + src_mask=mask_x, + tgt_mask=mask_y, + **kwargs, + ) + prob = linear_problem.LinearProblem( + geom, a=padded_weight_x, b=padded_weight_y + ) + solver = sinkhorn.Sinkhorn(**sinkhorn_kwargs) + return solver(prob).reg_ot_cost + + return segment._segment_interface( + x, + y, + eval_fn, + num_segments=num_segments, + max_measure_size=max_measure_size, + segment_ids_x=segment_ids_x, + segment_ids_y=segment_ids_y, + indices_are_sorted=indices_are_sorted, + num_per_segment_x=num_per_segment_x, + num_per_segment_y=num_per_segment_y, + weights_x=weights_x, + weights_y=weights_y, + padding_vector=padding_vector + ) diff --git a/ott/build/lib/ott/tools/sinkhorn_divergence.py b/ott/build/lib/ott/tools/sinkhorn_divergence.py new file mode 100644 index 0000000..572075e --- /dev/null +++ b/ott/build/lib/ott/tools/sinkhorn_divergence.py @@ -0,0 +1,354 @@ +# 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. +from types import MappingProxyType +from typing import Any, Mapping, Optional, Tuple, Type + +import jax.numpy as jnp + +from ott import utils +from ott.geometry import costs, geometry, pointcloud, segment +from ott.problems.linear import linear_problem, potentials +from ott.solvers import linear +from ott.solvers.linear import acceleration, sinkhorn + +__all__ = [ + "sinkhorn_divergence", "segment_sinkhorn_divergence", + "SinkhornDivergenceOutput" +] + +Potentials_t = Tuple[jnp.ndarray, jnp.ndarray] + + +@utils.register_pytree_node +class SinkhornDivergenceOutput: # noqa: D101 + divergence: float + potentials: Tuple[Potentials_t, Potentials_t, Potentials_t] + geoms: Tuple[geometry.Geometry, geometry.Geometry, geometry.Geometry] + errors: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray], + Optional[jnp.ndarray]] + converged: Tuple[bool, bool, bool] + a: jnp.ndarray + b: jnp.ndarray + n_iters: Tuple[int, int, int] + + def to_dual_potentials(self) -> "potentials.EntropicPotentials": + """Return dual estimators :cite:`pooladian:22`, eq. 8.""" + geom_xy, *_ = self.geoms + prob_xy = linear_problem.LinearProblem(geom_xy, a=self.a, b=self.b) + (f_xy, g_xy), (f_x, _), (_, g_y) = self.potentials + return potentials.EntropicPotentials( + f_xy, g_xy, prob_xy, f_xx=f_x, g_yy=g_y + ) + + def tree_flatten(self): # noqa: D102 + return [ + self.divergence, + self.potentials, + self.geoms, + self.a, + self.b, + ], { + "n_iters": self.n_iters, + "converged": self.converged, + "errors": self.errors + } + + @classmethod + def tree_unflatten(cls, aux_data, children): # noqa: D102 + div, pots, geoms, a, b = children + return cls(div, pots, geoms, a=a, b=b, **aux_data) + + +def sinkhorn_divergence( + geom: Type[geometry.Geometry], + *args: Any, + a: Optional[jnp.ndarray] = None, + b: Optional[jnp.ndarray] = None, + sinkhorn_kwargs: Mapping[str, Any] = MappingProxyType({}), + static_b: bool = False, + share_epsilon: bool = True, + symmetric_sinkhorn: bool = True, + **kwargs: Any, +) -> SinkhornDivergenceOutput: + """Compute Sinkhorn divergence defined by a geometry, weights, parameters. + + Args: + geom: Type of the geometry. + args: Positional arguments to + :meth:`~ott.geometry.geometry.Geometry.prepare_divergences` that are + specific to each geometry. + a: the weight of each input point. The sum of all elements of `a` must + match that of `b` to converge. + b: the weight of each target point. The sum of all elements of `b` must + match that of `a` to converge. + sinkhorn_kwargs: keywords arguments for + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` that is called twice + if ``static_b = True`` else 3 times. + static_b: if True, divergence of measure `b` against itself is **not** + computed. + share_epsilon: if True, enforces that the same epsilon regularizer is shared + for all 2 or 3 terms of the Sinkhorn divergence. In that case, the epsilon + will be by default that used when comparing x to y (contained in the first + geometry). This flag is set to True by default, because in the default + setting, the epsilon regularization is a function of the mean of the cost + matrix. + symmetric_sinkhorn: Use Sinkhorn updates in Eq. 25 of :cite:`feydy:19` for + symmetric terms comparing x/x and y/y. + kwargs: keywords arguments to the generic class. This is specific to each + geometry. + + Returns: + Sinkhorn divergence value, three pairs of potentials, three costs. + """ + geoms = geom.prepare_divergences(*args, static_b=static_b, **kwargs) + geom_xy, geom_x, geom_y, *_ = geoms + (None,) * 3 + num_a, num_b = geom_xy.shape + + if share_epsilon: + if isinstance(geom_x, geometry.Geometry): + geom_x = geom_x.copy_epsilon(geom_xy) + if isinstance(geom_y, geometry.Geometry): + geom_y = geom_y.copy_epsilon(geom_xy) + + a = jnp.ones(num_a) / num_a if a is None else a + b = jnp.ones(num_b) / num_b if b is None else b + return _sinkhorn_divergence( + geom_xy, + geom_x, + geom_y, + a=a, + b=b, + symmetric_sinkhorn=symmetric_sinkhorn, + **sinkhorn_kwargs + ) + + +def _sinkhorn_divergence( + geometry_xy: geometry.Geometry, + geometry_xx: geometry.Geometry, + geometry_yy: Optional[geometry.Geometry], + a: jnp.ndarray, + b: jnp.ndarray, + symmetric_sinkhorn: bool, + **kwargs: Any, +) -> SinkhornDivergenceOutput: + """Compute the (unbalanced) Sinkhorn divergence for the wrapper function. + + This definition includes a correction depending on the total masses of each + measure, as defined in :sejourne:19:, eq. 15. + + Args: + geometry_xy: a Cost object able to apply kernels with a certain epsilon, + between the views X and Y. + geometry_xx: a Cost object able to apply kernels with a certain epsilon, + between elements of the view X. + geometry_yy: a Cost object able to apply kernels with a certain epsilon, + between elements of the view Y. + a: jnp.ndarray[n]: the weight of each input point. The sum of + all elements of ``b`` must match that of ``a`` to converge. + b: jnp.ndarray[m]: the weight of each target point. The sum of + all elements of ``b`` must match that of ``a`` to converge. + symmetric_sinkhorn: Use Sinkhorn updates in Eq. 25 of :cite:`feydy:19` for + symmetric terms comparing x/x and y/y. + kwargs: Keyword arguments to :func:`~ott.solvers.linear.sinkhorn.Sinkhorn`. + + Returns: + SinkhornDivergenceOutput named tuple. + """ + # When computing a Sinkhorn divergence, the (x,y) terms and (x,x) / (y,y) + # terms are computed independently. The user might want to pass some + # sinkhorn_kwargs to parameterize Sinkhorn's behavior, but those should + # only apply to the (x,y) part. For the (x,x) / (y,y) part we fall back + # on a simpler choice (parallel_dual_updates + momentum 0.5) that is known + # to work well in such settings. In the future we might want to give some + # freedom on setting parameters for the (x,x)/(y,y) part. + # Since symmetric terms are computed assuming a = b, the linear systems + # arising in implicit differentiation (if used) of the potentials computed for + # the symmetric parts should be marked as symmetric. + kwargs_symmetric = kwargs.copy() + if symmetric_sinkhorn: + kwargs_symmetric.update( + parallel_dual_updates=True, + momentum=acceleration.Momentum(start=0, value=0.5), + anderson=None, + ) + implicit_diff = kwargs.get("implicit_diff", None) + if implicit_diff is not None: + kwargs_symmetric["implicit_diff"] = implicit_diff.replace(symmetric=True) + + out_xy = linear.solve(geometry_xy, a, b, **kwargs) + out_xx = linear.solve(geometry_xx, a, a, **kwargs_symmetric) + if geometry_yy is None: + # Create dummy output, corresponds to scenario where static_b is True. + # This choice ensures that `converged`` of this dummy output is True. + out_yy = sinkhorn.SinkhornOutput( + (None, None), + errors=jnp.array([-jnp.inf]), + reg_ot_cost=0.0, + threshold=0.0, + inner_iterations=0, + ) + else: + out_yy = linear.solve(geometry_yy, b, b, **kwargs_symmetric) + + div = ( + out_xy.reg_ot_cost - 0.5 * (out_xx.reg_ot_cost + out_yy.reg_ot_cost) + + 0.5 * geometry_xy.epsilon * (jnp.sum(a) - jnp.sum(b)) ** 2 + ) + return SinkhornDivergenceOutput( + divergence=div, + potentials=((out_xy.f, out_xy.g), (out_xx.f, out_xx.g), + (out_yy.f, out_yy.g)), + geoms=(geometry_xy, geometry_xx, geometry_yy), + errors=(out_xy.errors, out_xx.errors, out_yy.errors), + converged=(out_xy.converged, out_xx.converged, out_yy.converged), + a=a, + b=b, + n_iters=(out_xy.n_iters, out_xx.n_iters, out_yy.n_iters), + ) + + +def segment_sinkhorn_divergence( + x: jnp.ndarray, + y: jnp.ndarray, + num_segments: Optional[int] = None, + max_measure_size: Optional[int] = None, + cost_fn: Optional[costs.CostFn] = None, + segment_ids_x: Optional[jnp.ndarray] = None, + segment_ids_y: Optional[jnp.ndarray] = None, + indices_are_sorted: bool = False, + num_per_segment_x: Optional[Tuple[int, ...]] = None, + num_per_segment_y: Optional[Tuple[int, ...]] = None, + weights_x: Optional[jnp.ndarray] = None, + weights_y: Optional[jnp.ndarray] = None, + sinkhorn_kwargs: Mapping[str, Any] = MappingProxyType({}), + static_b: bool = False, + share_epsilon: bool = True, + symmetric_sinkhorn: bool = False, + **kwargs: Any +) -> jnp.ndarray: + """Compute Sinkhorn divergence between subsets of vectors in `x` and `y`. + + Helper function designed to compute Sinkhorn divergences between several point + clouds of varying size, in parallel, using padding for efficiency. + In practice, The inputs `x` and `y` (and their weight vectors `weights_x` and + `weights_y`) are assumed to be large weighted point clouds, that describe + points taken from multiple measures. To extract several subsets of points, we + provide two interfaces. The first interface assumes that a vector of id's is + passed, describing for each point of `x` (resp. `y`) to which measure the + point belongs to. The second interface assumes that `x` and `y` were simply + formed by concatenating several measures contiguously, and that only indices + that segment these groups are needed to recover them. + + For both interfaces, both `x` and `y` should contain the same total number of + segments. Each segment will be padded as necessary, all segments rearranged as + a tensor, and `vmap` used to evaluate Sinkhorn divergences in parallel. + + Args: + x: Array of input points, of shape `[num_x, feature]`. + Multiple segments are held in this single array. + y: Array of target points, of shape `[num_y, feature]`. + num_segments: Number of segments contained in `x` and `y`. + Providing this is required for JIT compilation to work, + see also :func:`~ott.geometry.segment.segment_point_cloud`. + max_measure_size: Total size of measures after padding. Should ideally be + set to an upper bound on points clouds processed with the segment + interface. Should also be smaller than total length of `x` or `y`. + Providing this is required for JIT compilation to work. + cost_fn: Cost function, + defaults to :class:`~ott.geometry.costs.SqEuclidean`. + segment_ids_x: **1st interface** The segment ID for which each row of `x` + belongs. This is a similar interface to :func:`jax.ops.segment_sum`. + segment_ids_y: **1st interface** The segment ID for which each row of `y` + belongs. + indices_are_sorted: **1st interface** Whether `segment_ids_x` and + `segment_ids_y` are sorted. + num_per_segment_x: **2nd interface** Number of points in each segment in + `x`. For example, [100, 20, 30] would imply that `x` is segmented into + three arrays of length `[100]`, `[20]`, and `[30]` respectively. + num_per_segment_y: **2nd interface** Number of points in each segment in + `y`. + weights_x: Weights of each input points, arranged in the same segmented + order as `x`. + weights_y: Weights of each input points, arranged in the same segmented + order as `y`. + sinkhorn_kwargs: Optionally a dict containing the keywords arguments for + calls to the `sinkhorn` function, called three times to evaluate for each + segment the Sinkhorn regularized OT cost between `x`/`y`, `x`/`x`, and + `y`/`y` (except when `static_b` is `True`, in which case `y`/`y` is not + evaluated) + static_b: if True, divergence of measure b against itself is NOT computed + share_epsilon: if True, enforces that the same epsilon regularizer is shared + for all 2 or 3 terms of the Sinkhorn divergence. In that case, the epsilon + will be by default that used when comparing x to y (contained in the first + geometry). This flag is set to True by default, because in the default + setting, the epsilon regularization is a function of the mean of the cost + matrix. + symmetric_sinkhorn: Use Sinkhorn updates in Eq. 25 of :cite:`feydy:19` for + symmetric terms comparing x/x and y/y. + kwargs: keywords arguments passed to form + :class:`~ott.geometry.pointcloud.PointCloud` geometry objects from the + subsets of points and masses selected in `x` and `y`, this could be for + instance entropy regularization float, scheduler or normalization. + + Returns: + An array of Sinkhorn divergences for each segment. + """ + # instantiate padding vector + dim = x.shape[1] + if cost_fn is None: + # default padder + padding_vector = costs.CostFn._padder(dim=dim) + else: + padding_vector = cost_fn._padder(dim=dim) + + def eval_fn( + padded_x: jnp.ndarray, + padded_y: jnp.ndarray, + padded_weight_x: jnp.ndarray, + padded_weight_y: jnp.ndarray, + ) -> float: + mask_x = padded_weight_x > 0.0 + mask_y = padded_weight_y > 0.0 + return sinkhorn_divergence( + pointcloud.PointCloud, + padded_x, + padded_y, + a=padded_weight_x, + b=padded_weight_y, + sinkhorn_kwargs=sinkhorn_kwargs, + static_b=static_b, + share_epsilon=share_epsilon, + symmetric_sinkhorn=symmetric_sinkhorn, + cost_fn=cost_fn, + src_mask=mask_x, + tgt_mask=mask_y, + **kwargs + ).divergence + + return segment._segment_interface( + x, + y, + eval_fn, + num_segments=num_segments, + max_measure_size=max_measure_size, + segment_ids_x=segment_ids_x, + segment_ids_y=segment_ids_y, + indices_are_sorted=indices_are_sorted, + num_per_segment_x=num_per_segment_x, + num_per_segment_y=num_per_segment_y, + weights_x=weights_x, + weights_y=weights_y, + padding_vector=padding_vector + ) diff --git a/ott/build/lib/ott/tools/soft_sort.py b/ott/build/lib/ott/tools/soft_sort.py new file mode 100644 index 0000000..faf52f2 --- /dev/null +++ b/ott/build/lib/ott/tools/soft_sort.py @@ -0,0 +1,706 @@ +# 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 functools +from typing import Any, Callable, Optional, Tuple, Union + +import jax +import jax.numpy as jnp +import jax.tree_util as jtu +import numpy as np + +from ott import utils +from ott.geometry import costs, pointcloud +from ott.problems.linear import linear_problem +from ott.solvers import linear +from ott.solvers.linear import sinkhorn + +__all__ = [ + "sort", "ranks", "sort_with", "quantile", "quantile_normalization", + "quantize", "topk_mask", "multivariate_cdf_quantile_maps" +] + +Func_t = Callable[[jnp.ndarray], jnp.ndarray] + + +def transport_for_sort( + inputs: jnp.ndarray, + weights: Optional[jnp.ndarray] = None, + target_weights: Optional[jnp.ndarray] = None, + squashing_fun: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, + epsilon: float = 1e-2, + **kwargs: Any, +) -> sinkhorn.SinkhornOutput: + r"""Solve reg. OT, from inputs to a weighted family of increasing values. + + Args: + inputs: Array[num_points]. Must be one dimensional. + weights: Array[num_points]. Weight vector `a` for input values. + target_weights: Array[num_targets]: Weight vector of the target + measure. It may be of different size than `weights`. + squashing_fun: function taking an array to squash all its entries in [0,1]. + sigmoid of whitened values by default. Can be set to be the identity by + passing ``squashing_fun = lambda x : x`` instead. + epsilon: the regularization parameter. + kwargs: keyword arguments for + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. + + Returns: + A :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput` object. + """ + shape = inputs.shape + if len(shape) > 2 or (len(shape) == 2 and shape[1] != 1): + raise ValueError( + f"Shape ({shape}) not supported. The input should be one-dimensional." + ) + + x = jnp.expand_dims(jnp.squeeze(inputs), axis=1) + if squashing_fun is None: + squashing_fun = lambda z: jax.nn.sigmoid((z - jnp.mean(z)) / + (jnp.std(z) + 1e-10)) + x = squashing_fun(x) + + a = jnp.squeeze(weights) + b = jnp.squeeze(target_weights) + num_targets = inputs.shape[0] if b is None else b.shape[0] + y = jnp.linspace(0.0, 1.0, num_targets)[:, jnp.newaxis] + + geom = pointcloud.PointCloud(x, y, epsilon=epsilon) + prob = linear_problem.LinearProblem(geom, a=a, b=b) + + solver = sinkhorn.Sinkhorn(**kwargs) + + return solver(prob) + + +def apply_on_axis(op, inputs, axis, *args, **kwargs: Any) -> jnp.ndarray: + """Apply a differentiable operator on a given axis of the input. + + Args: + op: a differentiable operator (can be ranks, quantile, etc.) + inputs: Array of any shape. + axis: the axis (int) or tuple of ints on which to apply the operator. If + several axes are passed to the operator, those are merged as a single + dimension. + args: other positional arguments to the operator. + kwargs: other positional arguments to the operator. + + Returns: + An Array holding the output of the differentiable operator on the given + axis. + """ + op_inner = functools.partial(op, **kwargs) + axis = (axis,) if isinstance(axis, int) else axis + num_points = np.prod(np.array(inputs.shape)[(axis,)]) + permutation = np.arange(len(inputs.shape)) + axis = tuple(permutation[a] for a in axis) + permutation = tuple(sorted(set(permutation) - set(axis)) + sorted(axis)) + inputs = jnp.transpose(inputs, permutation) + + batch_fn = jax.vmap(op_inner, in_axes=(0,) + (None,) * len(args)) + result = batch_fn(jnp.reshape(inputs, (-1, num_points)), *args) + shrink = len(axis) + result = jnp.reshape(result, inputs.shape[:-shrink] + result.shape[-1:]) + + permutation = tuple(range(len(result.shape))) + rank = len(result.shape) - 1 + axis = min(axis) + permutation = permutation[:axis] + (rank,) + permutation[axis:-1] + return jnp.transpose(result, permutation) + + +def _sort( + inputs: jnp.ndarray, topk: int, num_targets: Optional[int], **kwargs: Any +) -> jnp.ndarray: + """Apply the soft sort operator on a one dimensional array.""" + num_points = inputs.shape[0] + a = jnp.ones((num_points,)) / num_points + if 0 < topk < num_points: + start_index = 1 + b = jnp.concatenate([ + jnp.array([(num_points - topk) / num_points]), + jnp.ones(topk) / num_points + ]) + else: + # Use the "sorting" initializer if default uniform weights of same size. + if num_targets is None or num_targets == num_points: + num_targets = num_points + # use sorting initializer in this case. + kwargs.setdefault("initializer", "sorting") + start_index = 0 + b = jnp.ones((num_targets,)) / num_targets + ot = transport_for_sort(inputs, a, b, **kwargs) + out = 1.0 / b * ot.apply(inputs, axis=0) + return out[start_index:] + + +def sort( + inputs: jnp.ndarray, + axis: int = -1, + topk: int = -1, + num_targets: Optional[int] = None, + **kwargs: Any, +) -> jnp.ndarray: + r"""Apply the soft sort operator on a given axis of the input. + + For instance: + + .. code-block:: python + + x = jax.random.uniform(rng, (100,)) + x_sorted = sort(x) + + will output sorted convex-combinations of values contained in ``x``, that are + differentiable approximations to the sorted vector of entries in ``x``. + These can be compared with the values produced by :func:`jax.numpy.sort`, + + .. code-block:: python + + x_sorted = jax.numpy.sort(x) + + Args: + inputs: Array of any shape. + axis: the axis on which to apply the soft-sorting operator. + topk: if set to a positive value, the returned vector will only contain + the top-k values. This also reduces the complexity of soft-sorting, since + the number of target points to which the slice of the ``inputs`` tensor + will be mapped to will be equal to ``topk + 1``. + num_targets: if ``topk`` is not specified, a vector of size``num_targets`` + is returned. This defines the number of (composite) sorted values computed + from the inputs (each value is a convex combination of values recorded in + the inputs, provided in increasing order). If neither ``topk`` nor + ``num_targets`` are specified, ``num_targets`` defaults to the size of the + slices of the input that are sorted, i.e. ``inputs.shape[axis]``, and the + number of composite sorted values is equal to the slice of the inputs that + are sorted. As a result, the output is of the same size as ``inputs``. + kwargs: keyword arguments passed on to lower level functions. Of interest + to the user are ``squashing_fun``, which will redistribute the values in + ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) + to solve the optimal transport problem; + :class:`cost_fn ` object of + :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground + 1D cost function to transport from ``inputs`` to the ``num_targets`` + target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` + are passed on to parameterize the + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. + + Returns: + An Array of the same shape as the input, except on ``axis``, where that size + will be equal to ``topk`` or ``num_targets``, with soft-sorted values on the + given axis. Same size as ``inputs`` if both these parameters are ``None``. + """ + return apply_on_axis(_sort, inputs, axis, topk, num_targets, **kwargs) + + +def _ranks( + inputs: jnp.ndarray, num_targets, target_weights, **kwargs: Any +) -> jnp.ndarray: + """Apply the soft ranks operator on a one dimensional array.""" + num_points = inputs.shape[0] + if target_weights is None: + num_targets = num_points if num_targets is None else num_targets + target_weights = jnp.ones((num_targets,)) / num_targets + else: + num_targets = target_weights.shape[0] + a = jnp.ones((num_points,)) / num_points + ot = transport_for_sort(inputs, a, target_weights, **kwargs) + out = 1.0 / a * ot.apply(jnp.arange(num_targets), axis=1) + out *= (num_points - 1.0) / (num_targets - 1.0) + return jnp.reshape(out, inputs.shape) + + +def ranks( + inputs: jnp.ndarray, + axis: int = -1, + num_targets: Optional[int] = None, + target_weights: Optional[jnp.ndarray] = None, + **kwargs: Any, +) -> jnp.ndarray: + r"""Apply the soft rank operator on input tensor. + + For instance: + + .. code-block:: python + + x = jax.random.uniform(rng, (100,)) + x_ranks = ranks(x) + + will output values that are differentiable approximations to the ranks of + entries in ``x``. These should be compared to the non-differentiable rank + vectors, namely the normalized inverse permutation produced by + :func:`jax.numpy.argsort`, which can be obtained as: + + .. code-block:: python + + x_ranks = jax.numpy.argsort(jax.numpy.argsort(x)) + + Args: + inputs: Array of any shape. + axis: the axis on which to apply the soft-sorting operator. + target_weights: This vector contains weights (summing to 1) that describe + amount of mass shipped to targets. + num_targets: If ``target_weights` is ``None``, ``num_targets`` is + considered to define the number of targets used to rank inputs. Each + rank in the output will be a convex combination of + ``{1, .., num_targets}``. The weight of each of these points + is assumed to be uniform. If neither ``num_targets`` nor + ``target_weights`` are specified, ``num_targets`` defaults to the size + of the slices of the input that are sorted, i.e. ``inputs.shape[axis]``. + kwargs: keyword arguments passed on to lower level functions. Of interest + to the user are ``squashing_fun``, which will redistribute the values in + ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) + to solve the optimal transport problem; + :class:`cost_fn ` object of + :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground + 1D cost function to transport from ``inputs`` to the ``num_targets`` + target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` + are passed on to parameterize the + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. + + Returns: + An Array of the same shape as the input with soft-rank values + normalized to be in :math:`[0, n-1]` where :math:`n` is + ``inputs.shape[axis]``. + """ + return apply_on_axis( + _ranks, inputs, axis, num_targets, target_weights, **kwargs + ) + + +def topk_mask( + inputs: jnp.ndarray, + axis: int = -1, + k: int = 1, + **kwargs: Any, +) -> jnp.ndarray: + r"""Soft :math:`\text{top-}k` selection mask. + + For instance: + + .. code-block:: python + + k = 5 + x = jax.random.uniform(rng, (100,)) + mask = topk_mask(x, k=k) + + will output a vector of shape ``x.shape``, with values in :math:`[0,1]`, that + are differentiable approximations to the binary mask selecting the top $k$ + entries in ``x``. These should be compared to the non-differentiable mask + obtained with :func:`jax.numpy.sort`, which can be obtained as: + + .. code-block:: python + + mask = x >= jax.numpy.sort(x).flip()[k-1] + + Args: + inputs: Array of any shape. + axis: the axis on which to apply the soft-sorting operator. + k: topk parameter. Should be smaller than ``inputs.shape[axis]``. + kwargs: keyword arguments passed on to lower level functions. Of interest + to the user are ``squashing_fun``, which will redistribute the values in + ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) + to solve the optimal transport problem; + :class:`cost_fn ` object of + :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground + 1D cost function to transport from ``inputs`` to the ``num_targets`` + target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` + are passed on to parameterize the + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. + + Returns: + The soft :math:`\text{top-}k` selection mask. + """ + num_points = inputs.shape[axis] + assert k < num_points, ( + f"`k` must be smaller than `inputs.shape[axis]`, yet {k} >= {num_points}." + ) + target_weights = jnp.array([1.0 - k / num_points, k / num_points]) + out = apply_on_axis( + _ranks, + inputs, + axis, + num_targets=None, + target_weights=target_weights, + **kwargs + ) + return out / (num_points - 1) + + +def quantile( + inputs: jnp.ndarray, + q: Optional[Union[float, jnp.ndarray]], + axis: Union[int, Tuple[int, ...]] = -1, + weight: Optional[Union[float, jnp.ndarray]] = None, + **kwargs: Any, +) -> jnp.ndarray: + r"""Apply the soft quantiles operator on the input tensor. + + For instance: + + .. code-block:: python + + x = jax.random.uniform(rng, (100,)) + x_quantiles = quantile(x, q=jnp.array([0.2, 0.8])) + + ``x_quantiles`` will hold an approximation to the 20 and 80 percentiles in + ``x``, computed as a convex combination (a weighted mean, with weights summing + to 1) of all values in ``x`` (and not, as for standard quantiles, the + values ``x_sorted[20]`` and ``x_sorted[80]`` if ``x_sorted=jnp.sort(x)``). + These values offer a trade-off between accuracy (closeness to the true + percentiles) and gradient (the Jacobian of ``x_quantiles`` w.r.t ``x`` will + impact all values listed in ``x``, not just those indexed at 20 and 80). + + The non-differentiable version is given by :func:`jax.numpy.quantile`, e.g. + + .. code-block:: python + + x_quantiles = jax.numpy.quantile(x, q=jnp.array([0.2, 0.8])) + + Args: + inputs: an Array of any shape. + q: values of the quantile level to be computed, e.g. [0.5] for median. + These values should all lie within the segment :math:`]0,1[`, excluding + boundary values :math:`0` and :math:`1`. + axis: the axis on which to apply the operator. + weight: the weight assigned to each quantile target value in the OT problem. + This weight should be small, typically of the order of ``1/n``, where ``n`` + is the size of ``x``. Note: Since the size of ``q`` times ``weight`` + must be strictly smaller than ``1``, in order to leave enough mass to set + other target values in the transport problem, the algorithm might ensure + this by setting, when needed, a lower value. + kwargs: keyword arguments passed on to lower level functions. Of interest + to the user are ``squashing_fun``, which will redistribute the values in + ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) + to solve the optimal transport problem; + :class:`cost_fn ` object of + :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground + 1D cost function to transport from ``inputs`` to the ``num_targets`` + target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` + are passed on to parameterize the + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. + + Returns: + An Array, which has the same shape as ``inputs``, except on the ``axis`` + that is passed, which has size ``q.shape[0]``, to collect soft-quantile + values. + """ + + def _quantile( + inputs: jnp.ndarray, q: float, weight: float, **kwargs + ) -> jnp.ndarray: + num_points = inputs.shape[0] + q = jnp.array([0.2, 0.5, 0.8]) if q is None else jnp.atleast_1d(q) + num_quantiles = q.shape[0] + a = jnp.ones((num_points,)) / num_points + idx = jnp.argsort(q) + q = q[idx] + + extended_q = jnp.concatenate([jnp.array([0.0]), q, jnp.array([1.0])]) + filler_weights = extended_q[1:] - extended_q[:-1] + safe_weight = 0.5 * jnp.concatenate([ + jnp.array([1.0 / num_quantiles]), filler_weights + ]) + if weight is None: + # Populate with other options. + safe_weight = jnp.concatenate([ + safe_weight, + jnp.array( + [0.02] + ), # reasonable mass per quantile for a small number of points + jnp.array( + [1.5 / num_points] + ), # this means each quantile would be ~ assigned 1.5 points. + ]) + else: + safe_weight = jnp.concatenate([safe_weight, jnp.atleast_1d(weight)]) + weight = jnp.min(safe_weight) + weights = jnp.ones(filler_weights.shape) * weight + + # Takes into account quantile_width in the definition of weights + shift = -jnp.ones(filler_weights.shape) + shift = shift + 0.5 * ( + jax.nn.one_hot(0, num_quantiles + 1) + + jax.nn.one_hot(num_quantiles, num_quantiles + 1) + ) + filler_weights = filler_weights + weights * shift + + # Adds one more value to have tensors of the same shape to interleave them. + quantile_weights = jnp.ones(num_quantiles + 1) * weights + + # Interleaves the filler_weights with the quantile weights. + weights = jnp.reshape( + jnp.stack([filler_weights, quantile_weights], axis=1), (-1,) + )[:-1] + + ot = transport_for_sort(inputs, a, weights, **kwargs) + out = 1.0 / weights * ot.apply(jnp.squeeze(inputs), axis=0) + + # Recover odd indices corresponding to the desired quantiles. + odds = np.concatenate([ + np.zeros((num_quantiles + 1, 1), dtype=bool), + np.ones((num_quantiles + 1, 1), dtype=bool) + ], + axis=1).ravel()[:-1] + return out[odds][idx] + + return apply_on_axis(_quantile, inputs, axis, q, weight, **kwargs) + + +def multivariate_cdf_quantile_maps( + inputs: jnp.ndarray, + target_sampler: Optional[Callable[[jax.Array, Tuple[int, int]], + jnp.ndarray]] = None, + rng: Optional[jax.Array] = None, + num_target_samples: Optional[int] = None, + cost_fn: Optional[costs.CostFn] = None, + epsilon: Optional[float] = None, + input_weights: Optional[jnp.ndarray] = None, + target_weights: Optional[jnp.ndarray] = None, + **kwargs: Any +) -> Tuple[Func_t, Func_t]: + r"""Returns multivariate CDF and quantile maps, given input samples. + + Implements the multivariate generalizations for CDF and quantiles proposed in + :cite:`chernozhukov:17`. The reference measure is assumed to be the uniform + measure by default, but can be modified. For consistency, the reference + measure should be symmetrically centered around + :math:`(\tfrac{1}{2},\cdots,\tfrac{1}{2})` and supported on :math:`[0, 1]^d`. + + The implementation return two entropic map estimators, one for the CDF map, + the other for the quantiles map. + + Args: + inputs: 2D array of ``[n, d]`` vectors. + target_sampler: Callable that takes a ``rng`` and ``[m, d]`` shape. + ``m`` is passed on as ``target_num_samples``, dimension ``d`` is inferred + directly from the shape passed in ``inputs``. This is assumed by default + to be :func:`~jax.random.uniform`, and could be any other random sampler + properly wrapped to have the signature above. + rng: rng key used by ``target_sampler``. + num_target_samples: number ``m`` of points generated in the target + distribution. + cost_fn: Cost function, used to compare ``inputs`` and ``targets``. + Passed on to instantiate a + :class:`~ott.geometry.pointcloud.PointCloud` object. If :obj:`None`, + :class:`~ott.geometry.costs.SqEuclidean` is used. + epsilon: entropic regularization parameter used to instantiate the + :class:`~ott.geometry.pointcloud.PointCloud` object. + input_weights: ``[n,]`` vector of weights for input measure. Assumed to + be uniform by default. + target_weights: ``[m,]`` vector of weights for target measure. Assumed + to be uniform by default. + kwargs: keyword arguments passed on to the :func:`~ott.solvers.linear.solve` + function, which solves the OT problem between ``inputs`` and ``targets`` + using the :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` algorithm. + + Returns: + - The multivariate CDF map, taking a ``[b, d]`` batch of vectors in the + range of the ``inputs``, and mapping each vector within the range + of the reference measure (assumed by default to be :math:`[0, 1]^d`). + - The quantile map, mapping a batch ``[b, d]`` of multivariate quantile + vectors onto ``[b, d]`` vectors in :math:`[0, 1]^d`, the range of + the reference measure. + """ + n, d = inputs.shape + rng = utils.default_prng_key(rng) + + if num_target_samples is None: + num_target_samples = n + if target_sampler is None: + target_sampler = jax.random.uniform + + targets = target_sampler(rng, (num_target_samples, d)) + geom = pointcloud.PointCloud( + inputs, targets, cost_fn=cost_fn, epsilon=epsilon + ) + + out = linear.solve(geom, a=input_weights, b=target_weights, **kwargs) + potentials = out.to_dual_potentials() + + cdf_map = jtu.Partial(lambda x, p: p.transport(x), p=potentials) + quantile_map = jtu.Partial( + lambda x, p: p.transport(x, forward=False), p=potentials + ) + return cdf_map, quantile_map + + +def _quantile_normalization( + inputs: jnp.ndarray, targets: jnp.ndarray, weights: float, **kwargs: Any +) -> jnp.ndarray: + """Apply soft quantile normalization on a one dimensional array.""" + num_points = inputs.shape[0] + a = jnp.ones((num_points,)) / num_points + ot = transport_for_sort(inputs, a, weights, **kwargs) + return 1.0 / a * ot.apply(targets, axis=1) + + +def quantile_normalization( + inputs: jnp.ndarray, + targets: jnp.ndarray, + weights: Optional[jnp.ndarray] = None, + axis: int = -1, + **kwargs: Any, +) -> jnp.ndarray: + r"""Re-normalize inputs so that its quantiles match those of targets/weights. + + Quantile normalization rearranges the values in inputs to values that match + the distribution of values described in the discrete distribution ``targets`` + weighted by ``weights``. This transformation preserves the order of values + in ``inputs`` along the specified ``axis``. + + Args: + inputs: array of any shape whose values will be changed to match those in + ``targets``. + targets: sorted array (in ascending order) of dimension 1 describing a + discrete distribution. Note: the ``targets`` values must be provided as + a sorted vector. + weights: vector of non-negative weights, summing to :math:`1`, of the same + size as ``targets``. When not set, this defaults to the uniform + distribution. + axis: the axis along which the quantile transformation is applied. + kwargs: keyword arguments passed on to lower level functions. Of interest + to the user are ``squashing_fun``, which will redistribute the values in + ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) + to solve the optimal transport problem; + :class:`cost_fn ` object of + :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground + 1D cost function to transport from ``inputs`` to the ``num_targets`` + target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` + are passed on to parameterize the + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. + + Returns: + An array, which has the same shape as the input, except on the give axis on + which the dimension is 1. + + Raises: + A ValueError in case the weights and the targets are both set and not of + compatible shapes. + """ + if weights is not None and weights.shape != targets.shape: + raise ValueError( + "The target weights and targets values should have the " + f"same shape: {targets.shape} != {weights.shape}" + ) + if weights is None: + num_targets = targets.shape[0] + weights = jnp.ones((num_targets,)) / num_targets + + op = _quantile_normalization + return apply_on_axis(op, inputs, axis, targets, weights, **kwargs) + + +def sort_with( + inputs: jnp.ndarray, + criterion: jnp.ndarray, + topk: int = -1, + **kwargs: Any, +) -> jnp.ndarray: + r"""Sort a multidimensional array according to a real valued criterion. + + Given ``batch`` vectors of dimension `dim`, to which, for each, a real value + ``criterion`` is associated, this function produces ``topk`` (or + ``batch`` if ``topk`` is set to -1, as by default) composite vectors of size + ``dim`` that will be convex combinations of all vectors, ranked from smallest + to largest criterion. Composite vectors with the largest indices will contain + convex combinations of those vectors with highest criterion, vectors with + smaller indices will contain combinations of vectors with smaller criterion. + + Args: + inputs: Array of size [batch, dim]. + criterion: the values according to which to sort the inputs. It has shape + [batch, 1]. + topk: The number of outputs to keep. + kwargs: keyword arguments passed on to lower level functions. Of interest + to the user are ``squashing_fun``, which will redistribute the values in + ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) + to solve the optimal transport problem; + :class:`cost_fn ` object of + :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground + 1D cost function to transport from ``inputs`` to the ``num_targets`` + target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` + are passed on to parameterize the + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. + + Returns: + An Array of size [batch | topk, dim]. + """ + num_points = criterion.shape[0] + weights = jnp.ones(num_points) / num_points + if 0 < topk < num_points: + start_index = 1 + target_weights = jnp.concatenate([ + jnp.array([(num_points - topk) / num_points]), + jnp.ones(topk) / num_points + ]) + else: + start_index = 0 + target_weights = jnp.ones((num_points,)) / num_points + ot = transport_for_sort(criterion, weights, target_weights, **kwargs) + # Applies the topk on each of the dimensions of the inputs. + sort_fn = jax.vmap( + lambda x: (1.0 / target_weights * ot.apply(x, axis=0))[start_index:], + in_axes=(1,), + out_axes=1 + ) + return sort_fn(inputs) + + +def _quantize(inputs: jnp.ndarray, num_q: int, **kwargs: Any) -> jnp.ndarray: + """Apply the soft quantization operator on a one dimensional array.""" + num_points = inputs.shape[0] + a = jnp.ones((num_points,)) / num_points + b = jnp.ones((num_q,)) / num_q + ot = transport_for_sort(inputs, a, b, **kwargs) + return 1.0 / a * ot.apply(1.0 / b * ot.apply(inputs), axis=1) + + +def quantize( + inputs: jnp.ndarray, + num_levels: int = 10, + axis: int = -1, + **kwargs: Any, +) -> jnp.ndarray: + r"""Soft quantizes an input according using ``num_levels`` values along axis. + + The quantization operator consists in concentrating several values around + a few predefined ``num_levels``. The soft quantization operator proposed here + does so by carrying out a `soft` concentration that is differentiable. The + ``inputs`` values are first soft-sorted, resulting in ``num_levels`` values. + In a second step, the ``inputs`` values are then provided again a composite + value that is equal (for each) to a convex combination of those soft-sorted + values using the transportation matrix. As the regularization parameter + ``epsilon`` of regularized optimal transport goes to 0, this operator recovers + the expected behavior of quantization, namely each value in ``inputs`` is + assigned a single level. When using ``epsilon>0`` the behavior is similar but + differentiable. + + Args: + inputs: an Array of size [batch, dim]. + num_levels: number of quantiles available to quantize the signal. + axis: axis along which quantization is carried out. + kwargs: keyword arguments passed on to lower level functions. Of interest + to the user are ``squashing_fun``, which will redistribute the values in + ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) + to solve the optimal transport problem; + :class:`cost_fn ` object of + :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground + 1D cost function to transport from ``inputs`` to the ``num_targets`` + target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` + are passed on to parameterize the + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. + + Returns: + An Array of the same size as ``inputs``. + """ + return apply_on_axis(_quantize, inputs, axis, num_levels, **kwargs) diff --git a/ott/build/lib/ott/types.py b/ott/build/lib/ott/types.py new file mode 100644 index 0000000..f6eefb7 --- /dev/null +++ b/ott/build/lib/ott/types.py @@ -0,0 +1,42 @@ +# 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. +from typing import Protocol, Union + +import jax.experimental.sparse as jesp +import jax.numpy as jnp + +__all__ = ["Transport", "Array_g"] + +# TODO(michalk8): introduce additional types here + +# Either a dense or sparse array. +Array_g = Union[jnp.ndarray, jesp.BCOO] + + +class Transport(Protocol): + """Interface for the solution of a transport problem. + + Classes implementing those function do not have to inherit from it, the + class can however be used in type hints to support duck typing. + """ + + @property + def matrix(self) -> jnp.ndarray: + ... + + def apply(self, inputs: jnp.ndarray, axis: int) -> jnp.ndarray: + ... + + def marginal(self, axis: int = 0) -> jnp.ndarray: + ... diff --git a/ott/build/lib/ott/utils.py b/ott/build/lib/ott/utils.py new file mode 100644 index 0000000..e9aecb7 --- /dev/null +++ b/ott/build/lib/ott/utils.py @@ -0,0 +1,210 @@ +# 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 dataclasses +import functools +import io +import warnings +from typing import Any, Callable, NamedTuple, Optional, Tuple + +import jax +import numpy as np + +try: + from tqdm import tqdm +except ImportError: + tqdm = Any + +__all__ = [ + "register_pytree_node", + "deprecate", + "default_prng_key", + "default_progress_fn", + "tqdm_progress_fn", +] + +Status_t = Tuple[np.ndarray, np.ndarray, np.ndarray, NamedTuple] +IOCallback_t = Callable[[Status_t], None] + + +def register_pytree_node(cls: type) -> type: + """Register dataclasses as pytree_nodes.""" + cls = dataclasses.dataclass()(cls) + flatten = lambda obj: jax.tree_util.tree_flatten(dataclasses.asdict(obj)) + unflatten = lambda d, children: cls(**d.unflatten(children)) + jax.tree_util.register_pytree_node(cls, flatten, unflatten) + return cls + + +def deprecate( # noqa: D103 + *, + version: Optional[str] = None, + alt: Optional[str] = None, + func: Optional[Callable[[Any], Any]] = None +) -> Callable[[Any], Any]: + + def wrapper(*args: Any, **kwargs: Any) -> Any: + warnings.warn(msg, category=DeprecationWarning, stacklevel=2) + return func(*args, **kwargs) + + if func is None: + return lambda fn: deprecate(version=version, alt=alt, func=fn) + + msg = f"`{func.__name__}` will be removed in the " + msg += ("next" if version is None else f"`ott-jax=={version}`") + " release." + if alt: + msg += " " + alt + + return functools.wraps(func)(wrapper) + + +def default_prng_key(rng: Optional[jax.Array] = None) -> jax.Array: + """Get the default PRNG key. + + Args: + rng: PRNG key. + + Returns: + If ``rng = None``, returns the default PRNG key. + Otherwise, it returns the unmodified ``rng`` key. + """ + return jax.random.PRNGKey(0) if rng is None else rng + + +def default_progress_fn( + fmt: str = "{iter} / {max_iter} -- {error}", + stream: Optional[io.TextIOBase] = None, +) -> IOCallback_t: + """Return a callback that prints the progress when solving + :mod:`linear problems `. + + It prints the progress only when the error is computed, that is every + :attr:`~ott.solvers.linear.sinkhorn.Sinkhorn.inner_iterations`. + + Args: + fmt: Format used to print. It can format ``iter``, ``max_iter`` and + ``error`` values. + stream: Output IO stream. + + Returns: + A callback function accepting the following arguments + + - the current iteration number, + - the number of inner iterations after which the error is computed, + - the total number of iterations, and + - the current :class:`~ott.solvers.linear.sinkhorn.SinkhornState` or + :class:`~ott.solvers.linear.sinkhorn_lr.LRSinkhornState`. + + Examples: + .. code-block:: python + + import jax + import jax.numpy as jnp + + from ott import utils + from ott.geometry import pointcloud + from ott.solvers.linear import sinkhorn + + x = jax.random.normal(jax.random.PRNGKey(0), (100, 5)) + geom = pointcloud.PointCloud(x) + + progress_fn = utils.default_progress_fn() + solve_fn = jax.jit(sinkhorn.solve, static_argnames=["progress_fn"]) + out = solve_fn(geom, progress_fn=progress_fn) + """ # noqa: D205 + + def progress_callback(status: Status_t) -> None: + iteration, inner_iterations, total_iter, errors = _prepare_info(status) + # Avoid reporting error on each iteration, + # because errors are only computed every `inner_iterations`. + if iteration % inner_iterations == 0: + error_idx = max(0, iteration // inner_iterations - 1) + error = errors[error_idx] + + print( + fmt.format(iter=iteration, max_iter=total_iter, error=error), + file=stream + ) + + return progress_callback + + +def tqdm_progress_fn( + pbar: tqdm, + fmt: str = "error: {error:0.6e}", +) -> IOCallback_t: + """Return a callback that updates a progress bar when solving + :mod:`linear problems `. + + It updates the progress bar only when the error is computed, that is every + :attr:`~ott.solvers.linear.sinkhorn.Sinkhorn.inner_iterations`. + + Args: + pbar: `tqdm `_ progress bar. + fmt: Format used for the postfix. It can format ``iter``, ``max_iter`` and + ``error`` values. + + Returns: + A callback function accepting the following arguments + + - the current iteration number, + - the number of inner iterations after which the error is computed, + - the total number of iterations, and + - the current :class:`~ott.solvers.linear.sinkhorn.SinkhornState` or + :class:`~ott.solvers.linear.sinkhorn_lr.LRSinkhornState`. + + Examples: + .. code-block:: python + + import tqdm + + import jax + import jax.numpy as jnp + + from ott import utils + from ott.geometry import pointcloud + from ott.solvers.linear import sinkhorn + + x = jax.random.normal(jax.random.PRNGKey(0), (100, 5)) + geom = pointcloud.PointCloud(x) + + with tqdm.tqdm() as pbar: + progress_fn = utils.tqdm_progress_fn(pbar) + solve_fn = jax.jit(sinkhorn.solve, static_argnames=["progress_fn"]) + out = solve_fn(geom, progress_fn=progress_fn) + """ # noqa: D205 + + def progress_callback(status: Status_t) -> None: + iteration, inner_iterations, total_iter, errors = _prepare_info(status) + # Avoid reporting error on each iteration, + # because errors are only computed every `inner_iterations`. + if iteration % inner_iterations == 0: + error_idx = max(0, iteration // inner_iterations - 1) + error = errors[error_idx] + + postfix = fmt.format(iter=iteration, max_iter=total_iter, error=error) + pbar.set_postfix_str(postfix) + pbar.total = total_iter // inner_iterations + pbar.update() + + return progress_callback + + +def _prepare_info(status: Status_t) -> Tuple[int, int, int, np.ndarray]: + iteration, inner_iterations, total_iter, state = status + iteration = int(iteration) + 1 + inner_iterations = int(inner_iterations) + total_iter = int(total_iter) + errors = np.array(state.errors).ravel() + + return iteration, inner_iterations, total_iter, errors diff --git a/ott/src/ott_jax.egg-info/PKG-INFO b/ott/src/ott_jax.egg-info/PKG-INFO new file mode 100644 index 0000000..fa26352 --- /dev/null +++ b/ott/src/ott_jax.egg-info/PKG-INFO @@ -0,0 +1,356 @@ +Metadata-Version: 2.4 +Name: ott-jax +Version: 0.4.6a0 +Summary: Optimal Transport Tools in JAX. +Author-email: OTT team +License: Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. + +Project-URL: Source Code, https://github.com/ott-jax/ott +Project-URL: Documentation, https://ott-jax.readthedocs.io +Project-URL: Issue Tracker, https://github.com/ott-jax/ott/issues +Project-URL: Changelog, https://github.com/ott-jax/ott/releases +Keywords: optimal transport,gromov wasserstein,sinkhorn,low-rank sinkhorn,sinkhorn divergences,wasserstein,wasserstein barycenter,jax,autodiff,implicit differentiation +Classifier: Typing :: Typed +Classifier: Development Status :: 4 - Beta +Classifier: License :: OSI Approved :: Apache Software License +Classifier: Topic :: Scientific/Engineering :: Mathematics +Classifier: Natural Language :: English +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Science/Research +Classifier: Operating System :: POSIX :: Linux +Classifier: Operating System :: MacOS :: MacOS X +Classifier: Operating System :: Microsoft :: Windows +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.8 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Requires-Python: >=3.8 +Description-Content-Type: text/markdown +License-File: LICENSE +Requires-Dist: jax>=0.4.0 +Requires-Dist: jaxopt>=0.8 +Requires-Dist: lineax>=0.0.1; python_version >= "3.9" +Requires-Dist: numpy>=1.20.0 +Provides-Extra: neural +Requires-Dist: flax>=0.6.6; extra == "neural" +Requires-Dist: optax>=0.1.1; extra == "neural" +Provides-Extra: dev +Requires-Dist: pre-commit>=2.16.0; extra == "dev" +Requires-Dist: tox>=4; extra == "dev" +Provides-Extra: test +Requires-Dist: pytest; extra == "test" +Requires-Dist: pytest-xdist; extra == "test" +Requires-Dist: pytest-cov; extra == "test" +Requires-Dist: pytest-memray; extra == "test" +Requires-Dist: coverage[toml]; extra == "test" +Requires-Dist: chex; extra == "test" +Requires-Dist: networkx>=2.5; extra == "test" +Requires-Dist: scikit-learn>=1.0; extra == "test" +Requires-Dist: tqdm; extra == "test" +Requires-Dist: tslearn>=0.5; python_version < "3.12" and extra == "test" +Requires-Dist: lineax; python_version >= "3.9" and extra == "test" +Requires-Dist: matplotlib; extra == "test" +Provides-Extra: docs +Requires-Dist: sphinx>=4.0; extra == "docs" +Requires-Dist: sphinx-book-theme>=1.0.1; extra == "docs" +Requires-Dist: sphinx_autodoc_typehints>=1.12.0; extra == "docs" +Requires-Dist: sphinx-copybutton>=0.5.1; extra == "docs" +Requires-Dist: sphinxcontrib-bibtex>=2.5.0; extra == "docs" +Requires-Dist: sphinxcontrib-spelling>=7.7.0; extra == "docs" +Requires-Dist: myst-nb>=0.17.1; extra == "docs" +Dynamic: license-file + +logo + +# Optimal Transport Tools (OTT) +[![Downloads](https://static.pepy.tech/badge/ott-jax)](https://pypi.org/project/ott-jax/) +[![Tests](https://img.shields.io/github/actions/workflow/status/ott-jax/ott/tests.yml?branch=main)](https://github.com/ott-jax/ott/actions/workflows/tests.yml) +[![Docs](https://img.shields.io/readthedocs/ott-jax/latest)](https://ott-jax.readthedocs.io/en/latest/) +[![Coverage](https://img.shields.io/codecov/c/github/ott-jax/ott/main)](https://app.codecov.io/gh/ott-jax/ott) + +**See the [full documentation](https://ott-jax.readthedocs.io/en/latest/).** + +## What is OTT-JAX? +A ``JAX`` powered library to compute optimal transport at scale and on accelerators, ``OTT-JAX`` includes the fastest +implementation of the Sinkhorn algorithm you will find around. We have implemented all tweaks (scheduling, momentum, acceleration, initializations) and extensions (low-rank, entropic maps). They can be used directly between two datasets, or within more advanced problems +(Gromov-Wasserstein, barycenters). Some of ``JAX`` features, including +[JIT](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html#Using-jit-to-speed-up-functions), +[auto-vectorization](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html#Auto-vectorization-with-vmap) and +[implicit differentiation](https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html) +work towards the goal of having end-to-end differentiable outputs. ``OTT-JAX`` is led by a team of researchers at Apple, with contributions from Google and Meta researchers, as well as many academic partners, including TU München, Oxford, ENSAE/IP Paris, ENS Paris and the Hebrew University. + +## Installation +Install ``OTT-JAX`` from [PyPI](https://pypi.org/project/ott-jax/) as: +```bash +pip install ott-jax +``` +or with ``conda`` via [conda-forge](https://anaconda.org/conda-forge/ott-jax) as: +```bash +conda install -c conda-forge ott-jax +``` + +## What is optimal transport? +Optimal transport can be loosely described as the branch of mathematics and optimization that studies +*matching problems*: given two families of points, and a cost function on pairs of points, find a "good" (low cost) way +to associate bijectively to every point in the first family another in the second. + +Such problems appear in all areas of science, are easy to describe, yet hard to solve. Indeed, while matching optimally +two sets of $n$ points using a pairwise cost can be solved with the +[Hungarian algorithm](https://en.wikipedia.org/wiki/Hungarian_algorithm), solving it costs an order of $O(n^3)$ +operations, and lacks flexibility, since one may want to couple families of different sizes. + +Optimal transport extends all of this, through faster algorithms (in $n^2$ or even linear in $n$) along with numerous +generalizations that can help it handle weighted sets of different size, partial matchings, and even more evolved +so-called quadratic matching problems. + +In the simple toy example below, we compute the optimal coupling matrix between two point clouds sampled randomly +(2D vectors, compared with the squared Euclidean distance): + +## Example +```python +import jax +import jax.numpy as jnp + +from ott.geometry import pointcloud +from ott.problems.linear import linear_problem +from ott.solvers.linear import sinkhorn + +# sample two point clouds and their weights. +rngs = jax.random.split(jax.random.PRNGKey(0), 4) +n, m, d = 12, 14, 2 +x = jax.random.normal(rngs[0], (n,d)) + 1 +y = jax.random.uniform(rngs[1], (m,d)) +a = jax.random.uniform(rngs[2], (n,)) +b = jax.random.uniform(rngs[3], (m,)) +a, b = a / jnp.sum(a), b / jnp.sum(b) +# Computes the couplings using the Sinkhorn algorithm. +geom = pointcloud.PointCloud(x, y) +prob = linear_problem.LinearProblem(geom, a, b) + +solver = sinkhorn.Sinkhorn() +out = solver(prob) +``` + +The call to `solver(prob)` above works out the optimal transport solution. The `out` object contains a transport matrix +(here of size $12\times 14$) that quantifies the association strength between each point of the first point cloud, to one or +more points from the second, as illustrated in the plot below. We provide more flexibility to define custom cost +functions, objectives, and solvers, as detailed in the [full documentation](https://ott-jax.readthedocs.io/en/latest/). + +![obtained coupling](https://raw.githubusercontent.com/ott-jax/ott/main/docs/_static/images/couplings.png) + +## Citation +If you have found this work useful, please consider citing this reference: + +``` +@article{cuturi2022optimal, + title={Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein}, + author={Cuturi, Marco and Meng-Papaxanthos, Laetitia and Tian, Yingtao and Bunne, Charlotte and + Davis, Geoff and Teboul, Olivier}, + journal={arXiv preprint arXiv:2201.12324}, + year={2022} +} +``` +## See also +The [moscot](https://moscot.readthedocs.io/en/latest/index.html) package for OT analysis of multi-omics data also uses OTT as a backbone. diff --git a/ott/src/ott_jax.egg-info/SOURCES.txt b/ott/src/ott_jax.egg-info/SOURCES.txt new file mode 100644 index 0000000..cdce04a --- /dev/null +++ b/ott/src/ott_jax.egg-info/SOURCES.txt @@ -0,0 +1,92 @@ +CITATION.bib +CODE_OF_CONDUCT.md +CONTRIBUTING.md +LICENSE +MANIFEST.in +README.md +environment.yml +pyproject.toml +src/ott/__init__.py +src/ott/_version.py +src/ott/datasets.py +src/ott/py.typed +src/ott/types.py +src/ott/utils.py +src/ott/geometry/__init__.py +src/ott/geometry/costs.py +src/ott/geometry/distrib_costs.py +src/ott/geometry/epsilon_scheduler.py +src/ott/geometry/geodesic.py +src/ott/geometry/geometry.py +src/ott/geometry/graph.py +src/ott/geometry/grid.py +src/ott/geometry/low_rank.py +src/ott/geometry/pointcloud.py +src/ott/geometry/segment.py +src/ott/initializers/__init__.py +src/ott/initializers/linear/__init__.py +src/ott/initializers/linear/initializers.py +src/ott/initializers/linear/initializers_lr.py +src/ott/initializers/quadratic/__init__.py +src/ott/initializers/quadratic/initializers.py +src/ott/math/__init__.py +src/ott/math/fixed_point_loop.py +src/ott/math/matrix_square_root.py +src/ott/math/unbalanced_functions.py +src/ott/math/utils.py +src/ott/neural/__init__.py +src/ott/neural/layers.py +src/ott/neural/losses.py +src/ott/neural/models.py +src/ott/neural/solvers/__init__.py +src/ott/neural/solvers/conjugate.py +src/ott/neural/solvers/map_estimator.py +src/ott/neural/solvers/neuraldual.py +src/ott/problems/__init__.py +src/ott/problems/linear/__init__.py +src/ott/problems/linear/barycenter_problem.py +src/ott/problems/linear/linear_problem.py +src/ott/problems/linear/potentials.py +src/ott/problems/quadratic/__init__.py +src/ott/problems/quadratic/gw_barycenter.py +src/ott/problems/quadratic/quadratic_costs.py +src/ott/problems/quadratic/quadratic_problem.py +src/ott/solvers/__init__.py +src/ott/solvers/was_solver.py +src/ott/solvers/linear/__init__.py +src/ott/solvers/linear/_solve.py +src/ott/solvers/linear/acceleration.py +src/ott/solvers/linear/continuous_barycenter.py +src/ott/solvers/linear/discrete_barycenter.py +src/ott/solvers/linear/implicit_differentiation.py +src/ott/solvers/linear/lineax_implicit.py +src/ott/solvers/linear/lr_utils.py +src/ott/solvers/linear/sinkhorn.py +src/ott/solvers/linear/sinkhorn_lr.py +src/ott/solvers/linear/univariate.py +src/ott/solvers/quadratic/__init__.py +src/ott/solvers/quadratic/_solve.py +src/ott/solvers/quadratic/gromov_wasserstein.py +src/ott/solvers/quadratic/gromov_wasserstein_lr.py +src/ott/solvers/quadratic/gw_barycenter.py +src/ott/solvers/quadratic/lower_bound.py +src/ott/tools/__init__.py +src/ott/tools/k_means.py +src/ott/tools/plot.py +src/ott/tools/segment_sinkhorn.py +src/ott/tools/sinkhorn_divergence.py +src/ott/tools/soft_sort.py +src/ott/tools/gaussian_mixture/__init__.py +src/ott/tools/gaussian_mixture/fit_gmm.py +src/ott/tools/gaussian_mixture/fit_gmm_pair.py +src/ott/tools/gaussian_mixture/gaussian.py +src/ott/tools/gaussian_mixture/gaussian_mixture.py +src/ott/tools/gaussian_mixture/gaussian_mixture_pair.py +src/ott/tools/gaussian_mixture/linalg.py +src/ott/tools/gaussian_mixture/probabilities.py +src/ott/tools/gaussian_mixture/scale_tril.py +src/ott_jax.egg-info/PKG-INFO +src/ott_jax.egg-info/SOURCES.txt +src/ott_jax.egg-info/dependency_links.txt +src/ott_jax.egg-info/requires.txt +src/ott_jax.egg-info/top_level.txt \ No newline at end of file diff --git a/ott/src/ott_jax.egg-info/dependency_links.txt b/ott/src/ott_jax.egg-info/dependency_links.txt new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/ott/src/ott_jax.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/ott/src/ott_jax.egg-info/requires.txt b/ott/src/ott_jax.egg-info/requires.txt new file mode 100644 index 0000000..c93e9c5 --- /dev/null +++ b/ott/src/ott_jax.egg-info/requires.txt @@ -0,0 +1,41 @@ +jax>=0.4.0 +jaxopt>=0.8 +numpy>=1.20.0 + +[:python_version >= "3.9"] +lineax>=0.0.1 + +[dev] +pre-commit>=2.16.0 +tox>=4 + +[docs] +sphinx>=4.0 +sphinx-book-theme>=1.0.1 +sphinx_autodoc_typehints>=1.12.0 +sphinx-copybutton>=0.5.1 +sphinxcontrib-bibtex>=2.5.0 +sphinxcontrib-spelling>=7.7.0 +myst-nb>=0.17.1 + +[neural] +flax>=0.6.6 +optax>=0.1.1 + +[test] +pytest +pytest-xdist +pytest-cov +pytest-memray +coverage[toml] +chex +networkx>=2.5 +scikit-learn>=1.0 +tqdm +matplotlib + +[test:python_version < "3.12"] +tslearn>=0.5 + +[test:python_version >= "3.9"] +lineax diff --git a/ott/src/ott_jax.egg-info/top_level.txt b/ott/src/ott_jax.egg-info/top_level.txt new file mode 100644 index 0000000..6ec4bc0 --- /dev/null +++ b/ott/src/ott_jax.egg-info/top_level.txt @@ -0,0 +1 @@ +ott diff --git a/perturbot/build/lib/perturbot/__init__.py b/perturbot/build/lib/perturbot/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/perturbot/build/lib/perturbot/eval/.run_loo.py b/perturbot/build/lib/perturbot/eval/.run_loo.py new file mode 100755 index 0000000..ddb2f82 --- /dev/null +++ b/perturbot/build/lib/perturbot/eval/.run_loo.py @@ -0,0 +1,126 @@ +# Run leave-one-out prediction & evaluation +from typing import Optional +import numpy as np +import argparse +import mudata as md +from sklearn.model_selection import LeaveOneOut +from software.perturbot.perturbot.ot.cot_labels import ( + predict_with_cot, + evaluate_training, +) + +method_dict = {"cot_labels": predict_with_cot} +get_train_metrics_dict = {"cot_labels": evaluate_training} + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Run leave-one-out prediction & evaluation." + ) + parser.add_argument("input_h5mu", type=str, help="Path to the input .h5mu file.") + parser.add_argument( + "--label-col", + "-l", + type=str, + default="treatment", + help="Column label in .obs of MuData that encodes the perturbation labels.", + ) + parser.add_argument( + "--within-label-group", + "-g", + type=str, + default=None, + help="If provided, evaluation includes the concordance of the grouping.", + ) + parser.add_argument( + "--within-label-group", + "-g", + type=str, + default=None, + help="If provided, evaluation includes the concordance of the grouping.", + ) + parser.add_argument( + "--train-with-group", + "-tg", + type=bool, + action="store_true", + help="train with `within_label_group`.", + ) + parser.add_argument( + "--source-modality", + "-s", + type=str, + default="protein", + help="Source modality in MuData", + ) + parser.add_argument( + "--target-modality", + "-t", + type=str, + default="RNA", + help="Target modality in MuData", + ) + parser.add_argument( + "--rng-seed", + "-r", + type=int, + default=2024, + help="random number generator seed for NumPy", + ) + return parser.parse_args() + + +def get_data_with_group( + x, y, include_y, z: Optional[np.ndarray] = None, nbperclass: Optional[int] = None +): + # Randomly select nbperclass train datasets + xr = np.zeros((0, x.shape[1])) + yr = np.zeros((0)) + if z is not None: + zr = np.zeros((0)) + else: + zr = None + for i in include_y: + xi = x[y.ravel() == i, :] + xr = np.concatenate((xr, xi), 0) + yr = np.concatenate((yr, np.repeat(i, xi.shape[0]))) + if z is not None: + zi = z[y.ravel() == i] + zr = np.concatenate((zr, zi), 0) + return xr, yr, zr + + +def main(): + args = parse_args() + np.random.seed(args.rng_seed) + mdata = md.read_h5ad(args.input_h5mu) + Xtot1 = mdata[args.source_modality].X + Xtot2 = mdata[args.target_modality].X + labels = mdata.obs[args.label_col].values + within_label_group = mdata.obs[args.within_label_group].values + loo = LeaveOneOut() + train_data = {} + test_data = {} + for i, (train_idx, test_idx) in enumerate(loo.split(labels)): + Xs, Ys, Zs = get_data_with_group( + Xtot1, labels, labels[train_idx], within_label_group + ) + Xt, Yt, Zt = get_data_with_group( + Xtot2, labels, labels[train_idx], within_label_group + ) + Xs_test, Ys_test, Zs_test = get_data_with_group( + Xtot1, labels, labels[test_idx], within_label_group + ) + Xt_test, Yt_test, Zt_test = get_data_with_group( + Xtot1, labels, labels[test_idx], within_label_group + ) + for method_label, run_fn in method_dict.items(): + Xt_pred, log = run_fn(Xs, Ys, Xt, Yt, Xs_test) + train_eval_metrics = get_train_metrics_dict[method_label]( + Xs, Ys, Xt, Yt, Zt, log + ) + test_eval_metrics = get_test_metrics(Xt_pred, Xt_test, Zt_test) + + +if __name__ == "__main__": + main() diff --git a/perturbot/build/lib/perturbot/eval/__init__.py b/perturbot/build/lib/perturbot/eval/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/perturbot/build/lib/perturbot/eval/all.py b/perturbot/build/lib/perturbot/eval/all.py new file mode 100755 index 0000000..67dca36 --- /dev/null +++ b/perturbot/build/lib/perturbot/eval/all.py @@ -0,0 +1,190 @@ +"""For 5-fold CV, run the single inner loop of train-validation to select hyperparameter.""" + +import os +import argparse +import pickle as pkl +from functools import partial +import numpy as np +import torch +from perturbot.eval.match import get_FOSCTTM, get_diag_fracs +from perturbot.eval.utils import get_Ts_from_nn_multKs +from perturbot.match.cot_labels import get_coupling_cotl_sinkhorn + +from perturbot.match.ott_egwl import ( + get_coupling_egw_labels_ott, + get_coupling_egw_all_ott, + get_coupling_eot_ott, + get_coupling_leot_ott, + get_coupling_egw_ott, +) +from perturbot.match.cot import ( + get_coupling_cot_sinkhorn, + get_coupling_each_cot_sinkhorn, +) +from perturbot.match.gw_labels import get_coupling_egw_labels +from perturbot.predict.scvi_vae import train_vae_model, infer_from_Xs, infer_from_Ys + +ot_method_map = { + "ECOOTL": get_coupling_cotl_sinkhorn, + "ECOOT": get_coupling_cot_sinkhorn, + "ECOOT_each": get_coupling_each_cot_sinkhorn, + "EGWL": get_coupling_egw_labels, + "EOT_ott": get_coupling_eot_ott, + "LEOT_ott": get_coupling_leot_ott, + "EGW_ott": get_coupling_egw_ott, + "EGW_all_ott": get_coupling_egw_all_ott, + "EGWL_ott": get_coupling_egw_labels_ott, + "VAE_label": train_vae_model, + "VAE": partial(train_vae_model, use_label=False), +} + + +def parse_args(): + parser = argparse.ArgumentParser( + "Run inner validation loop of CV", + "Run sample-matching OT in parallel and fit prediction model in leave-one-out mannter", + ) + parser.add_argument("method", type=str) + parser.add_argument("filepath", type=str) + parser.add_argument("eps", type=float) + + parser.add_argument( + "-l", + "--log-filepath", + type=str, + default=None, + help="Path to log.pkl if transport was already run. If not None, reevaluates but does not calculates the transport again.", + ) + parser.add_argument( + "-b", + "--baseline", + type=str, + default=None, + help="One of perfect, random, by_conc", + ) + return parser.parse_args() + + +ot_method_all_to_all = ["GW_all", "EGW_all_ott", "EOT_all_ott"] + + +def main(args): + # if args.log_filepath is not None and args.log_filepath != "None": + # with open(args.log_filepath, "rb") as f: + # prev_logs = pkl.load(f) + logs = { + "matching_evals": {}, + "pred_evals": {}, + "T": {}, + "pred": {}, + "log": {}, + "eps": {}, + } + try: + match_eps = args.eps + logs["eps"] = match_eps + with open(args.filepath, "rb") as f: + data_dict = pkl.load(f) + + X_dict = data_dict["Xs_dict"] + Y_dict = data_dict["Xt_dict"] + Zs_dict = data_dict["Zs_dict"]["dosage"] + Zt_dict = data_dict["Zt_dict"]["dosage"] + + labels = list(X_dict.keys()) + + train_X = X_dict + train_Y = Y_dict + train_Z = Zs_dict + train_data = (train_X, train_Y) + print(f"Calculating matching with {match_eps}") + if args.log_filepath is not None: + with open(args.log_filepath, "rb") as f: + d = pkl.load(f) + Ts_matching = d["T"] + log_matching = "" + else: + Ts_matching, log_matching = ot_method_map[args.method]( + train_data, match_eps + ) + if "VAE" in args.method: + dim_X = data_dict["Xs_dict"][labels[0]].shape[1] + dim_Y = data_dict["Xt_dict"][labels[0]].shape[1] + latent_Y = infer_from_Ys(train_Y, Ts_matching, dim_X) + latent_X = infer_from_Xs(train_X, Ts_matching, dim_Y) + _, mean_foscttm = get_FOSCTTM( + Ts_matching, + latent_X, + latent_Y, + use_agg="mean", + use_barycenter=False, + ) + ks = [1, 5, 10, 50, 100] + k_to_Ts = get_Ts_from_nn_multKs(latent_X, latent_Y, ks) # k -> T + dfracs = {} + rel_dfracs = {} + for k, Ts in k_to_Ts.items(): + dfracs[k], rel_dfracs[k] = get_diag_fracs( + Ts, train_X, train_Y, train_Z, train_Z + ) + + else: + if isinstance(Ts_matching, dict): + total_sum = 0 + for k, v in Ts_matching.items(): + total_sum += v.sum() + Ts_matching = { + k: v.astype(np.double) / total_sum for k, v in Ts_matching.items() + } + else: + Ts_matching = Ts_matching.astype(np.double) / Ts_matching.sum() + _, mean_foscttm = get_FOSCTTM(Ts_matching, train_X, train_Y, use_agg="mean") + dfracs, rel_dfracs = get_diag_fracs( + Ts_matching, train_X, train_Y, train_Z, train_Z + ) + + # if not all_to_all: + mean_mean_foscttm = mean_foscttm.mean() + # else: + # mean_mean_foscttm = mean_foscttm + print("foscttm", mean_foscttm) + logs["matching_evals"] = ( + { + "foscttm": mean_foscttm, + "mean_foscttm": mean_mean_foscttm, + "dfracs": dfracs, + "rel_dfracs": rel_dfracs, + }, + ) + logs["T"] = Ts_matching + logs["log"] = (log_matching,) + except Exception as e: + with open(f"all_{args.method}.{args.eps}.tmp.pkl", "wb") as f: + pkl.dump({"log": logs, "T": Ts_matching}, f) + raise e + + with open(f"all_{args.method}.{args.eps}.pkl", "wb") as f: + pkl.dump(logs, f) + + +def submit_all_run(data_path, ot_method_label, load_existing=False): + epsilons = [1e-2, 1e-3, 1e-4, 1e-5] + for eps in epsilons: + run_label = f"all.{ot_method_label}.{eps}" + f = open(f"{run_label}.bsub", "w") + f.write("#!/bin/bash\n") + f.write("source ~/.bashrc\n") + f.write("pwd\n") + f.write("conda activate ot \n") + command = f"python /gpfs/scratchfs01/site/u/ryuj6/OT/software/perturbot/perturbot/eval/all.py {ot_method_label} {data_path} {eps}" + f.write(f"echo {command}\n") + f.write(f"{command}\n") + f.close() + command = f"bsub -M 10G -n 10 -q long -J {run_label} -o log/log.{run_label}.o%J -e log/log.{run_label}.e%J < {run_label}.bsub" + print(command) + os.system(command) + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/perturbot/build/lib/perturbot/eval/cv.py b/perturbot/build/lib/perturbot/eval/cv.py new file mode 100755 index 0000000..b035245 --- /dev/null +++ b/perturbot/build/lib/perturbot/eval/cv.py @@ -0,0 +1,251 @@ +import os +from typing import List, Dict, Any +from functools import partial +from itertools import product +import argparse +import os + +os.environ["OPENBLAS_NUM_THREADS"] = "1" +import numpy as np +import pandas as pd +from sklearn.model_selection import KFold +from perturbot.eval.prediction import get_evals_preds, get_evals +from multiprocessing import Pool +import pickle as pkl +from perturbot.eval.match import get_FOSCTTM, get_diag_fracs +from perturbot.eval.utils import get_Ts_from_nn_multKs, _pop_key +from perturbot.predict.scvi_vae import predict_from_model, infer_from_Xs, infer_from_Ys +from perturbot.preprocess.vae import ( + train_vae_rna, + train_vae_acc, + train_vae_prot, + SCVI_LATENT_KEY, +) +from perturbot.match.cot_labels import get_coupling_cotl, get_coupling_cotl_sinkhorn +from perturbot.match.gw import ( + get_coupling_gw_cg, + get_coupling_egw, + get_coupling_egw_all, + get_coupling_gw_all, +) +from perturbot.match.ott_egwl import ( + get_coupling_egw_labels_ott, + get_coupling_egw_all_ott, + get_coupling_eot_ott, + get_coupling_leot_ott, + get_coupling_egw_ott, +) +from perturbot.match.cot import ( + get_coupling_cot, + get_coupling_cot_sinkhorn, + get_coupling_each_cot_sinkhorn, +) +from perturbot.match.gw_labels import get_coupling_gw_labels, get_coupling_egw_labels +from perturbot.predict.scvi_vae import train_vae_model + +np.seterr(all="raise") + +vae_trainer_dict = { + "rna": train_vae_rna, + "accessibility": train_vae_acc, + "protein": train_vae_prot, +} + +ot_method_map = { + "COOTL": get_coupling_cotl, + "ECOOTL": get_coupling_cotl_sinkhorn, + "COOT": get_coupling_cot, + "ECOOT": get_coupling_cot_sinkhorn, + "ECOOT_each": get_coupling_each_cot_sinkhorn, + "GW_all": get_coupling_gw_all, + # "EGW_all": get_coupling_egw_all, + "GW": get_coupling_gw_cg, + # "EGW": get_coupling_egw, + "GWL": get_coupling_gw_labels, + "EGWL": get_coupling_egw_labels, + "EOT_ott": get_coupling_eot_ott, + "LEOT_ott": get_coupling_leot_ott, + "EGW_ott": get_coupling_egw_ott, + "EGW_all_ott": get_coupling_egw_all_ott, + "EGWL_ott": get_coupling_egw_labels_ott, + "VAE_label": train_vae_model, + "VAE": partial(train_vae_model, use_label=False), +} +ot_method_hyperparams = {} +for method in ["EGWL", "EOT_ott", "LEOT_ott", "EGW_ott", "EGW_all_ott", "EGWL_OTT"]: + ot_method_hyperparams[method] = [ + 1e-2, + 1e-3, + 1e-4, + 1e-5, + 1e-6, + ] +for method in ["ECOOTL", "ECOOT", "ECOOT_each"]: + ot_method_hyperparams[method] = [0.1, 0.05, 0.01, 0.005, 0.001] +ot_method_all_to_all = ["GW_all", "EGW_all_ott", "EOT_all_ott"] + + +def run_cv_models( + filepath, + ot_method_label, + data_label="k", + log_filepath=None, + full_filepath=None, + run_mlp=False, + mlp_baseline=False, + rerun=False, +): + if ot_method_label == "VAE": + assert full_filepath is not None + with open(filepath, "rb") as f: + data_dict = pkl.load(f) + labels = list(data_dict["Xs_dict"].keys()) + cv = KFold(n_splits=5) + # Outer loop will evaluate final performance. + for i, (train_val_idx, test_idx) in enumerate(cv.split(labels)): + # Inner loop will select the best-performing hyperparameter. + # Run CV for multiple hyperparameters (eps) + run_label = f"{'mlp.' if run_mlp else ''}{ot_method_label}.{i}" + + if ( + run_mlp + and os.path.isfile(f"val_mlp.{ot_method_label}.{i}.best_eps.pkl") + and not rerun + ): + print(f"val_mlp.{ot_method_label}.{i}.best_eps.pkl exists. Skipping") + continue + elif ( + (not run_mlp) + and os.path.isfile(f"val_CV_{ot_method_label}.{i}.best_eps.pkl") + and log_filepath is None + ): + print(f"val_CV_{ot_method_label}.{i}.best_eps.pkl exists. Skipping") + continue + # Write job script + f = open(f"{run_label}.bsub", "w") + f.write("#!/bin/bash\n") + f.write("source ~/.bashrc\n") + f.write("pwd\n") + f.write("conda activate ot \n") + if "VAE" in ot_method_label: + command = f"python /gpfs/scratchfs01/site/u/ryuj6/OT/software/perturbot/perturbot/eval/cv_inner_loop.py {ot_method_label} {i} {full_filepath} {'' if log_filepath is None else f'-l val_CV_{ot_method_label}.{i}.pkl'} {'--mlp' if run_mlp else ''} " + else: + command = f"python /gpfs/scratchfs01/site/u/ryuj6/OT/software/perturbot/perturbot/eval/cv_inner_loop.py {ot_method_label} {i} {full_filepath if run_mlp else filepath} {'' if log_filepath is None else f'-l val_CV_{ot_method_label}.{i}.pkl'} {'' if full_filepath is None else '-p ' + full_filepath} {'--mlp' if run_mlp else ''} {'--baseline' if mlp_baseline else ''}" + + f.write(f"echo {command}\n") + f.write(f"{command}\n") + f.close() + label = "mlp" if run_mlp else "cv" + if log_filepath is not None: + label = "reval" + command = f"bsub {'-M 20G -n 10' if 'VAE' in ot_method_label or run_mlp else '-M 10G -n 25'} -q {'short' if log_filepath is not None else 'long'} -J {data_label}{'mlp' if run_mlp else 'CV'}.{run_label} -o log/log_{label}.{run_label}.o%J -e log/log_{label}.{run_label}.e%J < {run_label}.bsub" + print(command) + os.system(command) + + +def get_test_eval( + filepath, + ot_method_label, + data_label="k", + log_filepath=None, + full_filepath=None, + run_mlp=False, + baseline=None, +): + if ot_method_label == "VAE": + assert full_filepath is not None + with open(filepath, "rb") as f: + data_dict = pkl.load(f) + labels = list(data_dict["Xs_dict"].keys()) + cv = KFold(n_splits=5) + # Outer loop will evaluate final performance. + if baseline is not None: + for i, (train_val_idx, test_idx) in enumerate(cv.split(labels)): + if log_filepath is not None: + log_filepath = f"test_{baseline}.{i}.pkl" + run_label = f"TM.{baseline}.{i}" + f = open(f"{run_label}.bsub", "w") + f.write("#!/bin/bash\n") + f.write("source ~/.bashrc\n") + f.write("pwd\n") + f.write("conda activate ot \n") + command = f"python /gpfs/scratchfs01/site/u/ryuj6/OT/software/perturbot/perturbot/eval/cv_outer_loop.py {baseline} {i} {filepath} 0,0,0 -p {full_filepath} -l {log_filepath} --baseline {baseline}" + f.write(f"echo {command}\n") + f.write(f"{command}\n") + f.close() + command = f"bsub -M 10G -n 10 -q long -J {data_label}{run_label} -o test_log/log.{run_label}.o%J -e test_log/log.{run_label}.e%J < {run_label}.bsub" + print(command) + os.system(command) + else: + for i, (train_val_idx, test_idx) in enumerate(cv.split(labels)): + # Evaluate matching + try: + with open(f"val_CV_{ot_method_label}.{i}.best_eps.pkl", "rb") as f: + log = pkl.load(f) + + except FileNotFoundError: + print(f"val_CV_{ot_method_label}.{i}.best_eps.pkl not found. Skipping") + continue + if "VAE" not in ot_method_label: + try: + with open(f"val_mlp.{ot_method_label}.{i}.best_eps.pkl", "rb") as f: + log2 = pkl.load(f) + eps_pred = log2["eps"] + except FileNotFoundError: + try: + with open(f"val_CV_MLP_{ot_method_label}.{i}.pkl", "rb") as f: + log2 = pkl.load(f) + eps_pred = log2["pred_evals"].T["Pearson_samples"].idxmax() + except FileNotFoundError: + print( + f"val_CV_MLP_{ot_method_label}.{i}.pkl and val_mlp.{ot_method_label}.{i}.best_eps.pkl not found. Skipping" + ) + continue + # elif os.path.exists(f"test_{ot_method_label}.{i}.pkl"): + # print(f"test_{ot_method_label}.{i}.pkl already exists. Skipping") + # elif os.path.exists(f"test_log/test_{ot_method_label}.{i}.tmp.pkl"): + # print( + # f"test_log/test_{ot_method_label}.{i}.tmp.pkl already exists. Skipping" + # ) + else: + eps_pred = np.nan + + allowed_eps = [1e-2, 1e-3, 1e-4, 1e-5] + eps_lin = log["pred"] + + eps_matching = log["matching"] + if isinstance(eps_lin, tuple): + eps_lin = eps_lin[0] + if isinstance(eps_matching, tuple): + eps_matching = eps_matching[0] + + if "VAE" not in ot_method_label and ( + eps_lin not in allowed_eps or eps_matching not in allowed_eps + ): + try: + with open(f"val_CV_{ot_method_label}.{i}.pkl", "rb") as f: + eval_log = pkl.load(f) + except FileNotFoundError: + print(f"val_CV_{ot_method_label}.{i}.pkl not found. Skipping") + eps_lin = eval_log["pred_evals"].loc["MSE", allowed_eps].idxmin() + eps_matching = eval_log["matching_evals"][allowed_eps].idxmin() + elif "VAE" in ot_method_label: + if isinstance(eps_lin, tuple): + eps_lin = eps_lin[0] + if isinstance(eps_matching, tuple): + eps_matching = eps_matching[0] + if log_filepath is not None: + log_filepath = f"test_{ot_method_label}.{i}.pkl" + run_label = f"TM.{ot_method_label}.{i}" + f = open(f"{run_label}.bsub", "w") + f.write("#!/bin/bash\n") + f.write("source ~/.bashrc\n") + f.write("pwd\n") + f.write("conda activate ot \n") + command = f"python /gpfs/scratchfs01/site/u/ryuj6/OT/software/perturbot/perturbot/eval/cv_outer_loop.py {ot_method_label} {i} {filepath} {eps_matching},{eps_lin},{eps_pred} -p {full_filepath} -l {log_filepath} {'' if baseline is None else '--baseline '+baseline}" + f.write(f"echo {command}\n") + f.write(f"{command}\n") + f.close() + command = f"bsub -M 10G -n 10 -q long -J {data_label}{run_label} -o test_log/log.{run_label}.o%J -e test_log/log.{run_label}.e%J < {run_label}.bsub" + print(command) + os.system(command) diff --git a/perturbot/build/lib/perturbot/eval/cv_inner_loop.py b/perturbot/build/lib/perturbot/eval/cv_inner_loop.py new file mode 100755 index 0000000..a38316e --- /dev/null +++ b/perturbot/build/lib/perturbot/eval/cv_inner_loop.py @@ -0,0 +1,673 @@ +"""For 5-fold CV, run the single inner loop of train-validation to select hyperparameter.""" + +from typing import Dict, Any +import sys +import argparse +import pickle as pkl +import jax +from functools import partial +from itertools import product +from multiprocessing import Pool +import numpy as np +import pandas as pd +from perturbot.eval.utils import get_Ts_from_nn_multKs +from sklearn.model_selection import KFold +import torch +from tqdm import tqdm +import psutil +import gc +from perturbot.eval.utils import _pop_keys, _pop_key +from perturbot.eval.match import get_FOSCTTM, get_diag_fracs +from perturbot.eval.prediction import get_evals_preds, get_evals +from perturbot.preprocess.vae import ( + train_vae_rna, + train_vae_acc, + train_vae_prot, + SCVI_LATENT_KEY, +) +from perturbot.match.cot_labels import get_coupling_cotl, get_coupling_cotl_sinkhorn +from perturbot.match.gw import ( + get_coupling_gw_cg, + get_coupling_egw, + get_coupling_egw_all, + get_coupling_gw_all, +) +from perturbot.match.ott_egwl import ( + get_coupling_egw_labels_ott, + get_coupling_egw_all_ott, + get_coupling_eot_ott, + get_coupling_leot_ott, + get_coupling_egw_ott, +) +from perturbot.match.cot import ( + get_coupling_cot, + get_coupling_cot_sinkhorn, + get_coupling_each_cot_sinkhorn, +) +from perturbot.match.gw_labels import get_coupling_gw_labels, get_coupling_egw_labels +from perturbot.predict.scvi_vae import train_vae_model +from perturbot.predict.linear_regression import ( + ols_normed, + weight_conc_normed, + weighted_ols_normed, + weight_1_ols_normed, + predict, +) +from perturbot.predict.mlp import train_mlp +from perturbot.predict.scvi_vae import predict_from_model, infer_from_Xs, infer_from_Ys + +ot_method_map = { + "ECOOTL": get_coupling_cotl_sinkhorn, + "ECOOT_each": get_coupling_each_cot_sinkhorn, + "ECOOT": get_coupling_cot_sinkhorn, + "EGWL": get_coupling_egw_labels, + "EOT_ott": get_coupling_eot_ott, + "LEOT_ott": get_coupling_leot_ott, + "EGW_ott": get_coupling_egw_ott, + "EGW_all_ott": get_coupling_egw_all_ott, + "EGWL_ott": get_coupling_egw_labels_ott, + "VAE_label": train_vae_model, + "VAE": partial(train_vae_model, use_label=False), +} + + +def parse_args(): + parser = argparse.ArgumentParser( + "Run inner validation loop of CV", + "Run sample-matching OT in parallel and fit prediction model in leave-one-out mannter", + ) + parser.add_argument("method", type=str) + parser.add_argument("test_idx", type=int) + parser.add_argument("filepath", type=str) + parser.add_argument("-m", "--mlp", action="store_true") + parser.add_argument("--baseline", action="store_true") + + parser.add_argument( + "-p", + "--pred-filepath", + type=str, + default=None, + help="Path to data.pkl file with full features.", + ) + parser.add_argument( + "-l", + "--log-filepath", + type=str, + default=None, + help="Path to log.pkl if transport was already run. If not None, reevaluates but does not calculates the transport again.", + ) + return parser.parse_args() + + +ot_method_hyperparams = {} +for method in [ + "EGWL", + "EOT_ott", + "LEOT_ott", + "EGW_ott", + "EGW_all_ott", + "EGWL_ott", + "ECOOT", + "ECOOT_each", + "ECOOTL", +]: + ot_method_hyperparams[method] = [ + 0.1, + 1e-2, + 1e-3, + 1e-4, + 1e-5, + ] +vae_adv_loss = [1, 5, 10, 50, 100] +vae_latent_dims = [128] +vae_learning_rates = [1e-4] +# vae_latent_dims = [32] +# vae_learning_rates = [1e-4] +for method in ["VAE", "VAE_label"]: + ot_method_hyperparams[method] = list( + product(vae_adv_loss, vae_latent_dims, vae_learning_rates) + ) + +ot_method_all_to_all = ["GW_all", "EGW_all_ott", "EOT_all_ott"] +pred_method = weighted_ols_normed +baseline_pred_methods = [ols_normed, weight_1_ols_normed, weight_conc_normed] +baseline_pred_method_labels = ["perfect", "random", "by_conc"] +pred_from_param = predict + + +def main(args): + epsilons = ot_method_hyperparams[args.method] + all_to_all = args.method in ot_method_all_to_all + with open(args.filepath, "rb") as f: + data_dict = pkl.load(f) + if args.pred_filepath is not None: + with open(args.pred_filepath, "rb") as f: + pred_data_dict = pkl.load(f) + + X_dict = data_dict["Xs_dict"] + Y_dict = data_dict["Xt_dict"] + Zs_dict = data_dict["Zs_dict"]["dosage"] + Zt_dict = data_dict["Zt_dict"]["dosage"] + + labels = list(X_dict.keys()) + dim_X = X_dict[labels[0]].shape[1] + dim_Y = Y_dict[labels[0]].shape[1] + cv = KFold(n_splits=5) + train_val_idx, test_idx = list(cv.split(labels))[args.test_idx] + cv_inner = KFold(n_splits=5) + + test_labels = [labels[tidx] for tidx in test_idx] + train_val_X = _pop_keys(X_dict, test_labels) + train_val_Y = _pop_keys(Y_dict, test_labels) + if args.pred_filepath is not None: + train_val_X_full = _pop_keys(pred_data_dict["Xs_dict"], test_labels) + train_val_Y_full = _pop_keys(pred_data_dict["Xt_dict"], test_labels) + train_val_Z = _pop_keys(Zs_dict, test_labels) + + train_val_idx_split = cv_inner.split(train_val_idx) + val_labels_all = [] + train_Zs = [] + train_data = [] + train_val_labels = [labels[tvidx] for tvidx in train_val_idx] + for train_idx, val_idx in train_val_idx_split: + val_labels = tuple([train_val_labels[vidx] for vidx in val_idx]) + val_labels_all.append(val_labels) + train_X = _pop_keys(train_val_X, val_labels) + train_Y = _pop_keys(train_val_Y, val_labels) + train_Z = _pop_keys(train_val_Z, val_labels) + train_Zs.append(train_Z) + train_data.append((train_X, train_Y)) + + # Run 5-fold inner CV for all eps + eps_prod, train_data_prod = zip(*product(epsilons, train_data)) + val_labels_prod = val_labels_all * len(epsilons) + train_Z_prod = train_Zs * len(epsilons) + print("Len eps", eps_prod) + if args.log_filepath is None and "VAE" not in args.method: + Ts_list = [] + logs = [] + # for i, (_train_data, _eps) in tqdm(enumerate(zip(train_data_prod, eps_prod))): + # print(f"{i}th iteration with {_eps}") + # _T, _log = ot_method_map[args.method](_train_data, _eps) + # Ts_list.append(_T) + # logs.append(_log) + # jax.clear_caches() + try: + with Pool(5 * len(epsilons)) as p: + # with Pool(5) as p: + Ts_list, logs = zip( + *p.starmap( + ot_method_map[args.method], zip(train_data_prod, eps_prod) + ) + ) + + except KeyboardInterrupt: + p.terminate() + finally: + p.join() + elif args.log_filepath is not None: + with open(args.log_filepath, "rb") as f: + prev_log = pkl.load(f) + Ts_list = [None] * len(train_data_prod) + logs = [None] * len(train_data_prod) + if "VAE" not in args.method: + epsilons = [ + e for e in epsilons if e in np.unique(list(prev_log["T"].keys())) + ] + eps_prod, train_data_prod = zip(*product(epsilons, train_data)) + val_labels_prod = val_labels_all * len(epsilons) + else: + Ts_list = [] + logs = [] + for train_data, eps in zip(train_data_prod, eps_prod): + t, l = ot_method_map[args.method](train_data, eps) + Ts_list.append(t) + logs.append(l) + + # Evaluate & select best eps + eps_to_pred_evals = {eps: [] for eps in epsilons} + eps_to_pred_full_evals = {eps: [] for eps in epsilons} + eps_to_matching_evals = {eps: [] for eps in epsilons} + eps_to_dfracs_evals = {eps: [] for eps in epsilons} + eps_to_val_to_T = {eps: {} for eps in epsilons} + eps_to_val_to_log = {eps: {} for eps in epsilons} + iters = zip(eps_prod, val_labels_prod, train_data_prod, Ts_list, logs) + assert len(list(iters)) == len(eps_prod) + print(f"Ts_list: (len {len(Ts_list)})") + for eps, val_labels, train_pair, Ts, train_Z, log in zip( + eps_prod, val_labels_prod, train_data_prod, Ts_list, train_Z_prod, logs + ): + if args.log_filepath is not None: + Ts = prev_log["T"][eps][val_labels] + log = prev_log["log"][eps][val_labels] + # Ts: Dict[Union[float,int], np.ndarray] for each label, transport plan + val_X_dict = {tidx: train_val_X[tidx] for tidx in val_labels} + val_Y_dict = {tidx: train_val_Y[tidx] for tidx in val_labels} + if args.pred_filepath is not None: + val_X_dict_full = {tidx: train_val_X_full[tidx] for tidx in val_labels} + val_Y_dict_full = {tidx: train_val_Y_full[tidx] for tidx in val_labels} + train_X, train_Y = train_pair + print(f"in the loop; {eps}, {val_labels}") + + if isinstance(Ts, int): + # COOT underflow + print("Underflow", eps, val_labels) + eps_to_matching_evals[eps].append(100) + for tidx in val_labels: + eps_to_pred_evals[eps].append( + pd.Series( + [np.nan, np.nan, np.nan, np.nan, np.nan], + index=[ + "Pearson_corr", + "Spearman_corr", + "Pearson_samples", + "Spearman_samples", + "MSE", + ], + ).to_frame() + ) + if args.pred_filepath is not None: + for tidx in val_labels: + eps_to_pred_full_evals[eps].append( + pd.Series( + [np.nan, np.nan, np.nan, np.nan, np.nan], + index=[ + "Pearson_corr", + "Spearman_corr", + "Pearson_samples", + "Spearman_samples", + "MSE", + ], + ).to_frame() + ) + eps_to_val_to_T[eps][val_labels] = Ts + eps_to_val_to_log[eps][val_labels] = {} + continue + # Evaluate matching + if "VAE" in args.method: + ks = [5, 10, 25, 50] + latent_Y = infer_from_Ys(train_Y, Ts, dim_X) + latent_X = infer_from_Xs(train_X, Ts, dim_Y) + _, mean_foscttm = get_FOSCTTM( + Ts, latent_X, latent_Y, use_agg="mean", use_barycenter=False + ) + + Ts_multK = get_Ts_from_nn_multKs(latent_X, latent_Y, ks) # k -> T + diag_fracs_dict = {} + for k, T_k in Ts_multK.items(): + dfracs, rel_dfracs = get_diag_fracs( + T_k, train_X, train_Y, train_Z, train_Z + ) + diag_fracs_dict[k] = rel_dfracs + eps_to_dfracs_evals[eps].append(diag_fracs_dict) + else: + _, mean_foscttm = get_FOSCTTM(Ts, train_X, train_Y, use_agg="mean") + dfracs, rel_dfracs = get_diag_fracs(Ts, train_X, train_Y, train_Z, train_Z) + + eps_to_dfracs_evals[eps].append(rel_dfracs) + eps_to_matching_evals[eps].append(mean_foscttm.mean()) + print("map", eps_to_matching_evals[eps]) + print("foscttm", mean_foscttm) + + # Evaluate prediction + for tidx in val_labels: + val_X = val_X_dict[tidx] + val_Y = val_Y_dict[tidx] + if "VAE" in args.method: + pred = predict_from_model(val_X, Ts, dim_Y) + else: + param = pred_method(train_X, train_Y, Ts) + pred = pred_from_param(val_X, param) + try: + df = get_evals( + val_Y, + pred, + prediction_id=(eps, val_labels), + full=False, + agg_method="mean", + ) + except: + df = pd.Series( + [np.nan, np.nan, np.nan, np.nan, np.nan], + index=[ + "Pearson_corr", + "Spearman_corr", + "Pearson_samples", + "Spearman_samples", + "MSE", + ], + ).to_frame() + eps_to_pred_evals[eps].append(df) + if args.pred_filepath is not None: + for tidx in val_labels: + train_X_full = _pop_keys(train_val_X_full, val_labels) + train_Y_full = _pop_keys(train_val_Y_full, val_labels) + val_X_full = val_X_dict_full[tidx] + val_Y_full = val_Y_dict_full[tidx] + if "VAE" in args.method: + pred = predict_from_model(val_X, Ts, dim_Y) + else: + param = pred_method(train_X_full, train_Y_full, Ts) + pred = pred_from_param(val_X_full, param) + eps_to_pred_full_evals[eps].append( + get_evals( + val_Y_full, + pred, + prediction_id=(eps, val_labels), + full=False, + agg_method="mean", + ) + ) + eps_to_val_to_T[eps][val_labels] = Ts + eps_to_val_to_log[eps][val_labels] = log + print(eps_to_matching_evals) + eps_to_matching_evals = pd.Series( + {k: np.nanmean(v) for k, v in eps_to_matching_evals.items()}, + ) + max_matching_eps = eps_to_matching_evals.idxmin() + eps_to_pred_evals = pd.DataFrame( + {k: pd.concat(v, axis=1).mean(axis=1) for k, v in eps_to_pred_evals.items()} + ) + max_pred_eps = eps_to_pred_evals.loc["MSE"].idxmin() + best_eps = {"matching": max_matching_eps, "pred": max_pred_eps} + if args.pred_filepath is not None: + eps_to_pred_full_evals = pd.DataFrame( + { + k: pd.concat(v, axis=1).mean(axis=1) + for k, v in eps_to_pred_full_evals.items() + } + ) + max_pred_full_eps = eps_to_pred_full_evals.loc["Pearson_samples"].idxmax() + best_eps["pred_full"] = eps_to_pred_full_evals.loc["Pearson_samples"].idxmax() + val_logs = { + "matching_evals": eps_to_matching_evals, + "dfracs": eps_to_dfracs_evals, + "pred_evals": eps_to_pred_evals, + "T": eps_to_val_to_T, + "log": eps_to_val_to_log, + "best_eps": best_eps, + } + if args.pred_filepath is not None: + val_logs["pred_full_evals"] = eps_to_pred_full_evals + + edited_label = "e." if args.log_filepath is not None else "" + with open( + f"val_CV_{args.method}.{args.test_idx}.best_eps.{edited_label}pkl", + "wb", + ) as f: + pkl.dump(best_eps, f) + + with open( + f"val_CV_{args.method}.{args.test_idx}.{edited_label}pkl", + "wb", + ) as f: + pkl.dump(val_logs, f) + + +def run_mlp(args): + epsilons = [e for e in ot_method_hyperparams[args.method] if e != 0.1 and e != 1e-6] + with open(args.log_filepath, "rb") as f: + log = pkl.load(f) + with open(args.filepath, "rb") as f: + data_dict = pkl.load(f) + + X_dict = data_dict["Xs_dict"] + Y_dict = data_dict["Xt_dict"] + Z_dict = data_dict["Zs_dict"]["dosage"] + + labels = list(X_dict.keys()) + dim_X = X_dict[labels[0]].shape[1] + dim_Y = Y_dict[labels[0]].shape[1] + cv = KFold(n_splits=5) + train_val_idx, test_idx = list(cv.split(labels))[args.test_idx] + cv_inner = KFold(n_splits=5) + + test_labels = [labels[tidx] for tidx in test_idx] + train_val_X = _pop_keys(X_dict, test_labels) + train_val_Y = _pop_keys(Y_dict, test_labels) + train_val_Z = _pop_keys(Z_dict, test_labels) + + train_val_idx_split = cv_inner.split(train_val_idx) + val_labels_all = [] + train_data = [] + val_data = [] + train_val_labels = [labels[tvidx] for tvidx in train_val_idx] + for train_idx, val_idx in train_val_idx_split: + val_labels = tuple([train_val_labels[vidx] for vidx in val_idx]) + val_labels_all.append(val_labels) + train_X = _pop_keys(train_val_X, val_labels) + train_Y = _pop_keys(train_val_Y, val_labels) + val_X = np.concatenate([train_val_X[l] for l in val_labels], axis=0) + val_Y = np.concatenate([train_val_Y[l] for l in val_labels], axis=0) + val_Z = np.concatenate([train_val_Z[l] for l in val_labels], axis=0) + train_data.append((train_X, train_Y)) + val_data.append((val_X, val_Y, val_Z)) + + # Run 5-fold inner CV for all eps + eps_prod, train_data_prod = zip(*product(epsilons, train_data)) + val_labels_prod = val_labels_all * len(epsilons) + val_data_prod = val_data * len(epsilons) + Ts_list = [] + for eps, val_labels in zip(eps_prod, val_labels_prod): + Ts_list.append(log["T"][eps][val_labels]) + + trained_models = [] + logs = [] + for train_data, Ts in zip(train_data_prod, Ts_list): + t, l = train_mlp(train_data, Ts) + trained_models.append(t) + logs.append(l) + + # Evaluate & select best eps + eps_to_pred_evals = {eps: [] for eps in epsilons} + eps_to_val_to_preds = {eps: {} for eps in epsilons} + eps_to_matching_evals = {eps: [] for eps in epsilons} + eps_to_val_to_model = {eps: {} for eps in epsilons} + eps_to_val_to_log = {eps: {} for eps in epsilons} + iters = zip(val_data_prod, train_data_prod, trained_models) + assert len(list(iters)) == len(eps_prod) + print(f"Ts_list: (len {len(Ts_list)})") + for eps, val_labels, val_pair, model in zip( + eps_prod, val_labels_prod, val_data_prod, trained_models + ): + val_X, val_Y, val_Z = val_pair + print(f"in the loop; {val_labels}") + + # Evaluate prediction + Y_pred = model(torch.tensor(val_X)).detach().numpy() + eval_df = get_evals( + val_Y, + Y_pred, + prediction_id=(eps, val_labels), + full=False, + agg_method="mean", + ) + print(eval_df) + eps_to_pred_evals[eps] = eps_to_pred_evals[eps] + [eval_df] + eps_to_val_to_preds[eps][val_labels] = (Y_pred, val_Y, val_Z) + eps_to_val_to_model[eps][val_labels] = model + + def concat_if_possible(v): + if len(v) > 0: + return pd.concat(v, axis=1).mean(axis=1) + else: + return v + + eps_to_pred_evals = pd.DataFrame( + {k: concat_if_possible(v) for k, v in eps_to_pred_evals.items()} + ) + # max_pred_eps = eps_to_pred_evals.loc["Pearson_c"].idxmax() + val_logs = { + "pred": eps_to_val_to_preds, + "pred_evals": eps_to_pred_evals, + "model": eps_to_val_to_model, + } + # with open(f"val_CV_MLP_{args.method}.{args.test_idx}.best_eps.pkl", "wb") as f: + # pkl.dump({"eps": max_pred_eps}, f) + + with open(f"val_CV_MLP_{args.method}.{args.test_idx}.pkl", "wb") as f: + pkl.dump(val_logs, f) + + +def run_mlp_control(args): + with open(args.filepath, "rb") as f: + data_dict = pkl.load(f) + + X_dict = data_dict["Xs_dict"] + Y_dict = data_dict["Xt_dict"] + Z_dict = data_dict["Zs_dict"]["dosage"] + + labels = list(X_dict.keys()) + dim_X = X_dict[labels[0]].shape[1] + dim_Y = Y_dict[labels[0]].shape[1] + cv = KFold(n_splits=5) + train_val_idx, test_idx = list(cv.split(labels))[args.test_idx] + cv_inner = KFold(n_splits=5) + + test_labels = [labels[tidx] for tidx in test_idx] + train_val_X = _pop_keys(X_dict, test_labels) + train_val_Y = _pop_keys(Y_dict, test_labels) + train_val_Z = _pop_keys(Z_dict, test_labels) + + train_val_idx_split = cv_inner.split(train_val_idx) + val_labels_all = [] + train_data = [] + val_data = [] + train_Z = [] + train_val_labels = [labels[tvidx] for tvidx in train_val_idx] + for train_idx, val_idx in train_val_idx_split: + val_labels = tuple([train_val_labels[vidx] for vidx in val_idx]) + val_labels_all.append(val_labels) + train_X = _pop_keys(train_val_X, val_labels) + train_Y = _pop_keys(train_val_Y, val_labels) + val_X = np.concatenate([train_val_X[l] for l in val_labels], axis=0) + val_Y = np.concatenate([train_val_Y[l] for l in val_labels], axis=0) + val_Z = np.concatenate([train_val_Z[l] for l in val_labels], axis=0) + train_data.append((train_X, train_Y)) + train_Z.append(train_val_Z) + val_data.append((val_X, val_Y, val_Z)) + + # Make baseline Ts + perf_T = {k: np.diag(np.ones(v.shape[0])) for k, v in X_dict.items()} + + def make_G(size, label, k): + assert size == len(label), (k, size, len(label)) + G = np.zeros((size, size)) + for l in np.unique(label): + l_idx = np.where(label == l)[0] + for i in l_idx: + for j in l_idx: + G[i, j] = 1.0 + assert (G.sum(axis=0) > 0).all(), (k, (G.sum(axis=0) == 0).sum()) + return G + + cond_T = {k: make_G(X_dict[k].shape[0], Z_dict[k], k) for k in X_dict.keys()} + rand_T = {k: np.ones((v.shape[0], v.shape[0])) for k, v in X_dict.items()} + Ts_list = [perf_T, cond_T, rand_T] + Ts_prod, train_data_prod = zip(*product(Ts_list, train_data)) + Ts_label = ["perfect", "dosage", "random"] + Ts_label_prod = np.repeat(Ts_label, len(train_data)) + train_Z_prod = train_Z * len(Ts_label) + val_labels_prod = val_labels_all * len(Ts_label) + val_data_prod = val_data * len(Ts_label) + + trained_models = [] + logs = [] + for train_data, val_labels, Ts, tlab in zip( + train_data_prod, val_labels_prod, Ts_prod, Ts_label_prod + ): + T = _pop_keys(Ts, test_labels) + T = _pop_keys(T, val_labels) + try: + t, l = train_mlp(train_data, T) + except Exception as e: + print(f"error at {tlab}, {val_labels}") + raise e + trained_models.append(t) + logs.append(l) + + tlab_to_pred_evals = {tlab: [] for tlab in Ts_label} + tlab_to_matching_evals = {tlab: [] for tlab in Ts_label} + tlab_to_val_to_model = {tlab: {} for tlab in Ts_label} + tlab_to_val_to_log = {tlab: {} for tlab in Ts_label} + tlab_to_val_to_pred = {tlab: {} for tlab in Ts_label} + iters = zip(val_data_prod, train_data_prod, trained_models) + assert len(list(iters)) == len(Ts_label_prod) + print(f"Ts_list: (len {len(Ts_list)})") + for Ts, tlab, val_labels, val_pair, train_data, model, train_Z, log in zip( + Ts_prod, + Ts_label_prod, + val_labels_prod, + val_data_prod, + train_data_prod, + trained_models, + train_Z_prod, + logs, + ): + val_X, val_Y, val_Z = val_pair + train_X, train_Y = train_data + print(f"in the loop; {val_labels}") + # Evaluate matching + T = _pop_keys(Ts, test_labels) + T = _pop_keys(T, val_labels) + try: + _, mean_foscttm = get_FOSCTTM(T, train_X, train_Y, use_agg="mean") + except KeyError as e: + print(Ts.keys()) + print(T.keys()) + print(train_X.keys()) + raise e + dfracs, rel_dfracs = get_diag_fracs(T, train_X, train_Y, train_Z, train_Z) + tlab_to_matching_evals[tlab] = { + "foscttm": mean_foscttm, + "dfracs": dfracs, + "rel_dfracs": rel_dfracs, + } + # Evaluate prediction + Y_pred = model(torch.tensor(val_X)).detach().numpy() + eval_df = get_evals( + val_Y, + Y_pred, + prediction_id=(tlab, val_labels), + full=False, + agg_method="mean", + norm_Y=Y_dict[3].mean(axis=0), + ) + tlab_to_val_to_log[tlab] = log + tlab_to_val_to_pred[tlab][val_labels] = (val_Y, Y_pred, val_Z) + print(eval_df) + tlab_to_pred_evals[tlab] = tlab_to_pred_evals[tlab] + [eval_df] + tlab_to_val_to_model[tlab][val_labels] = model + + def concat_if_possible(v): + if len(v) > 0: + return pd.concat(v, axis=1).mean(axis=1) + else: + return v + + tlab_to_pred_evals = pd.DataFrame( + {k: concat_if_possible(v) for k, v in tlab_to_pred_evals.items()} + ) + val_logs = { + "pred_evals": tlab_to_pred_evals, + "match_evals": tlab_to_matching_evals, + "model": tlab_to_val_to_model, + "logs": tlab_to_val_to_log, + "preds": tlab_to_val_to_pred, + } + + with open(f"val_CV_MLP_baseline.{args.test_idx}.pkl", "wb") as f: + pkl.dump(val_logs, f) + + +if __name__ == "__main__": + args = parse_args() + if args.mlp: + if args.baseline: + run_mlp_control(args) + else: + if args.log_filepath is None: + args.log_filepath = f"val_CV_{args.method}.{args.test_idx}.pkl" + run_mlp(args) + else: + main(args) diff --git a/perturbot/build/lib/perturbot/eval/cv_outer_loop.py b/perturbot/build/lib/perturbot/eval/cv_outer_loop.py new file mode 100755 index 0000000..8acbe10 --- /dev/null +++ b/perturbot/build/lib/perturbot/eval/cv_outer_loop.py @@ -0,0 +1,325 @@ +"""For 5-fold CV, run the single inner loop of train-validation to select hyperparameter.""" + +from typing import Dict, Any +import sys +import argparse +import pickle as pkl +from functools import partial +from itertools import product +from multiprocessing import Pool +import numpy as np +import pandas as pd +from sklearn.model_selection import KFold +import torch + +from perturbot.eval.utils import _pop_keys, _pop_key, get_Ts_from_nn_multKs, make_G +from perturbot.eval.match import get_FOSCTTM, get_diag_fracs +from perturbot.eval.prediction import get_evals_preds, get_evals +from perturbot.preprocess.vae import ( + train_vae_rna, + train_vae_acc, + train_vae_prot, + SCVI_LATENT_KEY, +) +from perturbot.match.cot_labels import get_coupling_cotl, get_coupling_cotl_sinkhorn +from perturbot.match.gw import ( + get_coupling_gw_cg, + get_coupling_egw, + get_coupling_egw_all, + get_coupling_gw_all, +) +from perturbot.match.ott_egwl import ( + get_coupling_egw_labels_ott, + get_coupling_egw_all_ott, + get_coupling_eot_ott, + get_coupling_leot_ott, + get_coupling_egw_ott, +) +from perturbot.match.cot import get_coupling_cot, get_coupling_cot_sinkhorn +from perturbot.match.gw_labels import get_coupling_gw_labels, get_coupling_egw_labels +from perturbot.predict.scvi_vae import train_vae_model +from perturbot.predict.linear_regression import ( + ols_normed, + weight_conc_normed, + weighted_ols_normed, + weight_1_ols_normed, + predict, +) +from perturbot.predict.mlp import train_mlp +from perturbot.predict.scvi_vae import predict_from_model, infer_from_Xs, infer_from_Ys + +ot_method_map = { + "ECOOTL": get_coupling_cotl_sinkhorn, + "ECOOT": get_coupling_cot_sinkhorn, + "EGWL": get_coupling_egw_labels, + "EOT_ott": get_coupling_eot_ott, + "LEOT_ott": get_coupling_leot_ott, + "EGW_ott": get_coupling_egw_ott, + "EGW_all_ott": get_coupling_egw_all_ott, + "EGWL_ott": get_coupling_egw_labels_ott, + "VAE_label": train_vae_model, + "VAE": partial(train_vae_model, use_label=False), +} + + +def parse_args(): + parser = argparse.ArgumentParser( + "Run inner validation loop of CV", + "Run sample-matching OT in parallel and fit prediction model in leave-one-out mannter", + ) + parser.add_argument("method", type=str) + parser.add_argument("test_idx", type=int) + parser.add_argument("filepath", type=str) + parser.add_argument("eps", type=str) + parser.add_argument( + "-p", + "--pred-filepath", + type=str, + default=None, + help="Path to data.pkl file with full features.", + ) + parser.add_argument( + "-l", + "--log-filepath", + type=str, + default=None, + help="Path to log.pkl if transport was already run. If not None, reevaluates but does not calculates the transport again.", + ) + parser.add_argument( + "-b", + "--baseline", + type=str, + default=None, + help="One of perfect, random, by_conc", + ) + return parser.parse_args() + + +ot_method_all_to_all = ["GW_all", "EGW_all_ott", "EOT_all_ott"] +pred_from_param = predict + + +def main(args): + if args.log_filepath is not None and args.log_filepath != "None": + with open(args.log_filepath, "rb") as f: + prev_logs = pkl.load(f) + logs = { + "matching_evals": {}, + "pred_evals": {}, + "T": {}, + "pred": {}, + "log": {}, + "eps": {}, + } + try: + match_eps, lin_eps, pred_eps = tuple(map(float, args.eps.split(","))) + logs["eps"]["match"] = match_eps + logs["eps"]["lin"] = lin_eps + logs["eps"]["pred"] = pred_eps + all_to_all = args.method in ot_method_all_to_all + with open(args.filepath, "rb") as f: + data_dict = pkl.load(f) + if args.pred_filepath is not None: + with open(args.pred_filepath, "rb") as f: + pred_data_dict = pkl.load(f) + + X_dict = data_dict["Xs_dict"] + Y_dict = data_dict["Xt_dict"] + Zs_dict = data_dict["Zs_dict"]["dosage"] + Zt_dict = data_dict["Zt_dict"]["dosage"] + + labels = list(X_dict.keys()) + dim_X = pred_data_dict["Xs_dict"][labels[0]].shape[1] + dim_Y = pred_data_dict["Xt_dict"][labels[0]].shape[1] + cv = KFold(n_splits=5) + train_idx, test_idx = list(cv.split(labels))[args.test_idx] + + test_labels = [labels[tidx] for tidx in test_idx] + train_X = _pop_keys(X_dict, test_labels) + train_Y = _pop_keys(Y_dict, test_labels) + train_Z = _pop_keys(Zs_dict, test_labels) + if args.pred_filepath is not None: + train_X_full = _pop_keys(pred_data_dict["Xs_dict"], test_labels) + train_Y_full = _pop_keys(pred_data_dict["Xt_dict"], test_labels) + train_Z = _pop_keys(Zs_dict, test_labels) + train_data = (train_X, train_Y) + train_data_full = (train_X_full, train_Y_full) + print(f"Calculating matching with {match_eps}") + + # Get matching + if args.log_filepath is not None and args.log_filepath != "None": + Ts_matching = prev_logs["T"]["match"] + Ts_pred = prev_logs["T"]["pred"] + try: + log_matching = prev_logs["log"]["match"] + except: + log_matching = None + try: + log_matching_predeps = prev_logs["log"]["match_pred"] + except: + log_matching_predeps = None + elif args.baseline is not None: + if args.baseline == "perfect": + Ts_matching = Ts_pred = { + k: np.diag(np.ones(v.shape[0])) for k, v in train_X.items() + } + elif args.baseline == "random": + Ts_matching = Ts_pred = { + k: np.ones((v.shape[0], v.shape[0])) for k, v in train_X.items() + } + elif args.baseline == "by_conc": + Ts_matching = Ts_pred = { + k: make_G(train_X[k].shape[0], train_Z[k], k) + for k in train_X.keys() + } + log_matching = None + log_matching_predeps = None + else: + if "VAE" in args.method: + Ts_matching, log_matching = ot_method_map[args.method]( + train_data_full, match_eps + ) + else: + Ts_matching, log_matching = ot_method_map[args.method]( + train_data, match_eps + ) + + if "VAE" in args.method: + if match_eps == pred_eps: + Ts_pred = Ts_matching + log_matching_predeps = None + else: + Ts_pred, log_matching_predeps = ot_method_map[args.method]( + train_data_full, pred_eps + ) + else: + if match_eps != pred_eps: + print(f"Calculating pred matching with {pred_eps}") + Ts_pred, log_matching_predeps = ot_method_map[args.method]( + train_data, pred_eps + ) + else: + Ts_pred = Ts_matching + log_matching_predeps = None + logs["T"]["match"] = Ts_matching + logs["T"]["pred"] = Ts_pred + + if "VAE" in args.method: + latent_Y = infer_from_Ys(train_Y_full, Ts_matching, dim_X) + latent_X = infer_from_Xs(train_X_full, Ts_matching, dim_Y) + _, mean_foscttm = get_FOSCTTM( + Ts_matching, + latent_X, + latent_Y, + use_agg="mean", + use_barycenter=False, + ) + ks = [1, 5, 10, 50, 100] + k_to_Ts = get_Ts_from_nn_multKs(latent_X, latent_Y, ks) # k -> T + dfracs = {} + rel_dfracs = {} + for k, Ts in k_to_Ts.items(): + dfracs[k], rel_dfracs[k] = get_diag_fracs( + Ts, train_X, train_Y, train_Z, train_Z + ) + + else: + # normalize + if isinstance(Ts_matching, dict): + total_sum = 0 + for k, v in Ts_matching.items(): + total_sum += v.sum() + Ts_matching = { + k: v.astype(np.double) / total_sum for k, v in Ts_matching.items() + } + else: + Ts_matching = Ts_matching.astype(np.double) / Ts_matching.sum() + _, mean_foscttm = get_FOSCTTM(Ts_matching, train_X, train_Y, use_agg="mean") + dfracs, rel_dfracs = get_diag_fracs( + Ts_matching, train_X, train_Y, train_Z, train_Z + ) + + # if not all_to_all: + mean_mean_foscttm = mean_foscttm.mean() + # else: + # mean_mean_foscttm = mean_foscttm + print("foscttm", mean_foscttm) + logs["matching_evals"] = ( + { + "foscttm": mean_foscttm, + "mean_foscttm": mean_mean_foscttm, + "dfracs": dfracs, + "rel_dfracs": rel_dfracs, + }, + ) + # Eval prediction: full features + test_X_full = np.concatenate( + [pred_data_dict["Xs_dict"][l] for l in test_labels], axis=0 + ) + test_Y_full = np.concatenate( + [pred_data_dict["Xt_dict"][l] for l in test_labels], axis=0 + ) + if "VAE" not in args.method: + trained_model, log_pred = train_mlp((train_X_full, train_Y_full), Ts_pred) + Y_pred_full = trained_model(torch.tensor(test_X_full)).detach().numpy() + else: + Y_pred_full = predict_from_model(test_X_full, Ts_pred, dim_Y) + logs["pred"] = { + "model": trained_model if "VAE" not in args.method else None, + "Y_pred": Y_pred_full, + "Y_true": test_Y_full, + "train_Y": train_Y_full, + "test_Z": {k: Zs_dict[k] for k in test_labels}, + } + eval_df_full = get_evals( + test_Y_full, + Y_pred_full, + prediction_id="eval", + full=False, + agg_method="mean", + ) + print(eval_df_full) + logs["pred_evals"]["full"] = eval_df_full + + # # Eval prediction - PC + # test_X = np.concatenate([X_dict[l] for l in test_labels], axis=0) + # test_Y = np.concatenate([Y_dict[l] for l in test_labels], axis=0) + # if "VAE" in args.method: + # pred_Ys = predict_from_model(test_X, Ts_lin, dim_Y) + + # else: + # params = [ + # lin_method(train_X, train_Y, Ts_lin) + # for lin_method in [ols_normed, weight_1_ols_normed, weight_conc_normed] + # ] + # pred_Ys = [pred_from_param(test_X, param) for param in params] + # pred_labels = ["ot", "perfect", "random", "by_conc"] + # eval_df = get_evals_preds( + # test_Y, + # pred_Ys, + # pred_labels=pred_labels, + # full=False, + # ) + # logs["pred_evals"]["PC"] = eval_df + # print(eval_df) + + logs["log"] = ( + { + "match": log_matching, + "match_pred": log_matching_predeps, + # "match_lin": log_matching_lin, + "mlp": log_pred if "VAE" not in args.method else None, + }, + ) + except Exception as e: + with open(f"test_{args.method}.{args.test_idx}.tmp.pkl", "wb") as f: + pkl.dump(logs, f) + raise e + + with open(f"test_{args.method}.{args.test_idx}.e.pkl", "wb") as f: + pkl.dump(logs, f) + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/perturbot/build/lib/perturbot/eval/feature_matching.py b/perturbot/build/lib/perturbot/eval/feature_matching.py new file mode 100755 index 0000000..2bf61d7 --- /dev/null +++ b/perturbot/build/lib/perturbot/eval/feature_matching.py @@ -0,0 +1,160 @@ +"""For 5-fold CV, run the single inner loop of train-validation to select hyperparameter.""" + +import os +import argparse +import pickle as pkl +from functools import partial +import numpy as np +from perturbot.match.fot import get_coupling_fot +from perturbot.eval.utils import make_G +from perturbot.predict.scvi_vae import train_vae_model, infer_from_Xs, infer_from_Ys +from perturbot.eval.utils import get_Ts_from_nn_multKs + + +def parse_args(): + parser = argparse.ArgumentParser( + "Run inner validation loop of CV", + "Run sample-matching OT in parallel and fit prediction model in leave-one-out mannter", + ) + parser.add_argument("method", type=str) + parser.add_argument("filepath", type=str) + parser.add_argument("best_eps", type=float) + parser.add_argument("eps", type=float) + + parser.add_argument( + "-l", + "--log-filepath", + type=str, + default=None, + help="Path to log.pkl if transport was already run. If not None, reevaluates but does not calculates the transport again.", + ) + parser.add_argument( + "-b", + "--baseline", + type=str, + default=None, + help="One of perfect, random, by_conc", + ) + parser.add_argument( + "--best-k", + type=int, + default=None, + help="Best k to use for VAE", + ) + return parser.parse_args() + + +ot_method_all_to_all = ["GW_all", "EGW_all_ott", "EOT_all_ott"] + + +def main(args): + logs = { + "matching_evals": {}, + "pred_evals": {}, + "T": {}, + "pred": {}, + "log": {}, + "eps": {}, + } + try: + if args.best_eps != 0: + with open(f"all_{args.method}.{args.best_eps}.pkl", "rb") as f: + d = pkl.load(f) + Ts = d["T"] + logs["sample_eps"] = args.best_eps + # logs['min_foscttm'] = d[] + else: + Ts = None + with open(args.filepath, "rb") as f: + data_dict = pkl.load(f) # full features + + X_dict = data_dict["Xs_dict"] + Y_dict = data_dict["Xt_dict"] + Z_dict = data_dict["Zs_dict"]["dosage"] + + if "VAE" in args.method: + labels = list(X_dict.keys()) + dim_X = data_dict["Xs_dict"][labels[0]].shape[1] + dim_Y = data_dict["Xt_dict"][labels[0]].shape[1] + latent_Y = infer_from_Ys(Y_dict, Ts, dim_X) + latent_X = infer_from_Xs(X_dict, Ts, dim_Y) + Ts = get_Ts_from_nn_multKs(latent_X, latent_Y, [args.best_k])[args.best_k] + + if Ts is None: + if args.method == "random": + Ts = { + k: np.ones((X_dict[k].shape[0], Y_dict[k].shape[0])) + / (X_dict[k].shape[0] * Y_dict[k].shape[0]) + for k in X_dict.keys() + } + elif args.method == "perfect": + Ts = { + k: np.diag( + np.ones(X_dict[k].shape[0]), + ) + / X_dict[k].shape[0] + for k in X_dict.keys() + } + elif args.method == "by_conc": + Ts = { + k: make_G(X_dict[k].shape[0], Z_dict[k], k) for k in X_dict.keys() + } + Tv, log = get_coupling_fot((X_dict, Y_dict), Ts, args.eps) + logs["log"] = log + logs["Tv"] = Tv + except Exception as e: + with open(f"features_{args.method}.{args.eps}.tmp.pkl", "wb") as f: + pkl.dump(logs, f) + raise e + + with open(f"features_{args.method}.{args.eps}.pkl", "wb") as f: + pkl.dump(logs, f) + + +def submit_feature_run(data_path, ot_method_label): + epsilons = [1e-2, 1e-3, 1e-4, 1e-5] + rel_dfracs = [] + if ot_method_label in ["perfect", "random", "by_conc"]: + best_eps = 0 + else: + best_k_dict = {} + for eps in epsilons: + # try: + with open(f"all_{ot_method_label}.{eps}.pkl", "rb") as f: + d = pkl.load(f) + _rel_dfracs = d["matching_evals"][0]["rel_dfracs"] + if isinstance(_rel_dfracs, dict): + max_rel_dfracs = -10 + for k, v in _rel_dfracs.items(): + if max_rel_dfracs < v: + max_rel_dfracs = v + best_k_dict[eps] = k + _rel_dfracs = max_rel_dfracs + rel_dfracs.append(_rel_dfracs) + # except Exception as e: + # print(e) + # rel_dfracs.append(-1) + best_eps = epsilons[rel_dfracs.index(max(rel_dfracs))] + if "VAE" in ot_method_label: + best_k = best_k_dict[best_eps] + for eps in epsilons: + run_label = f"FM.{ot_method_label}.{eps}" + f = open(f"{run_label}.bsub", "w") + f.write("source ~/.bashrc\n") + f.write("pwd\n") + f.write("conda activate ot \n") + command = f"python /gpfs/scratchfs01/site/u/ryuj6/OT/software/perturbot/perturbot/eval/feature_matching.py {ot_method_label} {data_path} {best_eps} {eps}" + + if "VAE" in ot_method_label: + command += f" --best-k {best_k}" + f.write(f"echo {command}\n") + f.write(f"{command}\n") + f.close() + command = f"bsub -M 10G -n 1 -q short -J {run_label} -o log/log.{run_label}.o%J -e log/log.{run_label}.e%J < {run_label}.bsub" + print(command) + os.system(command) + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/perturbot/build/lib/perturbot/eval/loo.py b/perturbot/build/lib/perturbot/eval/loo.py new file mode 100755 index 0000000..a3d0e28 --- /dev/null +++ b/perturbot/build/lib/perturbot/eval/loo.py @@ -0,0 +1,432 @@ +from typing import List, Dict +import numpy as np +import pandas as pd +from sklearn.model_selection import LeaveOneOut +from .prediction import get_evals_preds +from multiprocessing import Pool +import pickle as pkl +from .match import get_FOSCTTM, get_diag_fracs +from .utils import get_Ts_from_nn_multKs +from perturbot.predict.scvi_vae import predict_from_model, infer_from_Xs, infer_from_Ys +from perturbot.preprocess.vae import ( + train_vae_rna, + train_vae_acc, + train_vae_prot, + SCVI_LATENT_KEY, +) +from .utils import _pop_key + +np.seterr(all="raise") + +vae_trainer_dict = { + "rna": train_vae_rna, + "accessibility": train_vae_acc, + "protein": train_vae_prot, +} + + +def run_models( + X_dict, + Y_dict, + ot_method, + pred_method, + pred_method_label: str, + baseline_pred_methods: List, + baseline_pred_method_labels: List[str], + pred_from_param, + Ts_dict=None, + Z_dict=None, + *args, + **kwargs, +): + if pred_method_label == "VAE": + return run_models_vae( + X_dict, + Y_dict, + ot_method, + Ts_dict=None, + *args, + **kwargs, + ) + labels = list(X_dict.keys()) + loo = LeaveOneOut() + log = {"ot_couplings": {}, "params": {}, "preds": {}, "logs": {}} + eval_dfs = [] + train_data = [] + test_data = {} + + for i, (train_idx, test_idx) in enumerate(loo.split(labels)): + test_idx = test_idx.item() + test_X = X_dict[test_idx] + test_Y = Y_dict[test_idx] + train_X = X_dict.copy() + train_Y = Y_dict.copy() + del train_X[test_idx] + del train_Y[test_idx] + train_data.append((train_X, train_Y)) + test_data[test_idx] = (test_X, test_Y) + if Ts_dict is None: + try: + with Pool(10) as p: + Ts_list, logs = zip(*p.map(ot_method, train_data)) + except KeyboardInterrupt: + p.terminate() + finally: + p.join() + + for i, (test_idx, (test_X, test_Y)) in enumerate(test_data.items()): + if Ts_dict is None: + Ts = Ts_list[i] + log["logs"][i] = logs[i] + else: + Ts = Ts_dict[test_idx] + log["ot_couplings"][test_idx] = Ts + train_X = X_dict.copy() + train_Y = Y_dict.copy() + if Z_dict is not None: + train_Z = Z_dict.copy() + del train_X[test_idx] + del train_Y[test_idx] + params = [pred_method(train_X, train_Y, Ts)] + if Z_dict is not None: + params += [ + baseline_method(train_X, train_Y, train_Z) + for baseline_method in baseline_pred_methods + ] + else: + params += [ + baseline_method(train_X, train_Y) + for baseline_method in baseline_pred_methods + ] + log["params"][test_idx] = params + + preds = [pred_from_param(test_X, param) for param in params] + log["preds"][test_idx] = preds + method_labels = [pred_method_label] + baseline_pred_method_labels + assert len(preds) == len(method_labels) + eval_df = get_evals_preds(test_Y, preds, method_labels) + eval_df["loo_test_idx"] = test_idx + eval_dfs.append(eval_df) + + return pd.concat(eval_dfs, axis=0), log + + +def run_models_vae( + X_dict, + Y_dict, + ot_method, + Ts_dict=None, + *args, + **kwargs, +): + labels = list(X_dict.keys()) + loo = LeaveOneOut() + log = { + "models": {}, + "params": {}, + "preds": {}, + "latent_X": {}, + "latent_Y": {}, + "logs": {}, + } + eval_dfs = [] + train_data = [] + test_data = {} + + for i, (train_idx, test_idx) in enumerate(loo.split(labels)): + test_idx = test_idx.item() + test_X = X_dict[test_idx] + test_Y = Y_dict[test_idx] + train_X = X_dict.copy() + train_Y = Y_dict.copy() + del train_X[test_idx] + del train_Y[test_idx] + train_data.append((train_X, train_Y)) + test_data[test_idx] = (test_X, test_Y) + if Ts_dict is None: + try: + with Pool(10) as p: + model_list, logs = zip( + *p.map( + ot_method, + train_data, + ) + ) + except KeyboardInterrupt: + p.terminate() + finally: + p.join() + ks = [5, 10, 50, 100] + for k in ks: + log[f"pred_T_k{k}"] = {} # test_idx -> train_idx -> T + for i, (test_idx, (test_X, test_Y)) in enumerate(test_data.items()): + if Ts_dict is None: + model = model_list[i] + log["models"][test_idx] = model + log["logs"][test_idx] = logs[i] + else: + model = Ts_dict[test_idx] + log["models"][test_idx] = model + dX = test_X.shape[1] + dY = test_Y.shape[1] + latent_Y = infer_from_Ys(_pop_key(Y_dict, test_idx), model, dX) + latent_X = infer_from_Xs(_pop_key(X_dict, test_idx), model, dY) + pred_Y = predict_from_model(test_X, model, test_Y.shape[1]) + log["preds"][test_idx] = pred_Y + log["latent_X"][test_idx] = latent_X + log["latent_Y"][test_idx] = latent_Y + Ts = get_Ts_from_nn_multKs(latent_X, latent_Y, ks) # k -> T + for k, v in Ts.items(): + log[f"pred_T_k{k}"][test_idx] = v + eval_df = get_evals_preds(test_Y, [pred_Y], ["VAE"]) + eval_df["loo_test_idx"] = test_idx + eval_dfs.append(eval_df) + + return pd.concat(eval_dfs, axis=0), log + + +def run_models_vae_then_ot( + source_adata, + target_adata, + source_modality, + target_modality, + label_col, + ot_method, + pred_method, + pred_method_label: str, + baseline_pred_methods: List, + baseline_pred_method_labels: List[str], + pred_from_param, + T_dict=None, + model_dict=None, + Z_dict=None, + n_threads=1, + *args, + **kwargs, +): + labels = source_adata.obs[label_col].unique().values + loo = LeaveOneOut() + log = {"ot_couplings": {}, "params": {}, "preds": {}, "logs": {}} + eval_dfs = [] + test_data = {} + source_adatas_train = [] + target_adatas_train = [] + for i, (train_idx, test_idx) in enumerate(loo.split(labels)): + test_idx = test_idx.item() + source_adata_test = source_adata.obs[label_col == test_idx].copy() + target_adata_test = target_adata.obs[label_col == test_idx].copy() + source_adata_train = source_adata.obs[label_col != test_idx].copy() + target_adata_train = target_adata.obs[label_col != test_idx].copy() + source_adatas_train.append(source_adata_train) + target_adatas_train.append(target_adata_train) + test_data[test_idx] = (source_adata_test, target_adata_test) + + # Train VAE + if model_dict is None: + try: + with Pool(n_threads) as p: + source_adatas, source_models = zip( + *p.map(vae_trainer_dict[source_modality], source_adatas_train) + ) + except KeyboardInterrupt: + p.terminate() + finally: + p.join() + try: + with Pool(n_threads) as p: + target_adatas, target_models = zip( + *p.map(vae_trainer_dict[target_modality], target_adatas_train) + ) + except KeyboardInterrupt: + p.terminate() + finally: + p.join() + + # Train OT + train_data = [] + for i, (test_idx, (test_X, test_Y)) in enumerate(test_data.items()): + train_X = source_adatas[i].obsm[SCVI_LATENT_KEY] + train_Y = target_adatas[i].obsm[SCVI_LATENT_KEY] + train_data.append((train_X, train_Y)) + + if T_dict is None: + try: + with Pool(1) as p: + T_list, logs = zip(*p.map(ot_method, train_data)) + except KeyboardInterrupt: + p.terminate() + finally: + p.join() + + for i, (test_idx, (test_X, test_Y)) in enumerate(test_data.items()): + if T_dict is None: + Ts = T_list[i] + log["logs"][i] = logs[i] + else: + Ts = T_dict[test_idx] + log["ot_couplings"][test_idx] = Ts + + # @TODO see from here + pred_Y = predict_from_model_with_OT( + test_X, source_models[i], target_models[i], Ts, test_Y.shape[1] + ) + + preds = [pred_from_param(test_X, param) for param in params] + log["preds"][test_idx] = preds + method_labels = [pred_method_label] + baseline_pred_method_labels + assert len(preds) == len(method_labels) + eval_df = get_evals_preds(test_Y, preds, method_labels) + eval_df["loo_test_idx"] = test_idx + eval_dfs.append(eval_df) + + return pd.concat(eval_dfs, axis=0), log + + +def evaluate_loo_vae( + data_dict, logfile_name, test_idx_file, use_z_key="dosage", ks=[5, 10, 50, 100] +): + Xs_dict = data_dict["Xs_dict"] + Xt_dict = data_dict["Xt_dict"] + Zs_dict = data_dict["Zs_dict"] + Zt_dict = data_dict["Zt_dict"] + with open(logfile_name, "rb") as f: + log = pkl.load(f) + dfracs = {k: [] for k in ks} + rel_dfracs = {k: [] for k in ks} + med_foscttms = [] + idx_name = pd.read_csv(test_idx_file, header=None) + for test_idx in log["preds"].keys(): + model = log["models"] + Xs_train = _pop_key(Xs_dict, test_idx) + Xt_train = _pop_key(Xt_dict, test_idx) + for k in ks: + dfrac, rel_dfrac = get_diag_fracs( + log[f"pred_T_k{k}"][test_idx], + Xs_train, + Xt_train, + Zs_dict[use_z_key], + Zt_dict[use_z_key], + ) + dfracs[k].append(dfrac) + rel_dfracs[k].append(rel_dfrac) + _, median_foscttm = get_FOSCTTM( + model, + log["latent_X"][test_idx], + log["latent_Y"][test_idx], + use_barycenter=False, + ) + med_foscttms.append(median_foscttm) + med_foscttms = pd.concat(med_foscttms, axis=1).mean(axis=1) + rel_dfracs_ks = {} + dfracs_ks = {} + for k in ks: + rel_dfracs_ks[k] = pd.concat(rel_dfracs[k], axis=1).mean(axis=1) + dfracs_ks[k] = pd.concat(dfracs[k], axis=1).mean(axis=1) + main_k = 10 + other_ks = [k for k in ks if k != 3] + metrics_list = [ + med_foscttms, + dfracs_ks[main_k], + rel_dfracs_ks[main_k], + ] + labels = ["treat_idx", "FOSCTTM", "DFracs", "Relative DFracs"] + for k in other_ks: + metrics_list = metrics_list + [dfracs_ks[k], rel_dfracs_ks[k]] + labels = labels + [f"DFracs_{k}", f"Relative DFracs_{k}"] + metrics = pd.concat( + metrics_list, + axis=1, + ).reset_index() + + metrics.columns = labels + metrics["treatment"] = idx_name.loc[metrics["treat_idx"].astype(int), 0].tolist() + return metrics.set_index("treatment", drop=True)[labels[1:]].to_dict("series") + + +def evaluate_loo_run( + data_dict, + logfile_name, + test_idx_file="/home/ryuj6/scratch/OT/data/chemical_screen/chemical_screen_pca_idx.txt", + use_z_key="dosage", +): + Xs_dict = data_dict["Xs_dict"] + Xt_dict = data_dict["Xt_dict"] + Zs_dict = data_dict["Zs_dict"] + Zt_dict = data_dict["Zt_dict"] + with open(logfile_name, "rb") as f: + log = pkl.load(f) + dfracs = [] + rel_dfracs = [] + med_foscttms = [] + idx_name = pd.read_csv(test_idx_file, header=None) + for test_idx, Ts in log["ot_couplings"].items(): + Xs_train = _pop_key(Xs_dict, test_idx) + Xt_train = _pop_key(Xt_dict, test_idx) + + dfrac, rel_dfrac = get_diag_fracs( + Ts, Xs_train, Xt_train, Zs_dict[use_z_key], Zt_dict[use_z_key] + ) + _, median_foscttm = get_FOSCTTM(Ts, Xs_train, Xt_train) + med_foscttms.append(median_foscttm) + dfracs.append(dfrac) + rel_dfracs.append(rel_dfrac) + if isinstance(Ts, dict): + med_foscttms = pd.concat(med_foscttms, axis=1).mean(axis=1) + rel_dfracs = pd.concat(rel_dfracs, axis=1).mean(axis=1) + dfracs_df = pd.concat(dfracs, axis=1).mean(axis=1) + metrics = pd.concat( + [ + med_foscttms, + dfracs_df, + rel_dfracs, + ], + axis=1, + ).reset_index() + else: + med_foscttms = pd.Series(med_foscttms, index=log["ot_couplings"].keys()) + rel_dfracs = pd.Series(rel_dfracs, index=log["ot_couplings"].keys()) + dfracs_df = pd.Series(dfracs_df, index=log["ot_couplings"].keys()) + metrics = pd.concat( + [ + med_foscttms, + dfracs_df, + rel_dfracs, + ], + axis=1, + ).reset_index() + + metrics.columns = ["treat_idx", "FOSCTTM", "DFracs", "Relative DFracs"] + metrics["treatment"] = idx_name.loc[metrics["treat_idx"].astype(int), 0].tolist() + return metrics.set_index("treatment", drop=True)[ + ["FOSCTTM", "DFracs", "Relative DFracs"] + ].to_dict("series") + + +def evaluate_loo_runs_OT( + data_dict, + logfile_paths: List[str], + labels: List[str], + test_idx_file="/home/ryuj6/scratch/OT/data/chemical_screen/chemical_screen_pca_idx.txt", +): + metric_names = ["FOSCTTM", "DFracs", "Relative DFracs"] + metrics = {m: [] for m in metric_names} + non_foscttm_labels = [] + for logfile_path, label in zip(logfile_paths, labels): + if "vae" in label: + evals = evaluate_loo_vae(data_dict, logfile_path, test_idx_file) + else: + non_foscttm_labels.append(label) + evals = evaluate_loo_run(data_dict, logfile_path, test_idx_file) + for m in metric_names: + if m in evals: + metrics[m].append(evals[m]) + else: + metrics[m].append(pd.Series(np.zeros(evals["FOSCTTM"].shape))) + + def collect_series(list_series): + df = pd.concat(list_series, axis=1) + df.columns = labels + return df + + metrics = {m: collect_series(metrics[m]) for m in metrics.keys()} + return metrics diff --git a/perturbot/build/lib/perturbot/eval/match.py b/perturbot/build/lib/perturbot/eval/match.py new file mode 100755 index 0000000..be475a2 --- /dev/null +++ b/perturbot/build/lib/perturbot/eval/match.py @@ -0,0 +1,242 @@ +"""Evaluate the quality of OT""" + +from typing import Dict, Tuple, Union, Sequence, List +import numpy as np +import pandas as pd +from perturbot.eval.utils import foscttm +from perturbot.utils import mdict_to_matrix + + +def get_rel_mse(T_dict): + rel_err = {} + for k, T in T_dict.items(): + T = T / T.sum() + perfect_match = np.identity( + T_dict[k].shape[0], + ) + perfect_match /= perfect_match.sum() + err = np.mean((np.diag(T_dict[k]) - np.diag(perfect_match)) ** 2) + + all2all = np.ones((T_dict[k].shape[0], T_dict[k].shape[1])) + all2all /= all2all.sum() + worst_err = np.mean((np.diag(all2all) - np.diag(perfect_match)) ** 2) + + rel_err[k] = err / worst_err + + return rel_err + + +def get_confusion_matrix( + T_dict, + Xs_dict, + Xt_dict, + Zs_dict, + Zt_dict, + norm=True, +) -> Tuple[Dict[Union[float, int], np.ndarray], pd.Series]: + """ + Ts: Dictionary of n_l x n'_l matrices for each label l + zs: m x m Confusion matrix is calculated where k is the number of unique values in z. + Must contain integer values (0, ..., m-1) + """ + labels = list(Xs_dict.keys()) + if not isinstance(T_dict, dict): + return get_confusion_matrix_single(T_dict, Xs_dict, Xt_dict, Zs_dict, Zt_dict) + classes = set(val for vals in Zs_dict.values() for val in vals) + Cmat_dict = np.zeros((len(classes), len(classes))) + diag_frac = {} + for k in labels: + # Cmat_dict[k] = np.zeros((len(classes), len(classes))) + try: + Zs = Zs_dict[k] + except KeyError as e: + print(e) + print(Zs_dict.keys()) + raise e + Zt = Zt_dict[k] + # if norm: + # T = T_dict[k] / T_dict[k].sum() + # else: + T = T_dict[k] + for i in range(T.shape[0]): + for j in range(T.shape[1]): + if T[i, j]: + Cmat_dict[int(Zs[i]), int(Zt[j])] += T[i, j] + diag_frac = np.diag(Cmat_dict).sum() + return Cmat_dict, diag_frac + + +def get_confusion_matrix_single( + T, + Xs_dict, + Xt_dict, + Zs_dict, + Zt_dict, +) -> Tuple[np.ndarray, float]: + classes = set(val for vals in Zs_dict.values() for val in vals) + Cmat = np.zeros((len(classes), len(classes))) + Zs = np.concatenate([Zs_dict[k] for k in Xs_dict.keys()]) + Zt = np.concatenate([Zt_dict[k] for k in Xs_dict.keys()]) + T = T / T.sum() + for i in range(T.shape[0]): + for j in range(T.shape[1]): + if T[i, j]: + Cmat[int(Zs[i]), int(Zt[j])] += T[i, j] + diag_frac = np.diag(Cmat).sum() + return Cmat, diag_frac + + +def get_diag_fracs( + T_dict: Union[np.ndarray, Dict[float, np.ndarray]], + Xs_dict, + Xt_dict, + Zs_dict, + Zt_dict, +) -> Union[Tuple[pd.Series, pd.Series], Tuple[float, float]]: + if not isinstance(T_dict, dict): + T = T_dict.copy() + sidx = 0 + tidx = 0 + T_dict = {} + for k, v in Xs_dict.items(): + T_dict[k] = T[ + sidx : (sidx + v.shape[0]), tidx : (tidx + Xt_dict[k].shape[0]) + ] + sidx += v.shape[0] + tidx += Xt_dict[k].shape[0] + # return get_diag_fracs_single(T_dict, Xs_dict, Xt_dict, Zs_dict, Zt_dict) + Cmat_dict, diag_fracs = get_confusion_matrix( + T_dict, Xs_dict, Xt_dict, Zs_dict, Zt_dict + ) + total_size_perfect = sum([T_dict[k].shape[0] for k in T_dict.keys()]) + T_perfect = {} + for k in T_dict.keys(): + T_perfect[k] = ( + np.identity( + T_dict[k].shape[0], + ) + / total_size_perfect + ) + Cmats_perfect, perfect_diag_fracs = get_confusion_matrix( + T_perfect, Xs_dict, Xt_dict, Zs_dict, Zt_dict + ) + total_size_random = sum([T_dict[k].size for k in T_dict.keys()]) + T_random = {k: np.ones(T_dict[k].shape) / total_size_random for k in T_dict.keys()} + Cmats_random, random_diag_fracs = get_confusion_matrix( + T_random, Xs_dict, Xt_dict, Zs_dict, Zt_dict + ) + + return diag_fracs, (diag_fracs - random_diag_fracs) / ( + perfect_diag_fracs - random_diag_fracs + ) + + +def get_diag_fracs_single( + T, + Xs_dict, + Xt_dict, + Zs_dict, + Zt_dict, +): + Cmat, diag_fracs = get_confusion_matrix_single( + T, Xs_dict, Xt_dict, Zs_dict, Zt_dict + ) + T_perfect = {} + for k in Xs_dict.keys(): + T_perfect[k] = ( + np.identity( + Xs_dict[k].shape[0], + ) + / Xs_dict[k].shape[0] + ) + T_perfect = mdict_to_matrix( + T_perfect, + np.concatenate([np.ones(Xs_dict[l].shape[0]) * l for l in Xs_dict.keys()]), + np.concatenate([np.ones(Xt_dict[l].shape[0]) * l for l in Xs_dict.keys()]), + ) + Cmats_perfect, perfect_diag_fracs = get_confusion_matrix_single( + T_perfect, Xs_dict, Xt_dict, Zs_dict, Zt_dict + ) + + T_random = { + k: np.ones((Xs_dict[k].shape[0], Xt_dict[k].shape[0])) for k in Xs_dict.keys() + } + T_random = mdict_to_matrix( + T_random, + np.concatenate([np.ones(Xs_dict[l].shape[0]) * l for l in Xs_dict.keys()]), + np.concatenate([np.ones(Xt_dict[l].shape[0]) * l for l in Xs_dict.keys()]), + ) + Cmats_random, random_diag_fracs = get_confusion_matrix_single( + T_random, Xs_dict, Xt_dict, Zs_dict, Zt_dict + ) + + return diag_fracs, (diag_fracs - random_diag_fracs) / ( + perfect_diag_fracs - random_diag_fracs + ) + + +def get_FOSCTTM( + T_dict, Xs_dict, Xt_dict, use_barycenter=True, use_agg="mean" +) -> Union[Tuple[Dict[float, Sequence], pd.Series], Tuple[List[float], float]]: + """Obtain the fraction of cells with higher assignment probability than the + true match (assume the diagonal is the true match.). Barycenter is used as the prediction. + + Returns + foscttm_dict: label to the list of FOSCTTM per cells + median_foscttm_dict: label to the aggregated FOSCTTM + """ + agg_fn = np.nanmedian if use_agg == "median" else np.nanmean + if isinstance(T_dict, dict): + T = mdict_to_matrix( + T_dict, + np.concatenate([np.ones(Xs_dict[l].shape[0]) * l for l in Xs_dict.keys()]), + np.concatenate([np.ones(Xt_dict[l].shape[0]) * l for l in Xs_dict.keys()]), + ) + else: + T = T_dict + + foscttm_dict = {} + median_foscttm_dict = {} + Xs_true = np.concatenate([Xs_dict[l] for l in Xs_dict.keys()]) + Xt_true = np.concatenate([Xt_dict[l] for l in Xt_dict.keys()]) + if use_barycenter: + marg = T.sum(axis=-1) + marg[marg == 0] = 1e-30 + Xt_pred = (T / marg[:, None]) @ Xt_true + foscttm_ = foscttm(Xt_pred, Xt_true) + else: + foscttm_ = foscttm(Xs_true, Xt_true) + return foscttm_, agg_fn(foscttm_) + for l, Xt_true in Xt_dict.items(): + Xs_true = Xs_dict[l] + if use_barycenter: + T = T_dict[l] + marg = T.sum(axis=-1) + marg[marg == 0] = 1e-30 + Xt_pred = (T / marg[:, None]) @ Xt_true + foscttm_ = foscttm(Xt_pred, Xt_true) + else: + foscttm_ = foscttm(Xs_true, Xt_true) + foscttm_dict[l] = foscttm_ + median_foscttm_dict[l] = agg_fn(foscttm_) + return foscttm_dict, pd.Series(median_foscttm_dict) + + +def get_FOSCTTM_single( + T, Xs_dict, Xt_dict, use_barycenter=True +) -> Tuple[List[float], float]: + """Obtain the fraction of cells with higher assignment probability than the + true match (assume the diagonal is the true match.). Barycenter is used as the prediction. + """ + Xs = np.concatenate([Xs_dict[l] for l in Xs_dict.keys()], axis=0) + Xt = np.concatenate([Xt_dict[l] for l in Xt_dict.keys()], axis=0) + + if use_barycenter: + marg = T.sum(axis=-1) + marg[marg == 0] = 1e-30 + Xt_pred = (T / marg[:, None]) @ Xt + foscttm_ = foscttm(Xt_pred, Xt) + else: + foscttm_ = foscttm(Xs, Xt) + + return foscttm_, np.nanmedian(foscttm_) diff --git a/perturbot/build/lib/perturbot/eval/prediction.py b/perturbot/build/lib/perturbot/eval/prediction.py new file mode 100755 index 0000000..87d4111 --- /dev/null +++ b/perturbot/build/lib/perturbot/eval/prediction.py @@ -0,0 +1,210 @@ +import numpy as np +import pandas as pd +from scipy.stats import spearmanr +from perturbot.eval.utils import foscttm + + +def _pearson_rowwise(A, B, eps=1e-8): + assert A.shape[0] == B.shape[0] + A_mA = A - A.mean(1)[:, None] + B_mB = B - B.mean(1)[:, None] + ssA = np.einsum("ij,ij->i", A_mA, A_mA) + ssB = np.einsum("ij,ij->i", B_mB, B_mB) + return np.einsum("ij,ij->i", A_mA, B_mB) / (np.sqrt(ssA * ssB) + eps) + + +def _spearman_rowwise(A, B): + assert A.shape[0] == B.shape[0] + scorr = [] + for i in range(A.shape[0]): + scorr.append( + spearmanr( + A[i, :], + B[i, :], + )[0] + ) + return scorr + + +def get_corrs(Y_pred, Y_true, idx=None): + if idx is not None: + Y_pred = Y_pred[:, idx] + Y_true = Y_true[:, idx] + pearson = _pearson_rowwise(Y_pred, Y_true) + spearman = _spearman_rowwise(Y_pred, Y_true) + return pearson, spearman + + +def mse(Y_pred, Y_true, idx=None): + if idx is None: + return (np.abs(Y_pred - Y_true) ** 2).mean(axis=1) + else: + return (np.abs(Y_pred[:, idx] - Y_true[:, idx]) ** 2).mean(axis=1) + + +def get_evals( + Y_pred, + Y_true, + idx=None, + idx_id="subset", + prediction_id="pred", + full=True, + agg_method="mean", + norm_Y: np.ndarray = None, +): + if agg_method == "median": + agg = np.median + elif agg_method == "mean": + agg = np.mean + else: + raise ValueError() + if norm_Y is not None: + pearson, spearman = get_corrs( + Y_pred / norm_Y[None, :], + Y_true / norm_Y[None, :], + ) + else: + pearson, spearman = get_corrs(Y_pred, Y_true) + pearson_c, spearman_c = get_corrs(Y_pred.T, Y_true.T) + _mse = mse(Y_pred, Y_true) + metrics = pd.Series( + [agg(pearson), agg(spearman), agg(pearson_c), agg(spearman_c), agg(_mse)], + index=[ + "Pearson_corr", + "Spearman_corr", + "Pearson_samples", + "Spearman_samples", + "MSE", + ], + ).to_frame() + if idx is not None: + if norm_Y is not None: + pearson_idx, spearman_idx = get_corrs( + Y_pred / norm_Y[None, :], Y_true / norm_Y[None, :], idx + ) + else: + pearson_idx, spearman_idx = get_corrs(Y_pred, Y_true, idx) + pearson_idx_c, spearman_idx_c = get_corrs(Y_pred, Y_true, idx) + _foscttm_idx = foscttm(Y_pred, Y_true, idx) + _mse_idx = mse(Y_pred, Y_true, idx) + metrics = pd.concat( + [ + metrics, + pd.Series( + [ + agg(pearson_idx), + agg(spearman_idx), + agg(pearson_idx_c), + agg(spearman_idx_c), + agg(_mse_idx), + ], + index=[ + f"Pearson_corr_{idx_id}", + f"Spearman_corr_{idx_id}", + f"Pearson_samples_{idx_id}", + f"Spearman_samples_{idx_id}", + f"MSE_{idx_id}", + ], + ).to_frame(), + ], + axis=0, + ) + metrics.columns = [prediction_id] + if full: + if idx is not None: + full_metrics = pd.DataFrame( + { + "metric": np.concatenate( + [ + pearson, + spearman, + pearson_c, + spearman_c, + _mse, + pearson_idx, + spearman_idx, + pearson_idx_c, + spearman_idx_c, + _mse_idx, + ] + ), + "group": np.repeat( + [ + "Pearson_corr", + "Spearman_corr", + "Pearson_samples", + "Spearman_samples", + "MSE", + f"Pearson_corr_{idx_id}", + f"Spearman_corr_{idx_id}", + f"Pearson_samples_{idx_id}", + f"Spearman_samples_{idx_id}", + f"MSE_{idx_id}", + ], + len(pearson), + ), + "pred_label": prediction_id, + } + ) + else: + full_metrics = pd.DataFrame( + { + "metric": np.concatenate( + [ + pearson, + spearman, + pearson_c, + spearman_c, + _mse, + ] + ), + "group": np.repeat( + [ + "Pearson_corr", + "Spearman_corr", + "Pearson_samples", + "Spearman_samples", + "MSE", + ], + len(pearson), + ), + "pred_label": prediction_id, + } + ) + + return metrics, full_metrics + return metrics + + +def get_evals_preds( + Y_true, Y_preds, pred_labels, idx=None, idx_id="subset", full=False +): + metric_dfs = [] + if full: + foscttm_dfs = [] + for Y_pred, pred_label in zip(Y_preds, pred_labels): + if full: + metrics, foscttm_df = get_evals( + Y_true, + Y_pred, + idx=idx, + idx_id=idx_id, + full=full, + prediction_id=pred_label, + ) + metric_dfs.append(metrics) + foscttm_dfs.append(foscttm_df) + else: + metric_dfs.append( + get_evals( + Y_true, + Y_pred, + idx=idx, + idx_id=idx_id, + full=full, + prediction_id=pred_label, + ) + ) + if full: + return pd.concat(metric_dfs, axis=1), pd.conat(foscttm_dfs, axis=1) + return pd.concat(metric_dfs, axis=1) diff --git a/perturbot/build/lib/perturbot/eval/utils.py b/perturbot/build/lib/perturbot/eval/utils.py new file mode 100755 index 0000000..9680f52 --- /dev/null +++ b/perturbot/build/lib/perturbot/eval/utils.py @@ -0,0 +1,105 @@ +from typing import Sequence, List +import numpy as np +from sklearn.neighbors import NearestNeighbors + + +def make_G(size, label, k): + assert size == len(label), (k, size, len(label)) + G = np.zeros((size, size)) + for l in np.unique(label): + l_idx = np.where(label == l)[0] + for i in l_idx: + for j in l_idx: + G[i, j] = 1.0 + assert (G.sum(axis=0) > 0).all(), (k, (G.sum(axis=0) == 0).sum()) + return G + + +def foscttm(Y_pred, Y_true, idx=None) -> List[float]: + """ + From https://github.com/rsinghlab/SCOT/blob/9787d0a6a1a494059cb3fb49e8ceda9318273c06/src/.ipynb_checkpoints/evals-checkpoint.py#L32 + Returns fraction closer than true match for each sample (as an array) + """ + if idx is not None: + Y_pred = Y_pred[:, idx] + Y_true = Y_true[:, idx] + fracs = [] + nsamp = Y_pred.shape[0] + rank = 0 + for row_idx in range(nsamp): + euc_dist = np.sqrt( + np.sum(np.square(np.subtract(Y_pred[row_idx, :], Y_true)), axis=1) + ) + true_nbr = euc_dist[row_idx] + sort_euc_dist = sorted(euc_dist) + try: + rank = np.where(sort_euc_dist == true_nbr)[0].mean() + except FloatingPointError: + print(true_nbr) + print(sort_euc_dist) + np.where(np.abs(sort_euc_dist - true_nbr < 1e-8))[0] + rank = np.where(np.abs(sort_euc_dist - true_nbr < 1e-8))[0].mean() + + frac = float(rank) / (nsamp - 1) + fracs.append(frac) + return fracs + + +def get_T_from_nn(X, Y, k): + """Obtain kNN label of observation in X from the k nearest neighbors in Y""" + nsamp = X.shape[0] + T = np.zeros((X.shape[0], Y.shape[0])) + for row_idx in range(nsamp): + euc_dist = np.sqrt(np.sum(np.square(np.subtract(X[row_idx, :], Y)), axis=1)) + smallest_k_idx = np.argpartition(euc_dist, k)[:k] + T[row_idx, smallest_k_idx] = 1.0 / (nsamp * k) + return T + + +def get_Ts_from_nn_multKs(X_dict, Y_dict, ks): + """Obtain kNN label of observation in X from the k nearest neighbors in Y""" + X = np.concatenate([X_dict[l] for l in X_dict.keys()]) + Y = np.concatenate([Y_dict[l] for l in X_dict.keys()]) + k_to_T = {k: np.zeros((X.shape[0], Y.shape[0])) for k in ks} + nsamp = X.shape[0] + for row_idx in range(nsamp): + euc_dist = np.sqrt(np.sum(np.square(np.subtract(X[row_idx, :], Y)), axis=1)) + for k in ks: + smallest_k_idx = np.argpartition(euc_dist, k)[:k] + k_to_T[k][row_idx, smallest_k_idx] = 1.0 / (nsamp * k) + k_to_Tdict = {} + for k, T in k_to_T.items(): + i = 0 + j = 0 + k_to_Tdict[k] = { + l: np.zeros((X_dict[l].shape[0], Y_dict[l].shape[0])) for l in X_dict.keys() + } + for l, X in X_dict.items(): + k_to_Tdict[k][l] = T[ + i : (i + X_dict[l].shape[0]), j : (j + Y_dict[l].shape[0]) + ] + i += X_dict[l].shape[0] + j += Y_dict[l].shape[0] + # k_to_Tdict[k][l] = k_to_Tdict[k][l] / k_to_Tdict[k][l].sum() + + return k_to_Tdict + + +def _pop_key(d, k, sub_key=None): + d = d.copy() + if sub_key is None: + del d[k] + else: + del d[sub_key][k] + return d + + +def _pop_keys(d, ks, sub_key=None): + d = d.copy() + if sub_key is None: + for k in ks: + del d[k] + else: + for k in ks: + del d[sub_key][k] + return d diff --git a/perturbot/build/lib/perturbot/match/__init__.py b/perturbot/build/lib/perturbot/match/__init__.py new file mode 100755 index 0000000..4284d03 --- /dev/null +++ b/perturbot/build/lib/perturbot/match/__init__.py @@ -0,0 +1,27 @@ +from .cot_labels import get_coupling_cotl, get_coupling_cotl_sinkhorn, cotl_numpy +from .ott_egwl import ( + get_coupling_egw_labels_ott, + get_coupling_egw_all_ott, + get_coupling_eot_ott, + get_coupling_leot_ott, + get_coupling_egw_ott, +) +from .cot import get_coupling_cot, get_coupling_cot_sinkhorn, cot_numpy +from .gw_labels import get_coupling_gw_labels +from .fot import get_coupling_fot + +__all__ = [ + "cotl_numpy", + "cot_numpy", + "get_coupling_eot_ott", + "get_coupling_leot_ott", + "get_coupling_cot", + "get_coupling_cot_sinkhorn", + "get_coupling_cotl", + "get_coupling_cotl_sinkhorn", + "get_coupling_egw_ott", + "get_coupling_gw_labels", + "get_coupling_egw_all_ott", + "get_coupling_egw_labels_ott", + "get_coupling_fot", +] diff --git a/perturbot/build/lib/perturbot/match/cot.py b/perturbot/build/lib/perturbot/match/cot.py new file mode 100755 index 0000000..900c525 --- /dev/null +++ b/perturbot/build/lib/perturbot/match/cot.py @@ -0,0 +1,389 @@ +from typing import Union, Dict, Tuple, Optional +from numbers import Number +import numpy as np +import time +import ot +from ott.solvers import linear +from ott.geometry import geometry +from .utils import random_gamma_init, init_matrix_np +import jax + + +def cot_numpy( + X1, + X2, + w1=None, + w2=None, + v1=None, + v2=None, + niter=10, + algo="emd", + reg=0, + algo2="emd", + reg2=0, + verbose=True, + log=False, + random_init=False, + C_lin=None, +): + """Returns COOT between two datasets X1,X2 (see [1]), Sinkhorn reimplemented with OTT + + The function solves the following optimization problem: + + .. math:: + + COOT = \min_{Ts,Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}*Tv_{k,l} + + Where : + - X1 : The source dataset + - X2 : The target dataset + - w1,w2 : weights (histograms) on the samples (rows) of resp. X1 and X2 + - v1,v2 : weights (histograms) on the features (columns) of resp. X1 and X2 + + Parameters + ---------- + X1 : numpy array, shape (n, d) + Source dataset + X2 : numpy array, shape (n', d') + Target dataset + w1 : numpy array, shape (n,) + Weight (histogram) on the samples of X1. If None uniform distribution is considered. + w2 : numpy array, shape (n',) + Weight (histogram) on the samples of X2. If None uniform distribution is considered. + v1 : numpy array, shape (d,) + Weight (histogram) on the features of X1. If None uniform distribution is considered. + v2 : numpy array, shape (d',) + Weight (histogram) on the features of X2. If None uniform distribution is considered. + niter : integer + Number max of iterations of the BCD for solving COOT. + algo : string + Choice of algorithm for solving OT problems on samples each iteration. Choice ['emd','sinkhorn']. + If 'emd' returns sparse solution + If 'sinkhorn' returns regularized solution + algo2 : string + Choice of algorithm for solving OT problems on features each iteration. Choice ['emd','sinkhorn']. + If 'emd' returns sparse solution + If 'sinkhorn' returns regularized solution + reg : float + Regularization parameter for samples coupling matrix. Ignored if algo='emd' + reg2 : float + Regularization parameter for features coupling matrix. Ignored if algo='emd' + eps : float + Threshold for the convergence + random_init : bool + Wether to use random initialization for the coupling matrices. If false identity couplings are considered. + log : bool, optional + record log if True + C_lin : numpy array, shape (n, n') + Prior on the sample correspondences. Added to the cost for the samples transport + + Returns + ------- + Ts : numpy array, shape (n,n') + Optimal Transport coupling between the samples + Tv : numpy array, shape (d,d') + Optimal Transport coupling between the features + cost : float + Optimization value after convergence + log : dict + convergence information and coupling marices + References + ---------- + .. [1] Redko Ievgen, Vayer Titouan, Flamary R{\'e}mi and Courty Nicolas + "CO-Optimal Transport" + + Examples + ---------- + .. code-block:: python + + import numpy as np + from cot import cot_numpy + + n_samples=300 + Xs=np.random.rand(n_samples,2) + Xt=np.random.rand(n_samples,1) + cot_numpy(Xs,Xt) + """ + if v1 is None: + v1 = np.ones(X1.shape[1]) / X1.shape[1] # is (d,) + if v2 is None: + v2 = np.ones(X2.shape[1]) / X2.shape[1] # is (d',) + if w1 is None: + w1 = np.ones(X1.shape[0]) / X1.shape[0] # is (n',) + if w2 is None: + w2 = np.ones(X2.shape[0]) / X2.shape[0] # is (n,) + + if not random_init: + Ts = np.ones((X1.shape[0], X2.shape[0])) / ( + X1.shape[0] * X2.shape[0] + ) # is (n,n') + Tv = np.ones((X1.shape[1], X2.shape[1])) / ( + X1.shape[1] * X2.shape[1] + ) # is (d,d') + else: + Ts = random_gamma_init(w1, w2) + Tv = random_gamma_init(v1, v2) + + constC_s, hC1_s, hC2_s = init_matrix_np(X1, X2, v1, v2) + + constC_v, hC1_v, hC2_v = init_matrix_np(X1.T, X2.T, w1, w2) + cost = np.inf + + log_out = {} + log_out["cost"] = [] + + for i in range(niter): + Tsold = Ts + Tvold = Tv + costold = cost + + M = constC_s - np.dot(hC1_s, Tv).dot(hC2_s.T) + if C_lin is not None: + M = M + C_lin + if algo == "emd": + Ts = ot.emd(w1, w2, M, numItermax=1e7) + elif algo == "sinkhorn": + Ts = np.array( + linear.solve( + geometry.Geometry( + cost_matrix=M, epsilon=reg, scale_cost="max_cost" + ), + max_iterations=2000, + ).matrix + ) + + M = constC_v - np.dot(hC1_v, Ts).dot(hC2_v.T) + + if algo2 == "emd": + Tv = ot.emd(v1, v2, M, numItermax=1e7) + elif algo2 == "sinkhorn": + Tv = np.array( + linear.solve( + geometry.Geometry( + cost_matrix=M, epsilon=reg2, scale_cost="max_cost" + ), + max_iterations=2000, + ).matrix + ) + + delta = np.linalg.norm(Ts - Tsold) + np.linalg.norm(Tv - Tvold) + cost = np.sum(M * Tv) + + if log: + log_out["cost"].append(cost) + + if verbose: + print("Delta: {0} Loss: {1}".format(delta, cost)) + + if delta < 1e-16 or np.abs(costold - cost) < 1e-7: + if verbose: + print("converged at iter ", i) + break + jax.clear_caches() + if log: + return Ts, Tv, cost, log_out + else: + return Ts, Tv, cost + + +def predict_with_cot(Xs, ys, Xt, yt, Xs_test, log=True): + """Learn mapping from Xs, Ys, Yt, Yt and predict Xt_test from Xs_test.""" + vs = Xs.sum(axis=0) # set the weights on the features + vs /= vs.sum() + vt = Xt.sum(axis=0) + vt /= vt.sum() + + Ts, Tv, _, log = cot_numpy(Xs, Xt, y1=ys, y2=yt, v1=vs, v2=vt, niter=100, log=True) + + log["Ts"] = Ts + log["Tv"] = Tv + Xt_pred = Xs_test @ (Tv / Tv.sum(axis=-1)[:, None]) + return Xt_pred, log + + +def get_coupling_cot( + data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], +) -> Tuple[Union[int, Dict[Number, np.array]], Union[int, Dict]]: + """Returns sample coupling between two datasets X, Y given the labels, disregarding label information. + + The function solves the following optimization problem: + + .. math:: + + COOT = \min_{Ts,Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + eps: + Regularization parameter, relative to the max cost. + + Returns + ------- + T_dict : + Optimal Transport coupling between the samples per label + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_cot + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,2) for k in labels} + get_coupling_cot((Xs_dict, Xt_dict)) + """ + X_dict = data[0] + Y_dict = data[1] + X = np.concatenate([X_dict[l] for l in X_dict.keys()]) + Y = np.concatenate([Y_dict[l] for l in X_dict.keys()]) + start = time.time() + try: + T, Tv, cost, log = cot_numpy(X, Y, log=True, niter=2000) + except FloatingPointError: + return -1, -1 + log["time"] = time.time() - start + return T, log + + +def get_coupling_cot_sinkhorn( + data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], + eps: float = 5e-3, + eps2: Optional[float] = None, +) -> Tuple[Union[int, Dict[Number, np.array]], Union[int, Dict]]: + """Returns sample coupling between two datasets X, Y given the labels, disregarding label information. + + The function solves the following optimization problem: + + .. math:: + + ECOOT = \min_{Ts,Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} - \epsilon H(T_s) + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + eps: + Regularization parameter, relative to the max cost. + + Returns + ------- + T_dict : + Optimal Transport coupling between the samples per label + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_cot_sinkhorn + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,2) for k in labels} + get_coupling_cot_sinkhorn((Xs_dict, Xt_dict), 0.05) + """ + print(f"calculating with eps {eps}") + X_dict = data[0] + Y_dict = data[1] + X = np.concatenate([X_dict[l] for l in X_dict.keys()]) + Y = np.concatenate([Y_dict[l] for l in X_dict.keys()]) + if eps2 is None: + eps2 = eps + start = time.time() + try: + T, Tv, cost, log = cot_numpy( + X, + Y, + algo="sinkhorn", + reg=eps, + algo2="sinkhorn", + reg2=eps2, + log=True, + niter=2000, + ) + except FloatingPointError: + return -1, -1 + log["time"] = time.time() - start + return T, log + + +def get_coupling_each_cot_sinkhorn( + data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], + eps: float = 5e-3, + eps2: Optional[float] = None, +) -> Tuple[Union[int, Dict[Number, np.array]], Union[int, Dict]]: + """Returns sample coupling between two datasets X, Y given the labels, disregarding label information. + + The function solves the following optimization problem: + + .. math:: + + ECOOT = \min_{Ts,Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} - \epsilon H(T_s) + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + eps: + Regularization parameter, relative to the max cost. + + Returns + ------- + T_dict : + Optimal Transport coupling between the samples per label + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_cot_sinkhorn + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,2) for k in labels} + get_coupling_cot_sinkhorn((Xs_dict, Xt_dict), 0.05) + """ + print(f"calculating with eps {eps}") + X_dict = data[0] + Y_dict = data[1] + X = np.concatenate([X_dict[l] for l in X_dict.keys()]) + Y = np.concatenate([Y_dict[l] for l in X_dict.keys()]) + if eps2 is None: + eps2 = eps + start = time.time() + try: + T_dict = {} + for l in X_dict.keys(): + T, Tv, cost, log = cot_numpy( + X_dict[l], + Y_dict[l], + algo="sinkhorn", + reg=eps, + algo2="sinkhorn", + reg2=eps2, + log=True, + niter=2000, + ) + T_dict[l] = T + print(f"Done calculating with eps {eps}") + except FloatingPointError: + return -1, -1 + log["time"] = time.time() - start + return T_dict, log diff --git a/perturbot/build/lib/perturbot/match/cot_labels.py b/perturbot/build/lib/perturbot/match/cot_labels.py new file mode 100755 index 0000000..43fbf15 --- /dev/null +++ b/perturbot/build/lib/perturbot/match/cot_labels.py @@ -0,0 +1,340 @@ +# Adapted from https://github.com/PythonOT/COOT +from typing import Dict, Optional, Tuple, Union +from numbers import Number +import pandas as pd +import numpy as np +import time +import ot as pot +import matplotlib.pyplot as plt +from ott.solvers import linear +from ott.geometry import geometry +from .utils import random_gamma_init, init_matrix_np + + +def cotl_numpy( + X_dict: Dict[Number, np.ndarray], + Y_dict: Dict[Number, np.ndarray], + w1: Dict[Number, np.ndarray] = None, + w2: Dict[Number, np.ndarray] = None, + v1: Optional[np.ndarray] = None, + v2: Optional[np.ndarray] = None, + niter: int = 100, + algo: str = "emd", + reg: float = 0.1, + algo2: str = "emd", + reg2: float = 10.0, + verbose: bool = True, + log: bool = False, + random_init: bool = False, + C_lin: bool = None, +): + r"""Returns COOT between two datasets X, Y given labels. + + The function solves the following optimization problem: + + .. math:: + + COOTL = \min_{Ts^1,..Ts^{L},Tv} \sum_{t=1}^L \sum_{i,k \in {i|l_{x_i}=t}, j,l \in {j|l_{y_j}=t}} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} + + + Parameters + ---------- + X1 : numpy array, shape (n, d) + Source dataset + X2 : numpy array, shape (n', d') + Target dataset + y1 : numpy array, shape (n,) + y2 : numpy array, shape (n',) + w1 : numpy array, shape (n,) + Weight (histogram) on the samples of X1. If None uniform distribution is considered. + w2 : Ditionary of numpy array, shape (n',) + Weight (histogram) on the samples of X2. If None uniform distribution is considered. + v1 : numpy array, shape (d,) + Weight (histogram) on the features of X1. If None uniform distribution is considered. + v2 : numpy array, shape (d',) + Weight (histogram) on the features of X2. If None uniform distribution is considered. + niter : integer + Number max of iterations of the BCD for solving COOT. + algo : string + Choice of algorithm for solving OT problems on samples each iteration. Choice ['emd','sinkhorn']. + If 'emd' returns sparse solution + If 'sinkhorn' returns regularized solution + algo2 : string + Choice of algorithm for solving OT problems on features each iteration. Choice ['emd','sinkhorn']. + If 'emd' returns sparse solution + If 'sinkhorn' returns regularized solution + reg : float + Regularization parameter for samples coupling matrix. Ignored if algo='emd' + reg2 : float + Regularization parameter for features coupling matrix. Ignored if algo='emd' + eps : float + Threshold for the convergence + random_init : bool + Wether to use random initialization for the coupling matrices. If false identity couplings are considered. + log : bool, optional + record log if True + C_lin : numpy array, shape (n, n') + Prior on the sample correspondences. Added to the cost for the samples transport + + Returns + ------- + Ts : numpy array, shape (n,n') + Optimal Transport coupling between the samples + Tv : numpy array, shape (d,d') + Optimal Transport coupling between the features + cost : float + Optimization value after convergence + log : dict + convergence information and coupling marices + References + ---------- + .. [1] Redko Ievgen, Vayer Titouan, Flamary R{\'e}mi and Courty Nicolas + "CO-Optimal Transport" + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import cotl_numpy + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} + cotl_numpy(Xs_dict, Xt_dict) + """ + assert sorted(X_dict.keys()) == sorted( + Y_dict.keys() + ), "Labels don't match in y1 & y2." + labels = list(X_dict.keys()) + if v1 is None: + X = np.concatenate([X_dict[k] for k in labels], axis=0) + if (X >= 0).all(): + v1 = X.sum(axis=0) / X.sum() + else: + v1 = np.ones(X.shape[1]) / X.shape[1] + + if v2 is None: + Y = np.concatenate([Y_dict[k] for k in labels], axis=0) + if (Y >= 0).all(): + v2 = Y.sum(axis=0) / Y.sum() + else: + v2 = np.ones(Y.shape[1]) / Y.shape[1] + + if w1 is None: + w1 = { + k: np.ones(X_dict[k].shape[0]) / X_dict[k].shape[0] for k in labels + } # is (n',) + if w2 is None: + w2 = { + k: np.ones(Y_dict[k].shape[0]) / Y_dict[k].shape[0] for k in labels + } # is (n,) + + if not random_init: + Ts = { + k: np.ones((X_dict[k].shape[0], Y_dict[k].shape[0])) + / (X_dict[k].shape[0] * Y_dict[k].shape[0]) + for k in labels + } # is (n,n') + Tv = np.ones((X_dict[labels[0]].shape[1], Y_dict[labels[0]].shape[1])) / ( + X_dict[labels[0]].shape[1] * Y_dict[labels[0]].shape[1] + ) # is (d,d') + else: + Ts = {k: random_gamma_init(w1[k], w2[k]) for k in labels} + Tv = random_gamma_init(v1, v2) + + constC_s_dict = {} + hC1_s_dict = {} + hC2_s_dict = {} + + constC_v_dict = {} + hC1_v_dict = {} + hC2_v_dict = {} + + for k in labels: + constC_s_dict[k], hC1_s_dict[k], hC2_s_dict[k] = init_matrix_np( + X_dict[k], Y_dict[k], v1, v2 + ) + constC_v_dict[k], hC1_v_dict[k], hC2_v_dict[k] = init_matrix_np( + X_dict[k].T, Y_dict[k].T, w1[k], w2[k] + ) + cost = np.inf + + log_out = {} + log_out["cost"] = [] + + for i in range(niter): + Tsold = Ts + Tvold = Tv + costold = cost + + # Sample OT for each label + for k in labels: + M_k = constC_s_dict[k] - np.dot(hC1_s_dict[k], Tv).dot(hC2_s_dict[k].T) + print(f"M_{k}:{M_k.min()} - {M_k.max()}") + if C_lin is not None: + M_k = M_k + C_lin + if algo == "emd": + Ts[k] = pot.emd(w1[k], w2[k], M_k, numItermax=1e7) + elif algo == "sinkhorn": + Ts[k] = np.array( + linear.solve( + geometry.Geometry( + cost_matrix=M_k, epsilon=reg, scale_cost="max_cost" + ), + max_iterations=2000, + ).matrix + ) + + # Global feature OT + M = 0 + for k in labels: + M += constC_v_dict[k] - np.dot(hC1_v_dict[k], Ts[k]).dot(hC2_v_dict[k].T) + print(f"M:{M.min()} - {M.max()}") + if algo2 == "emd": + Tv = pot.emd(v1, v2, M, numItermax=1e7) + elif algo2 == "sinkhorn": + Tv = np.array( + linear.solve( + geometry.Geometry( + cost_matrix=M, epsilon=reg, scale_cost="max_cost" + ), + max_iterations=2000, + ).matrix + ) + if not np.abs(Tv.sum() - 1.0) < 1e-8: + Tv = Tv / Tv.sum() + delta = sum( + [np.linalg.norm(Ts[k] - Tsold[k]) for k in labels] + ) + np.linalg.norm(Tv - Tvold) + cost = np.sum(M * Tv) + + if log: + log_out["cost"].append(cost) + + if verbose: + print(f"It {i} Delta: {delta} Loss: {cost}") + + if delta < 1e-16 or np.abs(costold - cost) < 1e-7: + if verbose: + print("converged at iter ", i) + break + if log: + return Ts, Tv, cost, log_out + else: + return Ts, Tv, cost + + +def get_coupling_cotl( + data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], +) -> Tuple[Union[int, Dict[Number, np.array]], Union[int, Dict]]: + """Returns sample coupling between two datasets X, Y given the labels, disregarding label information. + + The function solves the following optimization problem: + + .. math:: + + COOTL = \min_{Ts^1,..Ts^{L},Tv} \sum_{t=1}^L \sum_{i,k \in {i|l_{x_i}=t}, j,l \in {j|l_{y_j}=t}} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} + + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + eps: + Regularization parameter, relative to the max cost. + + Returns + ------- + T_dict : + Optimal Transport coupling between the samples per label + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_eot_ott + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} + get_coupling_eot_ott((Xs_dict, Xt_dict), 0.05) + """ + X_dict = data[0] + Y_dict = data[1] + start = time.time() + try: + Ts, Tv, cost, log = cotl_numpy(X_dict, Y_dict, log=True, niter=2000) + except FloatingPointError: + return -1, -1 + log["time"] = time.time() - start + return Ts, log + + +def get_coupling_cotl_sinkhorn( + data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], + eps: float = 5e-3, + eps2: float = None, +) -> Tuple[Dict[Number, np.array], Dict]: + """Returns sample coupling between two datasets X, Y given the labels. + + The function solves the following optimization problem: + + .. math:: + + COOTL = \min_{Ts^1,..Ts^{L},Tv} \sum_{t=1}^L \sum_{i,k \in {i|l_{x_i}=t}, j,l \in {j|l_{y_j}=t}} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} - \epsilon_1 H(Ts) -\epsilon_2 H(Tv) + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + eps: + Regularization parameter, relative to the max cost. + + Returns + ------- + T_dict : + Optimal Transport coupling between the samples per label + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_eot_ott + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} + get_coupling_eot_ott((Xs_dict, Xt_dict), 0.05) + """ + print(f"calculating with eps {eps}") + X_dict = data[0] + Y_dict = data[1] + start = time.time() + if eps2 is None: + eps2 = eps + try: + Ts, Tv, cost, log = cotl_numpy( + X_dict, + Y_dict, + algo="sinkhorn", + reg=eps, + algo2="sinkhorn", + reg2=eps2, + log=True, + niter=2000, + ) + except FloatingPointError: + return -1, -1 + log["time"] = time.time() - start + return Ts, log diff --git a/perturbot/build/lib/perturbot/match/fot.py b/perturbot/build/lib/perturbot/match/fot.py new file mode 100755 index 0000000..7f25929 --- /dev/null +++ b/perturbot/build/lib/perturbot/match/fot.py @@ -0,0 +1,220 @@ +from typing import Dict, Tuple, Union +from numbers import Number +import numpy as np +import time +import ot +from scipy import stats +from scipy.sparse import random +from ott.solvers import linear +from ott.geometry import geometry +from .utils import init_matrix_np, random_gamma_init +from perturbot.utils import mdict_to_matrix + + +def fot_numpy( + X1, + X2, + Ts, + v1=None, + v2=None, + niter=10, + algo="emd", + reg=0, + algo2="emd", + reg2=0, + verbose=True, + log=False, + random_init=False, + C_lin=None, +): + """Returns COOT between two datasets X1,X2 (see [1]) + + The function solves the following optimization problem: + .. math:: + + FOT = \min_{Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}*Tv_{k,l} + + Where : + - X1 : The source dataset + - X2 : The target dataset + - w1,w2 : weights (histograms) on the samples (rows) of resp. X1 and X2 + - v1,v2 : weights (histograms) on the features (columns) of resp. X1 and X2 + + Parameters + ---------- + X1 : numpy array, shape (n, d) + Source dataset + X2 : numpy array, shape (n', d') + Target dataset + w1 : numpy array, shape (n,) + Weight (histogram) on the samples of X1. If None uniform distribution is considered. + w2 : numpy array, shape (n',) + Weight (histogram) on the samples of X2. If None uniform distribution is considered. + v1 : numpy array, shape (d,) + Weight (histogram) on the features of X1. If None uniform distribution is considered. + v2 : numpy array, shape (d',) + Weight (histogram) on the features of X2. If None uniform distribution is considered. + niter : integer + Number max of iterations of the BCD for solving COOT. + algo : string + Choice of algorithm for solving OT problems on samples each iteration. Choice ['emd','sinkhorn']. + If 'emd' returns sparse solution + If 'sinkhorn' returns regularized solution + algo2 : string + Choice of algorithm for solving OT problems on features each iteration. Choice ['emd','sinkhorn']. + If 'emd' returns sparse solution + If 'sinkhorn' returns regularized solution + reg : float + Regularization parameter for samples coupling matrix. Ignored if algo='emd' + reg2 : float + Regularization parameter for features coupling matrix. Ignored if algo='emd' + eps : float + Threshold for the convergence + random_init : bool + Wether to use random initialization for the coupling matrices. If false identity couplings are considered. + log : bool, optional + record log if True + C_lin : numpy array, shape (n, n') + Prior on the sample correspondences. Added to the cost for the samples transport + + Returns + ------- + Ts : numpy array, shape (n,n') + Optimal Transport coupling between the samples + Tv : numpy array, shape (d,d') + Optimal Transport coupling between the features + cost : float + Optimization value after convergence + log : dict + convergence information and coupling marices + References + ---------- + .. [1] Redko Ievgen, Vayer Titouan, Flamary R{\'e}mi and Courty Nicolas + "CO-Optimal Transport" + Example + ---------- + import numpy as np + from cot import cot_numpy + + n_samples=300 + Xs=np.random.rand(n_samples,2) + Xt=np.random.rand(n_samples,1) + cot_numpy(Xs,Xt) + """ + if v1 is None: + v1 = np.ones(X1.shape[1]) / X1.shape[1] # is (d,) + if v2 is None: + v2 = np.ones(X2.shape[1]) / X2.shape[1] # is (d',) + Ts = Ts / Ts.sum() + w1 = Ts.sum(axis=0) + w2 = Ts.sum(axis=1) + if not random_init: + Tv = np.ones((X1.shape[1], X2.shape[1])) / ( + X1.shape[1] * X2.shape[1] + ) # is (d,d') + else: + Tv = random_gamma_init(v1, v2) + + constC_v, hC1_v, hC2_v = init_matrix_np(X1.T, X2.T, w1, w2) + cost = np.inf + + log_out = {} + log_out["cost"] = [] + + for i in range(niter): + Tvold = Tv + costold = cost + + M = constC_v - np.dot(hC1_v, Ts).dot(hC2_v.T) + Tv = np.array( + linear.solve( + geometry.Geometry(cost_matrix=M, epsilon=reg2, scale_cost="max_cost"), + max_iterations=2000, + ).matrix + ) + + delta = np.linalg.norm(Tv - Tvold) + cost = np.sum(M * Tv) + + if log: + log_out["cost"].append(cost) + + if verbose: + print("Delta: {0} Loss: {1}".format(delta, cost)) + + if delta < 1e-16 or np.abs(costold - cost) < 1e-7: + if verbose: + print("converged at iter ", i) + break + if log: + return Tv, cost, log_out + else: + return Tv, cost + + +def get_coupling_fot( + data: Tuple[Dict[Number, np.ndarray], Dict[Number, np.ndarray]], + Ts: Union[Dict[Number, np.ndarray], np.ndarray], + eps=5e-3, +): + r"""Returns GW coupling between features given two datasets X, Y and the sample coupling. + + The function solves the following optimization problem: + + .. math:: + + FOT = \min_{Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} - \epsilon H(T_v) + + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + Ts: + Sample-to-sample transport. + Per-label transport matched with source dataset, target dataset + or a global coupling matrix where the samples are concatenated by + the order of labels in data[0].keys(). + eps: + Regularization parameter, relative to the max cost. + + Returns + ------- + Tv : + Feature-to-feature coupling. + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_egw_labels_ott, get_coupling_fot + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} + Ts, log = get_coupling_egw_labels_ott((Xs_dict, Xt_dict), 0.05) + Tv, feature_matching_log = get_coupling_fot((Xs_dict, Xt_dict), Ts, 0.05) + """ + + X_dict = data[0] + Y_dict = data[1] + if isinstance(Ts, dict): + Ts = mdict_to_matrix( + Ts, + np.concatenate([np.ones(X_dict[l].shape[0]) * l for l in X_dict.keys()]), + np.concatenate([np.ones(Y_dict[l].shape[0]) * l for l in X_dict.keys()]), + ) + X = np.concatenate([X_dict[l] for l in X_dict.keys()]) + Y = np.concatenate([Y_dict[l] for l in X_dict.keys()]) + start = time.time() + try: + Tv, cost, log = fot_numpy(X, Y, Ts, log=True, reg=eps, reg2=eps, niter=2000) + except FloatingPointError: + return -1, -1 + log["time"] = time.time() - start + return Tv, log diff --git a/perturbot/build/lib/perturbot/match/gw.py b/perturbot/build/lib/perturbot/match/gw.py new file mode 100755 index 0000000..cbcec81 --- /dev/null +++ b/perturbot/build/lib/perturbot/match/gw.py @@ -0,0 +1,137 @@ +import numpy as np +import scipy as sp +import ot +import time + + +def gw_cg(X_dict, Y_dict): + """GW with conditional gradient algorithm""" + labels = X_dict.keys() + Ts = {} + log = {"cost_time": {}, "time": {}} + for l in labels: + start = time.time() + C1 = sp.spatial.distance.cdist(X_dict[l], X_dict[l]) + C2 = sp.spatial.distance.cdist(Y_dict[l], Y_dict[l]) + + C1 /= C1.max() + C2 /= C2.max() + log["cost_time"][l] = time.time() - start + p = ot.unif(C1.shape[0]) + q = ot.unif(C2.shape[0]) + + start = time.time() + Ts[l], log[l] = ot.gromov.gromov_wasserstein( + C1, C2, p, q, "square_loss", verbose=True, log=True + ) + log["time"][l] = time.time() - start + return Ts, log + + +def egw_pgd(X_dict, Y_dict, epsilon=5e-3): + """EGW with projected gradient descent algorithm""" + labels = X_dict.keys() + Ts = {} + log = {"cost_time": {}, "time": {}} + for l in labels: + start = time.time() + C1 = sp.spatial.distance.cdist(X_dict[l], X_dict[l]) + C2 = sp.spatial.distance.cdist(Y_dict[l], Y_dict[l]) + + C1 /= C1.max() + C2 /= C2.max() + log["cost_time"][l] = time.time() - start + p = ot.unif(C1.shape[0]) + q = ot.unif(C2.shape[0]) + start = time.time() + Ts[l], log[l] = ot.gromov.entropic_gromov_wasserstein( + C1, + C2, + p, + q, + "square_loss", + epsilon=epsilon, + solver="PGD", + log=True, + verbose=True, + ) + log["time"][l] = time.time() - start + return Ts, log + + +def gw_all(Xtot, Ytot): + start = time.time() + C1 = sp.spatial.distance.cdist(Xtot, Xtot) + C2 = sp.spatial.distance.cdist(Ytot, Ytot) + C1 /= C1.max() + C2 /= C2.max() + cost_time = time.time() - start + p = ot.unif(C1.shape[0]) + q = ot.unif(C2.shape[0]) + start = time.time() + Ts, log = ot.gromov.gromov_wasserstein( + C1, C2, p, q, "square_loss", verbose=True, log=True, max_iter=1e8 + ) + log["time"] = time.time() - start + log["cost_time"] = cost_time + return Ts, log + + +def egw_all(Xtot, Ytot, epsilon=5e-3): + start = time.time() + C1 = sp.spatial.distance.cdist(Xtot, Xtot) + C2 = sp.spatial.distance.cdist(Ytot, Ytot) + C1 /= C1.max() + C2 /= C2.max() + cost_time = time.time() - start + p = ot.unif(C1.shape[0]) + q = ot.unif(C2.shape[0]) + start = time.time() + Ts, log = ot.gromov.entropic_gromov_wasserstein( + C1, + C2, + p, + q, + "square_loss", + verbose=True, + log=True, + epsilon=epsilon, + ) + log["time"] = time.time() - start + log["cost_time"] = cost_time + return Ts, log + + +def get_coupling_gw_cg(data): + X_dict = data[0] + Y_dict = data[1] + Ts, log = gw_cg(X_dict, Y_dict) + return Ts, log + + +def get_coupling_egw(data, eps=5e-3): + X_dict = data[0] + Y_dict = data[1] + Ts, log = egw_pgd(X_dict, Y_dict, epsilon=eps) + return Ts, log + + +def get_coupling_gw_all(data): + """Run GW all-to-all, ignore labels.""" + X_dict = data[0] + Y_dict = data[1] + Xtot = np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0) + Ytot = np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0) + T, log = gw_all(Xtot, Ytot) + return T, log + + +def get_coupling_egw_all(data, eps=5e-3): + """Run GW all-to-all, ignore labels.""" + X_dict = data[0] + Y_dict = data[1] + Xtot = np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0) + Ytot = np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0) + T, log = egw_all(Xtot, Ytot, eps) + + return T, log diff --git a/perturbot/build/lib/perturbot/match/gw_labels.py b/perturbot/build/lib/perturbot/match/gw_labels.py new file mode 100755 index 0000000..5a59b92 --- /dev/null +++ b/perturbot/build/lib/perturbot/match/gw_labels.py @@ -0,0 +1,148 @@ +from typing import Tuple, Dict +from numbers import Number +import numpy as np +import scipy as sp +import ot +import time + + +def get_coupling_gw_labels( + data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], +) -> Tuple[Dict[Number, np.array], Dict]: + r"""Returns GW coupling between two datasets X, Y given the labels. + + The function solves the following optimization problem: + + .. math:: + + GWL = \min_{T\in C_{p,q}^\ell} \sum_{i,k \in \{i|l_{x_i}=t\}, j,l \in \{j|l_{y_j}=t}\} |(x_i-x_k)^2 - (y_j-y_l)^2|^{2}*T_{i,j}T_{k,l} \\ + C_{p,q}^\ell = \{T | T \in C{p,q}, T_{ij} > 0 \implies l_{x_i} = l_{y_j}\} + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + + Returns + ------- + T_dict : + Optimal Transport coupling between the samples per label + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_gw_labels + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} + get_coupling_gw_labels((Xs_dict, Xt_dict)) + """ + X_dict = data[0] + Y_dict = data[1] + Xs_tot = np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0) + Xt_tot = np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0) + source_labels = np.concatenate( + [np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()] + ) + target_labels = np.concatenate( + [np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()] + ) + start = time.time() + C1_tot = sp.spatial.distance.cdist(Xs_tot, Xs_tot) + C2_tot = sp.spatial.distance.cdist(Xt_tot, Xt_tot) + C1_tot /= C1_tot.max() + C2_tot /= C2_tot.max() + cost_time = time.time() - start + start = time.time() + T, log = ot.gromov.gromov_wasserstein_labeled( + C1_tot, C2_tot, source_labels, target_labels, log=True + ) + end = time.time() + log["time"] = end - start + log["cosst_time"] = cost_time + T_dict = {} + for l in np.unique(source_labels): + T_dict[l] = T[source_labels == l, :][:, target_labels == l] + return T_dict, log + + +def get_coupling_egw_labels(data, eps=5e-3): + """Returns GW coupling between two datasets X, Y given the labels. + + The function solves the following optimization problem: + + .. math:: + + GWL = \min_{T\in C_{p,q}^\ell} \sum_{i,k \in \{i|l_{x_i}=t\} j,l \in \{j|l_{y_j}=t}\} |(x_i-x_k)^2 - (y_j-y_l)^2|^{2}*T_{i,j}T_{k,l} - \epsilon H(T)\\ + C_{p,q}^\ell = \{T | T \in C{p,q}, T_{ij} > 0 \implies l_{x_i} = l_{y_j}\} + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + eps: + Regularization parameter, relative to the max cost. + + Returns + ------- + T_dict : + Optimal Transport coupling between the samples per label + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_egw_labels + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} + get_coupling_egw_labels((Xs_dict, Xt_dict), 0.05) + """ + X_dict = data[0] + Y_dict = data[1] + Xs_tot = np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0) + Xt_tot = np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0) + source_labels = np.concatenate( + [np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()] + ) + target_labels = np.concatenate( + [np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()] + ) + start = time.time() + C1_tot = sp.spatial.distance.cdist(Xs_tot, Xs_tot) + C2_tot = sp.spatial.distance.cdist(Xt_tot, Xt_tot) + C1_tot /= C1_tot.max() + C2_tot /= C2_tot.max() + cost_time = time.time() - start + start = time.time() + print("running LEOT") + T, log = ot.gromov.entropic_gromov_wasserstein_labeled( + C1_tot, + C2_tot, + source_labels, + target_labels, + epsilon=eps, + log=True, + verbose=True, + ) + end = time.time() + print("Done running LEOT") + log["time"] = end - start + log["cost_time"] = cost_time + T_dict = {} + for l in np.unique(source_labels): + T_dict[l] = T[source_labels == l, :][:, target_labels == l] + return T_dict, log diff --git a/perturbot/build/lib/perturbot/match/ot_labels.py b/perturbot/build/lib/perturbot/match/ot_labels.py new file mode 100755 index 0000000..4b474a1 --- /dev/null +++ b/perturbot/build/lib/perturbot/match/ot_labels.py @@ -0,0 +1,78 @@ +import numpy as np +import scipy as sp +import ot +import time + + +def get_coupling_ot_labels(data, eps=5e-3): + X_dict = data[0] + Y_dict = data[1] + Xs_tot = np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0) + Xt_tot = np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0) + source_labels = np.concatenate( + [np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()] + ) + target_labels = np.concatenate( + [np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()] + ) + start = time.time() + C = sp.spatial.distance.cdist(Xs_tot, Xs_tot) + C = C / C.max() + cost_time = time.time() - start + start = time.time() + p = np.ones(C.shape[0]) / C.shape[0] + q = np.ones(C.shape[1]) / C.shape[1] + T, log = ot.bergman.sinkhorn_labeled( + p, + q, + T, + source_labels, + target_labels, + reg=eps, + log=True, + verbose=True, + ) + end = time.time() + log["time"] = end - start + log["cost_time"] = cost_time + T_dict = {} + for l in np.unique(source_labels): + T_dict[l] = T[source_labels == l, :][:, target_labels == l] + return T_dict, log + + +def get_coupling_eot_labels(data, eps=5e-3): + X_dict = data[0] + Y_dict = data[1] + Xs_tot = np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0) + Xt_tot = np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0) + source_labels = np.concatenate( + [np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()] + ) + target_labels = np.concatenate( + [np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()] + ) + start = time.time() + C = sp.spatial.distance.cdist(Xs_tot, Xs_tot) + C = C / C.max() + cost_time = time.time() - start + start = time.time() + p = np.ones(C.shape[0]) / C.shape[0] + q = np.ones(C.shape[1]) / C.shape[1] + T, log = ot.bergman.sinkhorn_labeled( + p, + q, + T, + source_labels, + target_labels, + reg=eps, + log=True, + verbose=True, + ) + end = time.time() + log["time"] = end - start + log["cost_time"] = cost_time + T_dict = {} + for l in np.unique(source_labels): + T_dict[l] = T[source_labels == l, :][:, target_labels == l] + return T_dict, log diff --git a/perturbot/build/lib/perturbot/match/ott_egwl.py b/perturbot/build/lib/perturbot/match/ott_egwl.py new file mode 100755 index 0000000..7f9e4df --- /dev/null +++ b/perturbot/build/lib/perturbot/match/ott_egwl.py @@ -0,0 +1,450 @@ +# Modified OTT with labels +from typing import Dict, Tuple +from numbers import Number +import numpy as np +import jax.numpy as jnp +import time +import ott +from ott.geometry import pointcloud, geometry +from ott.problems.linear import linear_problem +from ott.problems.quadratic import quadratic_problem +from ott.solvers.quadratic import gromov_wasserstein +from ott.solvers import linear +from ott.solvers.linear import acceleration, sinkhorn + + +def create_block_diag_mat(labels_a, labels_b): + block_diag_mat = np.zeros((len(labels_a), len(labels_b))) + for l in np.unique(labels_a): + block_diag_mat[ + np.ix_(np.where(labels_a == l)[0], np.where(labels_b == l)[0]) + ] = 1.0 + return jnp.array(block_diag_mat) + + +def get_coupling_egw_labels_ott( + data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], eps: float = 5e-3 +) -> Tuple[Dict[Number, np.array], Dict]: + r"""Returns GW coupling between two datasets X, Y given the labels. + + The function solves the following optimization problem: + + .. math:: + + EGWL = \min_{T\in C_{p,q}^\ell} \sum_{i,k \in \{i|l_{x_i}=t\}, j,l \in \{j|l_{y_j}=t\}} |(x_i-x_k)^2 - (y_j-y_l)^2|^{2}*T_{i,j}T_{k,l} - \epsilon H(T)\\ + C_{p,q}^\ell = \{T | T \in C{p,q}, T_{ij} > 0 \implies l_{x_i} = l_{y_j}\} + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + eps: + Regularization parameter, relative to the max cost. + + Returns + ------- + T_dict : + Optimal Transport coupling between the samples per label + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_egw_labels_ott + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} + get_coupling_egw_labels_ott((Xs_dict, Xt_dict), 0.05) + """ + X_dict = data[0] + Y_dict = data[1] + Xs_tot = jnp.array(np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0)) + Xt_tot = jnp.array(np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0)) + source_labels = jnp.array( + np.concatenate([np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()]) + ) + target_labels = jnp.array( + np.concatenate([np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()]) + ) + start = time.time() + + geom_xx = pointcloud.PointCloud(x=Xs_tot, y=Xs_tot, scale_cost="max_cost") + geom_yy = pointcloud.PointCloud(x=Xt_tot, y=Xt_tot, scale_cost="max_cost") + bdm = create_block_diag_mat(source_labels, target_labels) + + cost_time = time.time() - start + start = time.time() + print("running EGWL with ott") + + prob = quadratic_problem.QuadraticProblem( + geom_xx, + geom_yy, + labels_a=source_labels, + labels_b=target_labels, + n_labels=len(np.unique(source_labels)), + block_diag_mat=bdm, + ) + + solver = gromov_wasserstein.GromovWasserstein( + epsilon=eps, + store_inner_errors=True, + max_iterations=2000, + kwargs_sinkhorn={"max_iterations": 2000}, + ) + + out = solver(prob) + + has_converged = bool(out.linear_convergence[out.n_iters - 1]) + log = {} + log["n_iters_outer"] = out.n_iters + log["converged_inner"] = has_converged + log["converged_outer"] = out.converged + log["GW cost"] = out.reg_gw_cost + T = np.array(out.matrix) + print(f"{out.n_iters} outer iterations were needed.") + print(f"The last Sinkhorn iteration has converged: {has_converged}") + print(f"The outer loop of Gromov Wasserstein has converged: {out.converged}") + print(f"The final regularized GW cost is: {out.reg_gw_cost:.3f}") + + end = time.time() + print("Done running LEGWOT with ott") + log["time"] = end - start + log["cost_time"] = cost_time + T_dict = {} + for l in np.unique(source_labels): + T_dict[l] = T[np.array(source_labels) == l, :][:, np.array(target_labels) == l] + return T_dict, log + + +def get_coupling_egw_ott( + data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], eps: float = 5e-3 +) -> Tuple[Dict[Number, np.array], Dict]: + r"""Returns GW coupling between two datasets X, Y per label. + + The function solves the following optimization problem: + + .. math:: + + GW^l = \min_{T^l} \sum_{i,j,k,l} |(x_i-x_k)^2 - (y_j-y_l)^2|^{2}*T^l_{i,j}T^l_{k,l} - \epsilon H(T^l)\\ + + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + eps: + Regularization parameter, relative to the max cost. + + Returns + ------- + T_dict : + Optimal Transport coupling between the samples per label + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_egw_ott + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} + get_coupling_egw_ott((Xs_dict, Xt_dict), 0.05) + """ + X_dict = data[0] + Y_dict = data[1] + labels = X_dict.keys() + Ts = {} + log = {} + for l in labels: + log[l] = {} + start = time.time() + + geom_xx = pointcloud.PointCloud(x=X_dict[l], y=X_dict[l], scale_cost="max_cost") + geom_yy = pointcloud.PointCloud(x=Y_dict[l], y=Y_dict[l], scale_cost="max_cost") + + cost_time = time.time() - start + start = time.time() + + prob = quadratic_problem.QuadraticProblem( + geom_xx, + geom_yy, + ) + + # Instantiate a jitt'ed Gromov-Wasserstein solver + solver = gromov_wasserstein.GromovWasserstein( + epsilon=eps, store_inner_errors=True, max_iterations=1000 + ) + + out = solver(prob) + + end = time.time() + Ts[l] = np.array(out.matrix) + + has_converged = bool(out.linear_convergence[out.n_iters - 1]) + log[l]["n_iters_outer"] = out.n_iters + log[l]["converged_inner"] = has_converged + log[l]["converged_outer"] = out.converged + log[l]["GW cost"] = out.reg_gw_cost + log[l]["time"] = end - start + log[l]["cost_time"] = cost_time + return Ts, log + + +def get_coupling_egw_all_ott( + data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], eps: float = 5e-3 +) -> Tuple[Dict[Number, np.array], Dict]: + r"""Returns GW coupling between two datasets X, Y, all-to-all manner disregarding labels. + + The function solves the following optimization problem: + + .. math:: + + GW = \min_{T\in C_{p,q}} \sum_{i,j,k,l} |(x_i-x_k)^2 - (y_j-y_l)^2|^{2}*T_{i,j}T_{k,l} - \epsilon H(T) + + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + eps: + Regularization parameter, relative to the max cost. + + Returns + ------- + T_dict : + Optimal Transport coupling between the samples per label + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_egw_all_ott + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} + get_coupling_egw_all_ott((Xs_dict, Xt_dict), 0.05) + """ + X_dict = data[0] + Y_dict = data[1] + Xs_tot = jnp.array(np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0)) + Xt_tot = jnp.array(np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0)) + source_labels = jnp.array( + np.concatenate([np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()]) + ) + target_labels = jnp.array( + np.concatenate([np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()]) + ) + start = time.time() + + geom_xx = pointcloud.PointCloud(x=Xs_tot, y=Xs_tot, scale_cost="max_cost") + geom_yy = pointcloud.PointCloud(x=Xt_tot, y=Xt_tot, scale_cost="max_cost") + + cost_time = time.time() - start + start = time.time() + print("running EGWOT with ott") + + prob = quadratic_problem.QuadraticProblem( + geom_xx, + geom_yy, + ) + + # Instantiate a jitt'ed Gromov-Wasserstein solver + solver = gromov_wasserstein.GromovWasserstein( + epsilon=eps, store_inner_errors=True, max_iterations=1000 + ) + + out = solver(prob) + + has_converged = bool(out.linear_convergence[out.n_iters - 1]) + log = {} + log["n_iters_outer"] = out.n_iters + log["converged_inner"] = has_converged + log["converged_outer"] = out.converged + log["GW cost"] = out.reg_gw_cost + T = np.array(out.matrix) + print(f"{out.n_iters} outer iterations were needed.") + print(f"The last Sinkhorn iteration has converged: {has_converged}") + print(f"The outer loop of Gromov Wasserstein has converged: {out.converged}") + print(f"The final regularized GW cost is: {out.reg_gw_cost:.3f}") + + end = time.time() + print("Done running EGWOT with ott") + log["time"] = end - start + log["cost_time"] = cost_time + return T, log + + +def get_coupling_eot_ott( + data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], eps: float = 5e-3 +) -> Tuple[Dict[Number, np.array], Dict]: + r"""Returns OT coupling between two datasets X, Y given the labels, disregarding label information. + + The function solves the following optimization problem: + + .. math:: + + EOT = \min_{T\in C_{p,q}} \sum_{i,j} (x_i-y_j)^2 T_{i,j} - \epsilon H(T)\\ + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + eps: + Regularization parameter, relative to the max cost. + + Returns + ------- + T_dict : + Optimal Transport coupling between the samples per label + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_eot_ott + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} + get_coupling_eot_ott((Xs_dict, Xt_dict), 0.05) + """ + X_dict = data[0] + Y_dict = data[1] + Xs_tot = jnp.array(np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0)) + Xt_tot = jnp.array(np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0)) + source_labels = jnp.array( + np.concatenate([np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()]) + ) + target_labels = jnp.array( + np.concatenate([np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()]) + ) + start = time.time() + + geom = pointcloud.PointCloud(x=Xs_tot, y=Xt_tot, scale_cost="max_cost") + + cost_time = time.time() - start + start = time.time() + print("running LEOT with ott") + + out = linear.solve(geometry.Geometry(cost_matrix=geom.cost_matrix, epsilon=eps)) + + log = {} + log["n_iters_outer"] = out.n_iters + log["converged"] = out.converged + log["OT cost"] = out.reg_ot_cost + T = np.array(out.matrix) + print(f"{out.n_iters} outer iterations were needed.") + print(f"The last Sinkhorn iteration has converged: {out.converged}") + print(f"The final regularized OT cost is: {out.reg_ot_cost:.3f}") + + end = time.time() + print("Done running EOT with ott") + log["time"] = end - start + log["cost_time"] = cost_time + + return T, log + + +def get_coupling_leot_ott( + data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], eps: float = 5e-3 +) -> Tuple[Dict[Number, np.array], Dict]: + r"""Returns OT coupling between two datasets X, Y per label. + + The function solves the following optimization problem: + + .. math:: + + EOT^l = \min_{T^l} \sum_{i,j} (x_i-y_j)^2 T^l_{i,j} - \epsilon H(T^l)\\ + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + eps: + Regularization parameter, relative to the max cost. + + Returns + ------- + T_dict : + Optimal Transport coupling between the samples per label + log : + Running log + + Example + ---------- + .. code-block:: python + + import numpy as np + from perturbot.match import get_coupling_leot_ott + + n_samples = 300 + labels = [0,1,2,3] + Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} + Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} + get_coupling_leot_ott((Xs_dict, Xt_dict), 0.05) + """ + X_dict = data[0] + Y_dict = data[1] + Xs_tot = jnp.array(np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0)) + Xt_tot = jnp.array(np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0)) + source_labels = jnp.array( + np.concatenate([np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()]) + ) + target_labels = jnp.array( + np.concatenate([np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()]) + ) + start = time.time() + + geom = pointcloud.PointCloud(x=Xs_tot, y=Xt_tot, scale_cost="max_cost") + geom = geometry.Geometry(cost_matrix=geom.cost_matrix, epsilon=eps) + + cost_time = time.time() - start + start = time.time() + print("running LEOT with ott") + problem = linear_problem.LinearProblem( + geom, labels_a=source_labels, labels_b=target_labels + ) + solver = sinkhorn.Sinkhorn() + out = solver(problem) + + log = {} + log["n_iters_outer"] = out.n_iters + log["converged"] = out.converged + log["OT cost"] = out.reg_ot_cost + T = np.array(out.matrix) + print(f"{out.n_iters} outer iterations were needed.") + print(f"The last Sinkhorn iteration has converged: {out.converged}") + print(f"The final regularized OT cost is: {out.reg_ot_cost:.3f}") + + end = time.time() + print("Done running LEOT with ott") + log["time"] = end - start + log["cost_time"] = cost_time + T_dict = {} + for l in np.unique(source_labels): + T_dict[l] = T[np.array(source_labels) == l, :][:, np.array(target_labels) == l] + return T_dict, log diff --git a/perturbot/build/lib/perturbot/match/utils.py b/perturbot/build/lib/perturbot/match/utils.py new file mode 100755 index 0000000..3ea2951 --- /dev/null +++ b/perturbot/build/lib/perturbot/match/utils.py @@ -0,0 +1,184 @@ +import numpy as np +from scipy import stats +from scipy.sparse import random + + +def sinkhorn_scaling( + a, + b, + K, + numItermax=1000, + stopThr=1e-9, + verbose=False, + log=False, + always_raise=False, + **kwargs, +): + a = np.asarray(a, dtype=np.float64) + b = np.asarray(b, dtype=np.float64) + K = np.asarray(K, dtype=np.float64) + + # init data + Nini = len(a) + Nfin = len(b) + + if len(b.shape) > 1: + nbb = b.shape[1] + else: + nbb = 0 + + if log: + log = {"err": []} + + # we assume that no distances are null except those of the diagonal of + # distances + if nbb: + u = np.ones((Nini, nbb)) / Nini + v = np.ones((Nfin, nbb)) / Nfin + else: + u = np.ones(Nini) / Nini + v = np.ones(Nfin) / Nfin + + # print(reg) + # print(np.min(K)) + + Kp = (1 / a).reshape(-1, 1) * K + cpt = 0 + err = 1 + while err > stopThr and cpt < numItermax: + uprev = u + vprev = v + KtransposeU = np.dot(K.T, u) + v = np.divide(b, KtransposeU) + u = 1.0 / np.dot(Kp, v) + + zero_in_transp = np.any(KtransposeU == 0) + nan_in_dual = np.any(np.isnan(u)) or np.any(np.isnan(v)) + inf_in_dual = np.any(np.isinf(u)) or np.any(np.isinf(v)) + if zero_in_transp or nan_in_dual or inf_in_dual: + # we have reached the machine precision + # come back to previous solution and quit loop + print("Warning: numerical errors at iteration in sinkhorn_scaling", cpt) + # if zero_in_transp: + # print('Zero in transp : ',KtransposeU) + # if nan_in_dual: + # print('Nan in dual') + # print('u : ',u) + # print('v : ',v) + # print('KtransposeU ',KtransposeU) + # print('K ',K) + # print('M ',M) + + # if always_raise: + # raise NanInDualError + # if inf_in_dual: + # print('Inf in dual') + u = uprev + v = vprev + + break + if cpt % 10 == 0: + # we can speed up the process by checking for the error only all + # the 10th iterations + if nbb: + err = np.sum((u - uprev) ** 2) / np.sum((u) ** 2) + np.sum( + (v - vprev) ** 2 + ) / np.sum((v) ** 2) + else: + transp = u.reshape(-1, 1) * (K * v) + err = np.linalg.norm((np.sum(transp, axis=0) - b)) ** 2 + if log: + log["err"].append(err) + + if verbose: + if cpt % 200 == 0: + print("{:5s}|{:12s}".format("It.", "Err") + "\n" + "-" * 19) + print("{:5d}|{:8e}|".format(cpt, err)) + cpt = cpt + 1 + if log: + log["u"] = u + log["v"] = v + + if nbb: # return only loss + res = np.zeros((nbb)) + for i in range(nbb): + res[i] = np.sum(u[:, i].reshape((-1, 1)) * K * v[:, i].reshape((1, -1)) * M) + if log: + return res, log + else: + return res + + else: # return OT matrix + if log: + return u.reshape((-1, 1)) * K * v.reshape((1, -1)), log + else: + return u.reshape((-1, 1)) * K * v.reshape((1, -1)) + + +def random_gamma_init(p, q, **kwargs): + """Returns random coupling matrix with marginal p,q""" + rvs = stats.beta(1e-1, 1e-1).rvs + S = random(len(p), len(q), density=1, data_rvs=rvs) + return sinkhorn_scaling(p, q, S.A, **kwargs) + + +def init_matrix_np(X1, X2, v1, v2): + """Return loss matrices and tensors for COOT fast computation + Returns the value of |X1-X2|^{2} \otimes T as done in [1] based on [2] for the Gromov-Wasserstein distance. + Where : + - X1 : The source dataset of shape (n,d) + - X2 : The target dataset of shape (n',d') + - v1 ,v2 : weights (histograms) on the columns of resp. X1 and X2 + - T : Coupling matrix of shape (n,n') + Parameters + ---------- + X1 : numpy array, shape (n, d) + Source dataset + X2 : numpy array, shape (n', d') + Target dataset + v1 : numpy array, shape (d,) + Weight (histogram) on the features of X1. + v2 : numpy array, shape (d',) + Weight (histogram) on the features of X2. + + Returns + ------- + constC : ndarray, shape (n, n') + Constant C matrix (see paragraph 1.2 of supplementary material in [1]) + hC1 : ndarray, shape (n, d) + h1(X1) matrix (see paragraph 1.2 of supplementary material in [1]) + hC2 : ndarray, shape (n', d') + h2(X2) matrix (see paragraph 1.2 of supplementary material in [1]) + References + ---------- + .. [1] Redko Ievgen, Vayer Titouan, Flamary R{\'e}mi and Courty Nicolas + "CO-Optimal Transport" + .. [2] Peyré, Gabriel, Marco Cuturi, and Justin Solomon, + "Gromov-Wasserstein averaging of kernel and distance matrices." + International Conference on Machine Learning (ICML). 2016. + """ + + def f1(a): + return a**2 + + def f2(b): + return b**2 + + def h1(a): + return a + + def h2(b): + return 2 * b + + constC1 = np.dot( + np.dot(f1(X1), v1.reshape(-1, 1)), np.ones(f1(X2).shape[0]).reshape(1, -1) + ) + constC2 = np.dot( + np.ones(f1(X1).shape[0]).reshape(-1, 1), np.dot(v2.reshape(1, -1), f2(X2).T) + ) + + constC = constC1 + constC2 + hX1 = h1(X1) + hX2 = h2(X2) + + return constC, hX1, hX2 diff --git a/perturbot/build/lib/perturbot/predict/__init__.py b/perturbot/build/lib/perturbot/predict/__init__.py new file mode 100755 index 0000000..07310cd --- /dev/null +++ b/perturbot/build/lib/perturbot/predict/__init__.py @@ -0,0 +1,3 @@ +from .mlp import train_mlp + +__all__ = ["train_mlp"] diff --git a/perturbot/build/lib/perturbot/predict/linear_regression.py b/perturbot/build/lib/perturbot/predict/linear_regression.py new file mode 100755 index 0000000..56c7fea --- /dev/null +++ b/perturbot/build/lib/perturbot/predict/linear_regression.py @@ -0,0 +1,155 @@ +import numpy as np +from typing import Dict, Union +from numpy.linalg import inv +from tqdm.auto import tqdm + + +def ols(X, Y): + """Regular OLS where we know item-to. Note that given different number of perturbation, + each data points are weighted as the same.""" + X = np.concatenate([np.ones((X.shape[0], 1)), X], axis=1) + return inv(X.T @ X) @ (X.T @ Y) + + +def ols_label(X_dict, Y_dict): + assert X_dict.keys() == Y_dict.keys() + xtx = 0 + xty = 0 + for l, X in X_dict.items(): + X = np.concatenate([np.ones((X.shape[0], 1)), X], axis=1) + xtx += X.T @ X + xty += X.T @ Y_dict[l] + """Equivalent to ols""" + return inv(xtx) @ (xty) + + +def weighted_ols( + X_dict: Dict[int, np.ndarray], + Y_dict: Dict[int, np.ndarray], + G_dict: Dict[int, np.ndarray], +): + """OLS for weighted matching between X and Y provided as G. + Return B = (\sum_l{(X^l)'Diag(G_x)(X^l)})^{-1}(\sum_l{\sum_i{(X_i^l)'(G_i)(Y^l)}}). + Each sample from the source X has the same weight. + where + l: label (keys of the dictionaries), + i: sample index in X^l, + G_i: i th row of G, + G_x: Marginal distribution of x from G + """ + assert X_dict.keys() == Y_dict.keys() + xtx = 0 + xty = 0 + for l, X in tqdm(X_dict.items()): + X = np.concatenate([np.ones((X.shape[0], 1)), X], axis=1) + Y = Y_dict[l] + G = G_dict[l] + G = G / G.sum(axis=-1)[:, None] + assert X.shape[0] == Y.shape[0] + assert X.shape[0] == G.shape[0] + assert Y.shape[0] == G.shape[1] + xtx += X.T @ np.diag(G.sum(axis=-1)) @ X + for i in range(X.shape[0]): + xty += X[[i], :].T @ (G[[i], :] @ Y) + return inv(xtx) @ xty + + +def weight_1_ols(X_dict: Dict[int, np.ndarray], Y_dict: Dict[int, np.ndarray]): + """OLS for all-to-all matching given label. + Return B = (\sum_l{(X^l)'X^l)})^{-1}(\sum_l (X^l)'Y^l). + Each sample from the source X has the same weight. + """ + assert X_dict.keys() == Y_dict.keys() + xtx = 0 + xty = 0 + for l, X in X_dict.items(): + X = np.concatenate([np.ones((X.shape[0], 1)), X], axis=1) + Y = Y_dict[l] + xtx += X.T @ X + xty += ( + X.sum(axis=0, keepdims=True).T @ Y.sum(axis=0, keepdims=True) / Y.shape[0] + ) + return inv(xtx) @ xty + + +def predict(X, params): + X = np.concatenate([np.ones((X.shape[0], 1)), X], axis=1) + return X @ params + + +def weighted_ols_normed( + X_dict: Dict[int, np.ndarray], + Y_dict: Dict[int, np.ndarray], + G_dict: Union[np.ndarray, Dict[int, np.ndarray]], +) -> np.ndarray: + """OLS for weighted matching between X and Y provided as G. + Return B = (\sum_l{(X^l)'Diag(G_x)(X^l)})^{-1}(\sum_l{\sum_i{(X_i^l)'(G_i)(Y^l)}}). + Each sample from the source X has the same weight. + where + l: label (keys of the dictionaries), + i: sample index in X^l, + G_i: i th row of G, + G_x: Marginal distribution of x from G + """ + assert X_dict.keys() == Y_dict.keys() + xtx = 0 + xty = 0 + if isinstance(G_dict, dict): + for l, X in tqdm(X_dict.items()): + X = np.concatenate([np.ones((X.shape[0], 1)), X], axis=1) + Y = Y_dict[l] + G = G_dict[l] + G = G / G.sum() + assert X.shape[0] == Y.shape[0] + assert X.shape[0] == G.shape[0] + assert Y.shape[0] == G.shape[1] + xtx += X.T @ np.diag(G.sum(axis=-1)) @ X + for i in range(X.shape[0]): + xty += X[[i], :].T @ (G[[i], :] @ Y) + else: + try: + G = G_dict / G_dict.sum() + except FloatingPointError: + G = (G_dict + 1e-6) / (G_dict.sum() + 1e-6) + labels = X_dict.keys() + X = np.concatenate( + [ + np.ones((sum([X_dict[l].shape[0] for l in labels]), 1)), + np.concatenate([X_dict[l] for l in labels], axis=0), + ], + axis=1, + ) + Y = np.concatenate([Y_dict[l] for l in labels], axis=0) + xtx += X.T @ np.diag(G.sum(axis=-1)) @ X + for i in range(X.shape[0]): + xty += X[[i], :].T @ (G[[i], :] @ Y) + xty += X[[i], :].T @ (G[[i], :] @ Y) + return inv(xtx) @ xty + + +def weight_1_ols_normed(X_dict, Y_dict, *args): + G_dict = { + k: np.ones((X_dict[k].shape[0], Y_dict[k].shape[0])) for k in X_dict.keys() + } + return weighted_ols_normed(X_dict, Y_dict, G_dict) + + +def ols_normed(X_dict, Y_dict, *args): + G_dict = {k: np.identity(X_dict[k].shape[0]) for k in X_dict.keys()} + return weighted_ols_normed(X_dict, Y_dict, G_dict) + + +def weight_conc_normed(X_dict, Y_dict, Z_dict, z_key="dosage"): + def make_G(size, label): + G = np.zeros((size, size)) + for l in np.unique(label): + l_idx = np.where(label == l)[0] + for i in l_idx: + for j in l_idx: + G[i, j] = 1 + return G + + if z_key in Z_dict: + Z_dict = Z_dict[z_key] + G_dict = {k: make_G(X_dict[k].shape[0], Z_dict[k]) for k in X_dict.keys()} + return weighted_ols_normed(X_dict, Y_dict, G_dict) diff --git a/perturbot/build/lib/perturbot/predict/mlp.py b/perturbot/build/lib/perturbot/predict/mlp.py new file mode 100755 index 0000000..e45cf42 --- /dev/null +++ b/perturbot/build/lib/perturbot/predict/mlp.py @@ -0,0 +1,250 @@ +from typing import Sequence, Dict, Union, Tuple +from numbers import Number +import sys +import time +import numpy as np +import torch + +from torch.utils.data import Dataset, DataLoader, random_split +from torch import optim, nn +import lightning as L +from lightning.pytorch.callbacks import TQDMProgressBar +from scvi.train._callbacks import LoudEarlyStopping +from scvi.train._progress import ProgressBar +from perturbot.utils import mdict_to_matrix + + +class MLP(L.LightningModule): + def __init__(self, n_input, hidden_layers: Sequence[int], n_output: int): + super().__init__() + model = nn.Sequential() + self.log_scale = nn.Parameter(torch.randn(n_output)) + prev_dim = n_input + for i, n_dim in enumerate(hidden_layers): + model.add_module(f"dense{i}", nn.Linear(prev_dim, n_dim)) + model.add_module(f"batchnorm{i}", nn.BatchNorm1d(n_dim)) + model.add_module(f"act{i}", nn.ReLU()) + prev_dim = n_dim + model.add_module(f"dense_out", nn.Linear(prev_dim, n_output)) + self.model = model + self.batch_losses = [] + + def nll(self, x, x_hat): + scale = torch.exp(self.log_scale) + mean = x_hat + dist = torch.distributions.Normal(mean, scale) + + # measure prob of seeing image under p(x|z) + log_pxz = dist.log_prob(x) + return -log_pxz.sum() + + def training_step(self, batch, batch_idx): + x = batch["source"] + y = batch["target"][:, 0, :] + y_hat = self.model(x).squeeze() + loss = nn.functional.mse_loss(y_hat, y) + # loss = self.nll(y, y_hat) + self.batch_losses.append(loss) + self.log("train_loss", loss, on_epoch=True) + return loss + + def validation_step(self, val_batch, val_batch_idx): + x = val_batch["source"] + y = val_batch["target"][:, 0, :].squeeze() + y_hat = self.model(x).squeeze() + # loss = self.nll(y, y_hat) + loss = nn.functional.mse_loss(y_hat, y) + self.log("val_loss", loss, on_epoch=True) + + def configure_optimizers(self): + optimizer = optim.Adam(self.parameters(), lr=1e-3) + return optimizer + + def on_train_epoch_end(self): + loss = torch.stack(self.batch_losses).mean() + self.log("train_loss_epoch", loss, on_step=False, on_epoch=True, prog_bar=True) + self.batch_losses.clear() + + +class LabeledDataset(Dataset): + def __init__( + self, + X_dict: Dict[Union[int, float], np.ndarray], + Y_dict: Dict[Union[int, float], np.ndarray], + ): + """ + For data (X, Y) and coupling (T) dictionaries (label -> np.ndarray), + for each X, sample Y according to the coupling probability. + Assume T has uniform marginal on X. + If T is not provided, uniform matrix is used. + """ + assert sorted(X_dict.keys()) == sorted(Y_dict.keys()) + + self.X_dict = X_dict + self.Y_dict = Y_dict + self.labels = sorted(X_dict.keys()) + self.X = np.concatenate([X_dict[l] for l in self.labels]) + self.Y = np.concatenate([Y_dict[l] for l in self.labels]) + self.label = np.concatenate( + [np.ones(X_dict[l].shape[0]) * l for l in self.labels] + ) + + def __len__(self): + return self.X.shape[0] + + def __getitem__(self, idx): + sample = {"source": self.X[idx, :], "target": self.Y[idx, :]} + return sample + + +class LabeledCouplingDataset(Dataset): + def __init__( + self, + X_dict: Dict[Union[int, float], np.ndarray], + Y_dict: Dict[Union[int, float], np.ndarray], + T_dict: Dict[Union[int, float], np.ndarray] = None, + ): + """ + For data (X, Y) and coupling (T) dictionaries (label -> np.ndarray), + for each X, sample Y according to the coupling probability. + Assume T has uniform marginal on X. + If T is not provided, uniform matrix is used. + """ + assert sorted(X_dict.keys()) == sorted(Y_dict.keys()) + + self.labels = sorted(X_dict.keys()) + self.X = torch.from_numpy(np.concatenate([X_dict[l] for l in self.labels])) + self.Y = torch.from_numpy(np.concatenate([Y_dict[l] for l in self.labels])) + + if T_dict is None: + T_dict = {} + for l in self.labels: + T = np.ones((X_dict[l].shape[0], Y_dict[l].shape[0])) + T_dict[l] = T / T.sum() + + else: + assert self.labels == sorted(T_dict.keys()) + self.T = torch.from_numpy( + mdict_to_matrix( + T_dict, + np.concatenate([np.ones(X_dict[l].shape[0]) * l for l in self.labels]), + np.concatenate([np.ones(Y_dict[l].shape[0]) * l for l in self.labels]), + ) + ) + self.T[self.T.sum(axis=-1) == 0, :] = 1e-8 + + def _sample_Y_flattened(self, x_idx): + try: + Y_idx = torch.multinomial(self.T[x_idx, :], 1) + except RuntimeError as e: + print(x_idx) + torch.set_printoptions(threshold=10_000) + print(self.T[x_idx, :]) + raise e + return self.Y[Y_idx, :] + + def __len__(self): + return self.X.shape[0] + + def __getitem__(self, idx): + sample = {"source": self.X[idx, :], "target": self._sample_Y_flattened(idx)} + return sample + + +def train( + model: L.LightningModule, + dataset: Dataset, + val_prop: float = 0.05, + batch_size: int = 256, + seed: int = 2024, + max_epochs: int = 2000, + checkpoint_dir: str = ".", + loader_kwargs={}, + log_dir="./", +): + if not torch.cuda.is_available(): + torch.set_num_threads(10) + # csv_logger = CSVLogger(save_dir=log_dir, name="csv_file") + torch_seed = torch.Generator().manual_seed(seed) + # seed_everything(seed, workers=True) + train_set_size = int(len(dataset) * (1 - val_prop)) + valid_set_size = len(dataset) - train_set_size + train_set, valid_set = random_split( + dataset, [train_set_size, valid_set_size], generator=torch_seed + ) + train_loader = DataLoader(train_set, batch_size=batch_size, **loader_kwargs) + valid_loader = DataLoader(valid_set, batch_size=batch_size, **loader_kwargs) + trainer = L.Trainer( + max_epochs=max_epochs, + accelerator="cpu", + callbacks=[ + ProgressBar(), + LoudEarlyStopping(monitor="val_loss", mode="min", patience=45), + ], # , EarlyStopping(monitor="val_loss", mode="min")], + default_root_dir=checkpoint_dir, + # logger=[csv_logger], + # deterministic=True, + ) + trainer.fit(model.double(), train_loader, valid_loader) + return model + + +def train_mlp( + train_data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], + T_dict: Dict[Number, np.array], +) -> Tuple[nn.Module, Dict]: + """Trains MLP predicting Y from X given the labeled sample-matching T_dict. + + Parameters + ---------- + data : + (source dataset, target dataset) where source and target datasets + are the dictionaries mapping label to np.ndarray with matched labels. + T_dict : + Optimal Transport coupling between the samples per label + + Returns + ------- + model : + Trained predictor + log : + Training log + """ + Xs_dict, Xt_dict = train_data + dim_X = Xs_dict[list(Xs_dict.keys())[0]].shape[1] + dim_Y = Xt_dict[list(Xt_dict.keys())[0]].shape[1] + if isinstance(T_dict, np.ndarray): + # T is all-to-all mapping, make it into a dictionary + dataset = LabeledCouplingDataset( + {0: np.concatenate([Xs_dict[l] for l in Xs_dict.keys()], axis=0)}, + {0: np.concatenate([Xt_dict[l] for l in Xs_dict.keys()], axis=0)}, + {0: T_dict}, + ) + else: + dataset = LabeledCouplingDataset(Xs_dict, Xt_dict, T_dict) + + mlp = MLP(dim_X, [min(dim_X, 256), min(dim_X, 256)], dim_Y) + start_time = time.time() + model = train(mlp, dataset).model + log = {"time": start_time - time.time()} + return model, log + + +class MyProgressBar(TQDMProgressBar): + def init_validation_tqdm(self): + bar = super().init_validation_tqdm() + if not sys.stdout.isatty(): + bar.disable = True + return bar + + def init_predict_tqdm(self): + bar = super().init_predict_tqdm() + if not sys.stdout.isatty(): + bar.disable = True + return bar + + def init_test_tqdm(self): + bar = super().init_test_tqdm() + if not sys.stdout.isatty(): + bar.disable = True + return bar diff --git a/perturbot/build/lib/perturbot/predict/scvi_vae.py b/perturbot/build/lib/perturbot/predict/scvi_vae.py new file mode 100755 index 0000000..a1bad76 --- /dev/null +++ b/perturbot/build/lib/perturbot/predict/scvi_vae.py @@ -0,0 +1,177 @@ +from typing import Dict, Union, Tuple +import numpy as np +import pandas as pd +import anndata as ad +import torch +import scvi +from tqdm import tqdm +import time + + +def make_adata(X_dict, Y_dict): + adata_X = ad.AnnData( + X=np.concatenate([v for k, v in X_dict.items()]), + obs=pd.DataFrame( + { + "modality": "source", + "labels": np.concatenate( + [np.repeat(k, v.shape[0]) for k, v in X_dict.items()] + ), + } + ), + var=pd.DataFrame( + index=[f"PC{n+1}_X" for n in range(X_dict[list(X_dict.keys())[0]].shape[1])] + ), + ) + adata_Y = ad.AnnData( + X=np.concatenate([v for k, v in Y_dict.items()]), + obs=pd.DataFrame( + { + "modality": "target", + "labels": np.concatenate( + [np.repeat(k, v.shape[0]) for k, v in Y_dict.items()] + ), + } + ), + var=pd.DataFrame( + index=[f"PC{n+1}_Y" for n in range(Y_dict[list(Y_dict.keys())[0]].shape[1])] + ), + ) + sim_multi = ad.concat([adata_X, adata_Y], axis=1)[:0, :].copy() + return sim_multi + + +def train_vae_model( + data_dict, + eps: Union[float, Tuple[float]] = None, + use_label=True, +): + """VAE with adversarial loss. Assumes normal likelihood for the observations.""" + if isinstance(eps, tuple): + eps, latent_dim, learning_rate = eps + else: + latent_dim = 50 + learning_rate = 1e-4 + if not torch.cuda.is_available(): + scvi.settings.num_threads = 10 + else: + torch.set_float32_matmul_precision("medium") + X_dict = data_dict[0] + Y_dict = data_dict[1] + dim_X = X_dict[list(X_dict.keys())[0]].shape[1] + dim_Y = Y_dict[list(Y_dict.keys())[0]].shape[1] + assert X_dict.keys() == Y_dict.keys() + adata_X = ad.AnnData( + X=np.concatenate([v for k, v in X_dict.items()]), + obs=pd.DataFrame( + { + "modality": "source", + "labels": np.concatenate( + [np.repeat(k, v.shape[0]) for k, v in X_dict.items()] + ), + } + ), + var=pd.DataFrame(index=[f"PC{n+1}_X" for n in range(dim_X)]), + ) + adata_Y = ad.AnnData( + X=np.concatenate([v for k, v in Y_dict.items()]), + obs=pd.DataFrame( + { + "modality": "target", + "labels": np.concatenate( + [np.repeat(k, v.shape[0]) for k, v in Y_dict.items()] + ), + } + ), + var=pd.DataFrame(index=[f"PC{n+1}_Y" for n in range(dim_Y)]), + ) + sim_multi = make_adata(X_dict, Y_dict) + if not use_label: + sim_multi.obs.labels = 0 + adata_mvi = scvi.data.organize_multiome_anndatas(sim_multi, adata_X, adata_Y) + scvi.model.MATCHVI.setup_anndata(adata_mvi, batch_key="modality") + model = scvi.model.MATCHVI( + adata_mvi, + n_genes=adata_X.n_vars, + n_regions=adata_Y.n_vars, + gene_likelihood="normal", + accessibility_likelihood="normal", + n_hidden=min(256, adata_X.n_vars), + n_latent=latent_dim, + ) + start_time = time.time() + model.train( + plan_kwargs={"match": True, "adversarial_classifier": True}, + scale_match_loss=0, + scale_adversarial_loss=eps, + max_epochs=2000, + lr=learning_rate, + ) + return model.module, {"time": time.time() - start_time, "history": model.history} + + +def infer_from_X(X, module, dY): + fake_acc = np.empty((X.shape[0], dY)) + fake_acc[:] = np.nan + X = torch.from_numpy(np.concatenate([X, fake_acc], axis=1)).float() + print("infer from X:") + print(X.shape) + print(module.n_input_genes, module.n_input_regions) + + inference_output = module.inference( + X, + y=torch.zeros((X.shape[0], 1), requires_grad=False), + batch_index=torch.zeros((X.shape[0], 1)), + cont_covs=None, + cat_covs=None, + label=None, + cell_idx=torch.range(0, X.shape[0] - 1).unsqueeze(-1), + ) + return inference_output + + +def infer_from_Y(Y, module, dX): + fake_gexp = np.empty((Y.shape[0], dX)) + fake_gexp[:] = np.nan + X = torch.from_numpy(np.concatenate([fake_gexp, Y], axis=1)).float() + print("infer from Y:") + print(f"input:{X.shape}={fake_gexp.shape} + {Y.shape}") + print(module.n_input_genes, module.n_input_regions) + inference_output = module.inference( + X, + y=torch.zeros((X.shape[0], 1), requires_grad=False), + batch_index=torch.zeros((X.shape[0], 1)), + cont_covs=None, + cat_covs=None, + label=None, + cell_idx=torch.range(0, X.shape[0] - 1).unsqueeze(-1), + ) + return inference_output + + +def infer_from_Xs(X_dict, model, dY): + res = {} + for k, X in X_dict.items(): + res[k] = infer_from_X(X, model, dY)["qz_m"].detach().cpu().numpy() + return res + + +def infer_from_Ys(Y_dict, model, dX): + res = {} + for k, Y in Y_dict.items(): + res[k] = infer_from_Y(Y, model, dX)["qz_m"].detach().cpu().numpy() + return res + + +def predict_from_model(X, module, dY): + # model = model.to("cpu") + inference_output = infer_from_X(X, module, dY) + generative_output = module.generative( + inference_output["z"], + inference_output["qz_m"], + batch_index=torch.ones((X.shape[0], 1)), + libsize_expr=inference_output["libsize_expr"], + libsize_acc=inference_output["libsize_acc"], + use_z_mean=True, + ) + return generative_output["p"].detach().numpy() diff --git a/perturbot/build/lib/perturbot/predict/vae.py b/perturbot/build/lib/perturbot/predict/vae.py new file mode 100755 index 0000000..aef119e --- /dev/null +++ b/perturbot/build/lib/perturbot/predict/vae.py @@ -0,0 +1,493 @@ +from typing import Sequence, Dict, Union, Tuple +import collections.abc +import numpy as np + +import torch +from torch import nn +from torch.nn.functional import one_hot +from torch.nn import CrossEntropyLoss +import lightning as L +from scvi.nn import FCLayers # , one_hot +from scvi.module import Classifier +from torch.utils.data import Dataset +import warnings + +warnings.filterwarnings( + "ignore", ".*Trying to infer the `batch_size` from an ambiguous collection.*" +) + + +class AbstractVAE(L.LightningModule): + # https://github.com/williamFalcon/pytorch-lightning-vae/blob/main/vae.py + def __init__( + self, + ): + super().__init__() + + self.save_hyperparameters() + + # for the gaussian likelihood + self.log_scale = nn.Parameter(torch.Tensor([0.0])) + + def configure_optimizers(self): + return torch.optim.Adam(self.parameters(), lr=1e-4) + + def gaussian_likelihood(self, x_hat, logscale, x): + scale = torch.exp(logscale) + mean = x_hat + dist = torch.distributions.Normal(mean, scale) + + # measure prob of seeing image under p(x|z) + log_pxz = dist.log_prob(x) + return log_pxz.sum(dim=1) + + def kl_divergence(self, z, mu, std): + # -------------------------- + # Monte carlo KL divergence + # -------------------------- + # 1. define the first two probabilities (in this case Normal for both) + p = torch.distributions.Normal(torch.zeros_like(mu), torch.ones_like(std)) + q = torch.distributions.Normal(mu, std) + + # 2. get the probabilities from the equation + log_qzx = q.log_prob(z) + log_pz = p.log_prob(z) + + # kl + kl = log_qzx - log_pz + kl = kl.sum(-1) + return kl + + def training_step(self, batch, batch_idx): + x, _ = batch + + # encode x to get the mu and variance parameters + x_encoded = self.encoder(x) + mu, log_var = self.fc_mu(x_encoded), self.fc_var(x_encoded) + + # sample z from q + std = torch.exp(log_var / 2) + q = torch.distributions.Normal(mu, std) + z = q.rsample() + + # decoded + x_hat = self.decoder(z) + + # reconstruction loss + recon_loss = self.gaussian_likelihood(x_hat, self.log_scale, x) + + # kl + kl = self.kl_divergence(z, mu, std) + + # elbo + elbo = kl - recon_loss + elbo = elbo.mean() + + self.log_dict( + { + "elbo": elbo, + "kl": kl.mean(), + "recon_loss": recon_loss.mean(), + "reconstruction": recon_loss.mean(), + "kl": kl.mean(), + } + ) + + return elbo + + +class MatchVAE(AbstractVAE): + def __init__( + self, + n_input1, + n_input2, + n_hidden, + n_latent, + n_layers, + n_labels, + mod2_loss_scale=None, + adv_loss_scale=1.0, + weight_decay=1e-6, + ): + super().__init__() + self.automatic_optimization = False + self.save_hyperparameters() + self.log_scale1 = nn.Parameter(torch.zeros(n_input1)) + self.log_scale2 = nn.Parameter(torch.zeros(n_input2)) + # encoder, decoder + self.encoder1 = FCLayers( + n_in=n_input1, + n_out=n_latent, + n_layers=n_layers, + n_hidden=n_hidden, + ) + self.decoder1 = FCLayers( + n_in=n_latent, + n_out=n_input1, + n_layers=n_layers, + n_hidden=n_hidden, + ) + self.encoder2 = FCLayers( + n_in=n_input2, + n_out=n_latent, + n_layers=n_layers, + n_hidden=n_hidden, + ) + self.decoder2 = FCLayers( + n_in=n_latent, + n_out=n_input2, + n_layers=n_layers, + n_hidden=n_hidden, + ) + + # distribution parameters + self.fc_mu1 = nn.Linear(n_latent, n_latent) + self.fc_var1 = nn.Linear(n_latent, n_latent) + self.fc_mu2 = nn.Linear(n_latent, n_latent) + self.fc_var2 = nn.Linear(n_latent, n_latent) + + self.adversarial_classifier = Classifier( + n_input=n_latent + n_labels, + n_hidden=32, + n_labels=n_labels, + n_layers=2, + logits=True, + ) + if mod2_loss_scale is None: + mod2_loss_scale = n_input1 / n_input2 + self.mod2_loss_scale = mod2_loss_scale + self.adv_loss_scale = adv_loss_scale + self.weight_decay = weight_decay + self.n_labels = n_labels + + def encode(self, x1, x2): + z1 = self.fc_mu1(self.encoder1(x1)) + z2 = self.fc_mu2(self.encoder2(x2)) + return z1, z2 + + def configure_optimizers(self): + """Configure optimizers for adversarial training.""" + # params1 = filter(lambda p: p.requires_grad, self.parameters()) + generators = [ + self.encoder1.parameters(), + self.decoder1.parameters(), + self.encoder2.parameters(), + self.decoder2.parameters(), + self.fc_mu1.parameters(), + self.fc_var1.parameters(), + self.fc_mu2.parameters(), + self.fc_var2.parameters(), + ] + param1 = [] + for g in generators: + param1 += list(g) + param1 + [self.log_scale1, self.log_scale2] + optimizer1 = torch.optim.Adam( + param1, + lr=1e-3, + eps=0.01, + weight_decay=self.weight_decay, + ) + config1 = {"optimizer": optimizer1} + + params2 = filter( + lambda p: p.requires_grad, self.adversarial_classifier.parameters() + ) + optimizer2 = torch.optim.Adam( + params2, lr=1e-3, eps=0.01, weight_decay=self.weight_decay + ) + config2 = {"optimizer": optimizer2} + + # pytorch lightning requires this way to return + opts = [config1.pop("optimizer"), config2["optimizer"]] + if "lr_scheduler" in config1: + scheds = [config1["lr_scheduler"]] + return opts, scheds + else: + return opts + + def loss_adversarial_classifier( + self, z, batch_index, label_index, predict_true_class=True + ): + """Loss for adversarial classifier.""" + n_classes = self.n_labels + onehot_labels = one_hot(label_index, n_classes) + logits = self.adversarial_classifier(torch.concat([z, onehot_labels], axis=-1)) + + if predict_true_class: + cls_target = one_hot(batch_index, n_classes).float() + else: + one_hot_batch = one_hot(batch_index, n_classes) + # place zeroes where true label is + cls_target = (~one_hot_batch.bool()).float() + cls_target = cls_target / (n_classes - 1) + + ce_loss = CrossEntropyLoss() + loss = ce_loss(logits, cls_target) + return loss + + def training_step(self, batch, batch_idx): + opt1, opt2 = self.optimizers() + x1, x2, x1_label, x2_label = batch + + # encode x to get the mu and variance parameters + z1 = self.encoder1(x1) + mu1, log_var1 = self.fc_mu1(z1), self.fc_var1(z1) + + # sample z from q + std1 = torch.exp(log_var1 / 2) + q1 = torch.distributions.Normal(mu1, std1) + z1 = q1.rsample() + + # decoded + x_hat1 = self.decoder1(z1) + + # reconstruction loss + recon_loss1 = self.gaussian_likelihood(x_hat1, self.log_scale1, x1) + + # kl + kl1 = self.kl_divergence(z1, mu1, std1) + + # elbo + elbo1 = kl1 - recon_loss1 + elbo1 = elbo1.sum() / (x1.shape[0] + x2.shape[0]) + + # encode x to get the mu and variance parameters + z2 = self.encoder1(x2) + mu2, log_var2 = self.fc_mu1(z2), self.fc_var1(z2) + + # sample z from q + std2 = torch.exp(log_var2 / 2) + q2 = torch.distributions.Normal(mu2, std2) + z2 = q2.rsample() + + # decoded + x_hat2 = self.decoder1(z2) + + # reconstruction loss + recon_loss2 = self.gaussian_likelihood(x_hat2, self.log_scale2, x2) + + # kl + kl2 = self.kl_divergence(z2, mu2, std2) + + # elbo + elbo2 = kl2 - recon_loss2 + elbo2 = elbo2.sum() / (x1.shape[0] + x2.shape[0]) + + z = torch.concat([z1, z2], axis=0) + modality = ( + torch.concat([torch.zeros(z1.shape[0]), torch.zeros(z2.shape[0])], axis=0) + .long() + .to(z.device) + ) + label = torch.concat([x1_label, x2_label], axis=0) + fool_loss = self.loss_adversarial_classifier(z, modality, label, False) + + loss = elbo1 + elbo2 * self.mod2_loss_scale + fool_loss * self.adv_loss_scale + opt1.zero_grad() + self.manual_backward(loss) + opt1.step() + + adv_classifier_loss = ( + self.loss_adversarial_classifier(z.detach(), modality, label, True) + * self.adv_loss_scale + ) + opt2.zero_grad() + self.manual_backward(adv_classifier_loss) + opt2.step() + + self.log_dict( + { + "elbo1": elbo1, + "kl1": kl1.mean(), + "recon_loss1": recon_loss1.mean(), + "kl1": kl1.mean(), + "elbo2": elbo2, + "kl2": kl2.mean(), + "recon_loss2": recon_loss2.mean(), + "kl2": kl2.mean(), + "adv_loss": fool_loss, + "train_loss": loss, + "train_adv_classifier_loss": adv_classifier_loss, + "batch_size": x1.shape[0] + x2.shape[0], + } + ) + # self.log(batch_size=x1.shape[0] + x2.shape[0]) + + def validation_step(self, batch, batch_idx): + opt1, opt2 = self.optimizers() + x1, x2, x1_label, x2_label = batch + + # encode x to get the mu and variance parameters + z1 = self.encoder1(x1) + mu1, log_var1 = self.fc_mu1(z1), self.fc_var1(z1) + + # sample z from q + std1 = torch.exp(log_var1 / 2) + q1 = torch.distributions.Normal(mu1, std1) + z1 = q1.rsample() + + # decoded + x_hat1 = self.decoder1(z1) + + # reconstruction loss + recon_loss1 = self.gaussian_likelihood(x_hat1, self.log_scale1, x1) + + # kl + kl1 = self.kl_divergence(z1, mu1, std1) + + # elbo + elbo1 = kl1 - recon_loss1 + elbo1 = elbo1.sum() / (x1.shape[0] + x2.shape[0]) + + # encode x to get the mu and variance parameters + z2 = self.encoder1(x2) + mu2, log_var2 = self.fc_mu1(z2), self.fc_var1(z2) + + # sample z from q + std2 = torch.exp(log_var2 / 2) + q2 = torch.distributions.Normal(mu2, std2) + z2 = q2.rsample() + + # decoded + x_hat2 = self.decoder1(z2) + + # reconstruction loss + recon_loss2 = self.gaussian_likelihood(x_hat2, self.log_scale2, x2) + + # kl + kl2 = self.kl_divergence(z2, mu2, std2) + + # elbo + elbo2 = kl2 - recon_loss2 + elbo2 = elbo2.sum() / (x1.shape[0] + x2.shape[0]) + + z = torch.concat([z1, z2], axis=0) + modality = ( + torch.concat([torch.zeros(z1.shape[0]), torch.zeros(z2.shape[0])], axis=0) + .long() + .to(z.device) + ) + label = torch.concat([x1_label, x2_label], axis=0) + fool_loss = self.loss_adversarial_classifier(z, modality, label, False) + + loss = elbo1 + elbo2 * self.mod2_loss_scale + fool_loss * self.adv_loss_scale + + adv_classifier_loss = ( + self.loss_adversarial_classifier(z.detach(), modality, label, True) + * self.adv_loss_scale + ) + self.log_dict( + { + "val_elbo1": elbo1, + "val_kl1": kl1.mean(), + "val_recon_loss1": recon_loss1.mean(), + "val_kl1": kl1.mean(), + "val_elbo2": elbo2, + "val_kl2": kl2.mean(), + "val_recon_loss2": recon_loss2.mean(), + "val_kl2": kl2.mean(), + "val_adv_loss": fool_loss, + "val_loss": loss, + "val_train_adv_classifier_loss": adv_classifier_loss, + } + ) + # self.log(batch_size=x1.shape[0] + x2.shape[0]) + + def transfer_batch_to_device(self, batch, device, dataloader_idx): + if isinstance(batch, Tuple): + # move all tensors in your custom data structure to the device + return tuple(item.to(device) for item in batch) + batch.samples = batch.samples.to(device) + batch.targets = batch.targets.to(device) + elif dataloader_idx == 0: + # skip device transfer for the first dataloader or anything you wish + pass + else: + print(type(batch)) + batch = super().transfer_batch_to_device(batch, device, dataloader_idx) + return batch + + +class LabeledBimodalDataset(Dataset): + def __init__( + self, + X_dict: Dict[Union[int, float], np.ndarray], + Y_dict: Dict[Union[int, float], np.ndarray], + **kwargs, + ): + """ + For data (X, Y) and coupling (T) dictionaries (label -> np.ndarray), + for each X, sample Y according to the coupling probability. + Assume T has uniform marginal on X. + If T is not provided, uniform matrix is used. + """ + super().__init__() + assert sorted(X_dict.keys()) == sorted(Y_dict.keys()) + self.labels = sorted(X_dict.keys()) + self.labels_to_id = {i: l for i, l in enumerate(self.labels)} + X = np.zeros((0, X_dict[self.labels[0]].shape[1])) + Y = np.zeros((0, Y_dict[self.labels[0]].shape[1])) + modality = np.zeros((0)) + labels_X = np.zeros((0)) + labels_Y = np.zeros((0)) + for l in self.labels: + lid = self.labels_to_id[l] + X = np.concatenate([X, X_dict[l]]) + Y = np.concatenate([Y, Y_dict[l]]) + modality = np.concatenate( + [ + modality, + np.zeros(X_dict[l].shape[0], dtype=int), + np.zeros(Y_dict[l].shape[0], dtype=int), + ] + ) + labels_X = np.concatenate( + [labels_X, np.ones(X_dict[l].shape[0], dtype=int) * lid] + ) + labels_Y = np.concatenate( + [labels_Y, np.ones(X_dict[l].shape[0], dtype=int) * lid] + ) + self.X = X + self.Y = Y + self.modality = modality + self.labels_X = labels_X + self.labels_Y = labels_Y + self.data_idx = np.arange(0, self.X.shape[0] + self.Y.shape[0]) + np.random.shuffle(self.data_idx) + self.idx_order = np.argsort(self.data_idx) + # data_idx = 3 4 0 | 1 2 + # idx_order= 2 3 4 | 0 1 + # idx = 0, 1, 2 + # idx_order[idx] = [2, 3, 4] + + def __len__(self): + return self.X.shape[0] + + def __getitem__(self, idx): + X_idx = self.idx_order[idx][self.idx_order[idx] < self.X.shape[0]] + Y_idx = ( + self.idx_order[idx][self.idx_order[idx] >= self.X.shape[0]] + - self.X.shape[0] + ) + X = self.X[X_idx, :] + Y = self.Y[Y_idx, :] + labels_X = self.labels_X[X_idx] + labels_Y = self.labels_Y[Y_idx] + return X, Y, labels_X, labels_Y + + +def collate_fn(data): + """ + data: is a list of tuples with (example, label, length) + where 'example' is a tensor of arbitrary shape + and label/length are scalars + """ + x, y, labels_x, labels_y = zip(*data) + + return ( + torch.from_numpy(np.vstack(x)), + torch.from_numpy(np.vstack(y)), + torch.from_numpy(np.concatenate(labels_x)).long(), + torch.from_numpy(np.concatenate(labels_y)).long(), + ) diff --git a/perturbot/build/lib/perturbot/preprocess/__init__.py b/perturbot/build/lib/perturbot/preprocess/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/perturbot/build/lib/perturbot/preprocess/vae.py b/perturbot/build/lib/perturbot/preprocess/vae.py new file mode 100755 index 0000000..f5da651 --- /dev/null +++ b/perturbot/build/lib/perturbot/preprocess/vae.py @@ -0,0 +1,65 @@ +"""Preprocess raw data to learn the latent embedding and projections.""" + +import os +import anndata as ad +import scanpy as sc +import scvi +import mudata as md + +SCVI_LATENT_KEY = "X_scVI" + + +def train_vae_rna(adata: ad.AnnData, save_dir="./"): + sc.pp.filter_genes(adata, min_counts=3) + adata.layers["counts"] = adata.X.copy() # preserve counts + sc.pp.normalize_total(adata, target_sum=1e4) + sc.pp.log1p(adata) + adata.raw = adata # freeze the state in `.raw` + sc.pp.highly_variable_genes( + adata, + n_top_genes=1200, + subset=True, + layer="counts", + flavor="seurat_v3", + batch_key="cell_source", + ) + model = scvi.model.SCVI(adata, n_latent=50) + model.train() + model_dir = os.path.join(save_dir, "scvi_model") + model.save(model_dir, overwrite=True) + + latent = model.get_latent_representation() + adata.obsm[SCVI_LATENT_KEY] = latent + return adata, model + + +def train_vae_acc(adata: ad.AnnData, save_dir="./"): + # compute the threshold: 5% of the cells + min_cells = int(adata.shape[0] * 0.05) + # in-place filtering of regions + sc.pp.filter_genes(adata, min_cells=min_cells) + scvi.model.PEAKVI.setup_anndata(adata) + model = scvi.model.PEAKVI(adata, n_hidden=50) + model.train() + model_dir = os.path.join(save_dir.name, "peakvi_pbmc") + model.save(model_dir, overwrite=True) + latent = model.get_latent_representation() + adata.obsm[SCVI_LATENT_KEY] = latent + return adata, model + + +def train_vae_prot(adata: ad.AnnData, save_dir="./"): + mdata = md.MuData({"rna": adata[:, :0].copy(), "protein": adata}) + scvi.model.TOTALVI.setup_mudata( + mdata, + rna_layer=None, + protein_layer=None, + modalities={ + "rna_layer": "rna", + "protein_layer": "protein", + }, + ) + model = scvi.model.TOTALVI(mdata, n_latent=50) + model.train() + adata.obsm[SCVI_LATENT_KEY] = model.get_latent_representation() + return adata, model diff --git a/perturbot/build/lib/perturbot/utils.py b/perturbot/build/lib/perturbot/utils.py new file mode 100644 index 0000000..e8abb3e --- /dev/null +++ b/perturbot/build/lib/perturbot/utils.py @@ -0,0 +1,10 @@ +import numpy as np + + +def mdict_to_matrix(M_dict, source_labels, target_labels): + Mtot = np.zeros((len(source_labels), len(target_labels))) + for l, M in M_dict.items(): + Mtot[ + np.ix_(np.where(source_labels == l)[0], np.where(target_labels == l)[0]) + ] = M + return Mtot diff --git a/perturbot/perturbot.egg-info/PKG-INFO b/perturbot/perturbot.egg-info/PKG-INFO new file mode 100644 index 0000000..629638f --- /dev/null +++ b/perturbot/perturbot.egg-info/PKG-INFO @@ -0,0 +1,17 @@ +Metadata-Version: 2.4 +Name: perturbot +Version: 0.0.1 +Summary: OT for aligning multimodal perturbation +Author-email: Jayoung Ryu +Project-URL: Homepage, https://github.com/Genentech/Perturb-OT +Project-URL: Issues, https://github.com/Genentech/Perturb-OT/issues +Classifier: Programming Language :: Python :: 3 +Classifier: License :: OSI Approved :: MIT License +Classifier: Operating System :: OS Independent +Requires-Python: >=3.8 +Description-Content-Type: text/markdown +Requires-Dist: scanpy>=1.6 +Requires-Dist: POT>=0.9 +Requires-Dist: lightning +Requires-Dist: ott-jax==0.4.6a +Requires-Dist: scvi-tools==1.1.0a diff --git a/perturbot/perturbot.egg-info/SOURCES.txt b/perturbot/perturbot.egg-info/SOURCES.txt new file mode 100644 index 0000000..1c08e58 --- /dev/null +++ b/perturbot/perturbot.egg-info/SOURCES.txt @@ -0,0 +1,8 @@ +pyproject.toml +perturbot/__init__.py +perturbot/utils.py +perturbot.egg-info/PKG-INFO +perturbot.egg-info/SOURCES.txt +perturbot.egg-info/dependency_links.txt +perturbot.egg-info/requires.txt +perturbot.egg-info/top_level.txt \ No newline at end of file diff --git a/perturbot/perturbot.egg-info/dependency_links.txt b/perturbot/perturbot.egg-info/dependency_links.txt new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/perturbot/perturbot.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/perturbot/perturbot.egg-info/requires.txt b/perturbot/perturbot.egg-info/requires.txt new file mode 100644 index 0000000..6d3a9eb --- /dev/null +++ b/perturbot/perturbot.egg-info/requires.txt @@ -0,0 +1,5 @@ +scanpy>=1.6 +POT>=0.9 +lightning +ott-jax==0.4.6a +scvi-tools==1.1.0a diff --git a/perturbot/perturbot.egg-info/top_level.txt b/perturbot/perturbot.egg-info/top_level.txt new file mode 100644 index 0000000..97cc772 --- /dev/null +++ b/perturbot/perturbot.egg-info/top_level.txt @@ -0,0 +1 @@ +perturbot diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..a499371 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,7 @@ +[project] +name = "perturb-ot" +version = "0.1.0" +description = "Add your description here" +readme = "README.md" +requires-python = ">=3.12" +dependencies = [] From 633c215ef7a4734852c987be12e836babffc930d Mon Sep 17 00:00:00 2001 From: raphael_rubrice <119074180+raphaelrubrice@users.noreply.github.com> Date: Sat, 29 Nov 2025 16:24:10 +0100 Subject: [PATCH 2/7] Update README.md --- README.md | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 27e8fc7..fc18086 100755 --- a/README.md +++ b/README.md @@ -23,7 +23,7 @@ Cross-modality matching and prediction of **perturb**ation response with labeled ``` ## Installation -`perturbot/` uses the modified `scvi-tools` and `ott` submodules which can be installed with `pip install`. +1. `perturbot/` uses the modified `scvi-tools` and `ott` submodules which can be installed with `pip install`. ```bash cd scvi-tools/ pip install . @@ -33,6 +33,36 @@ cd ../perturbot pip install . ``` +2. At that point `perturbot` is almost empty (only __init__ and utils), to fix that, go into Perturb-OT/perturbot/perturbot, copy all submodules and paste them inside the Perturb-OT/perturbot/build/lib/perturbot folder. +```bash +cp -r perturbot/perturbot/. perturbot/build/lib/perturbot/ +``` + +3. Now uninstall and reininstall `perturbot` specifically: +```bash +pip uninstall perturbot +cd perturbot/ +pip install . +``` + +At this point perturbot.match should work when imported. However you'll likely hit dependency issues regarding jax and anndata: + +### Fixing `jaxlib.xla_extension` error +1. What's happening is some good old mismatch between the code in some depedencies and the updated API of JAX 0.8.0 +```bash +pip install "jax[cpu]==0.4.36" "jax[cuda]==0.4.36" "jaxlib==0.4.36" +``` + +That should solve it. + +### Fixing `anndata` errors +1. The importing dynamics of anndata slightly changed in recent versions which makes it no longer possible to import read at the very top level so to keep dependency code as is we need to downgrade to an anndata version that supported it: +```bash +pip install "anndata==0.10.9" +``` + +That should be it ! Both perturbot.match and perturbot.predict should be importable now :D ! + ## Usage ```python import numpy as np From 9b4aaa584736f671f27bb29617d4d94c5b9f8d58 Mon Sep 17 00:00:00 2001 From: raphael_rubrice <119074180+raphaelrubrice@users.noreply.github.com> Date: Sat, 29 Nov 2025 16:36:11 +0100 Subject: [PATCH 3/7] Update README.md by describing the patch applied --- README.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index fc18086..ee167f2 100755 --- a/README.md +++ b/README.md @@ -23,7 +23,7 @@ Cross-modality matching and prediction of **perturb**ation response with labeled ``` ## Installation -1. `perturbot/` uses the modified `scvi-tools` and `ott` submodules which can be installed with `pip install`. +`perturbot/` uses the modified `scvi-tools` and `ott` submodules which can be installed with `pip install`. ```bash cd scvi-tools/ pip install . @@ -32,6 +32,8 @@ pip install . cd ../perturbot pip install . ``` +## Patch applied +1. Execute the commands in the Installation section of the repo 2. At that point `perturbot` is almost empty (only __init__ and utils), to fix that, go into Perturb-OT/perturbot/perturbot, copy all submodules and paste them inside the Perturb-OT/perturbot/build/lib/perturbot folder. ```bash @@ -45,7 +47,7 @@ cd perturbot/ pip install . ``` -At this point perturbot.match should work when imported. However you'll likely hit dependency issues regarding jax and anndata: +At this point `perturbot.match` **should work** when imported. **However** you'll likely hit dependency issues regarding **jax and anndata**: ### Fixing `jaxlib.xla_extension` error 1. What's happening is some good old mismatch between the code in some depedencies and the updated API of JAX 0.8.0 @@ -61,7 +63,7 @@ That should solve it. pip install "anndata==0.10.9" ``` -That should be it ! Both perturbot.match and perturbot.predict should be importable now :D ! +**That should be it ! Both perturbot.match and perturbot.predict should be importable now :D !** ## Usage ```python From dc6456194b3d2744d190ab9c9c934b7a938cca80 Mon Sep 17 00:00:00 2001 From: Aelrach Date: Sat, 29 Nov 2025 17:16:59 +0100 Subject: [PATCH 4/7] Added gitignore --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..fcd5172 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +.python-version +*build/ From 91e50a798807b967cdc69d419e94f24246c9a447 Mon Sep 17 00:00:00 2001 From: Aelrach Date: Sat, 29 Nov 2025 17:17:57 +0100 Subject: [PATCH 5/7] removed gitignore --- .gitignore | 2 -- 1 file changed, 2 deletions(-) delete mode 100644 .gitignore diff --git a/.gitignore b/.gitignore deleted file mode 100644 index fcd5172..0000000 --- a/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -.python-version -*build/ From 4ca9eef1d2fb6cb0219999ac2a9813d20f9f70d5 Mon Sep 17 00:00:00 2001 From: Aelrach Date: Sat, 29 Nov 2025 17:27:26 +0100 Subject: [PATCH 6/7] Removed the build content for ott and perturbot for a light and maintenance friendly PR --- ott/build/lib/ott/__init__.py | 33 - ott/build/lib/ott/_version.py | 21 - ott/build/lib/ott/datasets.py | 154 --- ott/build/lib/ott/geometry/__init__.py | 24 - ott/build/lib/ott/geometry/costs.py | 1141 --------------- ott/build/lib/ott/geometry/distrib_costs.py | 89 -- .../lib/ott/geometry/epsilon_scheduler.py | 102 -- ott/build/lib/ott/geometry/geodesic.py | 277 ---- ott/build/lib/ott/geometry/geometry.py | 921 ------------ ott/build/lib/ott/geometry/graph.py | 272 ---- ott/build/lib/ott/geometry/grid.py | 415 ------ ott/build/lib/ott/geometry/low_rank.py | 508 ------- ott/build/lib/ott/geometry/pointcloud.py | 792 ----------- ott/build/lib/ott/geometry/segment.py | 188 --- ott/build/lib/ott/initializers/__init__.py | 14 - .../lib/ott/initializers/linear/__init__.py | 14 - .../ott/initializers/linear/initializers.py | 408 ------ .../initializers/linear/initializers_lr.py | 654 --------- .../ott/initializers/quadratic/__init__.py | 14 - .../initializers/quadratic/initializers.py | 193 --- ott/build/lib/ott/math/__init__.py | 19 - ott/build/lib/ott/math/fixed_point_loop.py | 239 ---- ott/build/lib/ott/math/matrix_square_root.py | 337 ----- .../lib/ott/math/unbalanced_functions.py | 90 -- ott/build/lib/ott/math/utils.py | 298 ---- ott/build/lib/ott/neural/__init__.py | 14 - ott/build/lib/ott/neural/layers.py | 213 --- ott/build/lib/ott/neural/losses.py | 309 ----- ott/build/lib/ott/neural/models.py | 408 ------ ott/build/lib/ott/neural/solvers/__init__.py | 14 - ott/build/lib/ott/neural/solvers/conjugate.py | 121 -- .../lib/ott/neural/solvers/map_estimator.py | 289 ---- .../lib/ott/neural/solvers/neuraldual.py | 700 ---------- ott/build/lib/ott/problems/__init__.py | 14 - ott/build/lib/ott/problems/linear/__init__.py | 14 - .../ott/problems/linear/barycenter_problem.py | 200 --- .../lib/ott/problems/linear/linear_problem.py | 132 -- .../lib/ott/problems/linear/potentials.py | 437 ------ .../lib/ott/problems/quadratic/__init__.py | 14 - .../ott/problems/quadratic/gw_barycenter.py | 325 ----- .../ott/problems/quadratic/quadratic_costs.py | 78 -- .../problems/quadratic/quadratic_problem.py | 588 -------- ott/build/lib/ott/py.typed | 0 ott/build/lib/ott/solvers/__init__.py | 14 - ott/build/lib/ott/solvers/linear/__init__.py | 30 - ott/build/lib/ott/solvers/linear/_solve.py | 60 - .../lib/ott/solvers/linear/acceleration.py | 172 --- .../solvers/linear/continuous_barycenter.py | 233 ---- .../ott/solvers/linear/discrete_barycenter.py | 251 ---- .../linear/implicit_differentiation.py | 324 ----- .../lib/ott/solvers/linear/lineax_implicit.py | 98 -- ott/build/lib/ott/solvers/linear/lr_utils.py | 371 ----- ott/build/lib/ott/solvers/linear/sinkhorn.py | 1230 ----------------- .../lib/ott/solvers/linear/sinkhorn_lr.py | 822 ----------- .../lib/ott/solvers/linear/univariate.py | 346 ----- .../lib/ott/solvers/quadratic/__init__.py | 20 - ott/build/lib/ott/solvers/quadratic/_solve.py | 91 -- .../solvers/quadratic/gromov_wasserstein.py | 461 ------ .../quadratic/gromov_wasserstein_lr.py | 897 ------------ .../ott/solvers/quadratic/gw_barycenter.py | 339 ----- .../lib/ott/solvers/quadratic/lower_bound.py | 96 -- ott/build/lib/ott/solvers/was_solver.py | 131 -- ott/build/lib/ott/tools/__init__.py | 21 - .../ott/tools/gaussian_mixture/__init__.py | 14 - .../lib/ott/tools/gaussian_mixture/fit_gmm.py | 303 ---- .../tools/gaussian_mixture/fit_gmm_pair.py | 393 ------ .../ott/tools/gaussian_mixture/gaussian.py | 217 --- .../gaussian_mixture/gaussian_mixture.py | 329 ----- .../gaussian_mixture/gaussian_mixture_pair.py | 222 --- .../lib/ott/tools/gaussian_mixture/linalg.py | 138 -- .../tools/gaussian_mixture/probabilities.py | 100 -- .../ott/tools/gaussian_mixture/scale_tril.py | 214 --- ott/build/lib/ott/tools/k_means.py | 412 ------ ott/build/lib/ott/tools/plot.py | 247 ---- ott/build/lib/ott/tools/segment_sinkhorn.py | 143 -- .../lib/ott/tools/sinkhorn_divergence.py | 354 ----- ott/build/lib/ott/tools/soft_sort.py | 706 ---------- ott/build/lib/ott/types.py | 42 - ott/build/lib/ott/utils.py | 210 --- perturbot/build/lib/perturbot/__init__.py | 0 .../build/lib/perturbot/eval/.run_loo.py | 126 -- .../build/lib/perturbot/eval/__init__.py | 0 perturbot/build/lib/perturbot/eval/all.py | 190 --- perturbot/build/lib/perturbot/eval/cv.py | 251 ---- .../build/lib/perturbot/eval/cv_inner_loop.py | 673 --------- .../build/lib/perturbot/eval/cv_outer_loop.py | 325 ----- .../lib/perturbot/eval/feature_matching.py | 160 --- perturbot/build/lib/perturbot/eval/loo.py | 432 ------ perturbot/build/lib/perturbot/eval/match.py | 242 ---- .../build/lib/perturbot/eval/prediction.py | 210 --- perturbot/build/lib/perturbot/eval/utils.py | 105 -- .../build/lib/perturbot/match/__init__.py | 27 - perturbot/build/lib/perturbot/match/cot.py | 389 ------ .../build/lib/perturbot/match/cot_labels.py | 340 ----- perturbot/build/lib/perturbot/match/fot.py | 220 --- perturbot/build/lib/perturbot/match/gw.py | 137 -- .../build/lib/perturbot/match/gw_labels.py | 148 -- .../build/lib/perturbot/match/ot_labels.py | 78 -- .../build/lib/perturbot/match/ott_egwl.py | 450 ------ perturbot/build/lib/perturbot/match/utils.py | 184 --- .../build/lib/perturbot/predict/__init__.py | 3 - .../perturbot/predict/linear_regression.py | 155 --- perturbot/build/lib/perturbot/predict/mlp.py | 250 ---- .../build/lib/perturbot/predict/scvi_vae.py | 177 --- perturbot/build/lib/perturbot/predict/vae.py | 493 ------- .../lib/perturbot/preprocess/__init__.py | 0 .../build/lib/perturbot/preprocess/vae.py | 65 - perturbot/build/lib/perturbot/utils.py | 10 - 108 files changed, 26978 deletions(-) delete mode 100644 ott/build/lib/ott/__init__.py delete mode 100644 ott/build/lib/ott/_version.py delete mode 100644 ott/build/lib/ott/datasets.py delete mode 100644 ott/build/lib/ott/geometry/__init__.py delete mode 100644 ott/build/lib/ott/geometry/costs.py delete mode 100644 ott/build/lib/ott/geometry/distrib_costs.py delete mode 100644 ott/build/lib/ott/geometry/epsilon_scheduler.py delete mode 100644 ott/build/lib/ott/geometry/geodesic.py delete mode 100644 ott/build/lib/ott/geometry/geometry.py delete mode 100644 ott/build/lib/ott/geometry/graph.py delete mode 100644 ott/build/lib/ott/geometry/grid.py delete mode 100644 ott/build/lib/ott/geometry/low_rank.py delete mode 100644 ott/build/lib/ott/geometry/pointcloud.py delete mode 100644 ott/build/lib/ott/geometry/segment.py delete mode 100644 ott/build/lib/ott/initializers/__init__.py delete mode 100644 ott/build/lib/ott/initializers/linear/__init__.py delete mode 100644 ott/build/lib/ott/initializers/linear/initializers.py delete mode 100644 ott/build/lib/ott/initializers/linear/initializers_lr.py delete mode 100644 ott/build/lib/ott/initializers/quadratic/__init__.py delete mode 100644 ott/build/lib/ott/initializers/quadratic/initializers.py delete mode 100644 ott/build/lib/ott/math/__init__.py delete mode 100644 ott/build/lib/ott/math/fixed_point_loop.py delete mode 100644 ott/build/lib/ott/math/matrix_square_root.py delete mode 100644 ott/build/lib/ott/math/unbalanced_functions.py delete mode 100644 ott/build/lib/ott/math/utils.py delete mode 100644 ott/build/lib/ott/neural/__init__.py delete mode 100644 ott/build/lib/ott/neural/layers.py delete mode 100644 ott/build/lib/ott/neural/losses.py delete mode 100644 ott/build/lib/ott/neural/models.py delete mode 100644 ott/build/lib/ott/neural/solvers/__init__.py delete mode 100644 ott/build/lib/ott/neural/solvers/conjugate.py delete mode 100644 ott/build/lib/ott/neural/solvers/map_estimator.py delete mode 100644 ott/build/lib/ott/neural/solvers/neuraldual.py delete mode 100644 ott/build/lib/ott/problems/__init__.py delete mode 100644 ott/build/lib/ott/problems/linear/__init__.py delete mode 100644 ott/build/lib/ott/problems/linear/barycenter_problem.py delete mode 100644 ott/build/lib/ott/problems/linear/linear_problem.py delete mode 100644 ott/build/lib/ott/problems/linear/potentials.py delete mode 100644 ott/build/lib/ott/problems/quadratic/__init__.py delete mode 100644 ott/build/lib/ott/problems/quadratic/gw_barycenter.py delete mode 100644 ott/build/lib/ott/problems/quadratic/quadratic_costs.py delete mode 100644 ott/build/lib/ott/problems/quadratic/quadratic_problem.py delete mode 100755 ott/build/lib/ott/py.typed delete mode 100644 ott/build/lib/ott/solvers/__init__.py delete mode 100644 ott/build/lib/ott/solvers/linear/__init__.py delete mode 100644 ott/build/lib/ott/solvers/linear/_solve.py delete mode 100644 ott/build/lib/ott/solvers/linear/acceleration.py delete mode 100644 ott/build/lib/ott/solvers/linear/continuous_barycenter.py delete mode 100644 ott/build/lib/ott/solvers/linear/discrete_barycenter.py delete mode 100644 ott/build/lib/ott/solvers/linear/implicit_differentiation.py delete mode 100644 ott/build/lib/ott/solvers/linear/lineax_implicit.py delete mode 100644 ott/build/lib/ott/solvers/linear/lr_utils.py delete mode 100644 ott/build/lib/ott/solvers/linear/sinkhorn.py delete mode 100644 ott/build/lib/ott/solvers/linear/sinkhorn_lr.py delete mode 100644 ott/build/lib/ott/solvers/linear/univariate.py delete mode 100644 ott/build/lib/ott/solvers/quadratic/__init__.py delete mode 100644 ott/build/lib/ott/solvers/quadratic/_solve.py delete mode 100644 ott/build/lib/ott/solvers/quadratic/gromov_wasserstein.py delete mode 100644 ott/build/lib/ott/solvers/quadratic/gromov_wasserstein_lr.py delete mode 100644 ott/build/lib/ott/solvers/quadratic/gw_barycenter.py delete mode 100644 ott/build/lib/ott/solvers/quadratic/lower_bound.py delete mode 100644 ott/build/lib/ott/solvers/was_solver.py delete mode 100644 ott/build/lib/ott/tools/__init__.py delete mode 100644 ott/build/lib/ott/tools/gaussian_mixture/__init__.py delete mode 100644 ott/build/lib/ott/tools/gaussian_mixture/fit_gmm.py delete mode 100644 ott/build/lib/ott/tools/gaussian_mixture/fit_gmm_pair.py delete mode 100644 ott/build/lib/ott/tools/gaussian_mixture/gaussian.py delete mode 100644 ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture.py delete mode 100644 ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture_pair.py delete mode 100644 ott/build/lib/ott/tools/gaussian_mixture/linalg.py delete mode 100644 ott/build/lib/ott/tools/gaussian_mixture/probabilities.py delete mode 100644 ott/build/lib/ott/tools/gaussian_mixture/scale_tril.py delete mode 100644 ott/build/lib/ott/tools/k_means.py delete mode 100644 ott/build/lib/ott/tools/plot.py delete mode 100644 ott/build/lib/ott/tools/segment_sinkhorn.py delete mode 100644 ott/build/lib/ott/tools/sinkhorn_divergence.py delete mode 100644 ott/build/lib/ott/tools/soft_sort.py delete mode 100644 ott/build/lib/ott/types.py delete mode 100644 ott/build/lib/ott/utils.py delete mode 100644 perturbot/build/lib/perturbot/__init__.py delete mode 100755 perturbot/build/lib/perturbot/eval/.run_loo.py delete mode 100755 perturbot/build/lib/perturbot/eval/__init__.py delete mode 100755 perturbot/build/lib/perturbot/eval/all.py delete mode 100755 perturbot/build/lib/perturbot/eval/cv.py delete mode 100755 perturbot/build/lib/perturbot/eval/cv_inner_loop.py delete mode 100755 perturbot/build/lib/perturbot/eval/cv_outer_loop.py delete mode 100755 perturbot/build/lib/perturbot/eval/feature_matching.py delete mode 100755 perturbot/build/lib/perturbot/eval/loo.py delete mode 100755 perturbot/build/lib/perturbot/eval/match.py delete mode 100755 perturbot/build/lib/perturbot/eval/prediction.py delete mode 100755 perturbot/build/lib/perturbot/eval/utils.py delete mode 100755 perturbot/build/lib/perturbot/match/__init__.py delete mode 100755 perturbot/build/lib/perturbot/match/cot.py delete mode 100755 perturbot/build/lib/perturbot/match/cot_labels.py delete mode 100755 perturbot/build/lib/perturbot/match/fot.py delete mode 100755 perturbot/build/lib/perturbot/match/gw.py delete mode 100755 perturbot/build/lib/perturbot/match/gw_labels.py delete mode 100755 perturbot/build/lib/perturbot/match/ot_labels.py delete mode 100755 perturbot/build/lib/perturbot/match/ott_egwl.py delete mode 100755 perturbot/build/lib/perturbot/match/utils.py delete mode 100755 perturbot/build/lib/perturbot/predict/__init__.py delete mode 100755 perturbot/build/lib/perturbot/predict/linear_regression.py delete mode 100755 perturbot/build/lib/perturbot/predict/mlp.py delete mode 100755 perturbot/build/lib/perturbot/predict/scvi_vae.py delete mode 100755 perturbot/build/lib/perturbot/predict/vae.py delete mode 100755 perturbot/build/lib/perturbot/preprocess/__init__.py delete mode 100755 perturbot/build/lib/perturbot/preprocess/vae.py delete mode 100644 perturbot/build/lib/perturbot/utils.py diff --git a/ott/build/lib/ott/__init__.py b/ott/build/lib/ott/__init__.py deleted file mode 100644 index dac0eb8..0000000 --- a/ott/build/lib/ott/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -# 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 contextlib - -from . import ( - datasets, - geometry, - initializers, - math, - problems, - solvers, - tools, - utils, -) - -with contextlib.suppress(ImportError): - # TODO(michalk8): add warning that neural module is not imported - from . import neural - -from ._version import __version__ - -del contextlib diff --git a/ott/build/lib/ott/_version.py b/ott/build/lib/ott/_version.py deleted file mode 100644 index 972f66c..0000000 --- a/ott/build/lib/ott/_version.py +++ /dev/null @@ -1,21 +0,0 @@ -# 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. -from importlib.metadata import PackageNotFoundError, version - -try: - __version__ = version("ott-jax") -except PackageNotFoundError: - __version__ = "" - -del version, PackageNotFoundError diff --git a/ott/build/lib/ott/datasets.py b/ott/build/lib/ott/datasets.py deleted file mode 100644 index 2a8d353..0000000 --- a/ott/build/lib/ott/datasets.py +++ /dev/null @@ -1,154 +0,0 @@ -# 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 dataclasses -from typing import Iterator, Literal, NamedTuple, Optional, Tuple - -import jax -import jax.numpy as jnp -import numpy as np - -__all__ = ["create_gaussian_mixture_samplers", "Dataset", "GaussianMixture"] - -from ott import utils - -Name_t = Literal["simple", "circle", "square_five", "square_four"] - - -class Dataset(NamedTuple): - r"""Samplers from source and target measures. - - Args: - source_iter: loader for the source measure - target_iter: loader for the target measure - """ - source_iter: Iterator[jnp.ndarray] - target_iter: Iterator[jnp.ndarray] - - -@dataclasses.dataclass -class GaussianMixture: - """A mixture of Gaussians. - - Args: - name: the name specifying the centers of the mixture components: - - - ``simple`` - data clustered in one center, - - ``circle`` - two-dimensional Gaussians arranged on a circle, - - ``square_five`` - two-dimensional Gaussians on a square with - one Gaussian in the center, and - - ``square_four`` - two-dimensional Gaussians in the corners of a - rectangle - - batch_size: batch size of the samples - init_rng: initial PRNG key - scale: scale of the Gaussian means - std: the standard deviation of the individual Gaussian samples - """ - name: Name_t - batch_size: int - init_rng: jax.Array - scale: float = 5.0 - std: float = 0.5 - - def __post_init__(self): - gaussian_centers = { - "simple": - np.array([[0, 0]]), - "circle": - np.array([ - (1, 0), - (-1, 0), - (0, 1), - (0, -1), - (1.0 / np.sqrt(2), 1.0 / np.sqrt(2)), - (1.0 / np.sqrt(2), -1.0 / np.sqrt(2)), - (-1.0 / np.sqrt(2), 1.0 / np.sqrt(2)), - (-1.0 / np.sqrt(2), -1.0 / np.sqrt(2)), - ]), - "square_five": - np.array([[0, 0], [1, 1], [-1, 1], [-1, -1], [1, -1]]), - "square_four": - np.array([[1, 0], [0, 1], [-1, 0], [0, -1]]), - } - if self.name not in gaussian_centers: - raise ValueError( - f"{self.name} is not a valid dataset for GaussianMixture" - ) - self.centers = gaussian_centers[self.name] - - def __iter__(self) -> Iterator[jnp.array]: - """Random sample generator from Gaussian mixture. - - Returns: - A generator of samples from the Gaussian mixture. - """ - return self._create_sample_generators() - - def _create_sample_generators(self) -> Iterator[jnp.array]: - rng = self.init_rng - while True: - rng1, rng2, rng = jax.random.split(rng, 3) - means = jax.random.choice(rng1, self.centers, (self.batch_size,)) - normal_samples = jax.random.normal(rng2, (self.batch_size, 2)) - samples = self.scale * means + (self.std ** 2) * normal_samples - yield samples - - -def create_gaussian_mixture_samplers( - name_source: Name_t, - name_target: Name_t, - train_batch_size: int = 2048, - valid_batch_size: int = 2048, - rng: Optional[jax.Array] = None, -) -> Tuple[Dataset, Dataset, int]: - """Gaussian samplers. - - Args: - name_source: name of the source sampler - name_target: name of the target sampler - train_batch_size: the training batch size - valid_batch_size: the validation batch size - rng: initial PRNG key - - Returns: - The dataset and dimension of the data. - """ - rng = utils.default_prng_key(rng) - rng1, rng2, rng3, rng4 = jax.random.split(rng, 4) - train_dataset = Dataset( - source_iter=iter( - GaussianMixture( - name_source, batch_size=train_batch_size, init_rng=rng1 - ) - ), - target_iter=iter( - GaussianMixture( - name_target, batch_size=train_batch_size, init_rng=rng2 - ) - ) - ) - valid_dataset = Dataset( - source_iter=iter( - GaussianMixture( - name_source, batch_size=valid_batch_size, init_rng=rng3 - ) - ), - target_iter=iter( - GaussianMixture( - name_target, batch_size=valid_batch_size, init_rng=rng4 - ) - ) - ) - dim_data = 2 - return train_dataset, valid_dataset, dim_data diff --git a/ott/build/lib/ott/geometry/__init__.py b/ott/build/lib/ott/geometry/__init__.py deleted file mode 100644 index 32d9d96..0000000 --- a/ott/build/lib/ott/geometry/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -# 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. -from . import ( - costs, - distrib_costs, - epsilon_scheduler, - geodesic, - geometry, - graph, - grid, - pointcloud, - segment, -) diff --git a/ott/build/lib/ott/geometry/costs.py b/ott/build/lib/ott/geometry/costs.py deleted file mode 100644 index de76be3..0000000 --- a/ott/build/lib/ott/geometry/costs.py +++ /dev/null @@ -1,1141 +0,0 @@ -# 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 abc -import functools -import math -from typing import Any, Callable, Dict, Optional, Tuple, Union - -import jax -import jax.numpy as jnp -import jaxopt -import numpy as np - -from ott.math import fixed_point_loop, matrix_square_root -from ott.math import utils as mu - -__all__ = [ - "PNormP", - "SqPNorm", - "Euclidean", - "SqEuclidean", - "Cosine", - "Arccos", - "ElasticL1", - "ElasticL2", - "ElasticSTVS", - "ElasticSqKOverlap", - "Bures", - "UnbalancedBures", - "SoftDTW", -] - - -@jax.tree_util.register_pytree_node_class -class CostFn(abc.ABC): - """Base class for all costs. - - Cost functions evaluate a function on a pair of inputs. For convenience, - that function is split into two norms -- evaluated on each input separately -- - followed by a pairwise cost that involves both inputs, as in: - - .. math:: - c(x, y) = norm(x) + norm(y) + pairwise(x, y) - - If the :attr:`norm` function is not implemented, that value is handled as - :math:`0`, and only :func:`pairwise` is used. - """ - - # no norm function created by default. - norm: Optional[Callable[[jnp.ndarray], Union[float, jnp.ndarray]]] = None - - @abc.abstractmethod - def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: - """Compute cost between :math:`x` and :math:`y`. - - Args: - x: Array. - y: Array. - - Returns: - The cost. - """ - - def barycenter(self, weights: jnp.ndarray, - xs: jnp.ndarray) -> Tuple[jnp.ndarray, Any]: - """Barycentric operator. - - Args: - weights: Convex set of weights. - xs: Points. - - Returns: - A list, whose first element is the barycenter of `xs` using `weights` - coefficients, followed by auxiliary information on the convergence of - the algorithm. - """ - raise NotImplementedError("Barycenter is not implemented.") - - @classmethod - def _padder(cls, dim: int) -> jnp.ndarray: - """Create a padding vector of adequate dimension, well-suited to a cost. - - Args: - dim: Dimensionality of the data. - - Returns: - The padding vector. - """ - return jnp.zeros((1, dim)) - - def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: - """Compute cost between :math:`x` and :math:`y`. - - Args: - x: Array. - y: Array. - - Returns: - The cost, optionally including the :attr:`norms ` of - :math:`x`/:math:`y`. - """ - cost = self.pairwise(x, y) - if self.norm is None: - return cost - return cost + self.norm(x) + self.norm(y) - - def all_pairs(self, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: - """Compute matrix of all pairwise costs, including the :attr:`norms `. - - Args: - x: Array of shape ``[n, ...]``. - y: Array of shape ``[m, ...]``. - - Returns: - Array of shape ``[n, m]`` of cost evaluations. - """ - return jax.vmap(lambda x_: jax.vmap(lambda y_: self(x_, y_))(y))(x) - - def all_pairs_pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: - """Compute matrix of all pairwise costs, excluding the :attr:`norms `. - - Args: - x: Array of shape ``[n, ...]``. - y: Array of shape ``[m, ...]``. - - Returns: - Array of shape ``[n, m]`` of cost evaluations. - """ - return jax.vmap(lambda x_: jax.vmap(lambda y_: self.pairwise(x_, y_))(y))(x) - - def tree_flatten(self): # noqa: D102 - return (), None - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - del aux_data - return cls(*children) - - -@jax.tree_util.register_pytree_node_class -class TICost(CostFn): - """Base class for translation invariant (TI) costs. - - Such costs are defined using a function :math:`h`, mapping vectors to - real-values, to be used as: - - .. math:: - c(x, y) = h(z), z := x - y. - - If that cost function is used to form an Entropic map using the - :cite:`brenier:91` theorem, then the user should ensure :math:`h` is - strictly convex, as well as provide the Legendre transform of :math:`h`, - whose gradient is necessarily the inverse of the gradient of :math:`h`. - """ - - @abc.abstractmethod - def h(self, z: jnp.ndarray) -> float: - """TI function acting on difference of :math:`x-y` to output cost. - - Args: - z: Array of shape ``[d,]``. - - Returns: - The cost. - """ - - def h_legendre(self, z: jnp.ndarray) -> float: - """Legendre transform of :func:`h` when it is convex.""" - raise NotImplementedError("Legendre transform of `h` is not implemented.") - - def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: - """Compute cost as evaluation of :func:`h` on :math:`x-y`.""" - return self.h(x - y) - - -@jax.tree_util.register_pytree_node_class -class SqPNorm(TICost): - r"""Squared p-norm of the difference of two vectors. - - Uses custom implementation of `norm` to avoid `NaN` values when - differentiating the norm of `x-x`. - - Args: - p: Power of the p-norm, :math:`\ge 1`. - """ - - def __init__(self, p: float): - super().__init__() - self.p = p - self.q = 1.0 / (1.0 - (1.0 / p)) if p > 1.0 else jnp.inf - - def h(self, z: jnp.ndarray) -> float: # noqa: D102 - return 0.5 * mu.norm(z, self.p) ** 2 - - def h_legendre(self, z: jnp.ndarray) -> float: - """Legendre transform of :func:`h`. - - For details on the derivation, see e.g., :cite:`boyd:04`, p. 93/94. - """ - return 0.5 * mu.norm(z, self.q) ** 2 - - def tree_flatten(self): # noqa: D102 - return (), (self.p,) - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - del children - return cls(*aux_data) - - -@jax.tree_util.register_pytree_node_class -class PNormP(TICost): - r"""p-norm to the power p (and divided by p) of the difference of two vectors. - - Uses custom implementation of `norm` to avoid `NaN` values when - differentiating the norm of `x-x`. - - Args: - p: Power of the p-norm in :math:`[1, +\infty)`. - Note that :func:`h_legendre` is not defined for ``p = 1``. - """ - - def __init__(self, p: float): - super().__init__() - self.p = p - self.q = 1.0 / (1.0 - (1.0 / p)) if p > 1.0 else jnp.inf - - def h(self, z: jnp.ndarray) -> float: # noqa: D102 - return mu.norm(z, self.p) ** self.p / self.p - - def h_legendre(self, z: jnp.ndarray) -> float: # noqa: D102 - # not defined for `p=1` - return mu.norm(z, self.q) ** self.q / self.q - - def tree_flatten(self): # noqa: D102 - return (), (self.p,) - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - del children - return cls(*aux_data) - - -@jax.tree_util.register_pytree_node_class -class Euclidean(CostFn): - """Euclidean distance. - - Note that the Euclidean distance is not cast as a - :class:`~ott.geometry.costs.TICost`, since this would correspond to :math:`h` - being :func:`jax.numpy.linalg.norm`, whose gradient is not invertible, - because the function is not strictly convex (it is linear on rays). - """ - - def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: - """Compute Euclidean norm using custom jvp implementation. - - Here we use a custom jvp implementation for the norm that does not yield - `NaN` gradients when differentiating the norm of `(x-x)`, but defaults - instead to zero, using a `custom_jvp` rule. - """ - return mu.norm(x - y) - - -@jax.tree_util.register_pytree_node_class -class SqEuclidean(TICost): - r"""Squared Euclidean distance. - - Implemented as a translation invariant cost, :math:`h(z) = \|z\|^2`. - """ - - def norm(self, x: jnp.ndarray) -> Union[float, jnp.ndarray]: - """Compute squared Euclidean norm for vector.""" - return jnp.sum(x ** 2, axis=-1) - - def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: - """Compute minus twice the dot-product between vectors.""" - return -2.0 * jnp.vdot(x, y) - - def h(self, z: jnp.ndarray) -> float: # noqa: D102 - return jnp.sum(z ** 2) - - def h_legendre(self, z: jnp.ndarray) -> float: # noqa: D102 - return 0.25 * jnp.sum(z ** 2) - - def barycenter(self, weights: jnp.ndarray, - xs: jnp.ndarray) -> Tuple[jnp.ndarray, Any]: - """Output barycenter of vectors when using squared-Euclidean distance.""" - return jnp.average(xs, weights=weights, axis=0), None - - -@jax.tree_util.register_pytree_node_class -class Cosine(CostFn): - """Cosine distance cost function. - - Args: - ridge: Ridge regularization. - """ - - def __init__(self, ridge: float = 1e-8): - super().__init__() - self._ridge = ridge - - def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: - """Cosine distance between vectors, denominator regularized with ridge.""" - x_norm = jnp.linalg.norm(x, axis=-1) - y_norm = jnp.linalg.norm(y, axis=-1) - cosine_similarity = jnp.vdot(x, y) / (x_norm * y_norm + self._ridge) - return 1.0 - cosine_similarity - - @classmethod - def _padder(cls, dim: int) -> jnp.ndarray: - return jnp.ones((1, dim)) - - -@jax.tree_util.register_pytree_node_class -class Arccos(CostFn): - r"""Arc-cosine cost function :cite:`cho:09`. - - The cost is implemented as: - - .. math:: - c_n(x, y) = -\log(\frac{1}{\pi} \|x\|^n \|y\|^n J_n(\theta)) - - where :math:`\theta := \arccos(\frac{x \cdot y}{\|x\| \|y\|})` and - :math:`J_n(\theta) := (-1)^n (\sin \theta)^{2n + 1} - (\frac{1}{\sin \theta}\frac{\partial}{\partial \theta})^n - (\frac{\pi - \theta}{\sin \theta})`. - - Args: - n: Order of the kernel. For :math:`n > 2`, successive applications of - :func:`~jax.grad` are used to compute the :math:`J_n(\theta)`. - ridge: Ridge regularization. - """ - - def __init__(self, n: int, ridge: float = 1e-8): - self.n = n - self._ridge = ridge - - def pairwise(self, x: jnp.ndarray, y: jnp.ndarray): # noqa: D102 - x_norm = jnp.linalg.norm(x, axis=-1) - y_norm = jnp.linalg.norm(y, axis=-1) - cosine_similarity = jnp.vdot(x, y) / (x_norm * y_norm + self._ridge) - theta = jnp.arccos(cosine_similarity) - - if self.n == 0: - m = 1.0 - theta / jnp.pi - elif self.n == 1: - j = jnp.sin(theta) + (jnp.pi - theta) * jnp.cos(theta) - m = (x_norm * y_norm) * (j / jnp.pi) - elif self.n == 2: - j = 3.0 * jnp.sin(theta) * jnp.cos(theta) + (jnp.pi - theta) * ( - 1.0 + 2.0 * jnp.cos(theta) ** 2 - ) - m = (x_norm * y_norm) ** 2 * (j / jnp.pi) - else: - j = self._j(theta) # less optimized version using autodiff - m = (x_norm * y_norm) ** self.n * (j / jnp.pi) - - return -jnp.log(m + self._ridge) - - @jax.jit - def _j(self, theta: float) -> float: - - def f(t: float, i: int) -> float: - if i == 0: - return (jnp.pi - t) / jnp.sin(t) - return jax.grad(f)(t, i - 1) / jnp.sin(t) - - n = self.n - return (-1) ** n * jnp.sin(theta) ** (2.0 * n + 1.0) * f(theta, n) - - def tree_flatten(self): # noqa: D102 - return [], {"n": self.n, "ridge": self._ridge} - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - del children - return cls(**aux_data) - - -class RegTICost(TICost, abc.ABC): - r"""Base class for regularized translation-invariant costs. - - .. math:: - \frac{1}{2} \|\cdot\|_2^2 + \text{scaling_reg} reg\left(matrix \cdot\right) - - where :func:`reg` is the regularization function. - - Args: - scaling_reg: Strength of the :meth:`regularization `. - matrix: :math:`p \times d` projection matrix in the Stiefel manifold, - namely with **orthonormalized rows**. - orthogonal: Whether to regularize in the orthogonal complement - to promote displacements in the span of ``matrix``. - """ - - def __init__( - self, - scaling_reg: float = 1.0, - matrix: Optional[jnp.ndarray] = None, - orthogonal: bool = False, - ): - super().__init__() - self.scaling_reg = scaling_reg - self.matrix = matrix - self.orthogonal = orthogonal - - @abc.abstractmethod - def _reg(self, z: jnp.ndarray) -> float: - """Regularization function.""" - - def _reg_stiefel_orth(self, z: jnp.ndarray) -> float: - raise NotImplementedError( - "Regularization in the orthogonal " - "subspace is not implemented." - ) - - def reg(self, z: jnp.ndarray) -> float: - """Regularization function. - - Args: - z: Array of shape ``[d,]``. - - Returns: - The regularization value. - """ - if self.matrix is None: - return self._reg(z) - if self.orthogonal: - return self._reg_stiefel_orth(z) - return self._reg(self.matrix @ z) - - def prox_reg(self, z: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: - """Proximal operator of :meth:`reg`. - - Args: - z: Array of shape ``[d,]``. - tau: Positive weight. - - Returns: - The prox of ``z``. - """ - if self.matrix is None: - return self._prox_reg(z, tau) - if self.orthogonal: - # regularization in the orthogonal subspace - return self._prox_reg_stiefel_orth(z, tau) - return self._prox_reg_stiefel(z, tau) - - def _prox_reg(self, z: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: - raise NotImplementedError("Proximal operator is not implemented.") - - def _prox_reg_stiefel_orth( - self, z: jnp.ndarray, tau: float = 1.0 - ) -> jnp.ndarray: - - def orth(x: jnp.ndarray) -> jnp.ndarray: - return x - self.matrix.T @ (self.matrix @ x) - - # assumes `matrix` has orthogonal rows - tmp = orth(z) - return z - orth(tmp - self._prox_reg(tmp, tau)) - - def _prox_reg_stiefel(self, z: jnp.ndarray, tau: float) -> jnp.ndarray: - # assumes `matrix` has orthogonal rows - tmp = self.matrix @ z - return z - self.matrix.T @ (tmp - self._prox_reg(tmp, tau)) - - def prox_legendre_reg(self, z: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: - r"""Proximal operator of the Legendre transform of :meth:`reg`. - - Uses Moreau's decomposition: - - .. math:: - x = \text{prox}_{\tau f} \left(x\right) + - \tau \text{prox}_{\frac{1}{\tau} f^*} \left(\frac{x}{\tau}\right) - - Args: - z: Array of shape ``[d,]``. - tau: Positive weight. - - Returns: - The prox of ``z``. - """ - return z - tau * self.prox_reg(z / tau, 1.0 / tau) - - def h(self, z: jnp.ndarray) -> float: # noqa: D102 - out = 0.5 * jnp.sum(z ** 2) - return out + self.scaling_reg * self.reg(z) - - def h_legendre(self, z: jnp.ndarray) -> float: # noqa: D102 - q = jax.lax.stop_gradient(self.prox_reg(z)) - return jnp.sum(q * z) - self.h(q) - - def h_transform(self, f: Callable[[jnp.ndarray], float], - **kwargs: Any) -> Callable[[jnp.ndarray], float]: - r"""Compute the h-transform of a concave function. - - Return a callable :math:`f_h` defined as: - - .. math:: - f_h(x) = \min_y h(x - y) - f(y) - - This is equivalent, up to a change of variables, :math:`z = x - y`, to - define - - .. math:: - \min_z h(z) - f(x - z). \\ - \min_z h(z) + \tilde{f}(z, x). - - where :math:`\tilde{f}(z, x) := -f(x - z)`. - - This is solved using proximal gradient descent, which requires having - access to the prox of :math:`\text{scaling_h} \cdot h` and not only to that - of :meth:`h`. Given the properties of :meth:`h`, the prox is obtained by - rescaling the output of the prox of a suitable scaling of :meth:`prox_reg`. - - Args: - f: Concave function. - kwargs: Keyword arguments for :class:`~jaxopt.ProximalGradient`. - - Returns: - The h-transform of ``f``. - """ - - def minus_f(z: jnp.ndarray, x: jnp.ndarray) -> float: - return -f(x - z) - - def prox( - x: jnp.ndarray, scaling_reg: float, scaling_h: float - ) -> jnp.ndarray: - # https://web.stanford.edu/~boyd/papers/pdf/prox_algs.pdf 2.2. - tmp = 1.0 / (1.0 + scaling_h) - tau = scaling_reg * scaling_h * tmp - return self.prox_reg(x * tmp, tau) - - def f_h(x: jnp.ndarray) -> float: - pg = jaxopt.ProximalGradient(fun=minus_f, prox=prox, **kwargs) - pg_run = pg.run(x, self.scaling_reg, x=x) - pg_sol = jax.lax.stop_gradient(pg_run.params) - return self.h(pg_sol) + minus_f(pg_sol, x) - - return f_h - - def tree_flatten(self): # noqa: D102 - return (self.scaling_reg, self.matrix), {"orthogonal": self.orthogonal} - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - return cls(*children, **aux_data) - - -@jax.tree_util.register_pytree_node_class -class ElasticL1(RegTICost): - r"""Cost inspired by elastic net :cite:`zou:05` regularization. - - .. math:: - \frac{1}{2} \|\cdot\|_2^2 + \text{scaling_reg} \|\text{matrix} \cdot\|_1 - - Args: - scaling_reg: Strength of the :meth:`regularization `. - matrix: :math:`p \times d` projection matrix with **orthogonal rows**. - orthogonal: Whether to regularize in the orthogonal complement - to promote displacements in the span of ``matrix``. - """ - - def _reg(self, z: jnp.ndarray) -> float: # noqa: D102 - return jnp.linalg.norm(z, ord=1) - - def _prox_reg(self, z: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: - return jnp.sign(z) * jax.nn.relu(jnp.abs(z) - tau * self.scaling_reg) - - -@jax.tree_util.register_pytree_node_class -class ElasticL2(RegTICost): - r"""Cost with L2 regularization. - - .. math:: - \frac{1}{2} \|\cdot\|_2^2 + \text{scaling_reg} \|\text{matrix} \cdot\|_2^2 - - Args: - scaling_reg: Strength of the :meth:`regularization `. - matrix: :math:`p \times d` projection matrix with **orthogonal rows**. - orthogonal: Whether to regularize in the orthogonal complement - to promote displacements in the span of ``matrix``. - """ - - def _reg(self, z: jnp.ndarray) -> float: # noqa: D102 - return 0.5 * jnp.sum(z ** 2) - - def _reg_stiefel_orth(self, z: jnp.ndarray) -> float: - # Pythagorean identity - return self._reg(z) - self._reg(self.matrix @ z) - - def _prox_reg(self, z: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: - return z / (1.0 + tau * self.scaling_reg) - - def _prox_reg_stiefel_orth( - self, z: jnp.ndarray, tau: float = 1.0 - ) -> jnp.ndarray: - out = z + tau * self.scaling_reg * self.matrix.T @ (self.matrix @ z) - return self._prox_reg(out, tau) - - -@jax.tree_util.register_pytree_node_class -class ElasticSTVS(RegTICost): - r"""Cost with soft thresholding operator with vanishing shrinkage (STVS) - :cite:`schreck:15` regularization. - - .. math:: - \frac{1}{2} \|\cdot\|_2^2 + \text{scaling_reg}^2\mathbf{1}_d^T\left(\sigma(\cdot) - - \frac{1}{2} \exp\left(-2\sigma(\cdot)\right) + \frac{1}{2}\right) - - where :math:`\sigma(\cdot) := \text{asinh}\left(\frac{\cdot} - {2\text{scaling_reg}}\right)` - - Args: - scaling_reg: Strength of the :meth:`regularization `. - matrix: :math:`p \times d` projection matrix with **orthogonal rows**. - orthogonal: Whether to regularize in the orthogonal complement - to promote displacements in the span of ``matrix``. - """ # noqa: D205,E501 - - def _reg(self, z: jnp.ndarray) -> float: # noqa: D102 - u = jnp.arcsinh(jnp.abs(z) / (2 * self.scaling_reg)) - out = u - 0.5 * jnp.exp(-2.0 * u) - # Lemma 2.1 of `schreck:15`; - # don't use `self.scaling_reg ** 2` because it's included in `h` - return self.scaling_reg * jnp.sum(out + 0.5) # make positive - - def _prox_reg( # noqa: D102 - self, z: jnp.ndarray, tau: float = 1.0 - ) -> jnp.ndarray: - tmp = 1.0 - (self.scaling_reg * tau / (jnp.abs(z) + 1e-12)) ** 2 - return jax.nn.relu(tmp) * z - - -@jax.tree_util.register_pytree_node_class -class ElasticSqKOverlap(RegTICost): - r"""Cost with squared k-overlap norm regularization :cite:`argyriou:12`. - - .. math:: - \frac{1}{2} \|\cdot\|_2^2 + \frac{1}{2} \text{scaling_reg} \|\cdot\|_{ovk}^2 - - where :math:`\|\cdot\|_{ovk}^2` is the squared k-overlap norm, - see def. 2.1 of :cite:`argyriou:12`. - - Args: - k: Number of groups. Must be in ``[0, d)`` where :math:`d` is the - dimensionality of the data. - args: Positional arguments for :class:`~ott.geometry.costs.RegTICost`. - kwargs: Keyword arguments for :class:`~ott.geometry.costs.RegTICost`. - """ - - def __init__(self, k: int, *args, **kwargs: Any): - super().__init__(*args, **kwargs) - self.k = k - - def _reg(self, z: jnp.ndarray) -> float: # noqa: D102 - # Prop 2.1 in :cite:`argyriou:12` - k = self.k - top_w = jax.lax.top_k(jnp.abs(z), k)[0] # Fetch largest k values - top_w = jnp.flip(top_w) # Sort k-largest from smallest to largest - # sum (dim - k) smallest values - sum_bottom = jnp.sum(jnp.abs(z)) - jnp.sum(top_w) - cumsum_top = jnp.cumsum(top_w) - # Cesaro mean of top_w (each term offset with sum_bottom). - cesaro = sum_bottom + cumsum_top - cesaro /= jnp.arange(k) + 1 - # Choose first index satisfying constraint in Prop 2.1 - lower_bound = cesaro - top_w >= 0 - # Last upper bound is always True. - upper_bound = jnp.concatenate(((top_w[1:] - cesaro[:-1] - > 0), jnp.array((True,)))) - r = jnp.argmax(lower_bound * upper_bound) - s = jnp.sum(jnp.where(jnp.arange(k) < k - r - 1, jnp.flip(top_w) ** 2, 0)) - - return 0.5 * (s + (r + 1) * cesaro[r] ** 2) - - def prox_reg(self, z: jnp.ndarray, tau: float = 1.0) -> float: # noqa: D102 - - @functools.partial(jax.vmap, in_axes=[0, None, None]) - def find_indices(r: int, l: jnp.ndarray, - z: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]: - - @functools.partial(jax.vmap, in_axes=[None, 0, None]) - def inner(r: int, l: int, - z: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]: - i = k - r - 1 - res = jnp.sum(z * ((i <= ixs) & (ixs < l))) - res /= l - k + (beta + 1) * r + beta + 1 - - cond1_left = jnp.logical_or(i == 0, (z[i - 1] / beta + 1) > res) - cond1_right = res >= (z[i] / (beta + 1)) - cond1 = jnp.logical_and(cond1_left, cond1_right) - - cond2_left = z[l - 1] > res - cond2_right = jnp.logical_or(l == d, res >= z[l]) - cond2 = jnp.logical_and(cond2_left, cond2_right) - - return res, cond1 & cond2 - - return inner(r, l, z) - - del tau # this case is not handled and currently not needed - # Alg. 1 of :cite:`argyriou:12` - k, d, beta = self.k, z.shape[-1], 1.0 / self.scaling_reg - - ixs = jnp.arange(d) - z, sgn = jnp.abs(z), jnp.sign(z) - z_ixs = jnp.argsort(z)[::-1] - z_sorted = z[z_ixs] - - # (k, d - k + 1) - T, mask = find_indices(jnp.arange(k), jnp.arange(k, d + 1), z_sorted) - (r,), (l,) = jnp.where(mask, size=1) # size=1 for jitting - T = T[r, l] - - q1 = (beta / (beta + 1)) * z_sorted * (ixs < (k - r - 1)) - q2 = (z_sorted - T) * jnp.logical_and((k - r - 1) <= ixs, ixs < (l + k)) - q = q1 + q2 - - # change sign and reorder - return sgn * q[jnp.argsort(z_ixs.astype(float))] - - def tree_flatten(self): # noqa: D102 - children, aux_data = super().tree_flatten() - return children, (self.k, aux_data) - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - k, aux_data = aux_data - return cls(k, *children, **aux_data) - - -@jax.tree_util.register_pytree_node_class -class Bures(CostFn): - """Bures distance between a pair of (mean, covariance matrix). - - Args: - dimension: Dimensionality of the data. - sqrtm_kw: Dictionary of keyword arguments to control the - behavior of inner calls to :func:`~ott.math.matrix_square_root.sqrtm`. - """ - - def __init__(self, dimension: int, sqrtm_kw: Optional[Dict[str, Any]] = None): - super().__init__() - self._dimension = dimension - self._sqrtm_kw = {} if sqrtm_kw is None else sqrtm_kw - - def norm(self, x: jnp.ndarray) -> jnp.ndarray: - """Compute norm of Gaussian, sq. 2-norm of mean + trace of covariance.""" - mean, cov = x_to_means_and_covs(x, self._dimension) - norm = jnp.sum(mean ** 2, axis=-1) - norm += jnp.trace(cov, axis1=-2, axis2=-1) - return norm - - def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: - """Compute - 2 x Bures dot-product.""" - mean_x, cov_x = x_to_means_and_covs(x, self._dimension) - mean_y, cov_y = x_to_means_and_covs(y, self._dimension) - mean_dot_prod = jnp.vdot(mean_x, mean_y) - sq_x = matrix_square_root.sqrtm(cov_x, self._dimension, **self._sqrtm_kw)[0] - sq_x_y_sq_x = jnp.matmul(sq_x, jnp.matmul(cov_y, sq_x)) - sq__sq_x_y_sq_x = matrix_square_root.sqrtm( - sq_x_y_sq_x, self._dimension, **self._sqrtm_kw - )[0] - return -2 * (mean_dot_prod + jnp.trace(sq__sq_x_y_sq_x, axis1=-2, axis2=-1)) - - def covariance_fixpoint_iter( - self, - covs: jnp.ndarray, - weights: jnp.ndarray, - tolerance: float = 1e-4, - sqrtm_kw: Optional[Dict[str, Any]] = None, - **kwargs: Any - ) -> jnp.ndarray: - """Iterate fix-point updates to compute barycenter of Gaussians. - - Args: - covs: [batch, d^2] covariance matrices - weights: simplicial weights (non-negative, sum to 1) - tolerance: tolerance of the fixed-point procedure. That tolerance is - applied to the Frobenius norm (normalized by total size) - of two successive iterations of the algorithm - sqrtm_kw: keyword arguments for :func:`~ott.math.matrix_square_root.sqrtm` - kwargs: keyword arguments for the outer fixed-point iteration - - Returns: - List containing Weighted Bures average of the covariance matrices, and - vector of (normalized) 2-norms of successive differences between iterates, - to monitor convergence. - """ - sqrtm_kw = {} if sqrtm_kw is None else sqrtm_kw - # Pop values or set defaults for fixed-point loop. - min_iterations = kwargs.pop("min_iterations", 1) - max_iterations = kwargs.pop("max_iterations", 100) - inner_iterations = kwargs.pop("inner_iterations", 5) - - @functools.partial(jax.vmap, in_axes=[None, 0, 0]) - def scale_covariances( - cov_sqrt: jnp.ndarray, cov: jnp.ndarray, weight: jnp.ndarray - ) -> jnp.ndarray: - """Rescale covariance in barycenter step.""" - return weight * matrix_square_root.sqrtm_only((cov_sqrt @ cov) @ cov_sqrt, - **sqrtm_kw) - - def cond_fn(iteration: int, constants: Tuple[Any, ...], state) -> bool: - del constants - _, diffs = state - return diffs[iteration // inner_iterations] > tolerance - - def body_fn( - iteration: int, constants: Tuple[Any, ...], - state: Tuple[jnp.ndarray, float], compute_error: bool - ) -> Tuple[jnp.ndarray, float]: - del constants, compute_error - cov, diffs = state - cov_sqrt, cov_inv_sqrt, _ = matrix_square_root.sqrtm(cov, **sqrtm_kw) - scaled_cov = jnp.linalg.matrix_power( - jnp.sum(scale_covariances(cov_sqrt, covs, weights), axis=0), 2 - ) - next_cov = (cov_inv_sqrt @ scaled_cov) @ cov_inv_sqrt - diff = jnp.sum((next_cov - cov) ** 2) / jnp.prod(jnp.array(cov.shape)) - diffs = diffs.at[iteration // inner_iterations].set(diff) - return next_cov, diffs - - def init_state() -> Tuple[jnp.ndarray, float]: - cov_init = jnp.eye(self._dimension) - diffs = -jnp.ones(math.ceil(max_iterations / inner_iterations)) - return cov_init, diffs - - cov, diffs = fixed_point_loop.fixpoint_iter( - cond_fn=cond_fn, - body_fn=body_fn, - min_iterations=min_iterations, - max_iterations=max_iterations, - inner_iterations=inner_iterations, - constants=(), - state=init_state(), - ) - return cov, diffs - - def barycenter( - self, - weights: jnp.ndarray, - xs: jnp.ndarray, - tolerance: float = 1e-4, - sqrtm_kw: Optional[Dict[str, Any]] = None, - **kwargs: Any - ) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Compute the Bures barycenter of weighted Gaussian distributions. - - Implements the fixed point approach proposed in :cite:`alvarez-esteban:16` - for the computation of the mean and the covariance of the barycenter of - weighted Gaussian distributions. - - Args: - weights: The barycentric weights. - xs: The points to be used in the computation of the barycenter, where - each point is described by a concatenation of the mean and the - covariance (raveled). - tolerance: convergence tolerance to control the termination of the - algorithm. - sqrtm_kw: Arguments passed on to the - :func:`~ott.math.matrix_square_root.sqrtm` function used within - :meth:`covariance_fixpoint_iter`. This defines the precision - (in terms of convergence threshold, and number of iterations) of the - matrix square root calls that are used at each outer iteration of - the computation of Gaussian barycenters. These values are, by default, - the same as those used to define the Bures cost object itself. - kwargs: Passed on to :meth:`covariance_fixpoint_iter`, to specify the - number of iterations and tolerance of the fixed-point iteration of the - barycenter routine, by parameterizing `tolerance` and other relevant - arguments passed on to :func:`~ott.math.fixed_point_loop.fixpoint_iter`, - namely `min_iterations`, `max_iterations` and `inner_iterations`. - - Returns: - A list holding a concatenation of the mean and the raveled covariance - of the barycenter as its first element, followed by a vector of - norms of successive differences in iterates. - """ - # Ensure that barycentric weights sum to 1. - weights = weights / jnp.sum(weights) - mus, covs = x_to_means_and_covs(xs, self._dimension) - mu_bary = jnp.sum(weights[:, None] * mus, axis=0) - cov_bary, diffs = self.covariance_fixpoint_iter( - covs=covs, - weights=weights, - tolerance=tolerance, - sqrtm_kw=sqrtm_kw if sqrtm_kw is not None else self._sqrtm_kw, - **kwargs - ) - return mean_and_cov_to_x(mu_bary, cov_bary, self._dimension), diffs - - @classmethod - def _padder(cls, dim: int) -> jnp.ndarray: - dimension = int((-1 + math.sqrt(1 + 4 * dim)) / 2) - padding = mean_and_cov_to_x( - jnp.zeros((dimension,)), jnp.eye(dimension), dimension - ) - return padding[jnp.newaxis, :] - - def tree_flatten(self): # noqa: D102 - return (), (self._dimension, self._sqrtm_kw) - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - del children - return cls(*aux_data) - - -@jax.tree_util.register_pytree_node_class -class UnbalancedBures(CostFn): - """Unbalanced Bures distance between two triplets of `(mass, mean, cov)`. - - This cost uses the notation defined in :cite:`janati:20`, eq. 37, 39, 40. - - Args: - dimension: Dimensionality of the data. - sigma: Entropic regularization. - gamma: KL-divergence regularization for the marginals. - kwargs: Keyword arguments for :func:`~ott.math.matrix_square_root.sqrtm`. - """ - - def __init__( - self, - dimension: int, - *, - sigma: float = 1.0, - gamma: float = 1.0, - **kwargs: Any, - ): - super().__init__() - self._dimension = dimension - self._sigma = sigma - self._gamma = gamma - self._sqrtm_kw = kwargs - - def norm(self, x: jnp.ndarray) -> jnp.ndarray: - """Compute norm of Gaussian for unbalanced Bures. - - Args: - x: Array of shape ``[n_points + n_points + n_dim ** 2,]``, potentially - batched, corresponding to the raveled mass, means and the covariance - matrix. - - Returns: - The norm, array of shape ``[]`` or ``[batch,]`` in the batched case. - """ - return self._gamma * x[..., 0] - - def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: - """Compute dot-product for unbalanced Bures. - - Args: - x: Array of shape ``[n_points + n_points + n_dim ** 2,]`` - corresponding to the raveled mass, means and the covariance matrix. - y: Array of shape ``[n_points + n_points + n_dim ** 2,]`` - corresponding to the raveled mass, means and the covariance matrix. - - Returns: - The cost. - """ - # Sets a few constants - gam = self._gamma - sig2 = self._sigma ** 2 - lam = sig2 + gam / 2.0 - tau = gam / (2.0 * lam) - - # Extracts mass, mean vector, covariance matrices - mass_x, mass_y = x[0], y[0] - mean_x, cov_x = x_to_means_and_covs(x[1:], self._dimension) - mean_y, cov_y = x_to_means_and_covs(y[1:], self._dimension) - - diff_means = mean_x - mean_y - - # Identity matrix of suitable size - iden = jnp.eye(self._dimension) - - # Creates matrices needed in the computation - tilde_a = 0.5 * gam * (iden - lam * jnp.linalg.inv(cov_x + lam * iden)) - tilde_b = 0.5 * gam * (iden - lam * jnp.linalg.inv(cov_y + lam * iden)) - - tilde_a_b = jnp.matmul(tilde_a, tilde_b) - c_mat = matrix_square_root.sqrtm( - 1 / tau * tilde_a_b + 0.25 * (sig2 ** 2) * iden, **self._sqrtm_kw - )[0] - c_mat -= 0.5 * sig2 * iden - - # Computes log determinants (their sign should be >0). - sldet_c, ldet_c = jnp.linalg.slogdet(c_mat) - sldet_t_ab, ldet_t_ab = jnp.linalg.slogdet(tilde_a_b) - sldet_ab, ldet_ab = jnp.linalg.slogdet(jnp.matmul(cov_x, cov_y)) - sldet_c_ab, ldet_c_ab = jnp.linalg.slogdet(c_mat - 2.0 * tilde_a_b / gam) - - # Gathers all these results to compute log total mass of transport - log_m_pi = (0.5 * self._dimension * sig2 / (gam + sig2)) * jnp.log(sig2) - log_m_pi += (1.0 / (tau + 1.0)) * ( - jnp.log(mass_x) + jnp.log(mass_y) + ldet_c + 0.5 * - (tau * ldet_t_ab - ldet_ab) - ) - log_m_pi += -jnp.sum( - diff_means * jnp.linalg.solve(cov_x + cov_y + lam * iden, diff_means) - ) / (2.0 * (tau + 1.0)) - log_m_pi += -0.5 * ldet_c_ab - - # if all logdet signs are 1, output value, nan otherwise - pos_signs = (sldet_c + sldet_c_ab + sldet_t_ab + sldet_t_ab) == 4 - - return jax.lax.cond( - pos_signs, lambda: 2 * sig2 * mass_x * mass_y - 2 * - (sig2 + gam) * jnp.exp(log_m_pi), lambda: jnp.nan - ) - - def tree_flatten(self): # noqa: D102 - return (), (self._dimension, self._sigma, self._gamma, self._sqrtm_kw) - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - del children - dim, sigma, gamma, kwargs = aux_data - return cls(dim, sigma=sigma, gamma=gamma, **kwargs) - - -@jax.tree_util.register_pytree_node_class -class SoftDTW(CostFn): - """Soft dynamic time warping (DTW) cost :cite:`cuturi:17`. - - Args: - gamma: Smoothing parameter :math:`> 0` for the soft-min operator. - ground_cost: Ground cost function. If ``None``, - use :class:`~ott.geometry.costs.SqEuclidean`. - debiased: Whether to compute the debiased soft-DTW :cite:`blondel:21`. - """ - - def __init__( - self, - gamma: float, - ground_cost: Optional[CostFn] = None, - debiased: bool = False - ): - self.gamma = gamma - self.ground_cost = SqEuclidean() if ground_cost is None else ground_cost - self.debiased = debiased - - def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: # noqa: D102 - c_xy = self._soft_dtw(x, y) - if self.debiased: - return c_xy - 0.5 * (self._soft_dtw(x, x) + self._soft_dtw(y, y)) - return c_xy - - def _soft_dtw(self, t1: jnp.ndarray, t2: jnp.ndarray) -> float: - - def body( - carry: Tuple[jnp.ndarray, jnp.ndarray], - current_antidiagonal: jnp.ndarray - ) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], jnp.ndarray]: - # modified from: https://github.com/khdlr/softdtw_jax - two_ago, one_ago = carry - - diagonal, right, down = two_ago[:-1], one_ago[:-1], one_ago[1:] - best = mu.softmin( - jnp.stack([diagonal, right, down], axis=-1), self.gamma, axis=-1 - ) - - next_row = best + current_antidiagonal - next_row = jnp.pad(next_row, (1, 0), constant_values=jnp.inf) - - return (one_ago, next_row), next_row - - t1 = t1[:, None] if t1.ndim == 1 else t1 - t2 = t2[:, None] if t2.ndim == 1 else t2 - dist = self.ground_cost.all_pairs(t1, t2) - - n, m = dist.shape - if n < m: - dist = dist.T - n, m = m, n - - model_matrix = jnp.full((n + m - 1, n), fill_value=jnp.inf) - mask = np.tri(n + m - 1, n, k=0, dtype=bool) - mask = mask & mask[::-1, ::-1] - model_matrix = model_matrix.T.at[mask.T].set(dist.ravel()).T - - init = ( - jnp.pad(model_matrix[0], (1, 0), constant_values=jnp.inf), - jnp.pad( - model_matrix[1] + model_matrix[0, 0], (1, 0), - constant_values=jnp.inf - ) - ) - - (_, carry), _ = jax.lax.scan(body, init, model_matrix[2:]) - return carry[-1] - - def tree_flatten(self): # noqa: D102 - return (self.gamma, self.ground_cost), {"debiased": self.debiased} - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - return cls(*children, **aux_data) - - -def x_to_means_and_covs(x: jnp.ndarray, - dimension: int) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Extract means and covariance matrices of Gaussians from raveled vector. - - Args: - x: [num_gaussians, dimension, (1 + dimension)] array of concatenated means - and covariances (raveled) dimension: the dimension of the Gaussians. - dimension: Dimensionality of the Gaussians. - - Returns: - Means and covariances of shape ``[num_gaussian, dimension]``. - """ - x = jnp.atleast_2d(x) - means = x[:, :dimension] - covariances = jnp.reshape( - x[:, dimension:dimension + dimension ** 2], (-1, dimension, dimension) - ) - return jnp.squeeze(means), jnp.squeeze(covariances) - - -def mean_and_cov_to_x( - mean: jnp.ndarray, covariance: jnp.ndarray, dimension: int -) -> jnp.ndarray: - """Ravel a Gaussian's mean and covariance matrix to d(1 + d) vector.""" - return jnp.concatenate( - (mean, jnp.reshape(covariance, (dimension * dimension))) - ) diff --git a/ott/build/lib/ott/geometry/distrib_costs.py b/ott/build/lib/ott/geometry/distrib_costs.py deleted file mode 100644 index 3ddc01a..0000000 --- a/ott/build/lib/ott/geometry/distrib_costs.py +++ /dev/null @@ -1,89 +0,0 @@ -# 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. -from typing import Any, Optional - -import jax.numpy as jnp -import jax.tree_util as jtu - -from ott.geometry import costs, pointcloud -from ott.problems.linear import linear_problem -from ott.solvers.linear import univariate - -__all__ = [ - "UnivariateWasserstein", -] - - -@jtu.register_pytree_node_class -class UnivariateWasserstein(costs.CostFn): - """1D Wasserstein cost for two 1D distributions. - - This ground cost between considers vectors as a family of values. - The Wasserstein distance between them is the 1D OT cost, using a user-defined - ground cost. - - Args: - ground_cost: Cost used to compute the 1D optimal transport between vector, - should be a translation-invariant (TI) cost for correctness. - If :obj:`None`, defaults to :class:`~ott.geometry.costs.SqEuclidean`. - solver: 1D optimal transport solver. - kwargs: Arguments passed on when calling the - :class:`~ott.solvers.linear.univariate.UnivariateSolver`. May include - random key, or specific instructions to subsample or compute using - quantiles. - """ - - def __init__( - self, - ground_cost: Optional[costs.TICost] = None, - solver: Optional[univariate.UnivariateSolver] = None, - **kwargs: Any - ): - super().__init__() - - self.ground_cost = ( - costs.SqEuclidean() if ground_cost is None else ground_cost - ) - self._solver = univariate.UnivariateSolver() if solver is None else solver - self._kwargs_solve = kwargs - # ensure transport solutions are neither computed nor stored - self._kwargs_solve["return_transport"] = False - - def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float: - """Wasserstein distance between :math:`x` and :math:`y` seen as a 1D dist. - - Args: - x: Array of shape ``[n,]``. - y: Array of shape ``[m,]``. - - Returns: - The transport cost. - """ - out = self._solver( - linear_problem.LinearProblem( - pointcloud.PointCloud( - x[:, None], y[:, None], cost_fn=self.ground_cost - ) - ), **self._kwargs_solve - ) - return jnp.squeeze(out.ot_costs) - - def tree_flatten(self): # noqa: D102 - return (self.ground_cost,), (self._solver, self._kwargs_solve) - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - ground_cost, = children - solver, solve_kwargs = aux_data - return cls(ground_cost, solver, **solve_kwargs) diff --git a/ott/build/lib/ott/geometry/epsilon_scheduler.py b/ott/build/lib/ott/geometry/epsilon_scheduler.py deleted file mode 100644 index a5cca8f..0000000 --- a/ott/build/lib/ott/geometry/epsilon_scheduler.py +++ /dev/null @@ -1,102 +0,0 @@ -# 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. -from typing import Any, Optional - -import jax -import jax.numpy as jnp - -__all__ = ["Epsilon"] - - -@jax.tree_util.register_pytree_node_class -class Epsilon: - """Scheduler class for the regularization parameter epsilon. - - An epsilon scheduler outputs a regularization strength, to be used by in a - Sinkhorn-type algorithm, at any iteration count. That value is either the - final, targeted regularization, or one that is larger, obtained by - geometric decay of an initial value that is larger than the intended target. - Concretely, the value returned by such a scheduler will consider first - the max between ``target`` and ``init * target * decay ** iteration``. - If the ``scale_epsilon`` parameter is provided, that value is used to - multiply the max computed previously by ``scale_epsilon``. - - Args: - target: the epsilon regularizer that is targeted. - If ``None``, use :math:`0.05`. - scale_epsilon: if passed, used to multiply the regularizer, to rescale it. - If ``None``, use :math:`1`. - init: initial value when using epsilon scheduling, understood as multiple - of target value. if passed, ``int * decay ** iteration`` will be used - to rescale target. - decay: geometric decay factor, :math:`<1`. - """ - - def __init__( - self, - target: Optional[float] = None, - scale_epsilon: Optional[float] = None, - init: float = 1.0, - decay: float = 1.0 - ): - self._target_init = target - self._scale_epsilon = scale_epsilon - self._init = init - self._decay = decay - - @property - def target(self) -> float: - """Return the final regularizer value of scheduler.""" - target = 5e-2 if self._target_init is None else self._target_init - scale = 1.0 if self._scale_epsilon is None else self._scale_epsilon - return scale * target - - def at(self, iteration: Optional[int] = 1) -> float: - """Return (intermediate) regularizer value at a given iteration.""" - if iteration is None: - return self.target - # check the decay is smaller than 1.0. - decay = jnp.minimum(self._decay, 1.0) - # the multiple is either 1.0 or a larger init value that is decayed. - multiple = jnp.maximum(self._init * (decay ** iteration), 1.0) - return multiple * self.target - - def done(self, eps: float) -> bool: - """Return whether the scheduler is done at a given value.""" - return eps == self.target - - def done_at(self, iteration: Optional[int]) -> bool: - """Return whether the scheduler is done at a given iteration.""" - return self.done(self.at(iteration)) - - def set(self, **kwargs: Any) -> "Epsilon": - """Return a copy of self, with potential overwrites.""" - kwargs = { - "target": self._target_init, - "scale_epsilon": self._scale_epsilon, - "init": self._init, - "decay": self._decay, - **kwargs - } - return Epsilon(**kwargs) - - def tree_flatten(self): # noqa: D102 - return ( - self._target_init, self._scale_epsilon, self._init, self._decay - ), None - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - del aux_data - return cls(*children) diff --git a/ott/build/lib/ott/geometry/geodesic.py b/ott/build/lib/ott/geometry/geodesic.py deleted file mode 100644 index 9375b23..0000000 --- a/ott/build/lib/ott/geometry/geodesic.py +++ /dev/null @@ -1,277 +0,0 @@ -# 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. -from typing import Any, Dict, Optional, Sequence, Tuple - -import jax -import jax.experimental.sparse as jesp -import jax.numpy as jnp -import numpy as np -from scipy.special import ive - -from ott import utils -from ott.geometry import geometry -from ott.math import utils as mu -from ott.types import Array_g - -__all__ = ["Geodesic"] - - -@jax.tree_util.register_pytree_node_class -class Geodesic(geometry.Geometry): - r"""Graph distance approximation using heat kernel :cite:`huguet:2023`. - - .. note:: - This constructor is not meant to be called by the user, - please use the :meth:`from_graph` method instead. - - Approximates the heat kernel using - `Chebyshev polynomials `_ - of the first kind of max order ``order``, which for small ``t`` - approximates the geodesic exponential kernel :math:`e^{\frac{-d(x, y)^2}{t}}`. - - Args: - scaled_laplacian: The Laplacian scaled by the largest eigenvalue. - eigval: The largest eigenvalue of the Laplacian. - chebyshev_coeffs: Coefficients of the Chebyshev polynomials. - t: Time parameter for the heat kernel. - kwargs: Keyword arguments for :class:`~ott.geometry.geometry.Geometry`. - """ - - def __init__( - self, - scaled_laplacian: Array_g, - eigval: jnp.ndarray, - chebyshev_coeffs: jnp.ndarray, - t: float = 1e-3, - **kwargs: Any - ): - super().__init__(epsilon=1.0, **kwargs) - self.scaled_laplacian = scaled_laplacian - self.eigval = eigval - self.chebyshev_coeffs = chebyshev_coeffs - self.t = t - - @classmethod - def from_graph( - cls, - G: Array_g, - t: Optional[float] = 1e-3, - eigval: Optional[jnp.ndarray] = None, - order: int = 100, - directed: bool = False, - normalize: bool = False, - rng: Optional[jax.Array] = None, - **kwargs: Any - ) -> "Geodesic": - r"""Construct a Geodesic geometry from an adjacency matrix. - - Args: - G: Adjacency matrix. - t: Time parameter for approximating the geodesic exponential kernel. - If `None`, it defaults to :math:`\frac{1}{|E|} \sum_{(u, v) \in E} - \text{weight}(u, v)` :cite:`crane:13`. In this case, the ``graph`` - must be specified and the edge weights are assumed to be positive. - eigval: Largest eigenvalue of the Laplacian. If :obj:`None`, it's - computed using :func:`jax.experimental.sparse.linalg.lobpcg_standard`. - order: Max order of Chebyshev polynomials. - directed: Whether the ``graph`` is directed. If :obj:`True`, it's made - undirected as :math:`G + G^T`. This parameter is ignored when passing - the Laplacian directly, assumed to be symmetric. - normalize: Whether to normalize the Laplacian as - :math:`L^{sym} = \left(D^+\right)^{\frac{1}{2}} L - \left(D^+\right)^{\frac{1}{2}}`, where :math:`L` is the - non-normalized Laplacian and :math:`D` is the degree matrix. - rng: Random key used when computing the largest eigenvalue. - kwargs: Keyword arguments for :class:`~ott.geometry.geodesic.Geodesic`. - - Returns: - The Geodesic geometry. - """ - assert G.shape[0] == G.shape[1], G.shape - rng = utils.default_prng_key(rng) - - if directed: - G = G + G.T - if t is None: - t = (jnp.sum(G) / jnp.sum(G > 0.0)) ** 2 - - degree = jnp.sum(G, axis=1) - laplacian = jnp.diag(degree) - G - if normalize: - inv_sqrt_deg = jnp.diag( - jnp.where(degree > 0.0, 1.0 / jnp.sqrt(degree), 0.0) - ) - laplacian = inv_sqrt_deg @ laplacian @ inv_sqrt_deg - - if eigval is None: - eigval = compute_largest_eigenvalue(laplacian, rng) - - scaled_laplacian, eigval = jax.lax.cond((eigval > 2.0), lambda l: - (2.0 * l / eigval, 2.0), lambda l: - (l, eigval), laplacian) - - # compute the coeffs of the Chebyshev pols approx using Bessel funcs - chebyshev_coeffs = compute_chebychev_coeff_all( - 0.5 * eigval, t, order, laplacian.dtype - ) - - return cls( - scaled_laplacian=scaled_laplacian, - eigval=eigval, - chebyshev_coeffs=chebyshev_coeffs, - t=t, - **kwargs - ) - - def apply_kernel( - self, - scaling: jnp.ndarray, - eps: Optional[float] = None, - axis: int = 0, - ) -> jnp.ndarray: - r"""Apply :attr:`kernel_matrix` on positive scaling vector. - - Args: - scaling: Scaling to apply the kernel to. - eps: passed for consistency, not used yet. - axis: passed for consistency, not used yet. - - Returns: - Kernel applied to ``scaling``. - """ - return expm_multiply( - self.scaled_laplacian, scaling, self.chebyshev_coeffs, 0.5 * self.eigval - ) - - @property - def kernel_matrix(self) -> jnp.ndarray: # noqa: D102 - n, _ = self.shape - kernel = self.apply_kernel(jnp.eye(n)) - return jax.lax.cond( - jnp.allclose(kernel, kernel.T, atol=1e-8, rtol=1e-8), lambda x: x, - lambda x: (x + x.T) / 2.0, kernel - ) - - @property - def cost_matrix(self) -> jnp.ndarray: # noqa: D102 - # Calculate the cost matrix using the formula (5) from the main reference - return -4.0 * self.t * mu.safe_log(self.kernel_matrix) - - @property - def shape(self) -> Tuple[int, int]: # noqa: D102 - return self.scaled_laplacian.shape - - @property - def is_symmetric(self) -> bool: # noqa: D102 - return True - - @property - def dtype(self) -> jnp.dtype: # noqa: D102 - return self.scaled_laplacian.dtype - - def transport_from_potentials( - self, f: jnp.ndarray, g: jnp.ndarray - ) -> jnp.ndarray: - """Not implemented.""" - raise ValueError("Not implemented.") - - def apply_transport_from_potentials( - self, - f: jnp.ndarray, - g: jnp.ndarray, - vec: jnp.ndarray, - axis: int = 0 - ) -> jnp.ndarray: - """Not implemented.""" - raise ValueError("Not implemented.") - - def marginal_from_potentials( - self, - f: jnp.ndarray, - g: jnp.ndarray, - axis: int = 0, - ) -> jnp.ndarray: - """Not implemented.""" - raise ValueError("Not implemented.") - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return [ - self.scaled_laplacian, - self.eigval, - self.chebyshev_coeffs, - self.t, - ], {} - - @classmethod - def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] - ) -> "Geodesic": - return cls(*children, **aux_data) - - -def compute_largest_eigenvalue( - laplacian_matrix: jnp.ndarray, - rng: jax.Array, -) -> float: - # Compute the largest eigenvalue of the Laplacian matrix. - n = laplacian_matrix.shape[0] - # Generate random initial directions for eigenvalue computation - initial_dirs = jax.random.normal(rng, (n, 1)) - - # Create a sparse matrix-vector product function using sparsify - # This function multiplies the sparse laplacian_matrix with a vector - lapl_vector_product = jesp.sparsify(lambda v: laplacian_matrix @ v) - - # Compute eigenvalues using the sparse matrix-vector product - eigvals, _, _ = jesp.linalg.lobpcg_standard( - lapl_vector_product, - initial_dirs, - ) - return eigvals[0] - - -def expm_multiply( - L: jnp.ndarray, X: jnp.ndarray, coeff: jnp.ndarray, eigval: float -) -> jnp.ndarray: - - def body(carry, c): - T0, T1, Y = carry - T2 = (2.0 / eigval) * L @ T1 - 2.0 * T1 - T0 - Y = Y + c * T2 - return (T1, T2, Y), None - - T0 = X - Y = 0.5 * coeff[0] * T0 - T1 = (1.0 / eigval) * L @ X - T0 - Y = Y + coeff[1] * T1 - - initial_state = (T0, T1, Y) - (_, _, Y), _ = jax.lax.scan(body, initial_state, coeff[2:]) - return Y - - -def compute_chebychev_coeff_all( - eigval: float, tau: float, K: int, dtype: np.dtype -) -> jnp.ndarray: - """Jax wrapper to compute the K+1 Chebychev coefficients.""" - result_shape_dtype = jax.ShapeDtypeStruct( - shape=(K + 1,), - dtype=dtype, - ) - - chebychev_coeff = lambda eigval, tau, K: ( - 2.0 * ive(np.arange(0, K + 1), -tau * eigval) - ).astype(dtype) - - return jax.pure_callback(chebychev_coeff, result_shape_dtype, eigval, tau, K) diff --git a/ott/build/lib/ott/geometry/geometry.py b/ott/build/lib/ott/geometry/geometry.py deleted file mode 100644 index bcbc724..0000000 --- a/ott/build/lib/ott/geometry/geometry.py +++ /dev/null @@ -1,921 +0,0 @@ -# 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 functools -from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Tuple, Union - -if TYPE_CHECKING: - from ott.geometry import low_rank - -import jax -import jax.numpy as jnp -import jax.scipy as jsp -import numpy as np - -from ott import utils -from ott.geometry import epsilon_scheduler -from ott.math import utils as mu - -__all__ = ["Geometry", "is_linear", "is_affine"] - - -@jax.tree_util.register_pytree_node_class -class Geometry: - r"""Base class to define ground costs/kernels used in optimal transport. - - Optimal transport problems are intrinsically geometric: they compute an - optimal way to transport mass from one configuration onto another. To define - what is meant by optimality of transport requires defining a cost, of moving - mass from one among several sources, towards one out of multiple targets. - These sources and targets can be provided as points in vectors spaces, grids, - or more generally exclusively described through a (dissimilarity) cost matrix, - or almost equivalently, a (similarity) kernel matrix. - - Once that cost or kernel matrix is set, the ``Geometry`` class provides a - basic operations to be run with the Sinkhorn algorithm. - - Args: - cost_matrix: Cost matrix of shape ``[n, m]``. - kernel_matrix: Kernel matrix of shape ``[n, m]``. - epsilon: Regularization parameter. If ``None`` and either - ``relative_epsilon = True`` or ``relative_epsilon = None``, this defaults - to the value computed in :attr:`mean_cost_matrix` / 20. If passed as a - ``float``, then the regularizer that is ultimately used is either that - ``float`` value (if ``relative_epsilon = False`` or ``None``) or that - ``float`` times the :attr:`mean_cost_matrix` - (if ``relative_epsilon = True``). Look for - :class:`~ott.geometry.epsilon_scheduler.Epsilon` when passed as a - scheduler. - relative_epsilon: when `False`, the parameter ``epsilon`` specifies the - value of the entropic regularization parameter. When `True`, ``epsilon`` - refers to a fraction of the :attr:`mean_cost_matrix`, which is computed - adaptively from data. - scale_cost: option to rescale the cost matrix. Implemented scalings are - 'median', 'mean' and 'max_cost'. Alternatively, a float factor can be - given to rescale the cost such that ``cost_matrix /= scale_cost``. - If `True`, use 'mean'. - src_mask: Mask specifying valid rows when computing some statistics of - :attr:`cost_matrix`, see :attr:`src_mask`. - tgt_mask: Mask specifying valid columns when computing some statistics of - :attr:`cost_matrix`, see :attr:`tgt_mask`. - - Note: - When defining a :class:`~ott.geometry.geometry.Geometry` through a - ``cost_matrix``, it is important to select an ``epsilon`` regularization - parameter that is meaningful. That parameter can be provided by the user, - or assigned a default value through a simple rule, - using the :attr:`mean_cost_matrix`. - """ - - def __init__( - self, - cost_matrix: Optional[jnp.ndarray] = None, - kernel_matrix: Optional[jnp.ndarray] = None, - epsilon: Optional[Union[float, epsilon_scheduler.Epsilon]] = None, - relative_epsilon: Optional[bool] = None, - scale_cost: Union[bool, int, float, Literal["mean", "max_cost", - "median"]] = 1.0, - src_mask: Optional[jnp.ndarray] = None, - tgt_mask: Optional[jnp.ndarray] = None, - ): - self._cost_matrix = cost_matrix - self._kernel_matrix = kernel_matrix - - # needed for `copy_epsilon`, because of the `isinstance` check - self._epsilon_init = epsilon if isinstance( - epsilon, epsilon_scheduler.Epsilon - ) else epsilon_scheduler.Epsilon(epsilon) - self._relative_epsilon = relative_epsilon - - self._scale_cost = "mean" if scale_cost is True else scale_cost - - self._src_mask = src_mask - self._tgt_mask = tgt_mask - - @property - def cost_rank(self) -> Optional[int]: - """Output rank of cost matrix, if any was provided.""" - - @property - def cost_matrix(self) -> jnp.ndarray: - """Cost matrix, recomputed from kernel if only kernel was specified.""" - if self._cost_matrix is None: - # If no epsilon was passed on to the geometry, then assume it is one by - # default. - eps = jnp.finfo(self._kernel_matrix.dtype).tiny - cost = -jnp.log(self._kernel_matrix + eps) - cost *= self.inv_scale_cost - return cost if self._epsilon_init is None else self.epsilon * cost - return self._cost_matrix * self.inv_scale_cost - - @property - def median_cost_matrix(self) -> float: - """Median of the :attr:`cost_matrix`.""" - geom = self._masked_geom(mask_value=jnp.nan) - return jnp.nanmedian(geom.cost_matrix) # will fail for online PC - - @property - def mean_cost_matrix(self) -> float: - """Mean of the :attr:`cost_matrix`.""" - tmp = self._masked_geom().apply_cost(self._n_normed_ones).squeeze() - return jnp.sum(tmp * self._m_normed_ones) - - @property - def kernel_matrix(self) -> jnp.ndarray: - """Kernel matrix. - - Either provided by user or recomputed from :attr:`cost_matrix`. - """ - if self._kernel_matrix is None: - return jnp.exp(-(self._cost_matrix * self.inv_scale_cost / self.epsilon)) - return self._kernel_matrix ** self.inv_scale_cost - - @property - def _epsilon(self) -> epsilon_scheduler.Epsilon: - (target, scale_eps, _, _), _ = self._epsilon_init.tree_flatten() - rel = self._relative_epsilon - - use_mean_scale = rel is True or (rel is None and target is None) - if scale_eps is None and use_mean_scale: - scale_eps = jax.lax.stop_gradient(self.mean_cost_matrix) - - if isinstance(self._epsilon_init, epsilon_scheduler.Epsilon): - return self._epsilon_init.set(scale_epsilon=scale_eps) - - return epsilon_scheduler.Epsilon( - target=5e-2 if target is None else target, scale_epsilon=scale_eps - ) - - @property - def epsilon(self) -> float: - """Epsilon regularization value.""" - return self._epsilon.target - - @property - def shape(self) -> Tuple[int, int]: - """Shape of the geometry.""" - mat = ( - self._kernel_matrix if self._cost_matrix is None else self._cost_matrix - ) - if mat is not None: - return mat.shape - return 0, 0 - - @property - def can_LRC(self) -> bool: - """Check quickly if casting geometry as LRC makes sense. - - This check is only carried out using basic considerations from the geometry, - not using a rigorous check involving, e.g., SVD. - """ - return False - - @property - def is_squared_euclidean(self) -> bool: - """Whether cost is computed by taking squared Euclidean distance.""" - return False - - @property - def is_online(self) -> bool: - """Whether geometry cost/kernel should be recomputed on the fly.""" - return False - - @property - def is_symmetric(self) -> bool: - """Whether geometry cost/kernel is a symmetric matrix.""" - mat = self.kernel_matrix if self.cost_matrix is None else self.cost_matrix - return ( - mat.shape[0] == mat.shape[1] and jnp.all(mat == mat.T) - ) if mat is not None else False - - @property - def inv_scale_cost(self) -> float: - """Compute and return inverse of scaling factor for cost matrix.""" - if isinstance(self._scale_cost, (int, float, np.number, jax.Array)): - return 1.0 / self._scale_cost - self = self._masked_geom(mask_value=jnp.nan) - if self._scale_cost == "max_cost": - return 1.0 / jnp.nanmax(self._cost_matrix) - if self._scale_cost == "mean": - return 1.0 / jnp.nanmean(self._cost_matrix) - if self._scale_cost == "median": - return 1.0 / jnp.nanmedian(self._cost_matrix) - raise ValueError(f"Scaling {self._scale_cost} not implemented.") - - def set_scale_cost(self, scale_cost: Union[bool, float, str]) -> "Geometry": - """Modify how to rescale of the :attr:`cost_matrix`.""" - # case when `geom` doesn't have `scale_cost` or doesn't need to be modified - # `False` retains the original scale - if scale_cost is False or scale_cost == self._scale_cost: - return self - children, aux_data = self.tree_flatten() - aux_data["scale_cost"] = scale_cost - return type(self).tree_unflatten(aux_data, children) - - def copy_epsilon(self, other: "Geometry") -> "Geometry": - """Copy the epsilon parameters from another geometry.""" - other_epsilon = other._epsilon - children, aux_data = self.tree_flatten() - - new_children = [] - for child in children: - if isinstance(child, epsilon_scheduler.Epsilon): - child = child.set( - target=other_epsilon._target_init, - scale_epsilon=other_epsilon._scale_epsilon - ) - new_children.append(child) - - aux_data["relative_epsilon"] = False - return type(self).tree_unflatten(aux_data, new_children) - - # The functions below are at the core of Sinkhorn iterations, they - # are implemented here in their default form, either in lse (using directly - # cost matrices in stabilized form) or kernel mode (using kernel matrices). - - def apply_lse_kernel( - self, - f: jnp.ndarray, - g: jnp.ndarray, - eps: float, - vec: jnp.ndarray = None, - axis: int = 0 - ) -> jnp.ndarray: - r"""Apply :attr:`kernel_matrix` in log domain. - - This function applies the ground geometry's kernel in log domain, using - a stabilized formulation. At a high level, this iteration performs either: - - - output = eps * log (K (exp(g / eps) * vec)) (1) - - output = eps * log (K'(exp(f / eps) * vec)) (2) - - K is implicitly exp(-cost_matrix/eps). - - To carry this out in a stabilized way, we take advantage of the fact that - the entries of the matrix ``f[:,*] + g[*,:] - C`` are all negative, and - therefore their exponential never overflows, to add (and subtract after) - f and g in iterations 1 & 2 respectively. - - Args: - f: jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix - g: jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix - eps: float, regularization strength - vec: jnp.ndarray [num_a or num_b,] , when not None, this has the effect of - doing log-Kernel computations with an addition elementwise - multiplication of exp(g / eps) by a vector. This is carried out by - adding weights to the log-sum-exp function, and needs to handle signs - separately. - axis: summing over axis 0 when doing (2), or over axis 1 when doing (1) - - Returns: - A jnp.ndarray corresponding to output above, depending on axis. - """ - w_res, w_sgn = self._softmax(f, g, eps, vec, axis) - remove = f if axis == 1 else g - return w_res - jnp.where(jnp.isfinite(remove), remove, 0), w_sgn - - def apply_kernel( - self, - scaling: jnp.ndarray, - eps: Optional[float] = None, - axis: int = 0, - ) -> jnp.ndarray: - """Apply :attr:`kernel_matrix` on positive scaling vector. - - Args: - scaling: jnp.ndarray [num_a or num_b] , scaling of size num_rows or - num_cols of kernel_matrix - eps: passed for consistency, not used yet. - axis: standard kernel product if axis is 1, transpose if 0. - - Returns: - a jnp.ndarray corresponding to output above, depending on axis. - """ - if eps is None: - kernel = self.kernel_matrix - else: - kernel = self.kernel_matrix ** (self.epsilon / eps) - kernel = kernel if axis == 1 else kernel.T - - return jnp.dot(kernel, scaling) - - def marginal_from_potentials( - self, - f: jnp.ndarray, - g: jnp.ndarray, - axis: int = 0, - ) -> jnp.ndarray: - """Output marginal of transportation matrix from potentials. - - This applies first lse kernel in the standard way, removes the - correction used to stabilize computations, and lifts this with an exp to - recover either of the marginals corresponding to the transport map induced - by potentials. - - Args: - f: jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix - g: jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix - axis: axis along which to integrate, returns marginal on other axis. - - Returns: - a vector of marginals of the transport matrix. - """ - h = (f if axis == 1 else g) - z = self.apply_lse_kernel(f, g, self.epsilon, axis=axis)[0] - return jnp.exp((z + h) / self.epsilon) - - def marginal_from_scalings( - self, - u: jnp.ndarray, - v: jnp.ndarray, - axis: int = 0, - ) -> jnp.ndarray: - """Output marginal of transportation matrix from scalings.""" - u, v = (v, u) if axis == 0 else (u, v) - return u * self.apply_kernel(v, eps=self.epsilon, axis=axis) - - def transport_from_potentials( - self, f: jnp.ndarray, g: jnp.ndarray - ) -> jnp.ndarray: - """Output transport matrix from potentials.""" - return jnp.exp(self._center(f, g) / self.epsilon) - - def transport_from_scalings( - self, u: jnp.ndarray, v: jnp.ndarray - ) -> jnp.ndarray: - """Output transport matrix from pair of scalings.""" - return self.kernel_matrix * u[:, jnp.newaxis] * v[jnp.newaxis, :] - - # Functions that are not supposed to be changed by inherited classes. - # These are the point of entry for Sinkhorn's algorithm to use a geometry. - - def update_potential( - self, - f: jnp.ndarray, - g: jnp.ndarray, - log_marginal: jnp.ndarray, - iteration: Optional[int] = None, - axis: int = 0, - ) -> jnp.ndarray: - """Carry out one Sinkhorn update for potentials, i.e. in log space. - - Args: - f: jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix - g: jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix - log_marginal: targeted marginal - iteration: used to compute epsilon from schedule, if provided. - axis: axis along which the update should be carried out. - - Returns: - new potential value, g if axis=0, f if axis is 1. - """ - eps = self._epsilon.at(iteration) - app_lse = self.apply_lse_kernel(f, g, eps, axis=axis)[0] - return eps * log_marginal - jnp.where(jnp.isfinite(app_lse), app_lse, 0) - - def update_scaling( - self, - scaling: jnp.ndarray, - marginal: jnp.ndarray, - iteration: Optional[int] = None, - axis: int = 0, - ) -> jnp.ndarray: - """Carry out one Sinkhorn update for scalings, using kernel directly. - - Args: - scaling: jnp.ndarray of num_a or num_b positive values. - marginal: targeted marginal - iteration: used to compute epsilon from schedule, if provided. - axis: axis along which the update should be carried out. - - Returns: - new scaling vector, of size num_b if axis=0, num_a if axis is 1. - """ - eps = self._epsilon.at(iteration) - app_kernel = self.apply_kernel(scaling, eps, axis=axis) - return marginal / jnp.where(app_kernel > 0, app_kernel, 1.0) - - # Helper functions - def _center(self, f: jnp.ndarray, g: jnp.ndarray) -> jnp.ndarray: - return f[:, jnp.newaxis] + g[jnp.newaxis, :] - self.cost_matrix - - def _softmax( - self, f: jnp.ndarray, g: jnp.ndarray, eps: float, - vec: Optional[jnp.ndarray], axis: int - ) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Apply softmax row or column wise, weighted by vec.""" - if vec is not None: - if axis == 0: - vec = vec.reshape((-1, 1)) - lse_output = mu.logsumexp( - self._center(f, g) / eps, b=vec, axis=axis, return_sign=True - ) - return eps * lse_output[0], lse_output[1] - lse_output = mu.logsumexp( - self._center(f, g) / eps, axis=axis, return_sign=False - ) - return eps * lse_output, jnp.array([1.0]) - - @functools.partial(jax.vmap, in_axes=[None, None, None, 0, None]) - def _apply_transport_from_potentials( - self, f: jnp.ndarray, g: jnp.ndarray, vec: jnp.ndarray, axis: int - ) -> jnp.ndarray: - """Apply lse_kernel to arbitrary vector while keeping track of signs.""" - lse_res, lse_sgn = self.apply_lse_kernel( - f, g, self.epsilon, vec=vec, axis=axis - ) - lse_res += f if axis == 1 else g - return lse_sgn * jnp.exp(lse_res / self.epsilon) - - # wrapper to allow default option for axis. - def apply_transport_from_potentials( - self, - f: jnp.ndarray, - g: jnp.ndarray, - vec: jnp.ndarray, - axis: int = 0 - ) -> jnp.ndarray: - """Apply transport matrix computed from potentials to a (batched) vec. - - This approach does not instantiate the transport matrix itself, but uses - instead potentials to apply the transport using apply_lse_kernel, therefore - guaranteeing stability and lower memory footprint. - - Computations are done in log space, and take advantage of the - (b=..., return_sign=True) optional parameters of logsumexp. - - Args: - f: jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix - g: jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix - vec: jnp.ndarray [batch, num_a or num_b], vector that will be multiplied - by transport matrix corresponding to potentials f, g, and geom. - axis: axis to differentiate left (0) or right (1) multiply. - - Returns: - ndarray of the size of vec. - """ - if vec.ndim == 1: - return self._apply_transport_from_potentials( - f, g, vec[jnp.newaxis, :], axis - )[0, :] - return self._apply_transport_from_potentials(f, g, vec, axis) - - @functools.partial(jax.vmap, in_axes=[None, None, None, 0, None]) - def _apply_transport_from_scalings( - self, u: jnp.ndarray, v: jnp.ndarray, vec: jnp.ndarray, axis: int - ): - u, v = (u, v * vec) if axis == 1 else (v, u * vec) - return u * self.apply_kernel(v, eps=self.epsilon, axis=axis) - - # wrapper to allow default option for axis - def apply_transport_from_scalings( - self, - u: jnp.ndarray, - v: jnp.ndarray, - vec: jnp.ndarray, - axis: int = 0 - ) -> jnp.ndarray: - """Apply transport matrix computed from scalings to a (batched) vec. - - This approach does not instantiate the transport matrix itself, but - relies instead on the apply_kernel function. - - Args: - u: jnp.ndarray [num_a,] , scaling of size num_rows of cost_matrix - v: jnp.ndarray [num_b,] , scaling of size num_cols of cost_matrix - vec: jnp.ndarray [batch, num_a or num_b], vector that will be multiplied - by transport matrix corresponding to scalings u, v, and geom. - axis: axis to differentiate left (0) or right (1) multiply. - - Returns: - ndarray of the size of vec. - """ - if vec.ndim == 1: - return self._apply_transport_from_scalings( - u, v, vec[jnp.newaxis, :], axis - )[0, :] - return self._apply_transport_from_scalings(u, v, vec, axis) - - def potential_from_scaling(self, scaling: jnp.ndarray) -> jnp.ndarray: - """Compute dual potential vector from scaling vector. - - Args: - scaling: vector. - - Returns: - a vector of the same size. - """ - return self.epsilon * jnp.log(scaling) - - def scaling_from_potential(self, potential: jnp.ndarray) -> jnp.ndarray: - """Compute scaling vector from dual potential. - - Args: - potential: vector. - - Returns: - a vector of the same size. - """ - finite = jnp.isfinite(potential) - return jnp.where( - finite, jnp.exp(jnp.where(finite, potential / self.epsilon, 0.0)), 0.0 - ) - - def apply_square_cost(self, arr: jnp.ndarray, axis: int = 0, sub_idx = None) -> jnp.ndarray: - """Apply elementwise-square of cost matrix to array (vector or matrix). - - This function applies the ground 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], when the cost matrix itself is computed as a - squared-Euclidean distance between vectors, and therefore admits an - explicit low-rank factorization. - - Args: - arr: array. - axis: axis of the array on which the cost matrix should be applied. - - Returns: - An array, [num_b, p] if axis=0 or [num_a, p] if axis=1. - """ - return self.apply_cost(arr, axis=axis, fn=lambda x: x ** 2, sub_idx = sub_idx) - - def apply_cost( - self, - arr: jnp.ndarray, - axis: int = 0, - fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, - sub_idx = None, - **kwargs: Any - ) -> jnp.ndarray: - """Apply :attr:`cost_matrix` to array (vector or matrix). - - This function applies the ground 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] - - 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, transpose if 0 - fn: function to apply to cost matrix element-wise before the dot product - kwargs: Keyword arguments for :meth:`_apply_cost_to_vec`. - - Returns: - An array, [num_b, p] if axis=0 or [num_a, p] if axis=1 - """ - if arr.ndim == 1: - arr = arr.reshape(-1, 1) - - app = functools.partial(self._apply_cost_to_vec, axis=axis, fn=fn, sub_idx=sub_idx,**kwargs) - return jax.vmap(app, in_axes=1, out_axes=1)(arr) - - def _apply_cost_to_vec( - self, - vec: jnp.ndarray, - axis: int = 0, - fn=None, - sub_idx = None, - **_: Any, - ) -> jnp.ndarray: - """Apply ``[num_a, num_b]`` fn(cost) (or transpose) to vector. - - Args: - vec: jnp.ndarray [num_a,] ([num_b,] if axis=1) vector - axis: axis on which the reduction is done. - fn: function optionally applied to cost matrix element-wise, before the - doc product - - Returns: - A jnp.ndarray corresponding to cost x vector - """ - if sub_idx is None: - matrix = self.cost_matrix.T if axis == 0 else self.cost_matrix - matrix = fn(matrix) if fn is not None else matrix - return jnp.dot(matrix, vec) - matrix = self.cost_matrix[sub_idx, :].T if axis == 0 else self.cost_matrix[:, sub_idx] - matrix = fn(matrix) if fn is not None else matrix - return jnp.dot(matrix, vec[sub_idx]) - - @classmethod - def prepare_divergences( - cls, - *args: Any, - static_b: bool = False, - **kwargs: Any - ) -> Tuple["Geometry", ...]: - """Instantiate 2 (or 3) geometries to compute a Sinkhorn divergence.""" - size = 2 if static_b else 3 - nones = [None, None, None] - cost_matrices = kwargs.pop("cost_matrix", args) - kernel_matrices = kwargs.pop("kernel_matrix", nones) - cost_matrices = cost_matrices if cost_matrices is not None else nones - return tuple( - cls(cost_matrix=arg1, kernel_matrix=arg2, **kwargs) - for arg1, arg2, _ in zip(cost_matrices, kernel_matrices, range(size)) - ) - - def to_LRCGeometry( - self, - rank: int = 0, - tol: float = 1e-2, - rng: Optional[jax.Array] = None, - scale: float = 1.0 - ) -> "low_rank.LRCGeometry": - r"""Factorize the cost matrix using either SVD (full) or :cite:`indyk:19`. - - When `rank=min(n,m)` or `0` (by default), use :func:`jax.numpy.linalg.svd`. - - For other values, use the routine in sublinear time :cite:`indyk:19`. - Uses the implementation of :cite:`scetbon:21`, algorithm 4. - - It holds that with probability *0.99*, - :math:`||A - UV||_F^2 \leq || A - A_k ||_F^2 + tol \cdot ||A||_F^2`, - where :math:`A` is ``n x m`` cost matrix, :math:`UV` the factorization - computed in sublinear time and :math:`A_k` the best rank-k approximation. - - Args: - rank: Target rank of the :attr:`cost_matrix`. - tol: Tolerance of the error. The total number of sampled points is - :math:`min(n, m,\frac{rank}{tol})`. - rng: The PRNG key to use for initializing the model. - 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. - - Returns: - Low-rank geometry. - """ - from ott.geometry import low_rank - assert rank >= 0, f"Rank must be non-negative, got {rank}." - n, m = self.shape - - if rank == 0 or rank >= min(n, m): - # TODO(marcocuturi): add hermitian=self.is_symmetric, currently bugging. - u, s, vh = jnp.linalg.svd( - self.cost_matrix, - full_matrices=False, - compute_uv=True, - ) - - cost_1 = u - cost_2 = (s[:, None] * vh).T - else: - rng = utils.default_prng_key(rng) - rng1, rng2, rng3, rng4, rng5 = jax.random.split(rng, 5) - n_subset = min(int(rank / tol), n, m) - - i_star = jax.random.randint(rng1, shape=(), minval=0, maxval=n) - j_star = jax.random.randint(rng2, shape=(), minval=0, maxval=m) - - ci_star = self.subset([i_star], None).cost_matrix.ravel() ** 2 # (m,) - cj_star = self.subset(None, [j_star]).cost_matrix.ravel() ** 2 # (n,) - - p_row = cj_star + ci_star[j_star] + jnp.mean(ci_star) # (n,) - p_row /= jnp.sum(p_row) - row_ixs = jax.random.choice(rng3, n, shape=(n_subset,), p=p_row) - # (n_subset, m) - s = self.subset(row_ixs, None).cost_matrix - s /= jnp.sqrt(n_subset * p_row[row_ixs][:, None]) - - p_col = jnp.sum(s ** 2, axis=0) # (m,) - p_col /= jnp.sum(p_col) - # (n_subset,) - col_ixs = jax.random.choice(rng4, m, shape=(n_subset,), p=p_col) - # (n_subset, n_subset) - w = s[:, col_ixs] / jnp.sqrt(n_subset * p_col[col_ixs][None, :]) - - U, _, V = jsp.linalg.svd(w) - U = U[:, :rank] # (n_subset, rank) - U = (s.T @ U) / jnp.linalg.norm(w.T @ U, axis=0) # (m, rank) - - _, d, v = jnp.linalg.svd(U.T @ U) # (k,), (k, k) - v = v.T / jnp.sqrt(d)[None, :] - - inv_scale = (1.0 / jnp.sqrt(n_subset)) - col_ixs = jax.random.choice(rng5, m, shape=(n_subset,)) # (n_subset,) - - # (n, n_subset) - A_trans = self.subset(None, col_ixs).cost_matrix * inv_scale - B = (U[col_ixs, :] @ v * inv_scale) # (n_subset, k) - M = jnp.linalg.inv(B.T @ B) # (k, k) - V = jnp.linalg.multi_dot([A_trans, B, M.T, v.T]) # (n, k) - cost_1 = V - cost_2 = U - - return low_rank.LRCGeometry( - cost_1=cost_1, - cost_2=cost_2, - epsilon=self._epsilon_init, - relative_epsilon=self._relative_epsilon, - scale_cost=self._scale_cost, - scale_factor=scale, - ) - - def subset( - self, src_ixs: Optional[jnp.ndarray], tgt_ixs: Optional[jnp.ndarray], - **kwargs: Any - ) -> "Geometry": - """Subset rows or columns of a geometry. - - Args: - src_ixs: Row indices. If ``None``, use all rows. - tgt_ixs: Column indices. If ``None``, use all columns. - kwargs: Keyword arguments to override the initialization. - - Returns: - The modified geometry. - """ - - def subset_fn( - arr: Optional[jnp.ndarray], - src_ixs: Optional[jnp.ndarray], - tgt_ixs: Optional[jnp.ndarray], - ) -> Optional[jnp.ndarray]: - if arr is None: - return None - if src_ixs is not None: - arr = arr[src_ixs, ...] - if tgt_ixs is not None: - arr = arr[:, tgt_ixs] - return arr # noqa: RET504 - - return self._mask_subset_helper( - src_ixs, - tgt_ixs, - fn=subset_fn, - propagate_mask=True, - **kwargs, - ) - - def mask( - self, - src_mask: Optional[jnp.ndarray], - tgt_mask: Optional[jnp.ndarray], - mask_value: float = 0.0, - ) -> "Geometry": - """Mask rows or columns of a geometry. - - The mask is used only when computing some statistics of the - :attr:`cost_matrix`. - - - :attr:`mean_cost_matrix` - - :attr:`median_cost_matrix` - - :attr:`inv_scale_cost` - - Args: - src_mask: Row mask. Can be specified either as a boolean array of shape - ``[num_a,]`` or as an array of indices. If ``None``, no mask is applied. - tgt_mask: Column mask. Can be specified either as a boolean array of shape - ``[num_b,]`` or as an array of indices. If ``None``, no mask is applied. - mask_value: Value to use for masking. - - Returns: - The masked geometry. - """ - - def mask_fn( - arr: Optional[jnp.ndarray], - src_mask: Optional[jnp.ndarray], - tgt_mask: Optional[jnp.ndarray], - ) -> Optional[jnp.ndarray]: - if arr is None: - return arr - assert arr.ndim == 2, arr.ndim - if src_mask is not None: - arr = jnp.where(src_mask[:, None], arr, mask_value) - if tgt_mask is not None: - arr = jnp.where(tgt_mask[None, :], arr, mask_value) - return arr # noqa: RET504 - - 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]], - Optional[jnp.ndarray]], - propagate_mask: bool, - **kwargs: Any, - ) -> "Geometry": - (cost, kernel, eps, src_mask, tgt_mask), aux_data = self.tree_flatten() - cost = fn(cost, src_ixs, tgt_ixs) - kernel = fn(kernel, src_ixs, 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, None) - tgt_mask = fn(tgt_mask, tgt_ixs, None) - - aux_data = {**aux_data, **kwargs} - return type(self).tree_unflatten( - aux_data, [cost, kernel, eps, src_mask, tgt_mask] - ) - - @property - def src_mask(self) -> Optional[jnp.ndarray]: - """Mask of shape ``[num_a,]`` to compute :attr:`cost_matrix` statistics. - - Specifically, it is used when computing: - - - :attr:`mean_cost_matrix` - - :attr:`median_cost_matrix` - - :attr:`inv_scale_cost` - """ - return self._normalize_mask(self._src_mask, self.shape[0]) - - @property - def tgt_mask(self) -> Optional[jnp.ndarray]: - """Mask of shape ``[num_b,]`` to compute :attr:`cost_matrix` statistics. - - Specifically, it is used when computing: - - - :attr:`mean_cost_matrix` - - :attr:`median_cost_matrix` - - :attr:`inv_scale_cost` - """ - return self._normalize_mask(self._tgt_mask, self.shape[1]) - - @property - def dtype(self) -> jnp.dtype: - """The data type.""" - return ( - self._kernel_matrix if self._cost_matrix is None else self._cost_matrix - ).dtype - - def _masked_geom(self, mask_value: float = 0.0) -> "Geometry": - """Mask geometry based on :attr:`src_mask` and :attr:`tgt_mask`.""" - src_mask, tgt_mask = self.src_mask, self.tgt_mask - if src_mask is None and tgt_mask is None: - return self - return self.mask(src_mask, tgt_mask, mask_value=mask_value) - - @property - def _n_normed_ones(self) -> jnp.ndarray: - """Normalized array of shape ``[num_a,]``.""" - mask = self.src_mask - arr = jnp.ones(self.shape[0]) if mask is None else mask - return arr / jnp.sum(arr) - - @property - def _m_normed_ones(self) -> jnp.ndarray: - """Normalized array of shape ``[num_b,]``.""" - mask = self.tgt_mask - arr = jnp.ones(self.shape[1]) if mask is None else mask - return arr / jnp.sum(arr) - - @staticmethod - def _normalize_mask(mask: Optional[jnp.ndarray], - size: int) -> Optional[jnp.ndarray]: - """Convert array of indices to a boolean mask.""" - if mask is None: - return None - if not jnp.issubdtype(mask, (bool, jnp.bool_)): - mask = jnp.isin(jnp.arange(size), mask) - assert mask.shape == (size,) - return mask - - def tree_flatten(self): # noqa: D102 - return ( - self._cost_matrix, self._kernel_matrix, self._epsilon_init, - self._src_mask, self._tgt_mask - ), { - "scale_cost": self._scale_cost, - "relative_epsilon": self._relative_epsilon - } - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - cost, kernel, eps, src_mask, tgt_mask = children - return cls( - cost, kernel, eps, src_mask=src_mask, tgt_mask=tgt_mask, **aux_data - ) - - -def is_affine(fn) -> bool: - """Test heuristically if a function is affine.""" - x = jnp.arange(10.0) - out = jax.vmap(jax.grad(fn))(x) - return jnp.sum(jnp.diff(jnp.abs(out))) == 0.0 - - -def is_linear(fn) -> bool: - """Test heuristically if a function is linear.""" - return jnp.logical_and(fn(0.0) == 0.0, is_affine(fn)) diff --git a/ott/build/lib/ott/geometry/graph.py b/ott/build/lib/ott/geometry/graph.py deleted file mode 100644 index 4a97707..0000000 --- a/ott/build/lib/ott/geometry/graph.py +++ /dev/null @@ -1,272 +0,0 @@ -# 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. -from typing import Any, Dict, Literal, Optional, Sequence, Tuple - -import jax -import jax.numpy as jnp -import jax.scipy as jsp - -from ott.geometry import geometry -from ott.math import fixed_point_loop -from ott.math import utils as mu - -__all__ = ["Graph"] - - -@jax.tree_util.register_pytree_node_class -class Graph(geometry.Geometry): - r"""Graph distance approximation using heat kernel :cite:`heitz:21,crane:13`. - - Approximates the heat kernel for large ``n_steps``, which for small ``t`` - approximates the geodesic exponential kernel :math:`e^{\frac{-d(x, y)^2}{t}}`. - - Args: - laplacian: Symmetric graph Laplacian. The check for symmetry is **NOT** - performed. See also :meth:`from_graph`. - n_steps: Maximum number of steps used to approximate the heat kernel. - numerical_scheme: Numerical scheme used to solve the heat diffusion. - normalize: Whether to normalize the Laplacian as - :math:`L^{sym} = \left(D^+\right)^{\frac{1}{2}} L - \left(D^+\right)^{\frac{1}{2}}`, where :math:`L` is the - non-normalized Laplacian and :math:`D` is the degree matrix. - tol: Relative tolerance with respect to the Hilbert metric, see - :cite:`peyre:19`, Remark 4.12. Used when iteratively updating scalings. - If negative, this option is ignored and only ``n_steps`` is used. - kwargs: Keyword arguments for :class:`~ott.geometry.geometry.Geometry`. - """ - - def __init__( - self, - laplacian: jnp.ndarray, - t: float = 1e-3, - n_steps: int = 100, - numerical_scheme: Literal["backward_euler", - "crank_nicolson"] = "backward_euler", - tol: float = -1.0, - **kwargs: Any - ): - super().__init__(epsilon=1.0, **kwargs) - self.laplacian = laplacian - self.t = t - self.n_steps = n_steps - self.numerical_scheme = numerical_scheme - self.tol = tol - - @classmethod - def from_graph( - cls, - G: jnp.ndarray, - t: Optional[float] = 1e-3, - directed: bool = False, - normalize: bool = False, - **kwargs: Any - ) -> "Graph": - r"""Construct :class:`~ott.geometry.graph.Graph` from an adjacency matrix. - - Args: - G: Adjacency matrix. - t: Constant used when approximating the geodesic exponential kernel. - If `None`, use :math:`\frac{1}{|E|} \sum_{(u, v) \in E} weight(u, v)` - :cite:`crane:13`. In this case, the ``graph`` must be specified - and the edge weights are all assumed to be positive. - directed: Whether the ``graph`` is directed. If not, it will be made - undirected as :math:`G + G^T`. This parameter is ignored when directly - passing the Laplacian, which is assumed to be symmetric. - normalize: Whether to normalize the Laplacian as - :math:`L^{sym} = \left(D^+\right)^{\frac{1}{2}} L - \left(D^+\right)^{\frac{1}{2}}`, where :math:`L` is the - non-normalized Laplacian and :math:`D` is the degree matrix. - kwargs: Keyword arguments for :class:`~ott.geometry.graph.Graph`. - - Returns: - The graph geometry. - """ - assert G.shape[0] == G.shape[1], G.shape - - if directed: - G = G + G.T - - degree = jnp.sum(G, axis=1) - laplacian = jnp.diag(degree) - G - - if normalize: - inv_sqrt_deg = jnp.diag( - jnp.where(degree > 0.0, 1.0 / jnp.sqrt(degree), 0.0) - ) - laplacian = inv_sqrt_deg @ laplacian @ inv_sqrt_deg - - if t is None: - t = (jnp.sum(G) / jnp.sum(G > 0.0)) ** 2 - - return cls(laplacian, t=t, **kwargs) - - def apply_kernel( - self, - scaling: jnp.ndarray, - eps: Optional[float] = None, - axis: int = 0, - ) -> jnp.ndarray: - r"""Apply :attr:`kernel_matrix` on positive scaling vector. - - Args: - scaling: Scaling to apply the kernel to. - eps: passed for consistency, not used yet. - axis: passed for consistency, not used yet. - - Returns: - Kernel applied to ``scaling``. - """ - - def conf_fn( - iteration: int, consts: Tuple[jnp.ndarray, Optional[jnp.ndarray]], - old_new: Tuple[jnp.ndarray, jnp.ndarray] - ) -> bool: - del iteration, consts - - x_old, x_new = old_new - x_old, x_new = mu.safe_log(x_old), mu.safe_log(x_new) - # center - x_old, x_new = x_old - jnp.nanmax(x_old), x_new - jnp.nanmax(x_new) - # Hilbert metric, see Remark 4.12 in `Computational Optimal Transport` - f = x_new - x_old - return (jnp.nanmax(f) - jnp.nanmin(f)) > self.tol - - def body_fn( - iteration: int, consts: Tuple[jnp.ndarray, Optional[jnp.ndarray]], - old_new: Tuple[jnp.ndarray, jnp.ndarray], compute_errors: bool - ) -> Tuple[jnp.ndarray, jnp.ndarray]: - del iteration, compute_errors - - L, scaled_lap = consts - _, b = old_new - - if self.numerical_scheme == "crank_nicolson": - # below is a preferred way of specifying the update (albeit more FLOPS), - # as CSR/CSC/COO matrices don't support adding a diagonal matrix now: - # b' = (2 * I - M) @ b = (2 * I - (I + c * L)) @ b = (I - c * L) @ b - b = b - scaled_lap @ b - return b, jsp.linalg.solve_triangular(L, b, lower=True) - - # eps we cannot use since it would require a re-solve - # axis we can ignore since the matrix is symmetric - del eps, axis - - force_scan = self.tol < 0.0 - fixpoint_fn = ( - fixed_point_loop.fixpoint_iter - if force_scan else fixed_point_loop.fixpoint_iter_backprop - ) - - state = (jnp.full_like(scaling, jnp.nan), scaling) - L = jsp.linalg.cholesky(self._M, lower=True) - if self.numerical_scheme == "crank_nicolson": - constants = L, self._scaled_laplacian - else: - constants = L, None - - return fixpoint_fn( - cond_fn=(lambda *_, **__: True) if force_scan else conf_fn, - body_fn=body_fn, - min_iterations=self.n_steps if force_scan else 1, - max_iterations=self.n_steps, - inner_iterations=1, - constants=constants, - state=state, - )[1] - - @property - def kernel_matrix(self) -> jnp.ndarray: # noqa: D102 - n, _ = self.shape - kernel = self.apply_kernel(jnp.eye(n)) - # Symmetrize the kernel if needed. Numerical imprecision - # happens when `numerical_scheme='backward_euler'` and small `t` - return jax.lax.cond( - jnp.allclose(kernel, kernel.T, atol=1e-8, rtol=1e-8), lambda x: x, - lambda x: (x + x.T) / 2.0, kernel - ) - - @property - def cost_matrix(self) -> jnp.ndarray: # noqa: D102 - return -self.t * mu.safe_log(self.kernel_matrix) - - @property - def _scale(self) -> float: - """Constant used to scale the Laplacian.""" - if self.numerical_scheme == "backward_euler": - return self.t / (4.0 * self.n_steps) - if self.numerical_scheme == "crank_nicolson": - return self.t / (2.0 * self.n_steps) - raise NotImplementedError( - f"Numerical scheme `{self.numerical_scheme}` is not implemented." - ) - - @property - def _scaled_laplacian(self) -> jnp.ndarray: - """Laplacian scaled by a constant, depending on the numerical scheme.""" - return self._scale * self.laplacian - - @property - def _M(self) -> jnp.ndarray: - n, _ = self.shape - return self._scaled_laplacian + jnp.eye(n) - - @property - def shape(self) -> Tuple[int, int]: # noqa: D102 - return self.laplacian.shape - - @property - def is_symmetric(self) -> bool: # noqa: D102 - return True - - @property - def dtype(self) -> jnp.dtype: # noqa: D102 - return self.laplacian.dtype - - def transport_from_potentials( - self, f: jnp.ndarray, g: jnp.ndarray - ) -> jnp.ndarray: - """Not implemented.""" - raise ValueError("Not implemented.") - - def apply_transport_from_potentials( - self, - f: jnp.ndarray, - g: jnp.ndarray, - vec: jnp.ndarray, - axis: int = 0 - ) -> jnp.ndarray: - """Not implemented.""" - raise ValueError("Not implemented.") - - def marginal_from_potentials( - self, - f: jnp.ndarray, - g: jnp.ndarray, - axis: int = 0, - ) -> jnp.ndarray: - """Not implemented.""" - raise ValueError("Not implemented.") - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return [self.laplacian, self.t], { - "n_steps": self.n_steps, - "numerical_scheme": self.numerical_scheme, - "tol": self.tol, - } - - @classmethod - def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] - ) -> "Graph": - return cls(*children, **aux_data) diff --git a/ott/build/lib/ott/geometry/grid.py b/ott/build/lib/ott/geometry/grid.py deleted file mode 100644 index 708753d..0000000 --- a/ott/build/lib/ott/geometry/grid.py +++ /dev/null @@ -1,415 +0,0 @@ -# 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 itertools -from typing import Any, List, NoReturn, Optional, Sequence, Tuple - -import jax -import jax.numpy as jnp -import numpy as np - -from ott.geometry import costs, geometry, low_rank, pointcloud -from ott.math import utils - -__all__ = ["Grid"] - - -@jax.tree_util.register_pytree_node_class -class Grid(geometry.Geometry): - r"""Class describing the geometry of points taken in a Cartesian product. - - This class implements a geometry in which probability measures are supported - on a :math:`d`-dimensional Cartesian grid, a Cartesian product of :math:`d` - lists of values, each list being itself of size :math:`n_i`. - - The transportation cost between points in the grid is assumed to be separable, - namely a sum of coordinate-wise cost functions, as in: - - .. math:: - - cost(x,y) = \sum_{i=1}^d cost_i(x_i, y_i) - - where :math:`cost_i`: R x R → R. - - In such a regime, and despite the fact that the total number :math:`n_{total}` - of points in the grid is exponential :math:`d` (namely :math:`\prod_i n_i`), - applying a kernel in the context of regularized optimal transport can be - carried out in time that is of the order of :math:`n_{total}^{(1+1/d)}` using - convolutions, either in the original domain or log-space domain. This class - precomputes :math:`d` :math:`n_i` x :math:`n_i` cost matrices (one per - dimension) and implements these two operations by carrying out these - convolutions one dimension at a time. - - Args: - x: list of arrays of varying sizes, describing the locations of the grid. - Locations are provided as a list of arrays, that is :math:`d` - vectors of (possibly varying) size :math:`n_i`. The resulting grid - is the Cartesian product of these vectors. - grid_size: tuple of integers describing grid sizes, namely - :math:`(n_1,...,n_d)`. This will only be used if x is None. - In that case the grid will be assumed to lie in the hypercube - :math:`[0,1]^d`, with the :math:`d` dimensions, described as points - regularly sampled in :math:`[0,1]`. - cost_fns: a sequence of :math:`d` cost functions, each being a cost taking - two reals as inputs to output a real number. - num_a: total size of grid. This parameters will be computed from other - inputs. - grid_dimension: dimension of grid. This parameters will be computed from - other inputs. - kwargs: keyword arguments for :class:`~ott.geometry.geometry.Geometry`. - """ - - def __init__( - self, - x: Optional[Sequence[jnp.ndarray]] = None, - grid_size: Optional[Sequence[int]] = None, - cost_fns: Optional[Sequence[costs.CostFn]] = None, - num_a: Optional[int] = None, - grid_dimension: Optional[int] = None, - **kwargs: Any, - ): - if ( - grid_size is not None and x is not None and num_a is not None and - grid_dimension is not None - ): - self.grid_size = tuple(map(int, grid_size)) - self.x = x - self.num_a = num_a - self.grid_dimension = grid_dimension - elif x is not None: - self.x = x - self.grid_size = tuple(xs.shape[0] for xs in x) - self.num_a = np.prod(np.array(self.grid_size)) - self.grid_dimension = len(self.x) - elif grid_size is not None: - self.grid_size = tuple(map(int, grid_size)) - self.x = tuple(jnp.linspace(0, 1, n) for n in self.grid_size) - self.num_a = np.prod(np.array(grid_size)) - self.grid_dimension = len(self.grid_size) - else: - raise ValueError("Input either grid_size tuple or grid locations x.") - - if cost_fns is None: - cost_fns = [costs.SqEuclidean()] - self.cost_fns = cost_fns - self.kwargs = { - "num_a": self.num_a, - "grid_size": self.grid_size, - "grid_dimension": self.grid_dimension - } - - super().__init__(**kwargs) - - @property - def geometries(self) -> List[geometry.Geometry]: - """Cost matrices along each dimension of the grid.""" - geometries = [] - for dimension, cost_fn in itertools.zip_longest( - range(self.grid_dimension), self.cost_fns, fillvalue=self.cost_fns[-1] - ): - x_values = self.x[dimension][:, jnp.newaxis] - geom = pointcloud.PointCloud( - x_values, - cost_fn=cost_fn, - epsilon=self._epsilon_init, - ) - geometries.append(geom) - return geometries - - @property - def median_cost_matrix(self) -> NoReturn: - """Not implemented.""" - raise NotImplementedError("Median cost not implemented for grids.") - - @property - def can_LRC(self) -> bool: # noqa: D102 - return True - - @property - def shape(self) -> Tuple[int, int]: # noqa: D102 - return self.num_a, self.num_a - - @property - def is_symmetric(self) -> bool: # noqa: D102 - return True - - # Reimplemented functions to be used in regularized OT - def apply_lse_kernel( - self, - f: jnp.ndarray, - g: jnp.ndarray, - eps: float, - vec: Optional[jnp.ndarray] = None, - axis: int = 0 - ) -> jnp.ndarray: - """Apply grid kernel in log space. See notes in parent class for use case. - - Reshapes vector inputs below as grids, applies kernels onto each slice, and - then expands the outputs as vectors. - - More implementation details in :cite:`schmitz:18`. - - Args: - f: jnp.ndarray, a vector of potentials - g: jnp.ndarray, a vector of potentials - eps: float, regularization strength - vec: jnp.ndarray, if needed, a vector onto which apply the kernel weighted - by f and g. - axis: axis (0 or 1) along which summation should be carried out. - - Returns: - a vector, the result of kernel applied in lse space onto vec. - """ - f, g = jnp.reshape(f, self.grid_size), jnp.reshape(g, self.grid_size) - - if vec is not None: - vec = jnp.reshape(vec, self.grid_size) - - if axis == 0: - f, g = g, f - - for dimension in range(self.grid_dimension): - g, vec = self._apply_lse_kernel_one_dimension(dimension, f, g, eps, vec) - g -= jnp.where(jnp.isfinite(f), f, 0) - - if vec is None: - vec = jnp.array(1.0) - return g.ravel(), vec.ravel() - - def _apply_lse_kernel_one_dimension(self, dimension, f, g, eps, vec=None): - """Permute axis & apply the kernel on a single slice.""" - indices = np.arange(self.grid_dimension) - indices[dimension], indices[0] = 0, dimension - f, g = jnp.transpose(f, indices), jnp.transpose(g, indices) - centered_cost = ( - f[:, jnp.newaxis, ...] + g[jnp.newaxis, ...] - jnp.expand_dims( - self.geometries[dimension].cost_matrix, - axis=tuple(range(2, 1 + self.grid_dimension)) - ) - ) / eps - - if vec is not None: - vec = jnp.transpose(vec, indices) - softmax_res, softmax_sgn = utils.logsumexp( - centered_cost, b=vec, axis=1, return_sign=True - ) - return eps * jnp.transpose(softmax_res, - indices), jnp.transpose(softmax_sgn, indices) - softmax_res = eps * utils.logsumexp(centered_cost, axis=1) - return jnp.transpose(softmax_res, indices), None - - def _apply_cost_to_vec( - self, vec: jnp.ndarray, axis: int = 0, fn=None - ) -> jnp.ndarray: - r"""Apply grid's cost matrix (without instantiating it) to a vector. - - The `apply_cost` operation on grids rests on the following identity. - If it were to be cast as a [num_a, num_a] matrix, the corresponding cost - matrix :math:`C` would be a sum of `grid_dimension` matrices, each of the - form (here for the j-th slice) - :math:`\tilde{C}_j : = 1_{n_1} \otimes \dots \otimes C_j \otimes 1_{n_d}` - where each :math:`1_{n}` is the :math:`n\times n` square matrix full of 1's. - - Applying :math:`\tilde{C}_j` on a vector grid consists in carrying a tensor - multiplication on the dimension of that vector reshaped as a grid, followed - by a summation on all other axis while keeping dimensions. That identity is - a generalization of the formula - :math:`(1_{n} \otimes A) vec(X) = vec( A X 1_n)` - where the last multiplication by the matrix of ones is equivalent to - summation while keeping dimensions. - - Args: - vec: jnp.ndarray, flat vector of total size prod(grid_size). - axis: axis 0 if applying transpose costs, 1 if using the original cost. - fn: function optionally applied to cost matrix element-wise, before the - dot product. - - Returns: - A jnp.ndarray corresponding to cost x matrix - """ - vec = jnp.reshape(vec, self.grid_size) - accum_vec = jnp.zeros_like(vec) - indices = list(range(1, self.grid_dimension)) - for dimension, geom in enumerate(self.geometries): - cost = geom.cost_matrix - ind = indices.copy() - ind.insert(dimension, 0) - if axis == 0: - cost = cost.T - accum_vec += jnp.sum( - jnp.tensordot(cost, vec, axes=([0], [dimension])), - axis=indices, - keepdims=True - ).transpose(ind) - return accum_vec.ravel() - - def apply_kernel( - self, - scaling: jnp.ndarray, - eps: Optional[float] = None, - axis: Optional[int] = None - ) -> jnp.ndarray: - """Apply grid kernel on scaling vector. - - See notes in parent class for use. - - Reshapes scaling vector as a grid, applies kernels onto each slice, and - then ravels backs the output as a vector. - - More implementation details in :cite:`schmitz:18`, - - Args: - scaling: jnp.ndarray, a vector of scaling (>0) values. - eps: float, regularization strength - axis: axis (0 or 1) along which summation should be carried out. - - Returns: - a vector, the result of kernel applied onto scaling. - """ - scaling = jnp.reshape(scaling, self.grid_size) - indices = list(range(1, self.grid_dimension)) - for dimension, geom in enumerate(self.geometries): - kernel = geom.kernel_matrix - kernel = kernel if eps is None else kernel ** (self.epsilon / eps) - ind = indices.copy() - ind.insert(dimension, 0) - scaling = jnp.tensordot( - kernel, scaling, axes=([0], [dimension]) - ).transpose(ind) - return scaling.ravel() - - def transport_from_potentials( - self, f: jnp.ndarray, g: jnp.ndarray, axis: int = 0 - ) -> NoReturn: - """Not implemented, use :meth:`apply_transport_from_potentials` instead.""" - raise ValueError( - "Grid geometry cannot instantiate a transport matrix, use", - " apply_transport_from_potentials(...) if you wish to ", - " apply the transport matrix to a vector, or use a point " - " cloud geometry instead" - ) - - def transport_from_scalings( - self, f: jnp.ndarray, g: jnp.ndarray, axis: int = 0 - ) -> NoReturn: - """Not implemented, use :meth:`apply_transport_from_scalings` instead.""" - raise ValueError( - "Grid geometry cannot instantiate a transport matrix, use ", - "apply_transport_from_scalings(...) if you wish to ", - "apply the transport matrix to a vector, or use a point " - "cloud geometry instead." - ) - - def subset( - self, src_ixs: Optional[jnp.ndarray], tgt_ixs: Optional[jnp.ndarray] - ) -> NoReturn: - """Not implemented.""" - raise NotImplementedError("Subsetting is not implemented for grids.") - - def mask( - self, - src_mask: Optional[jnp.ndarray], - tgt_mask: Optional[jnp.ndarray], - mask_value: float = 0.0, - ) -> NoReturn: - """Not implemented.""" - raise NotImplementedError("Masking is not implemented for grids.") - - @classmethod - def prepare_divergences( - cls, - *args: Any, - static_b: bool = False, - **kwargs: Any - ) -> Tuple["Grid", ...]: - """Instantiate the geometries used for a divergence computation.""" - grid_size = kwargs.pop("grid_size", None) - x = kwargs.pop("x", args) - - sep_grid = cls(x=x, grid_size=grid_size, **kwargs) - size = 2 if static_b else 3 - return tuple(sep_grid for _ in range(size)) - - @property - def dtype(self) -> jnp.dtype: # noqa: D102 - return self.x[0].dtype - - def tree_flatten(self): # noqa: D102 - return (self.x, self.cost_fns, self._epsilon_init), self.kwargs - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - return cls( - x=children[0], cost_fns=children[1], epsilon=children[2], **aux_data - ) - - def to_LRCGeometry( - self, - scale: float = 1.0, - **kwargs: Any, - ) -> low_rank.LRCGeometry: - """Converts grid to low-rank geometry. - - Conversion is carried out by taking advantage of the fact that the true cost - matrix of a grid geometry is a sum of Kronecker products of local cost - matrices (for each dimension) with matrices of 1's (both on left and right - sides) of varying dimension. Each of the matrices in that sum can be - factorized if each of these cost matrices can be factorized, which we do - by forcing a conversion to a low rank geometry object. - - 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 - geometries on each slice are not low-rank. - - Returns: - :class:`~ott.geometry.low_rank.LRCGeometry` object. - """ - cost_1, cost_2 = [], [] - for dimension, geom in enumerate(self.geometries): - # An overall low-rank conversion of the cost matrix on a grid, to an - # object of :class:`~ott.geometry.low_rank.LRCGeometry`, necesitates an - # exact low-rank matrix decompisition of the cost matrix of each slice - # of that grid, even if costs on such slices are not low-rank. - # The idea here is that even if the cost matrix on slice `i` is full rank - # `n_i`, we are better off doing 2 redundant `n_i x n_i` matrix products, - # because this is the only way to access to an overall low-rank - # factorization for the entire cost matrix. To get such an exact - # decomposition, the parameter `rank` is set to `0`, triggering a full - # singular value decomposition if needed. - geom = geom.to_LRCGeometry(rank=0, scale=scale, **kwargs) - c_1, c_2 = geom.cost_1, geom.cost_2 - l, r = self.grid_size[:dimension], self.grid_size[dimension + 1:] - l = int(np.prod(np.array(l))) - r = int(np.prod(np.array(r))) - cost_1.append( - jnp.kron(jnp.ones((l, 1)), jnp.kron(c_1, jnp.ones((r, 1),))) - ) - cost_2.append( - jnp.kron(jnp.ones((l, 1)), jnp.kron(c_2, jnp.ones((r, 1),))) - ) - cost_1 = jnp.concatenate(cost_1, axis=-1) - cost_2 = jnp.concatenate(cost_2, 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, - ) diff --git a/ott/build/lib/ott/geometry/low_rank.py b/ott/build/lib/ott/geometry/low_rank.py deleted file mode 100644 index c28c314..0000000 --- a/ott/build/lib/ott/geometry/low_rank.py +++ /dev/null @@ -1,508 +0,0 @@ -# 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. -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 -from ott.math import utils as mu - -__all__ = ["LRCGeometry", "LRKGeometry"] - - -@jax.tree_util.register_pytree_node_class -class LRCGeometry(geometry.Geometry): - """Geometry whose cost is defined by product of two low-rank matrices. - - Implements geometries that are defined as low rank products, i.e. for which - there exists two matrices :math:`A` and :math:`B` of :math:`r` columns such - that the cost of the geometry equals :math:`AB^T`. Apart from being faster to - apply to a vector, these geometries are characterized by the fact that adding - two such geometries should be carried out by concatenating factors, i.e. - if :math:`C = AB^T` and :math:`D = EF^T` then :math:`C + D = [A,E][B,F]^T` - - Args: - cost_1: Array of shape ``[num_a, r]``. - cost_2: Array of shape ``[num_b, r]``. - bias: constant added to entire cost matrix. - scale: Value used to rescale the factors of the low-rank geometry. - scale_cost: option to rescale the cost matrix. Implemented scalings are - 'max_bound', 'mean' and 'max_cost'. Alternatively, a float - factor can be given to rescale the cost such that - ``cost_matrix /= scale_cost``. If `True`, use 'mean'. - batch_size: optional size of the batch to compute online (without - instantiating the matrix) the scale factor ``scale_cost`` of the - :attr:`cost_matrix` when ``scale_cost = 'max_cost'``. If `None`, the batch - size is set to `1024` or to the largest number of samples between - :attr:`cost_1` and :attr:`cost_2` if smaller than `1024`. - kwargs: keyword arguments for :class:`~ott.geometry.geometry.Geometry`. - """ - - def __init__( - self, - cost_1: jnp.ndarray, - cost_2: jnp.ndarray, - bias: float = 0.0, - scale_factor: float = 1.0, - scale_cost: Union[bool, int, float, Literal["mean", "max_bound", - "max_cost"]] = 1.0, - batch_size: Optional[int] = None, - **kwargs: Any, - ): - super().__init__(**kwargs) - self._cost_1 = cost_1 - self._cost_2 = cost_2 - self._bias = bias - self._scale_factor = scale_factor - self._scale_cost = "mean" if scale_cost is True else scale_cost - self.batch_size = batch_size - - @property - def cost_1(self) -> jnp.ndarray: - """First factor of the :attr:`cost_matrix`.""" - scale_factor = jnp.sqrt(self._scale_factor * self.inv_scale_cost) - return scale_factor * self._cost_1 - - @property - def cost_2(self) -> jnp.ndarray: - """Second factor of the :attr:`cost_matrix`.""" - scale_factor = jnp.sqrt(self._scale_factor * self.inv_scale_cost) - return scale_factor * self._cost_2 - - @property - def bias(self) -> float: - """Constant offset added to the entire :attr:`cost_matrix`.""" - return self._bias * self.inv_scale_cost - - @property - def cost_rank(self) -> int: # noqa: D102 - return self._cost_1.shape[1] - - @property - def cost_matrix(self) -> jnp.ndarray: - """Materialize the cost matrix.""" - return jnp.matmul(self.cost_1, self.cost_2.T) + self.bias - - @property - def shape(self) -> Tuple[int, int]: # noqa: D102 - return self._cost_1.shape[0], self._cost_2.shape[0] - - @property - def is_symmetric(self) -> bool: # noqa: D102 - return ( - self._cost_1.shape[0] == self._cost_2.shape[0] and - jnp.all(self._cost_1 == self._cost_2) - ) - - @property - def inv_scale_cost(self) -> float: # noqa: D102 - if isinstance(self._scale_cost, (int, float, jax.Array)): - return 1.0 / self._scale_cost - self = self._masked_geom() - if self._scale_cost == "max_bound": - x_norm = self._cost_1[:, 0].max() - y_norm = self._cost_2[:, 1].max() - max_bound = x_norm + y_norm + 2 * jnp.sqrt(x_norm * y_norm) - return 1.0 / (max_bound + self._bias) - if self._scale_cost == "mean": - factor1 = jnp.dot(self._n_normed_ones, self._cost_1) - factor2 = jnp.dot(self._cost_2.T, self._m_normed_ones) - mean = jnp.dot(factor1, factor2) + self._bias - return 1.0 / mean - if self._scale_cost == "max_cost": - return 1.0 / self.compute_max_cost() - raise ValueError(f"Scaling {self._scale_cost} not implemented.") - - def apply_square_cost(self, arr: jnp.ndarray, axis: int = 0) -> jnp.ndarray: - """Apply elementwise-square of cost matrix to array (vector or matrix).""" - (n, m), r = self.shape, self.cost_rank - # When applying square of a LRCGeometry, one can either elementwise square - # the cost matrix, or instantiate an augmented (rank^2) LRCGeometry - # and apply it. First is O(nm), the other is O((n+m)r^2). - if n * m < (n + m) * r ** 2: # better use regular apply - return super().apply_square_cost(arr, axis) - - new_cost_1 = self.cost_1[:, :, None] * self.cost_1[:, None, :] - new_cost_2 = self.cost_2[:, :, None] * self.cost_2[:, None, :] - return LRCGeometry( - cost_1=new_cost_1.reshape((n, r ** 2)), - cost_2=new_cost_2.reshape((m, r ** 2)) - ).apply_cost(arr, axis) - - def _apply_cost_to_vec( - self, - vec: jnp.ndarray, - axis: int = 0, - fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, - is_linear: bool = False, - ) -> jnp.ndarray: - """Apply [num_a, num_b] fn(cost) (or transpose) to vector. - - Args: - vec: jnp.ndarray [num_a,] ([num_b,] if axis=1) vector - axis: axis on which the reduction is done. - fn: function optionally applied to cost matrix element-wise, before the - doc product - is_linear: Whether ``fn`` is a linear function to enable efficient - implementation. See :func:`ott.geometry.geometry.is_linear` - for a heuristic to help determine if a function is linear. - - Returns: - A jnp.ndarray corresponding to cost x vector - """ - - def linear_apply( - vec: jnp.ndarray, axis: int, fn: Callable[[jnp.ndarray], jnp.ndarray] - ) -> jnp.ndarray: - c1 = self.cost_1 if axis == 1 else self.cost_2 - c2 = self.cost_2 if axis == 1 else self.cost_1 - c2 = fn(c2) if fn is not None else c2 - bias = fn(self.bias) if fn is not None else self.bias - out = jnp.dot(c1, jnp.dot(c2.T, vec)) - return out + bias * jnp.sum(vec) * jnp.ones_like(out) - - if fn is None or is_linear: - return linear_apply(vec, axis, fn=fn) - return super()._apply_cost_to_vec(vec, axis, fn=fn) - - def compute_max_cost(self) -> float: - """Compute the maximum of the :attr:`cost_matrix`. - - Three cases are taken into account: - - - If the number of samples of ``cost_1`` and ``cost_2`` are both smaller - than 1024 and if ``batch_size`` is `None`, the ``cost_matrix`` is - computed to obtain its maximum entry. - - If one of the number of samples of ``cost_1`` or ``cost_2`` is larger - than 1024 and if ``batch_size`` is `None`, then the maximum of the - cost matrix is calculated by batch. The batches are created on the - longest axis of the cost matrix and their size is fixed to 1024. - - If ``batch_size`` is provided as a float, then the maximum of the cost - matrix is calculated by batch. The batches are created on the longest - axis of the cost matrix and their size if fixed by ``batch_size``. - - Returns: - Maximum of the cost matrix. - """ - batch_for_y = self.shape[1] > self.shape[0] - - n = self.shape[1] if batch_for_y else self.shape[0] - p = self._cost_2.shape[1] if batch_for_y else self._cost_1.shape[1] - carry = ((self._cost_1, self._cost_2) if batch_for_y else - (self._cost_2, self._cost_1)) - - if self.batch_size: - batch_size = min(self.batch_size, n) - else: - batch_size = min(1024, max(self.shape[0], self.shape[1])) - n_batch = n // batch_size - - def body(carry, slice_idx): - cost1, cost2 = carry - cost2_slice = jax.lax.dynamic_slice( - cost2, (slice_idx * batch_size, 0), (batch_size, p) - ) - out_slice = jnp.max(jnp.dot(cost2_slice, cost1.T)) - return carry, out_slice - - def finalize(carry): - cost1, cost2 = carry - return jnp.dot(cost2[n_batch * batch_size:], cost1.T) - - _, out = jax.lax.scan(body, carry, jnp.arange(n_batch)) - last_slice = finalize(carry) - max_value = jnp.max(jnp.concatenate((out, last_slice.reshape(-1)))) - return max_value + self._bias - - def to_LRCGeometry( - self, - rank: int = 0, - tol: float = 1e-2, - rng: Optional[jax.Array] = None, - scale: float = 1.0, - ) -> "LRCGeometry": - """Return self.""" - del rank, tol, rng, scale - return self - - @property - def can_LRC(self): # noqa: D102 - return True - - def subset( # noqa: D102 - self, src_ixs: Optional[jnp.ndarray], tgt_ixs: Optional[jnp.ndarray], - **kwargs: Any - ) -> "LRCGeometry": - - 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[ixs, ...] - - return self._mask_subset_helper( - src_ixs, tgt_ixs, fn=subset_fn, propagate_mask=True, **kwargs - ) - - def mask( # noqa: D102 - self, - src_mask: Optional[jnp.ndarray], - tgt_mask: Optional[jnp.ndarray], - mask_value: float = 0.0, - ) -> "LRCGeometry": - - 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, - ) -> "LRCGeometry": - (c1, c2, src_mask, tgt_mask, *children), aux_data = self.tree_flatten() - c1 = fn(c1, src_ixs) - c2 = fn(c2, 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, [c1, c2, src_mask, tgt_mask] + children - ) - - def __add__(self, other: "LRCGeometry") -> "LRCGeometry": - if not isinstance(other, LRCGeometry): - return NotImplemented - return LRCGeometry( - cost_1=jnp.concatenate((self.cost_1, other.cost_1), axis=1), - cost_2=jnp.concatenate((self.cost_2, other.cost_2), axis=1), - bias=self._bias + other._bias, - # already included in `cost_{1,2}` - scale_factor=1.0, - scale_cost=1.0, - ) - - @property - def dtype(self) -> jnp.dtype: # noqa: D102 - return self._cost_1.dtype - - def tree_flatten(self): # noqa: D102 - return ( - self._cost_1, - self._cost_2, - self._src_mask, - self._tgt_mask, - self._epsilon_init, - self._bias, - self._scale_factor, - ), { - "scale_cost": self._scale_cost, - "batch_size": self.batch_size - } - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - c1, c2, src_mask, tgt_mask, epsilon, bias, scale_factor = children - return cls( - c1, - c2, - bias=bias, - scale_factor=scale_factor, - epsilon=epsilon, - src_mask=src_mask, - tgt_mask=tgt_mask, - **aux_data - ) - - -@jax.tree_util.register_pytree_node_class -class LRKGeometry(geometry.Geometry): - """Low-rank kernel geometry. - - .. note:: - This constructor is not meant to be called by the user, - please use the :meth:`from_pointcloud` method instead. - - Args: - k1: Array of shape ``[num_a, r]`` with positive features. - k2: Array of shape ``[num_b, r]`` with positive features. - epsilon: Epsilon regularization. - kwargs: Keyword arguments for :class:`~ott.geometry.geometry.Geometry`. - """ - - def __init__( - self, - k1: jnp.ndarray, - k2: jnp.ndarray, - epsilon: Optional[float] = None, - **kwargs: Any - ): - super().__init__(epsilon=epsilon, relative_epsilon=False, **kwargs) - self.k1 = k1 - self.k2 = k2 - - @classmethod - def from_pointcloud( - cls, - x: jnp.ndarray, - y: jnp.ndarray, - *, - kernel: Literal["gaussian", "arccos"], - rank: int = 100, - std: float = 1.0, - n: int = 1, - rng: Optional[jax.Array] = None - ) -> "LRKGeometry": - r"""Low-rank kernel approximation :cite:`scetbon:20`. - - Args: - x: Array of shape ``[n, d]``. - y: Array of shape ``[m, d]``. - kernel: Type of the kernel to approximate. - rank: Rank of the approximation. - std: Depending on the ``kernel`` approximation: - - - ``'gaussian'`` - scale of the Gibbs kernel. - - ``'arccos'`` - standard deviation of the random projections. - n: Order of the arc-cosine kernel, see :cite:`cho:09` for reference. - rng: Random key used for seeding. - - Returns: - Low-rank kernel geometry. - """ - rng = utils.default_prng_key(rng) - if kernel == "gaussian": - r = jnp.maximum( - jnp.linalg.norm(x, axis=-1).max(), - jnp.linalg.norm(y, axis=-1).max() - ) - k1 = _gaussian_kernel(rng, x, rank, eps=std, R=r) - k2 = _gaussian_kernel(rng, y, rank, eps=std, R=r) - eps = std - elif kernel == "arccos": - k1 = _arccos_kernel(rng, x, rank, n=n, std=std) - k2 = _arccos_kernel(rng, y, rank, n=n, std=std) - eps = 1.0 - else: - raise NotImplementedError(kernel) - - return cls(k1, k2, epsilon=eps) - - def apply_kernel( # noqa: D102 - self, - scaling: jnp.ndarray, - eps: Optional[float] = None, - axis: int = 0, - ) -> jnp.ndarray: - if axis == 0: - return self.k2 @ (self.k1.T @ scaling) - return self.k1 @ (self.k2.T @ scaling) - - @property - def kernel_matrix(self) -> jnp.ndarray: # noqa: D102 - return self.k1 @ self.k2.T - - @property - def cost_matrix(self) -> jnp.ndarray: # noqa: D102 - eps = jnp.finfo(self.dtype).tiny - return -self.epsilon * jnp.log(self.kernel_matrix + eps) - - @property - def rank(self) -> int: # noqa: D102 - return self.k1.shape[1] - - @property - def shape(self) -> Tuple[int, int]: # noqa: D102 - return self.k1.shape[0], self.k2.shape[0] - - @property - def dtype(self) -> jnp.dtype: # noqa: D102 - return self.k1.dtype - - def transport_from_potentials( - self, f: jnp.ndarray, g: jnp.ndarray - ) -> jnp.ndarray: - """Not implemented.""" - raise ValueError("Not implemented.") - - def tree_flatten(self): # noqa: D102 - return [self.k1, self.k2, self._epsilon_init], {} - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - return cls(*children, **aux_data) - - -def _gaussian_kernel( - rng: jax.Array, - x: jnp.ndarray, - n_features: int, - eps: float, - R: jnp.ndarray, -) -> jnp.ndarray: - _, d = x.shape - cost_fn = costs.SqEuclidean() - - y = (R ** 2) / (eps * d) - q = y / (2.0 * mu.lambertw(y)) - sigma = jnp.sqrt(q * eps * 0.25) - - u = jax.random.normal(rng, shape=(n_features, d)) * sigma - cost = cost_fn.all_pairs(x, u) - norm_u = cost_fn.norm(u) - - tmp = -2.0 * (cost / eps) + (norm_u / (eps * q)) - phi = (2 * q) ** (d / 4) * jnp.exp(tmp) - - return (1.0 / jnp.sqrt(n_features)) * phi - - -def _arccos_kernel( - rng: jax.Array, - x: jnp.ndarray, - n_features: int, - n: int, - std: float = 1.0, - kappa: float = 1e-6, -) -> jnp.ndarray: - n_points, d = x.shape - c = jnp.sqrt(2) * (std ** (d / 2)) - - u = jax.random.normal(rng, shape=(n_features, d)) * std - tmp = -(1 / 4) * jnp.sum(u ** 2, axis=-1) * (1.0 - (1.0 / (std ** 2))) - phi = c * (jnp.maximum(0.0, (x @ u.T)) ** n) * jnp.exp(tmp) - - return jnp.c_[(1.0 / jnp.sqrt(n_features)) * phi, - jnp.full((n_points,), fill_value=kappa)] diff --git a/ott/build/lib/ott/geometry/pointcloud.py b/ott/build/lib/ott/geometry/pointcloud.py deleted file mode 100644 index 313ecb3..0000000 --- a/ott/build/lib/ott/geometry/pointcloud.py +++ /dev/null @@ -1,792 +0,0 @@ -# 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.geometry import costs, geometry, low_rank -from ott.math import utils as mu - -__all__ = ["PointCloud"] - - -@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.0 - - @property - def _norm_y(self) -> Union[float, jnp.ndarray]: - if self._axis_norm == 0: - return self.cost_fn.norm(self.y) - return 0.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, jax.Array)): - 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 - - 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 - - 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 - ) - - 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 - ) - - 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 - ) - - def apply_cost( - self, - arr: jnp.ndarray, - axis: int = 0, - fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, - is_linear: bool = False, - sub_idx=None, - ) -> 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, sub_idx=None) - - def _apply_cost( - self, arr: jnp.ndarray, axis: int = 0, fn=None, sub_idx=None - ) -> jnp.ndarray: - """See :meth:`apply_cost`.""" - if not self.is_online or sub_idx is not None: - return super().apply_cost(arr, axis, fn, sub_idx=sub_idx) - - 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 - ) - - def vec_apply_cost( - self, - arr: jnp.ndarray, - axis: int = 0, - fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, - sub_idx=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." - ) - - 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] - - @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.0 / self.inv_scale_cost - cost_fn = costs.SqEuclidean() - return type(self).tree_unflatten(aux_data, [x, y] + args + [cost_fn]) - - 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, - ) - - 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[ixs, ...] - - return self._mask_subset_helper( - src_ixs, tgt_ixs, fn=subset_fn, propagate_mask=True, **kwargs - ) - - def mask( # noqa: D102 - self, - src_mask: Optional[jnp.ndarray], - tgt_mask: Optional[jnp.ndarray], - mask_value: float = 0.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)) diff --git a/ott/build/lib/ott/geometry/segment.py b/ott/build/lib/ott/geometry/segment.py deleted file mode 100644 index 20a1ee9..0000000 --- a/ott/build/lib/ott/geometry/segment.py +++ /dev/null @@ -1,188 +0,0 @@ -# -# 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. -from typing import Callable, Optional, Tuple - -import jax -import jax.numpy as jnp - -__all__ = ["segment_point_cloud"] - - -def segment_point_cloud( - x: jnp.ndarray, - a: Optional[jnp.ndarray] = None, - num_segments: Optional[int] = None, - max_measure_size: Optional[int] = None, - segment_ids: Optional[jnp.ndarray] = None, - indices_are_sorted: bool = False, - num_per_segment: Optional[Tuple[int, ...]] = None, - padding_vector: Optional[jnp.ndarray] = None -) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Segment and pad as needed the entries of a point cloud. - - There are two interfaces: - - #. use ``segment_ids``, and optionally ``indices_are_sorted`` to describe - for each data point in the matrix to which segment it belongs to. - #. use ``num_per_segment`` which describes contiguous segments. - - If using the first interface, ``num_segments`` is required for jitting. - Assumes ``range(0, num_segments)`` are the segment ids. - - In both cases, jitting requires defining a ``max_measure_size``, the - upper bound on the maximal size of measures, which will be used for padding. - - Args: - x: Array of input points, of shape ``[num_x, ndim]``. - Multiple segments are held in this single array. - a: Array of shape ``[num_x,]`` containing the weights (within each measure) - of all the points. - num_segments: Number of segments. Required for jitting. - If `None` and using the second interface, it will be computed as - ``len(num_per_segment)``. - max_measure_size: Overall size of padding. Required for jitting. - If `None` and using the second interface, it will be computed as - ``max(num_per_segment)``. - segment_ids: **1st interface** The segment ids for which each row of ``x`` - belongs. This is a similar interface to :func:`jax.ops.segment_sum`. - indices_are_sorted: **1st interface** Whether ``segment_ids`` are sorted. - num_per_segment: **2nd interface** Number of points in each segment. - For example, `[100, 20, 30]` would imply that ``x`` is segmented into 3 - arrays of length `[100]`, `[20]`, and `[30]`, respectively. - Must be a tuple and not a :class:`jax.numpy.ndarray` to allow jitting. - This means changes in ``num_per_segment`` will re-trigger compilation. - padding_vector: vector to be used to pad point cloud matrices. Most likely - to be zero, but can be adjusted to be other values to avoid errors or - over/underflow in cost matrix that could be problematic (even these values - are not supposed to be taken given their corresponding masses are 0). - See also :func:`~ott.geometry.costs.CostFn._padder`. - If ``None``, vector of 0s of shape ``[1, ndim]`` is used. - - Returns: - Segmented ``x`` as an array of shape - ``[num_measures, max_measure_size, ndim]`` and ``a`` as an array of shape - ``[num_measures, max_measure_size]``. - """ - num, dim = x.shape - use_segment_ids = segment_ids is not None - if use_segment_ids: - assert num_segments is not None, "Please specify `num_segments`." - assert max_measure_size is not None, "Please specify `max_measure_size`." - num_per_segment = jax.ops.segment_sum( - jnp.ones_like(segment_ids), - segment_ids, - num_segments=num_segments, - indices_are_sorted=indices_are_sorted - ) - else: - assert num_per_segment is not None, "Please specify `num_per_segment`." - if max_measure_size is None: - max_measure_size = max(num_per_segment) - if num_segments is None: - num_segments = len(num_per_segment) - else: - assert num_segments == len(num_per_segment) - # conversion to facilitate computation of default weight below. - num_per_segment = jnp.array(num_per_segment) - segment_ids = jnp.arange(num_segments).repeat( - num_per_segment, total_repeat_length=num - ) - - if a is None: - a = jnp.array( - (1.0 / - num_per_segment).repeat(num_per_segment, total_repeat_length=num) - ) - - if padding_vector is None: - padding_vector = jnp.zeros((1, dim)) - - x = jnp.concatenate((x, padding_vector)) - a = jnp.concatenate((a, jnp.zeros((1,)))) - segmented_a, segmented_x = [], [] - - for i in range(num_segments): - idx = jnp.where(segment_ids == i, jnp.arange(num), num + 1) - idx = jax.lax.dynamic_slice(jnp.sort(idx), (0,), (max_measure_size,)) - - # segment the weights - segmented_a.append(a.at[idx].get()) - # segment the positions - segmented_x.append(x.at[idx].get()) - - segmented_a = jnp.stack(segmented_a) - segmented_x = jnp.stack(segmented_x) - - return segmented_x, segmented_a - - -def _segment_interface( - x: jnp.ndarray, - y: jnp.ndarray, - eval_fn: Callable[[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray], - jnp.ndarray], - num_segments: Optional[int] = None, - max_measure_size: Optional[int] = None, - segment_ids_x: Optional[jnp.ndarray] = None, - segment_ids_y: Optional[jnp.ndarray] = None, - indices_are_sorted: bool = False, - num_per_segment_x: Optional[jnp.ndarray] = None, - num_per_segment_y: Optional[jnp.ndarray] = None, - weights_x: Optional[jnp.ndarray] = None, - weights_y: Optional[jnp.ndarray] = None, - padding_vector: Optional[jnp.ndarray] = None, -) -> jnp.ndarray: - """Wrapper to segment two point clouds and return parallel evaluations. - - Utility function that segments two point clouds using the approach outlined - in `segment_point_cloud` and evaluates `eval_fn` on pairs of segmented point - clouds. - """ - use_segment_ids = segment_ids_x is not None - if use_segment_ids: - assert segment_ids_y is not None - else: - assert num_per_segment_x is not None - assert num_per_segment_y is not None - - segmented_x, segmented_weights_x = segment_point_cloud( - x, - a=weights_x, - num_segments=num_segments, - max_measure_size=max_measure_size, - segment_ids=segment_ids_x, - indices_are_sorted=indices_are_sorted, - num_per_segment=num_per_segment_x, - padding_vector=padding_vector - ) - - segmented_y, segmented_weights_y = segment_point_cloud( - y, - a=weights_y, - num_segments=num_segments, - max_measure_size=max_measure_size, - segment_ids=segment_ids_y, - indices_are_sorted=indices_are_sorted, - num_per_segment=num_per_segment_y, - padding_vector=padding_vector - ) - - v_eval = jax.vmap(eval_fn, in_axes=[0] * 4) - return v_eval( - segmented_x, - segmented_y, - segmented_weights_x, - segmented_weights_y, - ) diff --git a/ott/build/lib/ott/initializers/__init__.py b/ott/build/lib/ott/initializers/__init__.py deleted file mode 100644 index 5406247..0000000 --- a/ott/build/lib/ott/initializers/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. -from . import linear, quadratic diff --git a/ott/build/lib/ott/initializers/linear/__init__.py b/ott/build/lib/ott/initializers/linear/__init__.py deleted file mode 100644 index c7a7dc4..0000000 --- a/ott/build/lib/ott/initializers/linear/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. -from . import initializers, initializers_lr diff --git a/ott/build/lib/ott/initializers/linear/initializers.py b/ott/build/lib/ott/initializers/linear/initializers.py deleted file mode 100644 index e486349..0000000 --- a/ott/build/lib/ott/initializers/linear/initializers.py +++ /dev/null @@ -1,408 +0,0 @@ -# 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 abc -from typing import Any, Dict, Optional, Sequence, Tuple - -import jax -import jax.numpy as jnp - -from ott import utils -from ott.geometry import pointcloud -from ott.problems.linear import linear_problem - -__all__ = [ - "DefaultInitializer", "GaussianInitializer", "SortingInitializer", - "SubsampleInitializer" -] - - -@jax.tree_util.register_pytree_node_class -class SinkhornInitializer(abc.ABC): - """Base class for Sinkhorn initializers.""" - - @abc.abstractmethod - def init_dual_a( - self, - ot_prob: linear_problem.LinearProblem, - lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: - """Initialize Sinkhorn potential/scaling f_u. - - Args: - ot_prob: Linear OT problem. - lse_mode: Return potential if ``True``, scaling if ``False``. - rng: Random number generator for stochastic initializers. - - Returns: - potential/scaling, array of size ``[n,]``. - """ - - @abc.abstractmethod - def init_dual_b( - self, - ot_prob: linear_problem.LinearProblem, - lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: - """Initialize Sinkhorn potential/scaling g_v. - - Args: - ot_prob: Linear OT problem. - lse_mode: Return potential if ``True``, scaling if ``False``. - rng: Random number generator for stochastic initializers. - - Returns: - potential/scaling, array of size ``[m,]``. - """ - - def __call__( - self, - ot_prob: linear_problem.LinearProblem, - a: Optional[jnp.ndarray], - b: Optional[jnp.ndarray], - lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Initialize Sinkhorn potentials/scalings f_u and g_v. - - Args: - ot_prob: Linear OT problem. - a: Initial potential/scaling f_u. - If ``None``, it will be initialized using :meth:`init_dual_a`. - b: Initial potential/scaling g_v. - If ``None``, it will be initialized using :meth:`init_dual_b`. - lse_mode: Return potentials if ``True``, scalings if ``False``. - rng: Random number generator for stochastic initializers. - - Returns: - The initial potentials/scalings. - """ - rng = utils.default_prng_key(rng) - rng_x, rng_y = jax.random.split(rng, 2) - n, m = ot_prob.geom.shape - if a is None: - a = self.init_dual_a(ot_prob, lse_mode=lse_mode, rng=rng_x) - if b is None: - b = self.init_dual_b(ot_prob, lse_mode=lse_mode, rng=rng_y) - - assert a.shape == ( - n, - ), f"Expected `f_u` to have shape `{n,}`, found `{a.shape}`." - assert b.shape == ( - m, - ), f"Expected `g_v` to have shape `{m,}`, found `{b.shape}`." - - # cancel dual variables for zero weights - a = jnp.where(ot_prob.a > 0.0, a, -jnp.inf if lse_mode else 0.0) - b = jnp.where(ot_prob.b > 0.0, b, -jnp.inf if lse_mode else 0.0) - - return a, b - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return [], {} - - @classmethod - def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] - ) -> "SinkhornInitializer": - return cls(*children, **aux_data) - - -@jax.tree_util.register_pytree_node_class -class DefaultInitializer(SinkhornInitializer): - """Default initialization of Sinkhorn dual potentials/primal scalings.""" - - def init_dual_a( # noqa: D102 - self, - ot_prob: linear_problem.LinearProblem, - lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: - del rng - return jnp.zeros_like(ot_prob.a) if lse_mode else jnp.ones_like(ot_prob.a) - - def init_dual_b( # noqa: D102 - self, - ot_prob: linear_problem.LinearProblem, - lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: - del rng - return jnp.zeros_like(ot_prob.b) if lse_mode else jnp.ones_like(ot_prob.b) - - -@jax.tree_util.register_pytree_node_class -class GaussianInitializer(DefaultInitializer): - """Gaussian initializer :cite:`thornton2022rethinking:22`. - - Compute Gaussian approximations of each - :class:`~ott.geometry.pointcloud.PointCloud`, then compute closed from - Kantorovich potential between Gaussian approximations using Brenier's theorem - (adapt convex/Brenier potential to Kantorovich). Use this Gaussian potential - to initialize Sinkhorn potentials/scalings. - """ - - def init_dual_a( # noqa: D102 - self, - ot_prob: linear_problem.LinearProblem, - lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: - # import Gaussian here due to circular imports - from ott.tools.gaussian_mixture import gaussian - - del rng - assert isinstance( - ot_prob.geom, pointcloud.PointCloud - ), "Gaussian initializer valid only for pointcloud geoms." - - x, y = ot_prob.geom.x, ot_prob.geom.y - a, b = ot_prob.a, ot_prob.b - - gaussian_a = gaussian.Gaussian.from_samples(x, weights=a) - gaussian_b = gaussian.Gaussian.from_samples(y, weights=b) - # Brenier potential for cost ||x-y||^2/2, multiply by two for ||x-y||^2 - f_potential = 2 * gaussian_a.f_potential(dest=gaussian_b, points=x) - f_potential = f_potential - jnp.mean(f_potential) - return f_potential if lse_mode else ot_prob.geom.scaling_from_potential( - f_potential - ) - - -@jax.tree_util.register_pytree_node_class -class SortingInitializer(DefaultInitializer): - """Sorting initializer :cite:`thornton2022rethinking:22`. - - Solve non-regularized OT problem via sorting, then compute potential through - iterated minimum on C-transform and use this potential to initialize - regularized potential. - - Args: - vectorized_update: Whether to use vectorized loop. - tolerance: DualSort convergence threshold. - max_iter: Max DualSort steps. - """ - - def __init__( - self, - vectorized_update: bool = True, - tolerance: float = 1e-2, - max_iter: int = 100 - ): - super().__init__() - self.tolerance = tolerance - self.max_iter = max_iter - self.vectorized_update = vectorized_update - - def _init_sorting_dual( - self, modified_cost: jnp.ndarray, init_f: jnp.ndarray - ) -> jnp.ndarray: - """Run DualSort algorithm. - - Args: - modified_cost: cost matrix minus diagonal column-wise. - init_f: potential f, array of size n. This is the starting potential, - which is then updated to make the init potential, so an init of an init. - - Returns: - potential f, array of size n. - """ - - def body_fn( - state: Tuple[jnp.ndarray, float, int] - ) -> Tuple[jnp.ndarray, float, int]: - prev_f, _, it = state - new_f = fn(prev_f, modified_cost) - diff = jnp.sum((new_f - prev_f) ** 2) - it += 1 - return new_f, diff, it - - def cond_fn(state: Tuple[jnp.ndarray, float, int]) -> bool: - _, diff, it = state - return jnp.logical_and(diff > self.tolerance, it < self.max_iter) - - fn = _vectorized_update if self.vectorized_update else _coordinate_update - state = (init_f, jnp.inf, 0) # init, error, iter - f_potential, _, _ = jax.lax.while_loop( - cond_fun=cond_fn, body_fun=body_fn, init_val=state - ) - - return f_potential - - def init_dual_a( - self, - ot_prob: linear_problem.LinearProblem, - lse_mode: bool, - rng: Optional[jax.Array] = None, - init_f: Optional[jnp.ndarray] = None, - ) -> jnp.ndarray: - """Apply DualSort algorithm. - - Args: - ot_prob: OT problem between discrete distributions. - lse_mode: Return potential if ``True``, scaling if ``False``. - rng: Random number generator for stochastic initializers, unused. - init_f: potential f, array of size ``[n,]``. This is the starting - potential, which is then updated to make the init potential, - so an init of an init. - - Returns: - potential/scaling f_u, array of size ``[n,]``. - """ - del rng - assert not ot_prob.geom.is_online, \ - "Sorting initializer does not work for online geometry." - # check for sorted x, y requires point cloud and could slow initializer - cost_matrix = ot_prob.geom.cost_matrix - - assert cost_matrix.shape[0] == cost_matrix.shape[ - 1], "Requires square cost matrix." - - modified_cost = cost_matrix - jnp.diag(cost_matrix)[None, :] - - n = cost_matrix.shape[0] - init_f = jnp.zeros(n) if init_f is None else init_f - - f_potential = self._init_sorting_dual(modified_cost, init_f) - f_potential = f_potential - jnp.mean(f_potential) - - return f_potential if lse_mode else ot_prob.geom.scaling_from_potential( - f_potential - ) - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return ([], { - "tolerance": self.tolerance, - "max_iter": self.max_iter, - "vectorized_update": self.vectorized_update - }) - - -@jax.tree_util.register_pytree_node_class -class SubsampleInitializer(DefaultInitializer): - """Subsample initializer :cite:`thornton2022rethinking:22`. - - Subsample each :class:`~ott.geometry.pointcloud.PointCloud`, then compute - :class:`Sinkhorn potential ` - from the subsampled approximations and use this potential to initialize - Sinkhorn potentials/scalings for the original problem. - - Args: - subsample_n_x: number of points to subsample from the first measure in - :class:`~ott.geometry.pointcloud.PointCloud`. - subsample_n_y: number of points to subsample from the second measure in - :class:`~ott.geometry.pointcloud.PointCloud`. - If ``None``, use ``subsample_n_x``. - kwargs: Keyword arguments for - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. - """ - - def __init__( - self, - subsample_n_x: int, - subsample_n_y: Optional[int] = None, - **kwargs: Any, - ): - super().__init__() - self.subsample_n_x = subsample_n_x - self.subsample_n_y = subsample_n_y or subsample_n_x - self.sinkhorn_kwargs = kwargs - - def init_dual_a( # noqa: D102 - self, - ot_prob: linear_problem.LinearProblem, - lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: - from ott.solvers import linear - - assert isinstance( - ot_prob.geom, pointcloud.PointCloud - ), "Subsample initializer valid only for pointcloud geom." - rng = utils.default_prng_key(rng) - rng_x, rng_y = jax.random.split(rng, 2) - - x, y = ot_prob.geom.x, ot_prob.geom.y - a, b = ot_prob.a, ot_prob.b - - # subsample - sub_x = jax.random.choice( - rng_x, a=x, shape=(self.subsample_n_x,), replace=True, p=a, axis=0 - ) - sub_y = jax.random.choice( - rng_y, a=y, shape=(self.subsample_n_y,), replace=True, p=b, axis=0 - ) - - # create subsampled point cloud geometry - sub_geom = pointcloud.PointCloud( - sub_x, - sub_y, - epsilon=ot_prob.geom.epsilon, - scale_cost=ot_prob.geom._scale_cost, - cost_fn=ot_prob.geom.cost_fn - ) - - # run sinkhorn - subsample_sink_out = linear.solve(sub_geom, **self.sinkhorn_kwargs) - - # interpolate potentials - dual_potentials = subsample_sink_out.to_dual_potentials() - f_potential = jax.vmap(dual_potentials.f)(x) - - return f_potential if lse_mode else ot_prob.geom.scaling_from_potential( - f_potential - ) - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return ([], { - "subsample_n_x": self.subsample_n_x, - "subsample_n_y": self.subsample_n_y, - **self.sinkhorn_kwargs - }) - - -def _vectorized_update( - f: jnp.ndarray, modified_cost: jnp.ndarray -) -> jnp.ndarray: - """Inner loop DualSort Update. - - Args: - f: potential f, array of size n. - modified_cost: cost matrix minus diagonal column-wise. - - Returns: - updated potential vector, f. - """ - return jnp.min(modified_cost + f[None, :], axis=1) - - -def _coordinate_update( - f: jnp.ndarray, modified_cost: jnp.ndarray -) -> jnp.ndarray: - """Coordinate-wise updates within inner loop. - - Args: - f: potential f, array of size n. - modified_cost: cost matrix minus diagonal column-wise. - - Returns: - updated potential vector, f. - """ - - def body_fn(i: int, f: jnp.ndarray) -> jnp.ndarray: - new_f = jnp.min(modified_cost[i, :] + f) - return f.at[i].set(new_f) - - return jax.lax.fori_loop(0, len(f), body_fn, f) diff --git a/ott/build/lib/ott/initializers/linear/initializers_lr.py b/ott/build/lib/ott/initializers/linear/initializers_lr.py deleted file mode 100644 index 8206f29..0000000 --- a/ott/build/lib/ott/initializers/linear/initializers_lr.py +++ /dev/null @@ -1,654 +0,0 @@ -# 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 abc -import functools -from typing import ( - TYPE_CHECKING, - Any, - Dict, - Literal, - Mapping, - NamedTuple, - Optional, - Sequence, - Tuple, - Union, -) - -import jax -import jax.numpy as jnp -import numpy as np - -from ott import utils -from ott.geometry import geometry, low_rank, pointcloud -from ott.math import fixed_point_loop -from ott.math import utils as mu - -if TYPE_CHECKING: - from ott.problems.linear import linear_problem - from ott.problems.quadratic import quadratic_problem - from ott.solvers.linear import sinkhorn, sinkhorn_lr - from ott.solvers.quadratic import gromov_wasserstein_lr - -Problem_t = Union["linear_problem.LinearProblem", - "quadratic_problem.QuadraticProblem"] - -__all__ = [ - "RandomInitializer", "Rank2Initializer", "KMeansInitializer", - "GeneralizedKMeansInitializer" -] - - -@jax.tree_util.register_pytree_node_class -class LRInitializer(abc.ABC): - """Base class for low-rank initializers. - - Args: - rank: Rank of the factorization. - kwargs: Additional keyword arguments. - """ - - def __init__(self, rank: int, **kwargs: Any): - self._rank = rank - self._kwargs = kwargs - - @abc.abstractmethod - def init_q( - self, - ot_prob: Problem_t, - rng: jax.Array, - *, - init_g: jnp.ndarray, - **kwargs: Any, - ) -> jnp.ndarray: - """Initialize the low-rank factor :math:`Q`. - - Args: - ot_prob: OT problem. - rng: Random key for seeding. - init_g: Initial value for :math:`g` factor. - kwargs: Additional keyword arguments. - - Returns: - Array of shape ``[n, rank]``. - """ - - @abc.abstractmethod - def init_r( - self, - ot_prob: Problem_t, - rng: jax.Array, - *, - init_g: jnp.ndarray, - **kwargs: Any, - ) -> jnp.ndarray: - """Initialize the low-rank factor :math:`R`. - - Args: - ot_prob: Linear OT problem. - rng: Random key for seeding. - init_g: Initial value for :math:`g` factor. - kwargs: Additional keyword arguments. - - Returns: - Array of shape ``[m, rank]``. - """ - - @abc.abstractmethod - def init_g( - self, - ot_prob: Problem_t, - rng: jax.Array, - **kwargs: Any, - ) -> jnp.ndarray: - """Initialize the low-rank factor :math:`g`. - - Args: - ot_prob: OT problem. - rng: Random key for seeding. - kwargs: Additional keyword arguments. - - Returns: - Array of shape ``[rank,]``. - """ - - @classmethod - def from_solver( - cls, - solver: Union["sinkhorn_lr.LRSinkhorn", - "gromov_wasserstein_lr.LRGromovWasserstein"], - *, - kind: Literal["random", "rank2", "k-means", "generalized-k-means"], - **kwargs: Any, - ) -> "LRInitializer": - """Create a low-rank initializer from a linear or quadratic solver. - - Args: - solver: Low-rank linear or quadratic solver. - kind: Which initializer to instantiate. - kwargs: Keyword arguments when creating the initializer. - - Returns: - Low-rank initializer. - """ - rank = solver.rank - sinkhorn_kwargs = { - "norm_error": solver._norm_error, - "lse_mode": solver.lse_mode, - "implicit_diff": solver.implicit_diff, - "use_danskin": solver.use_danskin - } - - if kind == "random": - return RandomInitializer(rank, **kwargs) - if kind == "rank2": - return Rank2Initializer(rank, **kwargs) - if kind == "k-means": - return KMeansInitializer(rank, sinkhorn_kwargs=sinkhorn_kwargs, **kwargs) - if kind == "generalized-k-means": - return GeneralizedKMeansInitializer( - rank, sinkhorn_kwargs=sinkhorn_kwargs, **kwargs - ) - raise NotImplementedError(f"Initializer `{kind}` is not implemented.") - - def __call__( - self, - ot_prob: Problem_t, - q: Optional[jnp.ndarray] = None, - r: Optional[jnp.ndarray] = None, - g: Optional[jnp.ndarray] = None, - *, - rng: Optional[jax.Array] = None, - **kwargs: Any - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Initialize the factors :math:`Q`, :math:`R` and :math:`g`. - - Args: - ot_prob: OT problem. - q: Factor of shape ``[n, rank]``. If `None`, it will be initialized - using :meth:`init_q`. - r: Factor of shape ``[m, rank]``. If `None`, it will be initialized - using :meth:`init_r`. - g: Factor of shape ``[rank,]``. If `None`, it will be initialized - using :meth:`init_g`. - rng: Random key for seeding. - kwargs: Additional keyword arguments for :meth:`init_q`, :meth:`init_r` - and :meth:`init_g`. - - Returns: - The factors :math:`Q`, :math:`R` and :math:`g`, respectively. - """ - rng = utils.default_prng_key(rng) - rng1, rng2, rng3 = jax.random.split(rng, 3) - - if g is None: - g = self.init_g(ot_prob, rng1, **kwargs) - if q is None: - q = self.init_q(ot_prob, rng2, init_g=g, **kwargs) - if r is None: - r = self.init_r(ot_prob, rng3, init_g=g, **kwargs) - - assert g.shape == (self.rank,) - assert q.shape == (ot_prob.a.shape[0], self.rank) - assert r.shape == (ot_prob.b.shape[0], self.rank) - - return q, r, g - - @property - def rank(self) -> int: - """Rank of the transport matrix factorization.""" - return self._rank - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return [], {**self._kwargs, "rank": self.rank} - - @classmethod - def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] - ) -> "LRInitializer": - return cls(*children, **aux_data) - - -@jax.tree_util.register_pytree_node_class -class RandomInitializer(LRInitializer): - """Low-rank Sinkhorn factorization using random factors. - - Args: - rank: Rank of the factorization. - kwargs: Additional keyword arguments. - """ - - def init_q( # noqa: D102 - self, - ot_prob: Problem_t, - rng: jax.Array, - *, - init_g: jnp.ndarray, - **kwargs: Any, - ) -> jnp.ndarray: - del kwargs, init_g - a = ot_prob.a - init_q = jnp.abs(jax.random.normal(rng, (a.shape[0], self.rank))) - return a[:, None] * (init_q / jnp.sum(init_q, axis=1, keepdims=True)) - - def init_r( # noqa: D102 - self, - ot_prob: Problem_t, - rng: jax.Array, - *, - init_g: jnp.ndarray, - **kwargs: Any, - ) -> jnp.ndarray: - del kwargs, init_g - b = ot_prob.b - init_r = jnp.abs(jax.random.normal(rng, (b.shape[0], self.rank))) - return b[:, None] * (init_r / jnp.sum(init_r, axis=1, keepdims=True)) - - def init_g( # noqa: D102 - self, - ot_prob: Problem_t, - rng: jax.Array, - **kwargs: Any, - ) -> jnp.ndarray: - del kwargs - init_g = jnp.abs(jax.random.uniform(rng, (self.rank,))) + 1.0 - return init_g / jnp.sum(init_g) - - -@jax.tree_util.register_pytree_node_class -class Rank2Initializer(LRInitializer): - """Low-rank Sinkhorn factorization using rank-2 factors :cite:`scetbon:21`. - - Args: - rank: Rank of the factorization. - kwargs: Additional keyword arguments. - """ - - def _compute_factor( - self, - ot_prob: Problem_t, - init_g: jnp.ndarray, - *, - which: Literal["q", "r"], - ) -> jnp.ndarray: - a, b = ot_prob.a, ot_prob.b - marginal = a if which == "q" else b - n, r = marginal.shape[0], self.rank - - lambda_1 = jnp.min( - jnp.array([jnp.min(a), jnp.min(init_g), - jnp.min(b)]) - ) * 0.5 - - g1 = jnp.arange(1, r + 1) - g1 /= g1.astype(float).sum() - g2 = (init_g - lambda_1 * g1) / (1.0 - lambda_1) - - x = jnp.arange(1, n + 1) - x /= x.astype(float).sum() - y = (marginal - lambda_1 * x) / (1.0 - lambda_1) - - return ((lambda_1 * x[:, None] @ g1.reshape(1, -1)) + - ((1.0 - lambda_1) * y[:, None] @ g2.reshape(1, -1))) - - def init_q( # noqa: D102 - self, - ot_prob: Problem_t, - rng: jax.Array, - *, - init_g: jnp.ndarray, - **kwargs: Any, - ) -> jnp.ndarray: - del rng, kwargs - return self._compute_factor(ot_prob, init_g, which="q") - - def init_r( # noqa: D102 - self, - ot_prob: Problem_t, - rng: jax.Array, - *, - init_g: jnp.ndarray, - **kwargs: Any, - ) -> jnp.ndarray: - del rng, kwargs - return self._compute_factor(ot_prob, init_g, which="r") - - def init_g( # noqa: D102 - self, - ot_prob: Problem_t, - rng: jax.Array, - **kwargs: Any, - ) -> jnp.ndarray: - del rng, kwargs - return jnp.ones((self.rank,)) / self.rank - - -@jax.tree_util.register_pytree_node_class -class KMeansInitializer(LRInitializer): - """K-means initializer for low-rank Sinkhorn :cite:`scetbon:22b`. - - Applicable for :class:`~ott.geometry.pointcloud.PointCloud` and - :class:`~ott.geometry.low_rank.LRCGeometry`. - - Args: - rank: Rank of the factorization. - min_iterations: Minimum number of k-means iterations. - max_iterations: Maximum number of k-means iterations. - sinkhorn_kwargs: Keyword arguments for - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. - kwargs: Keyword arguments for :func:`~ott.tools.k_means.k_means`. - """ - - def __init__( - self, - rank: int, - min_iterations: int = 100, - max_iterations: int = 100, - sinkhorn_kwargs: Optional[Mapping[str, Any]] = None, - **kwargs: Any - ): - super().__init__(rank, **kwargs) - self._min_iter = min_iterations - self._max_iter = max_iterations - self._sinkhorn_kwargs = {} if sinkhorn_kwargs is None else sinkhorn_kwargs - - @staticmethod - def _extract_array(geom: geometry.Geometry, *, first: bool) -> jnp.ndarray: - if isinstance(geom, pointcloud.PointCloud): - return geom.x if first else geom.y - if isinstance(geom, low_rank.LRCGeometry): - return geom.cost_1 if first else geom.cost_2 - raise TypeError( - f"k-means initializer not implemented for `{type(geom).__name__}`." - ) - - def _compute_factor( - self, - ot_prob: Problem_t, - rng: jax.Array, - *, - init_g: jnp.ndarray, - which: Literal["q", "r"], - **kwargs: Any, - ) -> jnp.ndarray: - from ott.problems.linear import linear_problem - from ott.problems.quadratic import quadratic_problem - from ott.solvers.linear import sinkhorn - from ott.tools import k_means - - del kwargs - fn = functools.partial( - k_means.k_means, - min_iterations=self._min_iter, - max_iterations=self._max_iter, - **self._kwargs - ) - - if isinstance(ot_prob, quadratic_problem.QuadraticProblem): - if ot_prob.geom_xy is not None and ot_prob.fused_penalty >= 1.0: - # prefer the linear term if it has a higher weight - geom = ot_prob.geom_xy - else: - geom = ot_prob.geom_xx if which == "q" else ot_prob.geom_yy - else: - geom = ot_prob.geom - arr = self._extract_array(geom, first=which == "q") - marginals = ot_prob.a if which == "q" else ot_prob.b - - centroids = fn(arr, self.rank, rng=rng).centroids - geom = pointcloud.PointCloud( - arr, centroids, epsilon=1e-1, scale_cost="max_cost" - ) - - prob = linear_problem.LinearProblem(geom, marginals, init_g) - solver = sinkhorn.Sinkhorn(**self._sinkhorn_kwargs) - return solver(prob).matrix - - def init_q( # noqa: D102 - self, - ot_prob: Problem_t, - rng: jax.Array, - *, - init_g: jnp.ndarray, - **kwargs: Any, - ) -> jnp.ndarray: - return self._compute_factor( - ot_prob, rng, init_g=init_g, which="q", **kwargs - ) - - def init_r( # noqa: D102 - self, - ot_prob: Problem_t, - rng: jax.Array, - *, - init_g: jnp.ndarray, - **kwargs: Any, - ) -> jnp.ndarray: - return self._compute_factor( - ot_prob, rng, init_g=init_g, which="r", **kwargs - ) - - def init_g( # noqa: D102 - self, - ot_prob: Problem_t, - rng: jax.Array, - **kwargs: Any, - ) -> jnp.ndarray: - del rng, kwargs - return jnp.ones((self.rank,)) / self.rank - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - children, aux_data = super().tree_flatten() - aux_data["sinkhorn_kwargs"] = self._sinkhorn_kwargs - aux_data["min_iterations"] = self._min_iter - aux_data["max_iterations"] = self._max_iter - return children, aux_data - - -class GeneralizedKMeansInitializer(KMeansInitializer): - """Generalized k-means initializer :cite:`scetbon:22b`. - - Applicable for any :class:`~ott.geometry.geometry.Geometry` with a - square shape. - - Args: - rank: Rank of the factorization. - gamma: The (inverse of) gradient step size used by mirror descent. - min_iterations: Minimum number of iterations. - max_iterations: Maximum number of iterations. - inner_iterations: Number of iterations used by the algorithm before - re-evaluating progress. - threshold: Convergence threshold. - sinkhorn_kwargs: Keyword arguments for - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. - """ - - def __init__( - self, - rank: int, - gamma: float = 10.0, - min_iterations: int = 0, - max_iterations: int = 100, - inner_iterations: int = 10, - threshold: float = 1e-6, - sinkhorn_kwargs: Optional[Mapping[str, Any]] = None, - ): - super().__init__( - rank, - sinkhorn_kwargs=sinkhorn_kwargs, - # below argument are stored in `_kwargs` - gamma=gamma, - min_iterations=min_iterations, - max_iterations=max_iterations, - inner_iterations=inner_iterations, - threshold=threshold, - ) - - class Constants(NamedTuple): # noqa: D106 - solver: "sinkhorn.Sinkhorn" - geom: geometry.Geometry # (n, n) - marginal: jnp.ndarray # (n,) - g: jnp.ndarray # (r,) - gamma: float - threshold: float - - class State(NamedTuple): # noqa: D106 - factor: jnp.ndarray - criterions: jnp.ndarray - crossed_threshold: bool - - def _compute_factor( - self, - ot_prob: Problem_t, - rng: jax.Array, - *, - init_g: jnp.ndarray, - which: Literal["q", "r"], - **kwargs: Any, - ) -> jnp.ndarray: - from ott.problems.linear import linear_problem - from ott.problems.quadratic import quadratic_problem - from ott.solvers.linear import sinkhorn - - def init_fn() -> GeneralizedKMeansInitializer.State: - n = geom.shape[0] - factor = jnp.abs(jax.random.normal(rng, (n, self.rank))) + 1.0 # (n, r) - factor *= consts.marginal[:, None] / jnp.sum( - factor, axis=1, keepdims=True - ) - - return self.State( - factor, - criterions=-jnp.ones(outer_iterations), - crossed_threshold=False - ) - - # see the explanation in `ott.solvers.linear.sinkhorn_lr` - def converged( - state: GeneralizedKMeansInitializer.State, - consts: GeneralizedKMeansInitializer.Constants, iteration: int - ) -> bool: - - def conv_crossed(prev_err: float, curr_err: float) -> bool: - return jnp.logical_and( - prev_err < consts.threshold, curr_err < consts.threshold - ) - - def conv_not_crossed(prev_err: float, curr_err: float) -> bool: - return jnp.logical_and(curr_err < prev_err, curr_err < consts.threshold) - - it = iteration // inner_iterations - return jax.lax.cond( - state.crossed_threshold, conv_crossed, conv_not_crossed, - state.criterions[it - 2], state.criterions[it - 1] - ) - - def diverged( - state: GeneralizedKMeansInitializer.State, iteration: int - ) -> bool: - it = iteration // inner_iterations - return jnp.logical_not(jnp.isfinite(state.criterions[it - 1])) - - def cond_fn( - iteration: int, - consts: GeneralizedKMeansInitializer.Constants, - state: GeneralizedKMeansInitializer.State, - ) -> bool: - return jnp.logical_or( - iteration <= 2, - jnp.logical_and( - jnp.logical_not(diverged(state, iteration)), - jnp.logical_not(converged(state, consts, iteration)) - ) - ) - - def body_fn( - iteration: int, consts: GeneralizedKMeansInitializer.Constants, - state: GeneralizedKMeansInitializer.State, compute_error: bool - ) -> GeneralizedKMeansInitializer.State: - del compute_error - it = iteration // inner_iterations - - grad = consts.geom.apply_cost(state.factor, axis=1) # (n, r) - grad = grad + consts.geom.apply_cost(state.factor, axis=0) # (n, r) - grad = grad / consts.g - - norm = jnp.max(jnp.abs(grad)) ** 2 - gamma = consts.gamma / norm - eps = 1.0 / gamma - - cost = grad - eps * mu.safe_log(state.factor) # (n, r) - cost = geometry.Geometry( - cost_matrix=cost, - epsilon=eps, - ) - problem = linear_problem.LinearProblem( - cost, a=consts.marginal, b=consts.g - ) - - out = consts.solver(problem) - new_factor = out.matrix - - criterion = ((1 / gamma) ** 2) * ( - mu.kl(new_factor, state.factor) + mu.kl(state.factor, new_factor) - ) - crossed_threshold = jnp.logical_or( - state.crossed_threshold, - jnp.logical_and( - state.criterions[it - 1] >= consts.threshold, criterion - < consts.threshold - ) - ) - - return self.State( - factor=new_factor, - criterions=state.criterions.at[it].set(criterion), - crossed_threshold=crossed_threshold - ) - - del kwargs - - if isinstance(ot_prob, quadratic_problem.QuadraticProblem): - geom = ot_prob.geom_xx if which == "q" else ot_prob.geom_yy - else: - geom = ot_prob.geom - assert geom.shape[0] == geom.shape[ - 1], f"Expected the shape to be square, found `{geom.shape}`." - - inner_iterations = self._kwargs["inner_iterations"] - outer_iterations = np.ceil(self._max_iter / inner_iterations).astype(int) - force_scan = self._min_iter == self._max_iter - fixpoint_fn = ( - fixed_point_loop.fixpoint_iter - if force_scan else fixed_point_loop.fixpoint_iter_backprop - ) - - consts = self.Constants( - solver=sinkhorn.Sinkhorn(**self._sinkhorn_kwargs), - geom=geom.set_scale_cost("max_cost"), - marginal=ot_prob.a if which == "q" else ot_prob.b, - g=init_g, - gamma=self._kwargs["gamma"], - threshold=self._kwargs["threshold"], - ) - - return fixpoint_fn( - cond_fn, - body_fn, - min_iterations=self._min_iter, - max_iterations=self._max_iter, - inner_iterations=inner_iterations, - constants=consts, - state=init_fn(), - ).factor diff --git a/ott/build/lib/ott/initializers/quadratic/__init__.py b/ott/build/lib/ott/initializers/quadratic/__init__.py deleted file mode 100644 index f188cae..0000000 --- a/ott/build/lib/ott/initializers/quadratic/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. -from . import initializers diff --git a/ott/build/lib/ott/initializers/quadratic/initializers.py b/ott/build/lib/ott/initializers/quadratic/initializers.py deleted file mode 100644 index 795e81c..0000000 --- a/ott/build/lib/ott/initializers/quadratic/initializers.py +++ /dev/null @@ -1,193 +0,0 @@ -# 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 abc -from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple - -import jax -import jax.numpy as jnp - -from ott.geometry import geometry - -if TYPE_CHECKING: - from ott.problems.linear import linear_problem - from ott.problems.quadratic import quadratic_problem - -__all__ = ["BaseQuadraticInitializer", "QuadraticInitializer"] - - -@jax.tree_util.register_pytree_node_class -class BaseQuadraticInitializer(abc.ABC): - """Base class for quadratic initializers. - - Args: - kwargs: Keyword arguments. - """ - - def __init__(self, **kwargs: Any): - self._kwargs = kwargs - - def __call__( - self, quad_prob: "quadratic_problem.QuadraticProblem", **kwargs: Any - ) -> "linear_problem.LinearProblem": - """Compute the initial linearization of a quadratic problem. - - Args: - quad_prob: Quadratic problem to linearize. - kwargs: Additional keyword arguments. - - Returns: - Linear problem. - """ - from ott.problems.linear import linear_problem - - n, m = quad_prob.geom_xx.shape[0], quad_prob.geom_yy.shape[0] - geom = self._create_geometry(quad_prob, **kwargs) - assert geom.shape == (n, m), ( - f"Expected geometry of shape `{n, m}`, " - f"found `{geom.shape}`." - ) - return linear_problem.LinearProblem( - geom, - a=quad_prob.a, - b=quad_prob.b, - tau_a=quad_prob.tau_a, - tau_b=quad_prob.tau_b, - ) - - @abc.abstractmethod - def _create_geometry( - self, quad_prob: "quadratic_problem.QuadraticProblem", **kwargs: Any - ) -> geometry.Geometry: - """Compute initial geometry for linearization. - - Args: - quad_prob: Quadratic problem. - kwargs: Additional keyword arguments. - - Returns: - Geometry used to initialize the linearized problem. - """ - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return [], self._kwargs - - @classmethod - def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] - ) -> "BaseQuadraticInitializer": - return cls(*children, **aux_data) - - -class QuadraticInitializer(BaseQuadraticInitializer): - r"""Initialize a linear problem locally around a selected coupling. - - If the problem is balanced (``tau_a = 1`` and ``tau_b = 1``), - the equation of the cost follows eq. 6, p. 1 of :cite:`peyre:16`. - - If the problem is unbalanced (``tau_a < 1`` or ``tau_b < 1``), there are two - possible cases. A first possibility is to introduce a quadratic KL - divergence on the marginals in the objective as done in :cite:`sejourne:21` - (``gw_unbalanced_correction = True``), which in turns modifies the - local cost matrix. - - Alternatively, it could be possible to leave the formulation of the - local cost unchanged, i.e. follow eq. 6, p. 1 of :cite:`peyre:16` - (``gw_unbalanced_correction = False``) and include the unbalanced terms - at the level of the linear problem only. - - Let :math:`P` [num_a, num_b] be the transport matrix, `cost_xx` is the - cost matrix of `geom_xx` and `cost_yy` is the cost matrix of `geom_yy`. - `left_x` and `right_y` depend on the loss chosen for GW. - `gw_unbalanced_correction` is flag indicating whether the unbalanced - correction applies. The equation of the local cost can be written as: - - .. math:: - - \text{marginal_dep_term} + \text{left}_x(\text{cost_xx}) P - \text{right}_y(\text{cost_yy}) + \text{unbalanced_correction} - - When working with the fused problem, a linear term is added to the cost - matrix: `cost_matrix` += `fused_penalty` * `geom_xy.cost_matrix` - - Args: - init_coupling: The coupling to use for initialization. If :obj:`None`, - defaults to the product coupling :math:`ab^T`. - """ - - def __init__( - self, init_coupling: Optional[jnp.ndarray] = None, **kwargs: Any - ): - super().__init__(**kwargs) - self.init_coupling = init_coupling - - def _create_geometry( - self, - quad_prob: "quadratic_problem.QuadraticProblem", - *, - epsilon: float, - relative_epsilon: Optional[bool] = None, - **kwargs: Any, - ) -> geometry.Geometry: - """Compute initial geometry for linearization. - - Args: - quad_prob: Quadratic OT problem. - epsilon: Epsilon regularization. - relative_epsilon: Flag, use `relative_epsilon` or not in geometry. - kwargs: Keyword arguments for :class:`~ott.geometry.geometry.Geometry`. - - Returns: - The initial geometry used to initialize the linearized problem. - """ - from ott.problems.quadratic import quadratic_problem - - del kwargs - - marginal_cost = quad_prob.marginal_dependent_cost(quad_prob.a, quad_prob.b) - geom_xx, geom_yy = quad_prob.geom_xx, quad_prob.geom_yy - - h1, h2 = quad_prob.quad_loss - if self.init_coupling is None: - tmp1 = quadratic_problem.apply_cost(geom_xx, quad_prob.a, axis=1, fn=h1) - tmp2 = quadratic_problem.apply_cost(geom_yy, quad_prob.b, axis=1, fn=h2) - tmp = jnp.outer(tmp1, tmp2) - else: - tmp1 = h1.func(geom_xx.cost_matrix) - tmp2 = h2.func(geom_yy.cost_matrix) - tmp = tmp1 @ self.init_coupling @ tmp2.T - - if quad_prob.is_balanced: - cost_matrix = marginal_cost.cost_matrix - tmp - else: - # initialize epsilon for Unbalanced GW according to Sejourne et. al (2021) - init_transport = jnp.outer(quad_prob.a, quad_prob.b) - marginal_1, marginal_2 = init_transport.sum(1), init_transport.sum(0) - - epsilon = quadratic_problem.update_epsilon_unbalanced( - epsilon=epsilon, transport_mass=marginal_1.sum() - ) - unbalanced_correction = quad_prob.cost_unbalanced_correction( - init_transport, marginal_1, marginal_2, epsilon=epsilon - ) - cost_matrix = marginal_cost.cost_matrix - tmp + unbalanced_correction - - cost_matrix += quad_prob.fused_penalty * quad_prob._fused_cost_matrix - return geometry.Geometry( - cost_matrix=cost_matrix, - epsilon=epsilon, - relative_epsilon=relative_epsilon - ) - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return [self.init_coupling], self._kwargs diff --git a/ott/build/lib/ott/math/__init__.py b/ott/build/lib/ott/math/__init__.py deleted file mode 100644 index 64bc1c0..0000000 --- a/ott/build/lib/ott/math/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -# 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. -from . import ( - fixed_point_loop, - matrix_square_root, - unbalanced_functions, - utils, -) diff --git a/ott/build/lib/ott/math/fixed_point_loop.py b/ott/build/lib/ott/math/fixed_point_loop.py deleted file mode 100644 index 9034eba..0000000 --- a/ott/build/lib/ott/math/fixed_point_loop.py +++ /dev/null @@ -1,239 +0,0 @@ -# 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. -from typing import Any, Callable - -import jax -import jax.numpy as jnp -import numpy as np - -__all__ = ["fixpoint_iter", "fixpoint_iter_backprop"] - - -def fixpoint_iter( - cond_fn: Callable[[int, Any, Any], bool], - body_fn: Callable[[Any, Any, Any, Any], Any], min_iterations: int, - max_iterations: int, inner_iterations: int, constants: Any, state: Any -): - """Implementation of a fixed point loop. - - This fixed point loop iterator applies ``body_fn`` to a tuple - ``(iteration, constants, state, compute_error)`` to output a new state, using - context provided in iteration and constants. - - ``body_fn`` is iterated (inner_iterations -1) times, and one last time with - the ``compute_error`` flag to ``True``, indicating that additional - computational effort can be spent on recalculating the latest error - (``errors`` are stored as the first element of the state tuple). - - upon termination of these ``inner_iterations``, the loop is continued if - iteration is smaller than ``min_iterations``, stopped if equal/larger than - ``max_iterations``, and interrupted if ``cond_fn`` returns False. - - Args: - cond_fn : termination condition function - body_fn : body loop instructions - min_iterations : lower bound on the total amount of fixed point iterations - max_iterations : upper bound on the total amount of fixed point iterations - inner_iterations : number of iterations ``body_fn`` will be executed - successively before calling ``cond_fn``. - constants : constant (during loop) parameters passed on to body - state : state variable - - Returns: - outputs state returned by ``body_fn`` upon termination. - """ # noqa: D401 - # If number of minimal iterations matches maximal number, force a scan instead - # of a while loop. - - force_scan = (min_iterations == max_iterations) - - compute_error_flags = jnp.arange(inner_iterations) == inner_iterations - 1 - - def max_cond_fn(iteration_state): - iteration, state = iteration_state - return jnp.logical_and( - iteration < max_iterations, - jnp.logical_or( - iteration < min_iterations, cond_fn(iteration, constants, state) - ) - ) - - def unrolled_body_fn(iteration_state): - - def one_iteration(iteration_state, compute_error): - iteration, state = iteration_state - state = body_fn(iteration, constants, state, compute_error) - iteration += 1 - return (iteration, state), None - - iteration_state, _ = jax.lax.scan( - one_iteration, iteration_state, compute_error_flags - ) - return (iteration_state, None) if force_scan else iteration_state - - if force_scan: - (_, state), _ = jax.lax.scan( - lambda carry, x: unrolled_body_fn(carry), (0, state), - None, - length=max_iterations // inner_iterations - ) - else: - _, state = jax.lax.while_loop(max_cond_fn, unrolled_body_fn, (0, state)) - return state - - -def fixpoint_iter_fwd( - cond_fn, body_fn, min_iterations, max_iterations, inner_iterations, - constants, state -): - """Forward iteration of fixed point iteration to handle backpropagation. - - The main difference with fixpoint_iter is the checkpointing, in variable - states, of the state variables as they are recorded through iterations, every - inner_iterations. This sequence of states will be used in the backward loop. - - Args: - cond_fn : termination condition function - body_fn : body loop instructions - min_iterations : lower bound on the total amount of fixed point iterations - max_iterations : upper bound on the total amount of fixed point iterations - inner_iterations : number of iterations body_fn will be executed - successively before calling cond_fn. - constants : constant (during loop) parameters passed on to body - state : state variable - - Returns: - outputs state returned by body_fn upon termination. - """ - force_scan = min_iterations == max_iterations - compute_error_flags = jnp.arange(inner_iterations) == inner_iterations - 1 - states = jax.tree_util.tree_map( - lambda x: jnp.zeros( - (max_iterations // inner_iterations + 1,) + jnp.shape(x), - dtype=jax.dtypes.result_type(x) - ), state - ) - - def max_cond_fn(iteration_states_state): - iteration, _, state = iteration_states_state - return jnp.logical_and( - iteration < max_iterations, - jnp.logical_or( - iteration < min_iterations, cond_fn(iteration, constants, state) - ) - ) - - def unrolled_body_fn(iteration_states_state): - iteration, states, state = iteration_states_state - states = jax.tree_util.tree_map( - lambda states, state: jax.lax.dynamic_update_index_in_dim( - states, state, iteration // inner_iterations, 0 - ), states, state - ) - - def one_iteration(iteration_state, compute_error): - iteration, state = iteration_state - state = body_fn(iteration, constants, state, compute_error) - iteration += 1 - return (iteration, state), None - - iteration_state, _ = jax.lax.scan( - one_iteration, (iteration, state), compute_error_flags - ) - iteration, state = iteration_state - out = (iteration, states, state) - return (out, None) if force_scan else out - - if force_scan: - (iteration, states, state), _ = jax.lax.scan( - lambda carry, x: unrolled_body_fn(carry), (0, states, state), - None, - length=max_iterations // inner_iterations - ) - else: - iteration, states, state = jax.lax.while_loop( - max_cond_fn, unrolled_body_fn, (0, states, state) - ) - - return state, (constants, iteration, states) - - -def fixpoint_iter_bwd( - cond_fn, body_fn, min_iterations, max_iterations, inner_iterations, res, g -): - """Backward iteration of fixed point iteration, using checkpointed states.""" - del cond_fn - force_scan = (min_iterations == max_iterations) - constants, iteration, states = res - # The tree may contain some python floats - g_constants = jax.tree_util.tree_map( - lambda x: jnp.zeros_like(x, dtype=x.dtype) - if isinstance(x, (np.ndarray, jnp.ndarray)) else 0, constants - ) - - def bwd_cond_fn(iteration_g_gconst): - iteration, _, _ = iteration_g_gconst - return iteration >= 0 - - def unrolled_body_fn_no_errors(iteration, constants, state): - compute_error_flags = jnp.zeros((inner_iterations,), dtype=bool) - - def one_iteration(iteration_state, compute_error): - iteration, state = iteration_state - state = body_fn(iteration, constants, state, compute_error) - iteration += 1 - return (iteration, state), None - - iteration_state, _ = jax.lax.scan( - one_iteration, (iteration, state), compute_error_flags - ) - _, state = iteration_state - return state - - def unrolled_body_fn(iteration_g_gconst): - iteration, g, g_constants = iteration_g_gconst - state = jax.tree_util.tree_map( - lambda x: x[iteration // inner_iterations], states - ) - _, pullback = jax.vjp( - unrolled_body_fn_no_errors, iteration, constants, state - ) - _, gi_constants, g_state = pullback(g) - g_constants = jax.tree_util.tree_map( - lambda x, y: x + y, g_constants, gi_constants - ) - out = (iteration - inner_iterations, g_state, g_constants) - return (out, None) if force_scan else out - - if force_scan: - (_, g_state, g_constants), _ = jax.lax.scan( - lambda carry, x: unrolled_body_fn(carry), (0, g, g_constants), - None, - length=max_iterations // inner_iterations - ) - else: - _, g_state, g_constants = jax.lax.while_loop( - bwd_cond_fn, unrolled_body_fn, - (iteration - inner_iterations, g, g_constants) - ) - - return g_constants, g_state - - -# definition of backprop friendly variant of fixpoint_iter. -fixpoint_iter_backprop = jax.custom_vjp( - fixpoint_iter, nondiff_argnums=(0, 1, 2, 3, 4) -) - -fixpoint_iter_backprop.defvjp(fixpoint_iter_fwd, fixpoint_iter_bwd) diff --git a/ott/build/lib/ott/math/matrix_square_root.py b/ott/build/lib/ott/math/matrix_square_root.py deleted file mode 100644 index 324a5ea..0000000 --- a/ott/build/lib/ott/math/matrix_square_root.py +++ /dev/null @@ -1,337 +0,0 @@ -# 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 functools -import math -from typing import Tuple - -import jax -import jax.numpy as jnp - -from ott.math import fixed_point_loop - -__all__ = ["sqrtm", "sqrtm_only", "inv_sqrtm_only"] - - -@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2, 3, 4, 5)) -def sqrtm( - x: jnp.ndarray, - threshold: float = 1e-6, - min_iterations: int = 0, - inner_iterations: int = 10, - max_iterations: int = 1000, - regularization: float = 1e-6 -) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Higham algorithm to compute matrix square root of p.d. matrix. - - See :cite:`higham:97`, eq. 2.6b - - Args: - x: a (batch of) square p.s.d. matrices of the same size. - threshold: convergence tolerance threshold for Newton-Schulz iterations. - min_iterations: min number of iterations after which error is computed. - inner_iterations: error is re-evaluated every inner_iterations iterations. - max_iterations: max number of iterations. - regularization: small regularizer added to norm of x, before normalization. - - Returns: - Square root matrix of x (or x's if batch), its inverse, - errors along iterates. - """ - dimension = x.shape[-1] - norm_x = jnp.linalg.norm(x, axis=(-2, -1)) * (1 + regularization) - - if jnp.ndim(x) > 2: - norm_x = norm_x[..., jnp.newaxis, jnp.newaxis] - - def cond_fn(iteration, const, state): - """Stopping criterion. Checking decrease of objective is needed here.""" - _, threshold = const - errors, _, _ = state - err = errors[iteration // inner_iterations - 1] - - return jnp.logical_or( - iteration == 0, - jnp.logical_and( - jnp.logical_and(jnp.isfinite(err), err > threshold), - jnp.all(jnp.diff(errors) <= 0) - ) - ) # check decreasing obj, else stop - - def body_fn(iteration, const, state, compute_error): - """Carry out matrix updates on y and z, stores error if requested. - - Args: - iteration: iteration number - const: tuple of constant parameters that do not change throughout the - loop. - state: state variables currently updated in the loop. - compute_error: flag to indicate this iteration computes/stores an error - - Returns: - state variables. - """ - x, _ = const - errors, y, z = state - w = 0.5 * jnp.matmul(z, y) - y = 1.5 * y - jnp.matmul(y, w) - z = 1.5 * z - jnp.matmul(w, z) - - err = jnp.where(compute_error, new_err(x, norm_x, y), jnp.inf) - - errors = errors.at[iteration // inner_iterations].set(err) - - return errors, y, z - - def new_err(x, norm_x, y): - res = x - norm_x * jnp.matmul(y, y) - norm_fn = functools.partial(jnp.linalg.norm, axis=(-2, -1)) - return jnp.max(norm_fn(res) / norm_fn(x)) - - y = x / norm_x - z = jnp.eye(dimension) - if jnp.ndim(x) > 2: - z = jnp.tile(z, list(x.shape[:-2]) + [1, 1]) - errors = -jnp.ones(math.ceil(max_iterations / inner_iterations)) - state = (errors, y, z) - const = (x, threshold) - errors, y, z = fixed_point_loop.fixpoint_iter_backprop( - cond_fn, body_fn, min_iterations, max_iterations, inner_iterations, const, - state - ) - sqrt_x = jnp.sqrt(norm_x) * y - inv_sqrt_x = z / jnp.sqrt(norm_x) - - return sqrt_x, inv_sqrt_x, errors - - -def solve_sylvester_bartels_stewart( - a: jnp.ndarray, - b: jnp.ndarray, - c: jnp.ndarray, -) -> jnp.ndarray: - """Solve the real Sylvester equation AX - XB = C using Bartels-Stewart.""" - # See https://nhigham.com/2020/09/01/what-is-the-sylvester-equation/ for - # discussion of the algorithm (but note that in the derivation, the sign on - # the right hand side is flipped in the equation in which the columns are set - # to be equal). - m = a.shape[-1] - n = b.shape[-1] - # Cast a and b to complex to ensure we get the complex Schur decomposition - # (the real Schur decomposition may not give an upper triangular solution). - # For the decomposition below, a = u r u* and b = v s v* - r, u = jax.lax.linalg.schur(a + 0j) - s, v = jax.lax.linalg.schur(b + 0j) - d = jnp.matmul( - jnp.conjugate(jnp.swapaxes(u, axis1=-2, axis2=-1)), jnp.matmul(c, v) - ) - # The solution in the transformed space will in general be complex, too. - y = jnp.zeros(a.shape[:-2] + (m, n)) + 0j - idx = jnp.arange(m) - for j in range(n): - lhs = r.at[..., idx, idx].add(-s[..., j:j + 1, j]) - rhs = d[..., j] + jnp.matmul(y[..., :j], s[..., :j, j:j + 1])[..., 0] - y = y.at[..., j].set(jax.scipy.linalg.solve_triangular(lhs, rhs)) - - x = jnp.matmul( - u, jnp.matmul(y, jnp.conjugate(jnp.swapaxes(v, axis1=-2, axis2=-1))) - ) - # The end result should be real; remove the imaginary part of the solution. - return jnp.real(x) - - -def sqrtm_fwd( - x: jnp.ndarray, - threshold: float, - min_iterations: int, - inner_iterations: int, - max_iterations: int, - regularization: float, -) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], Tuple[jnp.ndarray, - jnp.ndarray]]: - """Forward pass of custom VJP.""" - sqrt_x, inv_sqrt_x, errors = sqrtm( - x=x, - threshold=threshold, - min_iterations=min_iterations, - inner_iterations=inner_iterations, - max_iterations=max_iterations, - regularization=regularization, - ) - return (sqrt_x, inv_sqrt_x, errors), (sqrt_x, inv_sqrt_x) - - -def sqrtm_bwd( - threshold: float, - min_iterations: int, - inner_iterations: int, - max_iterations: int, - regularization: float, - residual: Tuple[jnp.ndarray, jnp.ndarray], - cotangent: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], -) -> Tuple[jnp.ndarray]: - """Compute the derivative by solving a Sylvester equation.""" - del threshold, min_iterations, inner_iterations, \ - max_iterations, regularization - sqrt_x, inv_sqrt_x = residual - # ignores cotangent associated with errors - cot_sqrt, cot_inv_sqrt, _ = cotangent - - # Solve for d(X^{1/2}): - # Start with X^{1/2} X^{1/2} = X - # Differentiate to obtain - # d(X^{1/2}) X^{1/2} + X^{1/2} d(X^{1/2}) = dX - # The above is a Sylvester equation that we can solve using Bartels-Stewart. - # Below think of cot_sqrt as (dX)^T and vjp_cot_sqrt as d(X^{1/2})^T. - # See https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html - vjp_cot_sqrt = jnp.swapaxes( - solve_sylvester_bartels_stewart( - a=sqrt_x, b=-sqrt_x, c=jnp.swapaxes(cot_sqrt, axis1=-1, axis2=-2) - ), - axis1=-1, - axis2=-2 - ) - - # Now solve for d(X^{-1/2}): - # Start with X^{-1/2} X^{-1/2} = X^{-1} - # Use the product rule and the fact that d(X^{-1}) = -X^{-1} dX X^{-1} - # to obtain - # (See The Matrix Cookbook section on derivatives of an inverse - # https://www.math.uwaterloo.ca/~hwolkowi/matrixcookbook.pdf ) - # d(X^{-1/2}) X^{-1/2} + X^{-1/2} d(X^{-1/2}) = -X^{-1} dX X^{-1} - # Again we have a Sylvester equation that we solve as above, and again we - # think of cot_inv_sqrt as (dX)^T and vjp_cot_inv_sqrt as d(X^{-1/2})^T - inv_x = jnp.matmul(inv_sqrt_x, inv_sqrt_x) - vjp_cot_inv_sqrt = jnp.swapaxes( - solve_sylvester_bartels_stewart( - a=inv_sqrt_x, - b=-inv_sqrt_x, - c=-jnp.matmul( - inv_x, - jnp.matmul(jnp.swapaxes(cot_inv_sqrt, axis1=-2, axis2=-1), inv_x) - ) - ), - axis1=-1, - axis2=-2 - ) - return vjp_cot_sqrt + vjp_cot_inv_sqrt, - - -sqrtm.defvjp(sqrtm_fwd, sqrtm_bwd) - -# Specialized versions of sqrtm that compute only the square root or inverse. -# These functions have lower complexity gradients than sqrtm. - - -@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2, 3, 4, 5)) -def sqrtm_only( # noqa: D103 - x: jnp.ndarray, - threshold: float = 1e-6, - min_iterations: int = 0, - inner_iterations: int = 10, - max_iterations: int = 1000, - regularization: float = 1e-6 -) -> jnp.ndarray: - return sqrtm( - x, threshold, min_iterations, inner_iterations, max_iterations, - regularization - )[0] - - -def sqrtm_only_fwd( # noqa: D103 - x: jnp.ndarray, threshold: float, min_iterations: int, - inner_iterations: int, max_iterations: int, regularization: float -) -> Tuple[jnp.ndarray, jnp.ndarray]: - sqrt_x = sqrtm( - x, threshold, min_iterations, inner_iterations, max_iterations, - regularization - )[0] - return sqrt_x, sqrt_x - - -def sqrtm_only_bwd( # noqa: D103 - threshold: float, min_iterations: int, inner_iterations: int, - max_iterations: int, regularization: float, sqrt_x: jnp.ndarray, - cotangent: jnp.ndarray -) -> Tuple[jnp.ndarray]: - del threshold, min_iterations, inner_iterations, \ - max_iterations, regularization - vjp = jnp.swapaxes( - solve_sylvester_bartels_stewart( - a=sqrt_x, b=-sqrt_x, c=jnp.swapaxes(cotangent, axis1=-2, axis2=-1) - ), - axis1=-2, - axis2=-1 - ) - return vjp, - - -sqrtm_only.defvjp(sqrtm_only_fwd, sqrtm_only_bwd) - - -@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2, 3, 4, 5)) -def inv_sqrtm_only( # noqa: D103 - x: jnp.ndarray, - threshold: float = 1e-6, - min_iterations: int = 0, - inner_iterations: int = 10, - max_iterations: int = 1000, - regularization: float = 1e-6 -) -> jnp.ndarray: - return sqrtm( - x, threshold, min_iterations, inner_iterations, max_iterations, - regularization - )[1] - - -def inv_sqrtm_only_fwd( # noqa: D103 - x: jnp.ndarray, - threshold: float, - min_iterations: int, - inner_iterations: int, - max_iterations: int, - regularization: float, -) -> Tuple[jnp.ndarray, jnp.ndarray]: - inv_sqrt_x = sqrtm( - x, threshold, min_iterations, inner_iterations, max_iterations, - regularization - )[1] - return inv_sqrt_x, inv_sqrt_x - - -def inv_sqrtm_only_bwd( # noqa: D103 - threshold: float, min_iterations: int, inner_iterations: int, - max_iterations: int, regularization: float, residual: jnp.ndarray, - cotangent: jnp.ndarray -) -> Tuple[jnp.ndarray]: - del threshold, min_iterations, inner_iterations, \ - max_iterations, regularization - - inv_sqrt_x = residual - inv_x = jnp.matmul(inv_sqrt_x, inv_sqrt_x) - vjp = jnp.swapaxes( - solve_sylvester_bartels_stewart( - a=inv_sqrt_x, - b=-inv_sqrt_x, - c=-jnp.matmul( - inv_x, - jnp.matmul(jnp.swapaxes(cotangent, axis1=-2, axis2=-1), inv_x) - ) - ), - axis1=-1, - axis2=-2 - ) - return vjp, - - -inv_sqrtm_only.defvjp(inv_sqrtm_only_fwd, inv_sqrtm_only_bwd) diff --git a/ott/build/lib/ott/math/unbalanced_functions.py b/ott/build/lib/ott/math/unbalanced_functions.py deleted file mode 100644 index 9ba20ba..0000000 --- a/ott/build/lib/ott/math/unbalanced_functions.py +++ /dev/null @@ -1,90 +0,0 @@ -# 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. -from typing import Callable - -import jax.numpy as jnp - - -def phi_star(h: jnp.ndarray, rho: float) -> jnp.ndarray: - """Legendre transform of KL, :cite:`sejourne:19`, p. 9.""" - return rho * (jnp.exp(h / rho) - 1) - - -def derivative_phi_star(f: jnp.ndarray, rho: float) -> jnp.ndarray: - """Derivative of Legendre transform of phi_starKL, see phi_star.""" - # TODO(cuturi): use jax.grad directly. - return jnp.exp(f / rho) - - -def grad_of_marginal_fit( - c: jnp.ndarray, h: jnp.ndarray, tau: float, epsilon: float -) -> jnp.ndarray: - """Compute grad of terms linked to marginals in objective. - - Computes gradient w.r.t. f ( or g) of terms in :cite:`sejourne:19`, - left-hand-side of eq. 15 terms involving phi_star). - - Args: - c: jnp.ndarray, first target marginal (either a or b in practice) - h: jnp.ndarray, potential (either f or g in practice) - tau: float, strength (in ]0,1]) of regularizer w.r.t. marginal - epsilon: regularization - - Returns: - a vector of the same size as c or h - """ - if tau == 1.0: - return c - r = rho(epsilon, tau) - return jnp.where(c > 0, c * derivative_phi_star(-h, r), 0.0) - - -def second_derivative_phi_star(f: jnp.ndarray, rho: float) -> jnp.ndarray: - """Second Derivative of Legendre transform of KL, see phi_star.""" - return jnp.exp(f / rho) / rho - - -def diag_jacobian_of_marginal_fit( - c: jnp.ndarray, h: jnp.ndarray, tau: float, epsilon: float, - derivative: Callable[[jnp.ndarray, float], jnp.ndarray] -): - """Compute grad of terms linked to marginals in objective. - - Computes second derivative w.r.t. f ( or g) of terms in :cite:`sejourne:19`, - left-hand-side of eq. 32 (terms involving phi_star) - - Args: - c: jnp.ndarray, first target marginal (either a or b in practice) - h: jnp.ndarray, potential (either f or g in practice) - tau: float, strength (in ]0,1]) of regularizer w.r.t. marginal - epsilon: regularization - derivative: Callable - - Returns: - a vector of the same size as c or h. - """ - if tau == 1.0: - return 0.0 - - r = rho(epsilon, tau) - # here no minus sign because we are taking derivative w.r.t -h - return jnp.where( - c > 0, - c * second_derivative_phi_star(-h, r) * - derivative(c * derivative_phi_star(-h, r)), 0.0 - ) - - -def rho(epsilon: float, tau: float) -> float: # noqa: D103 - return (epsilon * tau) / (1.0 - tau) diff --git a/ott/build/lib/ott/math/utils.py b/ott/build/lib/ott/math/utils.py deleted file mode 100644 index afad18d..0000000 --- a/ott/build/lib/ott/math/utils.py +++ /dev/null @@ -1,298 +0,0 @@ -# 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 functools -from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Union - -import jax -import jax.numpy as jnp -import jax.scipy as jsp - -if TYPE_CHECKING: - from ott.geometry import costs - -__all__ = [ - "safe_log", - "norm", - "kl", - "gen_kl", - "gen_js", - "logsumexp", - "softmin", - "barycentric_projection", - "sort_and_argsort", - "lambertw", -] - - -def safe_log( # noqa: D103 - x: jnp.ndarray, - *, - eps: Optional[float] = None -) -> jnp.ndarray: - if eps is None: - eps = jnp.finfo(x.dtype).tiny - return jnp.where(x > 0.0, jnp.log(x), jnp.log(eps)) - - -@functools.partial(jax.custom_jvp, nondiff_argnums=[1, 2, 3]) -@functools.partial(jax.jit, static_argnames=("ord", "axis", "keepdims")) -def norm( - x: jnp.ndarray, - ord: Union[int, str, None] = None, - axis: Union[None, Sequence[int], int] = None, - keepdims: bool = False -) -> jnp.ndarray: - """Computes order ord norm of vector, using `jnp.linalg` in forward pass. - - Evaluations of distances between a vector and itself using translation - invariant costs, typically norms, result in functions of the form - ``lambda x : jnp.linalg.norm(x-x)``. Such functions output `NaN` gradients, - because they involve computing the derivative of a negative exponent of 0 - (e.g. when differentiating the Euclidean norm, one gets a 0-denominator in the - expression, see e.g. https://github.com/google/jax/issues/6484 for context). - - While this makes sense mathematically, in the context of optimal transport - such distances between a point and itself can be safely ignored when they - contribute to an OT cost (when, for instance, computing Sinkhorn divergences, - involving computing the OT cost of a point cloud with itself). - - To avoid such `NaN` values, this custom norm implementation uses the - double-where trick, to avoid having branches that output any `NaN`, and - safely output a 0 instead. - - Args: - x: Input array. If `axis` is None, `x` must be 1-D or 2-D, unless `ord` - is None. If both `axis` and `ord` are None, the 2-norm of ``x.ravel`` - will be returned. - ord: `{non-zero int, jnp.inf, -jnp.inf, 'fro', 'nuc'}`, Order of the norm. - The default is `None`, which is equivalent to `2.0` for vectors. - axis: `{None, int, 2-tuple of ints}`, optional. If `axis` is an integer, it - specifies the axis of `x` along which to compute the vector norms. - If `axis` is a 2-tuple, it specifies the axes that hold 2-D matrices, and - the matrix norms of these matrices are computed. If `axis` is None then - either a vector norm (when `x` is 1-D) or a matrix norm (when `x` is 2-D) - is returned. The default is None. - keepdims: If set to True, the axes which are normed over are left in the - result as dimensions with size one. With this option the result will - broadcast correctly against the original `x`. - - Returns: - float or ndarray, Norm of the matrix or vector(s). - """ - return jnp.linalg.norm(x, ord=ord, axis=axis, keepdims=keepdims) - - -@norm.defjvp -def norm_jvp(ord, axis, keepdims, primals, tangents): - """Custom_jvp for norm, that returns 0.0 when evaluated at 0.""" - x, = primals - x_is_zero = jnp.all(jnp.logical_not(x)) - clean_x = jnp.where(x_is_zero, jnp.ones_like(x), x) - primals, tangents = jax.jvp( - functools.partial(jnp.linalg.norm, ord=ord, axis=axis, keepdims=keepdims), - (clean_x,), tangents - ) - return primals, jnp.where(x_is_zero, 0.0, tangents) - - -# TODO(michalk8): add axis argument -def kl(p: jnp.ndarray, q: jnp.ndarray) -> float: - """Kullback-Leibler divergence.""" - return jnp.vdot(p, (safe_log(p) - safe_log(q))) - - -def gen_kl(p: jnp.ndarray, q: jnp.ndarray) -> float: - """Generalized Kullback-Leibler divergence.""" - return jnp.vdot(p, (safe_log(p) - safe_log(q))) + jnp.sum(q) - jnp.sum(p) - - -# TODO(michalk8): add axis argument -def gen_js(p: jnp.ndarray, q: jnp.ndarray, c: float = 0.5) -> float: - """Jensen-Shannon divergence.""" - return c * (gen_kl(p, q) + gen_kl(q, p)) - - -@functools.partial(jax.custom_jvp, nondiff_argnums=(1, 2, 4)) -def logsumexp( # noqa: D103 - mat, axis=None, keepdims=False, b=None, return_sign=False -): - return jax.scipy.special.logsumexp( - mat, axis=axis, keepdims=keepdims, b=b, return_sign=return_sign - ) - - -@logsumexp.defjvp -def logsumexp_jvp(axis, keepdims, return_sign, primals, tangents): - """Custom derivative rule for lse that does not blow up with -inf. - - This logsumexp implementation uses the standard jax one in forward mode but - implements a custom rule to differentiate. Given the preference of jax for - jvp over vjp, and the fact that this is a simple linear rule, jvp is used. - This custom differentiation address issues when the output of lse is - -inf (which corresponds to the case where all inputs in a slice are -inf, - which happens typically when ``a`` or ``b`` weight vectors have zeros.) - - Although both exp(lse) and its derivative should be 0, automatic - differentiation returns a NaN derivative because of a -inf - (-inf) operation - appearing in the definition of centered_exp below. This is corrected in the - implementation below. - - Args: - axis: argument from original logsumexp - keepdims: argument from original logsumexp - return_sign: argument from original logsumexp - primals: mat and b, the two arguments against which we differentiate. - tangents: of same size as mat and b. - - Returns: - original primal outputs + their tangent. - """ # noqa: D401 - mat, b = primals - tan_mat, tan_b = tangents - lse = logsumexp(mat, axis, keepdims, b, return_sign) - if return_sign: - lse, sign = lse - lse = jnp.where(jnp.isfinite(lse), lse, 0.0) - centered_exp = jnp.exp(mat - jnp.expand_dims(lse, axis=axis)) - - if b is None: - res = jnp.sum(centered_exp * tan_mat, axis=axis, keepdims=keepdims) - else: - res = jnp.sum(b * centered_exp * tan_mat, axis=axis, keepdims=keepdims) - res += jnp.sum(tan_b * centered_exp, axis=axis, keepdims=keepdims) - if return_sign: - return (lse, sign), (sign * res, jnp.zeros_like(sign)) - return lse, res - - -@functools.partial(jax.custom_vjp, nondiff_argnums=(2,)) -def softmin( - x: jnp.ndarray, gamma: float, axis: Optional[int] = None -) -> jnp.ndarray: - r"""Soft-min operator. - - Args: - x: Input data. - gamma: Smoothing parameter :math:`> 0`. - axis: Axis or axes over which to operate. If ``None``, use flattened input. - - Returns: - The soft minimum. - """ - return -gamma * jsp.special.logsumexp(x / -gamma, axis=axis) - - -softmin.defvjp( - lambda x, gamma, axis: (softmin(x, gamma, axis), (x / -gamma, axis)), - lambda axis, res, g: ( - jnp.where( - jnp.isinf(res[0]), 0.0, - jax.nn.softmax(res[0], axis=axis) * - (g if axis is None else jnp.expand_dims(g, axis=axis)) - ), None - ) -) - - -@functools.partial(jax.vmap, in_axes=[0, 0, None]) -def barycentric_projection( - matrix: jnp.ndarray, y: jnp.ndarray, cost_fn: "costs.CostFn" -) -> jnp.ndarray: - """Compute the barycentric projection of a matrix. - - Args: - matrix: a matrix of shape (n, m) - y: a vector of shape (m,) - cost_fn: a CostFn instance. - - Returns: - a vector of shape (n,) containing the barycentric projection of matrix. - """ - return jax.vmap( - lambda m, y: cost_fn.barycenter(m, y)[0], in_axes=[0, None] - )(matrix, y) - - -def sort_and_argsort( - x: jnp.array, - *, - argsort: bool = False -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: - """Unified function that returns both sort and argsort, if latter needed.""" - if argsort: - i_x = jnp.argsort(x) - return x[i_x], i_x - return jnp.sort(x), None - - -@functools.partial(jax.custom_jvp, nondiff_argnums=(1, 2)) -def lambertw( - z: jnp.ndarray, tol: float = 1e-8, max_iter: int = 100 -) -> jnp.ndarray: - """Principal branch of the - `Lambert W function `_. - - This implementation uses Halley's iteration and the global initialization - proposed in :cite:`iacono:17`, Eq. 20 . - - Args: - z: Array. - tol: Tolerance threshold. - max_iter: Maximum number of iterations. - - Returns: - The Lambert W evaluated at ``z``. - """ # noqa: D205 - - def initial_iacono(x: jnp.ndarray) -> jnp.ndarray: - y = jnp.sqrt(1.0 + jnp.e * x) - num = 1.0 + 1.14956131 * y - denom = 1.0 + 0.45495740 * jnp.log1p(y) - return -1.0 + 2.036 * jnp.log(num / denom) - - def cond_fun(container): - it, converged, _ = container - return jnp.logical_and(jnp.any(~converged), it < max_iter) - - def halley_iteration(container): - it, _, w = container - - # modified from `tensorflow_probability` - f = w - z * jnp.exp(-w) - delta = f / (w + 1.0 - 0.5 * (w + 2.0) * f / (w + 1.0)) - - w_next = w - delta - - not_converged = jnp.abs(delta) <= tol * jnp.abs(w_next) - return it + 1, not_converged, w_next - - w0 = initial_iacono(z) - converged = jnp.zeros_like(w0, dtype=bool) - - _, _, w = jax.lax.while_loop( - cond_fun=cond_fun, body_fun=halley_iteration, init_val=(0, converged, w0) - ) - return w - - -@lambertw.defjvp -def _lambertw_jvp( - tol: float, max_iter: int, primals: Tuple[jnp.ndarray, ...], - tangents: Tuple[jnp.ndarray, ...] -) -> Tuple[jnp.ndarray, jnp.ndarray]: - z, = primals - dz, = tangents - w = lambertw(z, tol=tol, max_iter=max_iter) - pz = jnp.where(z == 0.0, 1.0, w / ((1.0 + w) * z)) - return w, pz * dz diff --git a/ott/build/lib/ott/neural/__init__.py b/ott/build/lib/ott/neural/__init__.py deleted file mode 100644 index aa1ca23..0000000 --- a/ott/build/lib/ott/neural/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. -from . import layers, losses, models, solvers diff --git a/ott/build/lib/ott/neural/layers.py b/ott/build/lib/ott/neural/layers.py deleted file mode 100644 index 78c2ef3..0000000 --- a/ott/build/lib/ott/neural/layers.py +++ /dev/null @@ -1,213 +0,0 @@ -# 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. -from typing import Any, Callable, Optional, Tuple - -import jax -import jax.numpy as jnp -from flax import linen as nn - -__all__ = ["PositiveDense", "PosDefPotentials"] - -PRNGKey = jax.Array -Shape = Tuple[int, ...] -Dtype = Any -Array = jnp.ndarray - -# wrap to silence docs linter -DEFAULT_KERNEL_INIT = lambda *a, **k: nn.initializers.lecun_normal()(*a, **k) -DEFAULT_BIAS_INIT = nn.initializers.zeros -DEFAULT_RECTIFIER = nn.activation.relu - - -class PositiveDense(nn.Module): - """A linear transformation using a matrix with all entries non-negative. - - Args: - dim_hidden: Number of output dimensions. - rectifier_fn: Rectifier function. The default is - :func:`~flax.linen.activation.relu`. - use_bias: Whether to add bias to the output. - kernel_init: Initializer for the matrix. The default is - :func:`~flax.linen.initializers.lecun_normal`. - bias_init: Initializer for the bias. The default is - :func:`~flax.linen.initializers.zeros`. - precision: Numerical precision of the computation. - """ - - dim_hidden: int - rectifier_fn: Callable[[Array], Array] = DEFAULT_RECTIFIER - use_bias: bool = True - kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = DEFAULT_KERNEL_INIT - bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = DEFAULT_BIAS_INIT - precision: Optional[jax.lax.Precision] = None - - @nn.compact - def __call__(self, x: jnp.ndarray) -> jnp.ndarray: - """Applies a linear transformation to x along the last dimension. - - Args: - x: Array of shape ``[batch, ..., features]``. - - Returns: - Array of shape ``[batch, ..., dim_hidden]``. - """ - # TODO(michalk8): update when refactoring neuraldual - # assert x.ndim > 1, x.ndim - - kernel = self.param( - "kernel", self.kernel_init, (x.shape[-1], self.dim_hidden) - ) - kernel = self.rectifier_fn(kernel) - - x = jnp.tensordot(x, kernel, axes=(-1, 0), precision=self.precision) - if self.use_bias: - x = x + self.param("bias", self.bias_init, (self.dim_hidden,)) - - return x - - -class PosDefPotentials(nn.Module): - r""":math:`\frac{1}{2} x^T (A_i A_i^T + \text{Diag}(d_i)) x + b_i^T x^2 + c_i` potentials. - - This class implements a layer that takes (batched) ``d``-dimensional vectors - ``x`` in, to output a ``num_potentials``-dimensional vector. Each of the - entries in that output is a positive definite quadratic form evaluated at - ``x``; each of these quadratic terms is parameterized as a low-rank plus - diagonal matrix. The low-rank term is parameterized as :math:`A_i A_i^T`, - where each of these matrices is of size ``(rank, d)``. Taken together, - these matrices form a tensor ``(num_potentials, rank, d)``. - The diagonal terms :math:`d_i` form a ``(num_potentials, d)`` matrix of - positive values; the linear terms :math:`b_i` form a ``(num_potentials, d)`` - matrix. Finally, the :math:`c_i` are contained in a vector of size - ``(num_potentials,)``. - - Args: - num_potentials: Dimension of the output. - rank: Rank of the matrices :math:`A_i` used as low-rank factors - for the quadratic potentials. - rectifier_fn: Rectifier function to ensure non-negativity of the diagonals - :math:`d_i`. The default is :func:`~flax.linen.activation.relu`. - use_linear: Whether to add a linear layers :math:`b_i` to the outputs. - use_bias: Whether to add biases :math:`c_i` to the outputs. - kernel_lr_init: Initializer for the matrices :math:`A_i` - of the quadratic potentials when ``rank > 0``. - The default is :func:`~flax.linen.initializers.lecun_normal`. - kernel_diag_init: Initializer for the diagonals :math:`d_i`. - The default is :func:`~flax.linen.initializers.ones`. - kernel_linear_init: Initializer for the linear layers :math:`b_i`. - The default is :func:`~flax.linen.initializers.lecun_normal`. - bias_init: Initializer for the bias. The default is - :func:`~flax.linen.initializers.zeros`. - precision: Numerical precision of the computation. - """ # noqa: E501 - - num_potentials: int - rank: int = 0 - rectifier_fn: Callable[[Array], Array] = DEFAULT_RECTIFIER - use_linear: bool = True - use_bias: bool = True - kernel_lr_init: Callable[[PRNGKey, Shape, Dtype], Array] = DEFAULT_KERNEL_INIT - kernel_diag_init: Callable[[PRNGKey, Shape, Dtype], - Array] = nn.initializers.ones - kernel_linear_init: Callable[[PRNGKey, Shape, Dtype], - Array] = DEFAULT_KERNEL_INIT - bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = DEFAULT_BIAS_INIT - precision: Optional[jax.lax.Precision] = None - - @nn.compact - def __call__(self, x: jnp.ndarray) -> jnp.ndarray: - """Compute quadratic forms of the input. - - Args: - x: Array of shape ``[batch, ..., features]``. - - Returns: - Array of shape ``[batch, ..., num_potentials]``. - """ - # TODO(michalk8): update when refactoring neuraldual - # assert x.ndim > 1, x.ndim - - dim_data = x.shape[-1] - x = x.reshape((-1, dim_data)) - - diag_kernel = self.param( - "diag_kernel", self.kernel_diag_init, (dim_data, self.num_potentials) - ) - # ensures the diag_kernel parameter stays non negative - diag_kernel = self.rectifier_fn(diag_kernel) - - # (batch, dim_data, 1), (1, dim_data, num_potentials) - y = 0.5 * jnp.sum(((x ** 2)[..., None] * diag_kernel[None]), axis=1) - - if self.rank > 0: - quad_kernel = self.param( - "quad_kernel", self.kernel_lr_init, - (self.num_potentials, dim_data, self.rank) - ) - # (batch, num_potentials, rank) - quad = 0.5 * jnp.tensordot( - x, quad_kernel, axes=(-1, 1), precision=self.precision - ) ** 2 - y = y + jnp.sum(quad, axis=-1) - - if self.use_linear: - linear_kernel = self.param( - "lin_kernel", self.kernel_linear_init, - (dim_data, self.num_potentials) - ) - y = y + jnp.dot(x, linear_kernel, precision=self.precision) - - if self.use_bias: - y = y + self.param("bias", self.bias_init, (self.num_potentials,)) - - return y - - @classmethod - def init_from_samples( - cls, source: jnp.ndarray, target: jnp.ndarray, **kwargs: Any - ) -> "PosDefPotentials": - """Initialize the layer using Gaussian approximation :cite:`bunne:22`. - - Args: - source: Samples from the source distribution, array of shape ``[n, d]``. - target: Samples from the target distribution, array of shape ``[m, d]``. - kwargs: Keyword arguments for initialization. Note that ``use_linear`` - will be always set to :obj:`True`. - - Returns: - The layer with fixed linear and quadratic initialization. - """ - factor, mean = _compute_gaussian_map_params(source, target) - - kwargs["use_linear"] = True - return cls( - kernel_lr_init=lambda *_, **__: factor, - kernel_linear_init=lambda *_, **__: mean.T, - **kwargs, - ) - - -def _compute_gaussian_map_params( - source: jnp.ndarray, target: jnp.ndarray -) -> Tuple[jnp.ndarray, jnp.ndarray]: - from ott.math import matrix_square_root - from ott.tools.gaussian_mixture import gaussian - - g_s = gaussian.Gaussian.from_samples(source) - g_t = gaussian.Gaussian.from_samples(target) - lin_op = g_s.scale.gaussian_map(g_t.scale) - b = jnp.squeeze(g_t.loc) - lin_op @ jnp.squeeze(g_s.loc) - lin_op = matrix_square_root.sqrtm_only(lin_op) - - return jnp.expand_dims(lin_op, 0), jnp.expand_dims(b, 0) diff --git a/ott/build/lib/ott/neural/losses.py b/ott/build/lib/ott/neural/losses.py deleted file mode 100644 index ca5abc6..0000000 --- a/ott/build/lib/ott/neural/losses.py +++ /dev/null @@ -1,309 +0,0 @@ -# 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. -from typing import Any, Callable, Literal, Optional, Tuple, Union - -import jax -import jax.numpy as jnp -import numpy as np - -from ott.geometry import costs, pointcloud -from ott.solvers import linear -from ott.solvers.linear import sinkhorn -from ott.problems.quadratic import quadratic_problem -from ott.solvers.quadratic import gromov_wasserstein - -__all__ = ["monge_gap", "monge_gap_from_samples"] - - -def monge_gap( - map_fn: Callable[[jnp.ndarray], jnp.ndarray], - reference_points: jnp.ndarray, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, - relative_epsilon: Optional[bool] = None, - scale_cost: Union[ - bool, int, float, Literal["mean", "max_cost", "median"] - ] = 1.0, - return_output: bool = False, - **kwargs: Any, -) -> Union[float, Tuple[float, sinkhorn.SinkhornOutput]]: - r"""Monge gap regularizer :cite:`uscidda:23`. - - For a cost function :math:`c` and empirical reference measure - :math:`\hat{\rho}_n=\frac{1}{n}\sum_{i=1}^n \delta_{x_i}`, the - (entropic) Monge gap of a map function - :math:`T:\mathbb{R}^d\rightarrow\mathbb{R}^d` is defined as: - - .. math:: - \mathcal{M}^c_{\hat{\rho}_n, \varepsilon} (T) - = \frac{1}{n} \sum_{i=1}^n c(x_i, T(x_i)) - - W_{c, \varepsilon}(\hat{\rho}_n, T \sharp \hat{\rho}_n) - - See :cite:`uscidda:23` Eq. (8). This function is a thin wrapper that calls - :func:`~ott.neural.losses.monge_gap_from_samples`. - - Args: - map_fn: Callable corresponding to map :math:`T` in definition above. The - callable should be vectorized (e.g. using :func:`jax.vmap`), i.e, - able to process a *batch* of vectors of size `d`, namely - ``map_fn`` applied to an array returns an array of the same shape. - reference_points: Array of `[n,d]` points, :math:`\hat\rho_n` in paper - cost_fn: An object of class :class:`~ott.geometry.costs.CostFn`. - epsilon: Regularization parameter. See - :class:`~ott.geometry.pointcloud.PointCloud` - relative_epsilon: when `False`, the parameter ``epsilon`` specifies the - value of the entropic regularization parameter. When `True`, ``epsilon`` - refers to a fraction of the - :attr:`~ott.geometry.pointcloud.PointCloud.mean_cost_matrix`, which is - computed adaptively using ``source`` and ``target`` points. - scale_cost: option to rescale the cost matrix. Implemented scalings are - 'median', 'mean' and 'max_cost'. Alternatively, a float factor can be - given to rescale the cost such that ``cost_matrix /= scale_cost``. - If `True`, use 'mean'. - return_output: boolean to also return the - :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput`. - kwargs: holds the kwargs to instantiate the or - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver to - compute the regularized OT cost. - - Returns: - The Monge gap value and optionally the - :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput` - """ - target = map_fn(reference_points) - return monge_gap_from_samples( - source=reference_points, - target=target, - cost_fn=cost_fn, - epsilon=epsilon, - relative_epsilon=relative_epsilon, - scale_cost=scale_cost, - return_output=return_output, - **kwargs, - ) - - -def monge_gap_from_samples( - source: jnp.ndarray, - target: jnp.ndarray, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, - relative_epsilon: Optional[bool] = None, - scale_cost: Union[ - bool, int, float, Literal["mean", "max_cost", "median"] - ] = 1.0, - return_output: bool = False, - **kwargs: Any, -) -> Union[float, Tuple[float, sinkhorn.SinkhornOutput]]: - r"""Monge gap, instantiated in terms of samples before / after applying map. - - .. math:: - \frac{1}{n} \sum_{i=1}^n c(x_i, y_i)) - - W_{c, \varepsilon}(\frac{1}{n}\sum_i \delta_{x_i}, - \frac{1}{n}\sum_i \delta_{y_i}) - - where :math:`W_{c, \varepsilon}` is an entropy-regularized optimal transport - cost, the :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.ent_reg_cost`. - - Args: - source: samples from first measure, array of shape ``[n, d]``. - target: samples from second measure, array of shape ``[n, d]``. - cost_fn: a cost function between two points in dimension :math:`d`. - If :obj:`None`, :class:`~ott.geometry.costs.SqEuclidean` is used. - epsilon: Regularization parameter. See - :class:`~ott.geometry.pointcloud.PointCloud` - relative_epsilon: when `False`, the parameter ``epsilon`` specifies the - value of the entropic regularization parameter. When `True`, ``epsilon`` - refers to a fraction of the - :attr:`~ott.geometry.pointcloud.PointCloud.mean_cost_matrix`, which is - computed adaptively using ``source`` and ``target`` points. - scale_cost: option to rescale the cost matrix. Implemented scalings are - 'median', 'mean' and 'max_cost'. Alternatively, a float factor can be - given to rescale the cost such that ``cost_matrix /= scale_cost``. - If `True`, use 'mean'. - return_output: boolean to also return the - :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput`. - kwargs: holds the kwargs to instantiate the or - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver to - compute the regularized OT cost. - - Returns: - The Monge gap value and optionally the - :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput` - """ - cost_fn = costs.SqEuclidean() if cost_fn is None else cost_fn - geom = pointcloud.PointCloud( - x=source, - y=target, - cost_fn=cost_fn, - epsilon=epsilon, - relative_epsilon=relative_epsilon, - scale_cost=scale_cost, - ) - gt_displacement_cost = jnp.mean(jax.vmap(cost_fn)(source, target)) - out = linear.solve(geom=geom, **kwargs) - loss = gt_displacement_cost - out.ent_reg_cost - return (loss, out) if return_output else loss - - -def gw_monge_gap_from_samples( - source: jnp.ndarray, - target: jnp.ndarray, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, - scale_cost: Union[ - bool, int, float, Literal["mean", "max_cost", "median"] - ] = 1.0, - return_output: bool = False, - **kwargs: Any, -) -> Union[float, Tuple[float, sinkhorn.SinkhornOutput]]: - r"""Monge gap, instantiated in terms of samples before / after applying map. - - .. math:: - \frac{1}{n} \sum_{i=1}^n c(x_i, y_i)) - - W_{c, \varepsilon}(\frac{1}{n}\sum_i \delta_{x_i}, - \frac{1}{n}\sum_i \delta_{y_i}) - - where :math:`W_{c, \varepsilon}` is an entropy-regularized optimal transport - cost, the :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.ent_reg_cost`. - - Args: - source: samples from first measure, array of shape ``[n, d]``. - target: samples from second measure, array of shape ``[n, d]``. - cost_fn: a cost function between two points in dimension :math:`d`. - If :obj:`None`, :class:`~ott.geometry.costs.SqEuclidean` is used. - epsilon: Regularization parameter. See - :class:`~ott.geometry.pointcloud.PointCloud` - relative_epsilon: when `False`, the parameter ``epsilon`` specifies the - value of the entropic regularization parameter. When `True`, ``epsilon`` - refers to a fraction of the - :attr:`~ott.geometry.pointcloud.PointCloud.mean_cost_matrix`, which is - computed adaptively using ``source`` and ``target`` points. - scale_cost: option to rescale the cost matrix. Implemented scalings are - 'median', 'mean' and 'max_cost'. Alternatively, a float factor can be - given to rescale the cost such that ``cost_matrix /= scale_cost``. - If `True`, use 'mean'. - return_output: boolean to also return the - :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput`. - kwargs: holds the kwargs to instantiate the or - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver to - compute the regularized OT cost. - - Returns: - The Monge gap value and optionally the - :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput` - """ - cost_fn = costs.SqEuclidean() if cost_fn is None else cost_fn - geom_xx = pointcloud.PointCloud( - x=source, y=source, cost_fn=cost_fn, scale_cost=scale_cost - ) - geom_yy = pointcloud.PointCloud( - x=target, y=target, cost_fn=cost_fn, scale_cost=scale_cost - ) - prob = quadratic_problem.QuadraticProblem(geom_xx, geom_yy) - solver = gromov_wasserstein.GromovWasserstein(epsilon=epsilon) - out = solver(prob) - - gt_displacement_cost = ( - (geom_xx.cost_matrix - geom_yy.cost_matrix) ** 2 - ).mean() - - loss = gt_displacement_cost - out.reg_gw_cost - return (loss, out) if return_output else loss - - -def labeled_gw_monge_gap_from_samples( - source: jnp.ndarray, - target: jnp.ndarray, - labels: jnp.array, - n_labels: int, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, - scale_cost: Union[ - bool, int, float, Literal["mean", "max_cost", "median"] - ] = 1.0, - return_output: bool = False, - **kwargs: Any, -) -> Union[float, Tuple[float, sinkhorn.SinkhornOutput]]: - r"""Monge gap, instantiated in terms of samples before / after applying map. - - .. math:: - \frac{1}{n} \sum_{i=1}^n c(x_i, y_i)) - - W_{c, \varepsilon}(\frac{1}{n}\sum_i \delta_{x_i}, - \frac{1}{n}\sum_i \delta_{y_i}) - - where :math:`W_{c, \varepsilon}` is an entropy-regularized optimal transport - cost, the :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.ent_reg_cost`. - - Args: - source: samples from first measure, array of shape ``[n, d]``. - target: samples from second measure, array of shape ``[n, d]``. - cost_fn: a cost function between two points in dimension :math:`d`. - If :obj:`None`, :class:`~ott.geometry.costs.SqEuclidean` is used. - epsilon: Regularization parameter. See - :class:`~ott.geometry.pointcloud.PointCloud` - relative_epsilon: when `False`, the parameter ``epsilon`` specifies the - value of the entropic regularization parameter. When `True`, ``epsilon`` - refers to a fraction of the - :attr:`~ott.geometry.pointcloud.PointCloud.mean_cost_matrix`, which is - computed adaptively using ``source`` and ``target`` points. - scale_cost: option to rescale the cost matrix. Implemented scalings are - 'median', 'mean' and 'max_cost'. Alternatively, a float factor can be - given to rescale the cost such that ``cost_matrix /= scale_cost``. - If `True`, use 'mean'. - return_output: boolean to also return the - :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput`. - kwargs: holds the kwargs to instantiate the or - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver to - compute the regularized OT cost. - - Returns: - The Monge gap value and optionally the - :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput` - """ - cost_fn = costs.SqEuclidean() if cost_fn is None else cost_fn - geom_xx = pointcloud.PointCloud( - x=source, y=source, cost_fn=cost_fn, scale_cost=scale_cost - ) - geom_yy = pointcloud.PointCloud( - x=target, y=target, cost_fn=cost_fn, scale_cost=scale_cost - ) - bdm = create_block_diag_mat(labels, labels) - prob = quadratic_problem.QuadraticProblem( - geom_xx, - geom_yy, - labels_a=labels, - labels_b=labels, - n_labels=n_labels, - block_diag_mat=bdm, - ) - solver = gromov_wasserstein.GromovWasserstein(epsilon=epsilon) - out = solver(prob) - - gt_displacement_cost = ( - (geom_xx.cost_matrix - geom_yy.cost_matrix) ** 2 * bdm - ).mean() - - loss = gt_displacement_cost - out.reg_gw_cost - return (loss, out) if return_output else loss - - -def create_block_diag_mat(labels_a, labels_b): - """Creates block diagonal matrix that has 1 entry when label_a == label_b otherwise 0.""" - block_diag_mat = np.zeros((len(labels_a), len(labels_b))) - for l in np.unique(labels_a): - block_diag_mat[ - np.ix_(np.where(labels_a == l)[0], np.where(labels_b == l)[0]) - ] = 1.0 - return jnp.array(block_diag_mat) diff --git a/ott/build/lib/ott/neural/models.py b/ott/build/lib/ott/neural/models.py deleted file mode 100644 index c8fc502..0000000 --- a/ott/build/lib/ott/neural/models.py +++ /dev/null @@ -1,408 +0,0 @@ -# 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 functools -from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union - -import jax -import jax.numpy as jnp -import optax -from flax import linen as nn -from flax.core import frozen_dict -from flax.training import train_state - -from ott import utils -from ott.geometry import geometry -from ott.initializers.linear import initializers as lin_init -from ott.neural import layers -from ott.neural.solvers import neuraldual -from ott.problems.linear import linear_problem - -__all__ = ["ICNN", "MLP", "MetaInitializer"] - -# wrap to silence docs linter -DEFAULT_KERNEL_INIT = lambda *a, **k: nn.initializers.normal()(*a, **k) -DEFAULT_RECTIFIER = nn.activation.relu -DEFAULT_ACTIVATION = nn.activation.relu - - -class ICNN(neuraldual.BaseW2NeuralDual): - """Input convex neural network (ICNN). - - Implementation of input convex neural networks as introduced in - :cite:`amos:17` with initialization schemes proposed by :cite:`bunne:22`. - - Args: - dim_data: data dimensionality. - dim_hidden: sequence specifying size of hidden dimensions. The - output dimension of the last layer is 1 by default. - ranks: ranks of the matrices :math:`A_i` used as low-rank factors - for the quadratic potentials. If a sequence is passed, it must contain - ``len(dim_hidden) + 2`` elements, where the last 2 elements correspond - to the ranks of the final layer with dimension 1 and the potentials, - respectively. - init_fn: Initializer for the kernel weight matrices. - The default is :func:`~flax.linen.initializers.normal`. - act_fn: choice of activation function used in network architecture, - needs to be convex. The default is :func:`~flax.linen.activation.relu`. - pos_weights: Enforce positive weights with a projection. - If :obj:`False`, the positive weights should be enforced with clipping - or regularization in the loss. - rectifier_fn: function to ensure the non negativity of the weights. - The default is :func:`~flax.linen.activation.relu`. - gaussian_map_samples: Tuple of source and target points, used to initialize - the ICNN to mimic the linear Bures map that morphs the (Gaussian - approximation) of the input measure to that of the target measure. If - :obj:`None`, the identity initialization is used, and ICNN mimics half the - squared Euclidean norm. - """ - - dim_data: int - dim_hidden: Sequence[int] - ranks: Union[int, Tuple[int, ...]] = 1 - init_fn: Callable[[jax.Array, Tuple[int, ...], Any], jnp.ndarray] = ( - DEFAULT_KERNEL_INIT - ) - act_fn: Callable[[jnp.ndarray], jnp.ndarray] = DEFAULT_ACTIVATION - pos_weights: bool = False - rectifier_fn: Callable[[jnp.ndarray], jnp.ndarray] = DEFAULT_RECTIFIER - gaussian_map_samples: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None - - def setup(self) -> None: # noqa: D102 - dim_hidden = list(self.dim_hidden) + [1] - *ranks, pos_def_rank = self._normalize_ranks() - - # final layer computes average, still with normalized rescaling - self.w_zs = [self._get_wz(dim) for dim in dim_hidden[1:]] - # subsequent layers re-injected into convex functions - self.w_xs = [ - self._get_wx(dim, rank) for dim, rank in zip(dim_hidden, ranks) - ] - self.pos_def_potentials = self._get_pos_def_potentials(pos_def_rank) - - @nn.compact - def __call__(self, x: jnp.ndarray) -> float: # noqa: D102 - w_x, *w_xs = self.w_xs - assert len(self.w_zs) == len(w_xs), (len(self.w_zs), len(w_xs)) - - z = self.act_fn(w_x(x)) - for w_z, w_x in zip(self.w_zs, w_xs): - z = self.act_fn(w_z(z) + w_x(x)) - z = z + self.pos_def_potentials(x) - - return z.squeeze() - - def _get_wz(self, dim: int) -> nn.Module: - if self.pos_weights: - return layers.PositiveDense( - dim, - kernel_init=self.init_fn, - use_bias=False, - rectifier_fn=self.rectifier_fn, - ) - - return nn.Dense( - dim, - kernel_init=self.init_fn, - use_bias=False, - ) - - def _get_wx(self, dim: int, rank: int) -> nn.Module: - return layers.PosDefPotentials( - rank=rank, - num_potentials=dim, - use_linear=True, - use_bias=True, - kernel_diag_init=nn.initializers.zeros, - kernel_lr_init=self.init_fn, - kernel_linear_init=self.init_fn, - bias_init=nn.initializers.zeros, - ) - - def _get_pos_def_potentials(self, rank: int) -> layers.PosDefPotentials: - kwargs = { - "num_potentials": 1, - "use_linear": True, - "use_bias": True, - "bias_init": nn.initializers.zeros, - } - - if self.gaussian_map_samples is None: - return layers.PosDefPotentials( - rank=rank, - kernel_diag_init=nn.initializers.ones, - kernel_lr_init=nn.initializers.zeros, - kernel_linear_init=nn.initializers.zeros, - **kwargs, - ) - - source, target = self.gaussian_map_samples - return layers.PosDefPotentials.init_from_samples( - source, - target, - rank=self.dim_data, - kernel_diag_init=nn.initializers.zeros, - **kwargs, - ) - - def _normalize_ranks(self) -> Tuple[int, ...]: - # +2 for the newly added layer with 1 + the final potentials - n_ranks = len(self.dim_hidden) + 2 - if isinstance(self.ranks, int): - return (self.ranks,) * n_ranks - - assert len(self.ranks) == n_ranks, (len(self.ranks), n_ranks) - return tuple(self.ranks) - - @property - def is_potential(self) -> bool: # noqa: D102 - return True - - -class MLP(neuraldual.BaseW2NeuralDual): - """A generic, not-convex MLP. - - Args: - dim_hidden: sequence specifying size of hidden dimensions. The output - dimension of the last layer is automatically set to 1 if - :attr:`is_potential` is ``True``, or the dimension of the input otherwise - is_potential: Model the potential if ``True``, otherwise - model the gradient of the potential - act_fn: Activation function - """ - - dim_hidden: Sequence[int] - dim_out: int - is_potential: bool = True - act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.leaky_relu - - @nn.compact - def __call__(self, x: jnp.ndarray) -> jnp.ndarray: # noqa: D102 - squeeze = x.ndim == 1 - if squeeze: - x = jnp.expand_dims(x, 0) - assert x.ndim == 2, x.ndim - n_input = x.shape[-1] - - z = x - for n_hidden in self.dim_hidden: - Wx = nn.Dense(n_hidden, use_bias=True) - z = self.act_fn(Wx(z)) - - if self.is_potential: - Wx = nn.Dense(1, use_bias=True) - z = Wx(z).squeeze(-1) - - quad_term = 0.5 * jax.vmap(jnp.dot)(x, x) - z += quad_term - elif self.dim_out is None: - Wx = nn.Dense(n_input, use_bias=True) - z = x + Wx(z) - else: - Wx = nn.Dense(self.dim_out, use_bias=True) - z = Wx(z) - - return z.squeeze(0) if squeeze else z - - -@jax.tree_util.register_pytree_node_class -class MetaInitializer(lin_init.DefaultInitializer): - """Meta OT Initializer with a fixed geometry :cite:`amos:22`. - - This initializer consists of a predictive model that outputs the - :math:`f` duals to solve the entropy-regularized OT problem given - input probability weights ``a`` and ``b``, and a given (assumed to be - fixed) geometry ``geom``. - - The model's parameters are learned using a training set of OT - instances (multiple pairs of probability weights), that assume the - **same** geometry ``geom`` is used throughout, both for training and - evaluation. - - Args: - geom: The fixed geometry of the problem instances. - meta_model: The model to predict the potential :math:`f` from the measures. - opt: The optimizer to update the parameters. If :obj:`None`, use - :func:`optax.adam` with :math:`0.001` learning rate. - rng: The PRNG key to use for initializing the model. - state: The training state of the model to start from. - - Examples: - The following code shows a simple - example of using ``update`` to train the model, where - ``a`` and ``b`` are the weights of the measures and - ``geom`` is the fixed geometry. - - .. code-block:: python - - meta_initializer = init_lib.MetaInitializer(geom) - while training(): - a, b = sample_batch() - loss, init_f, meta_initializer.state = meta_initializer.update( - meta_initializer.state, a=a, b=b - ) - """ - - def __init__( - self, - geom: geometry.Geometry, - meta_model: nn.Module, - opt: Optional[optax.GradientTransformation] = optax.adam( - learning_rate=1e-3 - ), # noqa: B008 - rng: Optional[jax.Array] = None, - state: Optional[train_state.TrainState] = None, - ): - self.geom = geom - self.opt = opt - self.rng = utils.default_prng_key(rng) - - na, nb = geom.shape - self.meta_model = meta_model - - if state is None: - # Initialize the model's training state. - placeholder = jnp.concatenate( - [jnp.zeros(na), jnp.zeros(nb)], axis=-1 - ) - params = self.meta_model.init(self.rng, placeholder)["params"] - self.state = train_state.TrainState.create( - apply_fn=self.meta_model.apply, params=params, tx=opt - ) - else: - self.state = state - - self.update_impl = self._get_update_fn() - - def update( - self, state: train_state.TrainState, a: jnp.ndarray, b: jnp.ndarray - ) -> Tuple[jnp.ndarray, jnp.ndarray, train_state.TrainState]: - r"""Update the meta model with the dual objective. - - The goal is for the model to match the optimal duals, i.e., - :math:`\hat f_\theta \approx f^\star`. - This can be done by training the predictions of :math:`\hat f_\theta` - to optimize the dual objective, which :math:`f^\star` also optimizes for. - The overall learning setup can thus be written as: - - .. math:: - \min_\theta\; {\mathbb E}_{(\alpha,\beta)\sim{\mathcal{D}}}\; - J(\hat f_\theta(a, b); \alpha, \beta), - - where :math:`a,b` are the probabilities of the measures :math:`\alpha,\beta` - ,:math:`\mathcal{D}` is a meta distribution of optimal transport problems, - - .. math:: - -J(f; \alpha, \beta, c) := \langle f, a\rangle + \langle g, b \rangle - - \varepsilon\left\langle \exp\{f/\varepsilon\}, K\exp\{g/\varepsilon\} - \right\rangle - - is the entropic dual objective, - and :math:`K_{i,j} := -C_{i,j}/\varepsilon` is the *Gibbs kernel*. - - Args: - state: Optimizer state of the meta model. - a: Probabilities of the :math:`\alpha` measure's atoms. - b: Probabilities of the :math:`\beta` measure's atoms. - - Returns: - The training loss, :math:`f`, and updated state. - """ - return self.update_impl(state, a, b) - - def init_dual_a( # noqa: D102 - self, - ot_prob: "linear_problem.LinearProblem", - lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: - del rng - # Detect if the problem is batched. - assert ot_prob.a.ndim in (1, 2) - assert ot_prob.b.ndim in (1, 2) - vmap_a_val = 0 if ot_prob.a.ndim == 2 else None - vmap_b_val = 0 if ot_prob.b.ndim == 2 else None - - if vmap_a_val is not None or vmap_b_val is not None: - compute_f_maybe_batch = jax.vmap( - self._compute_f, in_axes=(vmap_a_val, vmap_b_val, None) - ) - else: - compute_f_maybe_batch = self._compute_f - - init_f = compute_f_maybe_batch(ot_prob.a, ot_prob.b, self.state.params) - return ( - init_f if lse_mode else ot_prob.geom.scaling_from_potential(init_f) - ) - - def _get_update_fn(self): - """Return the implementation (and jitted) update function.""" - from ott.problems.linear import linear_problem - from ott.solvers.linear import sinkhorn - - def dual_obj_loss_single(params, a, b): - f_pred = self._compute_f(a, b, params) - g_pred = self.geom.update_potential( - f_pred, jnp.zeros_like(b), jnp.log(b), 0, axis=0 - ) - g_pred = jnp.where(jnp.isfinite(g_pred), g_pred, 0.0) - - ot_prob = linear_problem.LinearProblem(geom=self.geom, a=a, b=b) - dual_obj = sinkhorn.compute_kl_reg_cost( - f_pred, g_pred, ot_prob, lse_mode=True - ) - loss = -dual_obj - return loss, f_pred - - def loss_batch(params, a, b): - loss_fn = functools.partial(dual_obj_loss_single, params=params) - loss, f_pred = jax.vmap(loss_fn)(a=a, b=b) - return jnp.mean(loss), f_pred - - @jax.jit - def update(state, a, b): - a = jnp.atleast_2d(a) - b = jnp.atleast_2d(b) - grad_fn = jax.value_and_grad(loss_batch, has_aux=True) - (loss, init_f), grads = grad_fn(state.params, a, b) - return loss, init_f, state.apply_gradients(grads=grads) - - return update - - def _compute_f( - self, - a: jnp.ndarray, - b: jnp.ndarray, - params: frozen_dict.FrozenDict[str, jnp.ndarray], - ) -> jnp.ndarray: - r"""Predict the optimal :math:`f` potential. - - Args: - a: Probabilities of the :math:`\alpha` measure's atoms. - b: Probabilities of the :math:`\beta` measure's atoms. - params: The parameters of the Meta model. - - Returns: - The :math:`f` potential. - """ - return self.meta_model.apply( - {"params": params}, jnp.concatenate([a, b], axis=-1) - ) - - def tree_flatten( - self, - ) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return [self.geom, self.meta_model, self.opt], { - "rng": self.rng, - "state": self.state, - } diff --git a/ott/build/lib/ott/neural/solvers/__init__.py b/ott/build/lib/ott/neural/solvers/__init__.py deleted file mode 100644 index b09d8c6..0000000 --- a/ott/build/lib/ott/neural/solvers/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. -from . import conjugate, map_estimator, neuraldual diff --git a/ott/build/lib/ott/neural/solvers/conjugate.py b/ott/build/lib/ott/neural/solvers/conjugate.py deleted file mode 100644 index 7d0098f..0000000 --- a/ott/build/lib/ott/neural/solvers/conjugate.py +++ /dev/null @@ -1,121 +0,0 @@ -# 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 abc -from typing import Callable, Literal, NamedTuple, Optional - -import jax.numpy as jnp -from jaxopt import LBFGS - -from ott import utils - -__all__ = [ - "ConjugateResults", - "FenchelConjugateSolver", - "FenchelConjugateLBFGS", - "DEFAULT_CONJUGATE_SOLVER", -] - - -class ConjugateResults(NamedTuple): - r"""Holds the results of numerically conjugating a function. - - Args: - val: the conjugate value, i.e., :math:`f^\star(y)` - grad: the gradient, i.e., :math:`\nabla f^\star(y)` - num_iter: the number of iterations taken by the solver - """ - val: float - grad: jnp.ndarray - num_iter: int - - -class FenchelConjugateSolver(abc.ABC): - r"""Abstract conjugate solver class. - - Given a function :math:`f`, numerically estimate the Fenchel conjugate - :math:`f^\star(y) := -\inf_{x\in\mathbb{R}^n} f(x)-\langle x, y\rangle`. - """ - - @abc.abstractmethod - def solve( - self, - f: Callable[[jnp.ndarray], jnp.ndarray], - y: jnp.ndarray, - x_init: Optional[jnp.ndarray] = None - ) -> ConjugateResults: - """Solve for the conjugate. - - Args: - f: function to conjugate - y: point to conjugate - x_init: initial point to search over - - Returns: - The solution to the conjugation. - """ - - -@utils.register_pytree_node -class FenchelConjugateLBFGS(FenchelConjugateSolver): - """Solve for the conjugate using :class:`~jaxopt.LBFGS`. - - Args: - gtol: gradient tolerance - max_iter: maximum number of iterations - max_linesearch_iter: maximum number of line search iterations - linesearch_type: type of line search - linesearch_init: strategy for line search initialization - increase_factor: factor by which to increase the step size during - the line search - """ - - gtol: float = 1e-3 - max_iter: int = 10 - max_linesearch_iter: int = 10 - linesearch_type: Literal["zoom", "backtracking", - "hager-zhang"] = "backtracking" - linesearch_init: Literal["increase", "max", "current"] = "increase" - increase_factor: float = 1.5 - - def solve( # noqa: D102 - self, - f: Callable[[jnp.ndarray], jnp.ndarray], - y: jnp.ndarray, - x_init: Optional[jnp.array] = None - ) -> ConjugateResults: - assert y.ndim == 1, y.ndim - - solver = LBFGS( - fun=lambda x: f(x) - x.dot(y), - tol=self.gtol, - maxiter=self.max_iter, - linesearch=self.linesearch_type, - linesearch_init=self.linesearch_init, - increase_factor=self.increase_factor, - implicit_diff=False, - unroll=False - ) - - out = solver.run(y if x_init is None else x_init) - return ConjugateResults( - val=-out.state.value, grad=out.params, num_iter=out.state.iter_num - ) - - -DEFAULT_CONJUGATE_SOLVER = FenchelConjugateLBFGS( - gtol=1e-5, - max_iter=20, - max_linesearch_iter=20, - linesearch_type="backtracking", -) diff --git a/ott/build/lib/ott/neural/solvers/map_estimator.py b/ott/build/lib/ott/neural/solvers/map_estimator.py deleted file mode 100644 index f5389f5..0000000 --- a/ott/build/lib/ott/neural/solvers/map_estimator.py +++ /dev/null @@ -1,289 +0,0 @@ -# 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 collections -import functools -from typing import ( - Any, - Callable, - Dict, - Iterator, - Optional, - Sequence, - Tuple, - Union, -) - -import jax -import jax.numpy as jnp -import optax -from flax.core import frozen_dict -from flax.training import train_state - -from ott import utils -from ott.neural.solvers import neuraldual - -__all__ = ["MapEstimator"] - - -class MapEstimator: - r"""Mapping estimator between probability measures. - - It estimates a map :math:`T` by minimizing the loss: - - .. math:: - \text{min}_{\theta}\; \Delta(T_\theta \sharp \mu, \theta) - + \lambda R(T_\theta \sharp \rho, \rho) - - where :math:`\Delta` is a fitting loss and :math:`R` is a regularizer. - :math:`\Delta` allows to fit the marginal constraint, i.e. transport - :math:`\mu` to :math:`\nu` via :math:`T`, while :math:`R` - is a regularizer imposing an inductive bias on the learned map. The - regularizer in this case is a function used to compute a metric between two - sets of points. - - For instance, :math:`\Delta` can be the - :func:`~ott.tools.sinkhorn_divergence.sinkhorn_divergence` - and :math:`R` the :func:`~ott.neural.losses.monge_gap_from_samples` - :cite:`uscidda:23` for a given cost function :math:`c`. - In that case, it estimates a :math:`c`-OT map, i.e. a map :math:`T` - optimal for the Monge problem induced by :math:`c`. - - Args: - dim_data: input dimensionality of data required for network init. - model: network architecture for map :math:`T`. - optimizer: optimizer function for map :math:`T`. - fitting_loss: function that outputs a fitting loss :math:`\Delta` between - two families of points, as well as any log object. - regularizer: function that outputs a score from two families of points, - here assumed to be of the same size, as well as any log object. - regularizer_strength: strength of the :attr:`regularizer`. - num_train_iters: number of total training iterations. - logging: option to return logs. - valid_freq: frequency with training and validation are logged. - rng: random key used for seeding for network initializations. - """ - - def __init__( - self, - dim_data: int, - model: neuraldual.BaseW2NeuralDual, - optimizer: Optional[optax.OptState] = None, - fitting_loss: Optional[Callable[[jnp.ndarray, jnp.ndarray], - Tuple[float, Optional[Any]]]] = None, - regularizer: Optional[Callable[[jnp.ndarray, jnp.ndarray], - Tuple[float, Optional[Any]]]] = None, - regularizer_strength: Union[float, Sequence[float]] = 1.0, - num_train_iters: int = 10_000, - logging: bool = False, - valid_freq: int = 500, - rng: Optional[jax.Array] = None, - ): - self._fitting_loss = fitting_loss - self._regularizer = regularizer - # Can use either a fixed strength, or generalize to a schedule. - self.regularizer_strength = jnp.repeat( - jnp.atleast_2d(regularizer_strength), - num_train_iters, - total_repeat_length=num_train_iters, - axis=0 - ).ravel() - self.num_train_iters = num_train_iters - self.logging = logging - self.valid_freq = valid_freq - self.rng = utils.default_prng_key(rng) - - # set default optimizer - if optimizer is None: - optimizer = optax.adam(learning_rate=0.001) - - # setup training - self.setup(dim_data, model, optimizer) - - def setup( - self, - dim_data: int, - neural_net: neuraldual.BaseW2NeuralDual, - optimizer: optax.OptState, - ): - """Setup all components required to train the network.""" - # neural network - self.state_neural_net = neural_net.create_train_state( - self.rng, optimizer, dim_data - ) - - # step function - self.step_fn = self._get_step_fn() - - @property - def regularizer(self) -> Callable[[jnp.ndarray, jnp.ndarray], float]: - """Regularizer added to the fitting loss. - - Can be, e.g. the :func:`~ott.neural.losses.monge_gap_from_samples`. - If no regularizer is passed for solver instantiation, - or regularization weight :attr:`regularizer_strength` is 0, - return 0 by default along with an empty set of log values. - """ - if self._regularizer is not None: - return self._regularizer - return lambda *_, **__: (0.0, None) - - @property - def fitting_loss(self) -> Callable[[jnp.ndarray, jnp.ndarray], float]: - """Fitting loss to fit the marginal constraint. - - Can be, e.g. :func:`~ott.tools.sinkhorn_divergence.sinkhorn_divergence`. - If no fitting_loss is passed for solver instantiation, return 0 by default, - and no log values. - """ - if self._fitting_loss is not None: - return self._fitting_loss - return lambda *_, **__: (0.0, None) - - @staticmethod - def _generate_batch( - loader_source: Iterator[jnp.ndarray], - loader_target: Iterator[jnp.ndarray], - ) -> Dict[str, jnp.ndarray]: - """Generate batches a batch of samples. - - ``loader_source`` and ``loader_target`` can be training or - validation dataloaders. - """ - return { - "source": next(loader_source), - "target": next(loader_target), - } - - def train_map_estimator( - self, - trainloader_source: Iterator[jnp.ndarray], - trainloader_target: Iterator[jnp.ndarray], - validloader_source: Iterator[jnp.ndarray], - validloader_target: Iterator[jnp.ndarray], - ) -> Tuple[train_state.TrainState, Dict[str, Any]]: - """Training loop.""" - # define logs - logs = collections.defaultdict(lambda: collections.defaultdict(list)) - - # try to display training progress with tqdm - try: - from tqdm import trange - tbar = trange(self.num_train_iters, leave=True) - except ImportError: - tbar = range(self.num_train_iters) - - for step in tbar: - # update step - is_logging_step = ( - self.logging and ((step % self.valid_freq == 0) or - (step == self.num_train_iters - 1)) - ) - train_batch = self._generate_batch( - loader_source=trainloader_source, - loader_target=trainloader_target, - ) - valid_batch = ( - None if not is_logging_step else self._generate_batch( - loader_source=validloader_source, - loader_target=validloader_target, - ) - ) - self.state_neural_net, current_logs = self.step_fn( - self.state_neural_net, train_batch, valid_batch, is_logging_step, step - ) - - # store and print metrics if logging step - if is_logging_step: - for log_key in current_logs: - for metric_key in current_logs[log_key]: - logs[log_key][metric_key].append(current_logs[log_key][metric_key]) - - # update the tqdm bar if tqdm is available - if not isinstance(tbar, range): - reg_msg = ( - "NA" if current_logs["eval"]["regularizer"] == 0.0 else - f"{current_logs['eval']['regularizer']:.4f}" - ) - postfix_str = ( - f"fitting_loss: {current_logs['eval']['fitting_loss']:.4f}, " - f"regularizer: {reg_msg} ," - f"total: {current_logs['eval']['total_loss']:.4f}" - ) - tbar.set_postfix_str(postfix_str) - - return self.state_neural_net, logs - - def _get_step_fn(self) -> Callable: - """Create a one step training and evaluation function.""" - - def loss_fn( - params: frozen_dict.FrozenDict, apply_fn: Callable, - batch: Dict[str, jnp.ndarray], step: int - ) -> Tuple[float, Dict[str, float]]: - """Loss function.""" - # map samples with the fitted map - mapped_samples = apply_fn({"params": params}, batch["source"]) - - # compute the loss - val_fitting_loss, log_fitting_loss = self.fitting_loss( - mapped_samples, batch["target"] - ) - val_regularizer, log_regularizer = self.regularizer( - batch["source"], mapped_samples - ) - val_tot_loss = ( - val_fitting_loss + self.regularizer_strength[step] * val_regularizer - ) - - # store training logs - loss_logs = { - "total_loss": val_tot_loss, - "fitting_loss": val_fitting_loss, - "regularizer": val_regularizer, - "log_regularizer": log_regularizer, - "log_fitting": log_fitting_loss, - } - - return val_tot_loss, loss_logs - - @functools.partial(jax.jit, static_argnums=3) - def step_fn( - state_neural_net: train_state.TrainState, - train_batch: Dict[str, jnp.ndarray], - valid_batch: Optional[Dict[str, jnp.ndarray]] = None, - is_logging_step: bool = False, - step: int = 0 - ) -> Tuple[train_state.TrainState, Dict[str, float]]: - """One step function.""" - # compute loss and gradients - grad_fn = jax.value_and_grad(loss_fn, argnums=0, has_aux=True) - (_, current_train_logs), grads = grad_fn( - state_neural_net.params, state_neural_net.apply_fn, train_batch, step - ) - - # logging step - current_logs = {"train": current_train_logs, "eval": {}} - if is_logging_step: - _, current_eval_logs = loss_fn( - params=state_neural_net.params, - apply_fn=state_neural_net.apply_fn, - batch=valid_batch, - step=step - ) - current_logs["eval"] = current_eval_logs - - # update state - return state_neural_net.apply_gradients(grads=grads), current_logs - - return step_fn diff --git a/ott/build/lib/ott/neural/solvers/neuraldual.py b/ott/build/lib/ott/neural/solvers/neuraldual.py deleted file mode 100644 index fffa927..0000000 --- a/ott/build/lib/ott/neural/solvers/neuraldual.py +++ /dev/null @@ -1,700 +0,0 @@ -# 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 abc -import warnings -from typing import ( - Any, - Callable, - Dict, - Iterator, - List, - Literal, - Optional, - Tuple, - Union, -) - -import jax -import jax.numpy as jnp -import optax -from flax import linen as nn -from flax import struct -from flax.core import frozen_dict -from flax.training import train_state - -from ott import utils -from ott.geometry import costs -from ott.neural import models -from ott.neural.solvers import conjugate -from ott.problems.linear import potentials - -__all__ = ["W2NeuralTrainState", "BaseW2NeuralDual", "W2NeuralDual"] - -Train_t = Dict[Literal["train_logs", "valid_logs"], Dict[str, List[float]]] -Callback_t = Callable[[int, potentials.DualPotentials], None] - -PotentialValueFn_t = Callable[[jnp.ndarray], jnp.ndarray] -PotentialGradientFn_t = Callable[[jnp.ndarray], jnp.ndarray] - - -class W2NeuralTrainState(train_state.TrainState): - """Adds information about the model's value and gradient to the state. - - This extends :class:`~flax.training.train_state.TrainState` to include - the potential methods from the - :class:`~ott.neural.solvers.neuraldual.BaseW2NeuralDual` used during training. - - Args: - potential_value_fn: the potential's value function - potential_gradient_fn: the potential's gradient function - """ - potential_value_fn: Callable[ - [frozen_dict.FrozenDict[str, jnp.ndarray], Optional[PotentialValueFn_t]], - PotentialValueFn_t] = struct.field(pytree_node=False) - potential_gradient_fn: Callable[[frozen_dict.FrozenDict[str, jnp.ndarray]], - PotentialGradientFn_t] = struct.field( - pytree_node=False - ) - - -class BaseW2NeuralDual(abc.ABC, nn.Module): - """Base class for the neural solver models.""" - - @property - @abc.abstractmethod - def is_potential(self) -> bool: - """Indicates if the module implements a potential value or a vector field. - - Returns: - ``True`` if the module defines a potential, ``False`` if it defines a - vector field. - """ - - def potential_value_fn( - self, - params: frozen_dict.FrozenDict[str, jnp.ndarray], - other_potential_value_fn: Optional[PotentialValueFn_t] = None, - ) -> PotentialValueFn_t: - r"""Return a function giving the value of the potential. - - Applies the module if :attr:`is_potential` is ``True``, otherwise - constructs the value of the potential from the gradient with - - .. math:: - - g(y) = -f(\nabla_y g(y)) + y^T \nabla_y g(y) - - where :math:`\nabla_y g(y)` is detached for the envelope theorem - :cite:`danskin:67,bertsekas:71` - to give the appropriate first derivatives of this construction. - - Args: - params: parameters of the module - other_potential_value_fn: function giving the value of the other - potential. Only needed when :attr:`is_potential` is ``False``. - - Returns: - A function that can be evaluated to obtain a potential value, or a linear - interpolation of a potential. - """ - if self.is_potential: - return lambda x: self.apply({"params": params}, x) - - assert other_potential_value_fn is not None, \ - "The value of the gradient-based potential depends " \ - "on the value of the other potential." - - def value_fn(x: jnp.ndarray) -> jnp.ndarray: - squeeze = x.ndim == 1 - if squeeze: - x = jnp.expand_dims(x, 0) - grad_g_x = jax.lax.stop_gradient(self.apply({"params": params}, x)) - value = -other_potential_value_fn(grad_g_x) + \ - jax.vmap(jnp.dot)(grad_g_x, x) - return value.squeeze(0) if squeeze else value - - return value_fn - - def potential_gradient_fn( - self, - params: frozen_dict.FrozenDict[str, jnp.ndarray], - ) -> PotentialGradientFn_t: - """Return a function returning a vector or the gradient of the potential. - - Args: - params: parameters of the module - - Returns: - A function that can be evaluated to obtain the potential's gradient - """ - if self.is_potential: - return jax.vmap(jax.grad(self.potential_value_fn(params))) - return lambda x: self.apply({"params": params}, x) - - def create_train_state( - self, - rng: jax.Array, - optimizer: optax.OptState, - input: Union[int, Tuple[int, ...]], - **kwargs: Any, - ) -> W2NeuralTrainState: - """Create initial training state.""" - params = self.init(rng, jnp.ones(input))["params"] - - return W2NeuralTrainState.create( - apply_fn=self.apply, - params=params, - tx=optimizer, - potential_value_fn=self.potential_value_fn, - potential_gradient_fn=self.potential_gradient_fn, - **kwargs, - ) - - -class W2NeuralDual: - r"""Solver for the Wasserstein-2 Kantorovich dual between Euclidean spaces. - - Learn the Wasserstein-2 optimal transport between two measures - :math:`\alpha` and :math:`\beta` in :math:`n`-dimensional Euclidean space, - denoted source and target, respectively. This is achieved by parameterizing - a Kantorovich potential :math:`f_\theta: \mathbb{R}^n\rightarrow\mathbb{R}` - associated with the :math:`\alpha` measure with an - :class:`~ott.neural.models.ICNN` or :class:`~ott.neural.models.MLP`, where - :math:`\nabla f` transports source to target cells. This potential is learned - by optimizing the dual form associated with the negative inner product cost - - .. math:: - - \text{argsup}_{\theta}\; -\mathbb{E}_{x\sim\alpha}[f_\theta(x)] - - \mathbb{E}_{y\sim\beta}[f^\star_\theta(y)], - - where :math:`f^\star(y) := -\inf_{x\in\mathbb{R}^n} f(x)-\langle x, y\rangle` - is the convex conjugate. - :math:`\nabla f^\star` transports from the target - to source cells and provides the inverse optimal - transport map from :math:`\beta` to :math:`\alpha`. - This solver estimates the conjugate :math:`f^\star` - with a neural approximation :math:`g` that is fine-tuned - with :class:`~ott.neural.solvers.conjugate.FenchelConjugateSolver`, - which is a combination further described in :cite:`amos:23`. - - The :class:`~ott.neural.solvers.neuraldual.BaseW2NeuralDual` potentials for - ``neural_f`` and ``neural_g`` can - - 1. both provide the values of the potentials :math:`f` and :math:`g`, or - 2. one of them can provide the gradient mapping e.g., :math:`\nabla f` - or :math:`\nabla g` where the potential's value can be obtained - via the Fenchel conjugate as discussed in :cite:`amos:23`. - - The potential's value or gradient mapping is specified via - :attr:`~ott.neural.solvers.neuraldual.BaseW2NeuralDual.is_potential`. - - Args: - dim_data: input dimensionality of data required for network init - neural_f: network architecture for potential :math:`f`. - neural_g: network architecture for the conjugate potential - :math:`g\approx f^\star` - optimizer_f: optimizer function for potential :math:`f` - optimizer_g: optimizer function for the conjugate potential :math:`g` - num_train_iters: number of total training iterations - num_inner_iters: number of training iterations of :math:`g` per iteration - of :math:`f` - back_and_forth: alternate between updating the forward and backward - directions. Inspired from :cite:`jacobs:20` - valid_freq: frequency with which model is validated - log_freq: frequency with training and validation are logged - logging: option to return logs - rng: random key used for seeding for network initializations - pos_weights: option to train networks with positive weights or regularizer - beta: regularization parameter when not training with positive weights - conjugate_solver: numerical solver for the Fenchel conjugate. - amortization_loss: amortization loss for the conjugate - :math:`g\approx f^\star`. Options are `'objective'` :cite:`makkuva:20` or - `'regression'` :cite:`amos:23`. - parallel_updates: Update :math:`f` and :math:`g` at the same time - """ - - def __init__( - self, - dim_data: int, - neural_f: Optional[BaseW2NeuralDual] = None, - neural_g: Optional[BaseW2NeuralDual] = None, - optimizer_f: Optional[optax.OptState] = None, - optimizer_g: Optional[optax.OptState] = None, - num_train_iters: int = 20000, - num_inner_iters: int = 1, - back_and_forth: Optional[bool] = None, - valid_freq: int = 1000, - log_freq: int = 1000, - logging: bool = False, - rng: Optional[jax.Array] = None, - pos_weights: bool = True, - beta: float = 1.0, - conjugate_solver: Optional[conjugate.FenchelConjugateSolver - ] = conjugate.DEFAULT_CONJUGATE_SOLVER, - amortization_loss: Literal["objective", "regression"] = "regression", - parallel_updates: bool = True, - ): - self.num_train_iters = num_train_iters - self.num_inner_iters = num_inner_iters - self.back_and_forth = back_and_forth - self.valid_freq = valid_freq - self.log_freq = log_freq - self.logging = logging - self.pos_weights = pos_weights - self.beta = beta - self.parallel_updates = parallel_updates - self.conjugate_solver = conjugate_solver - self.amortization_loss = amortization_loss - - # set default optimizers - if optimizer_f is None: - optimizer_f = optax.adam(learning_rate=0.0001, b1=0.5, b2=0.9, eps=1e-8) - if optimizer_g is None: - optimizer_g = optax.adam(learning_rate=0.0001, b1=0.5, b2=0.9, eps=1e-8) - - # set default neural architectures - if neural_f is None: - neural_f = models.ICNN(dim_data=dim_data, dim_hidden=[64, 64, 64, 64]) - if neural_g is None: - neural_g = models.ICNN(dim_data=dim_data, dim_hidden=[64, 64, 64, 64]) - self.neural_f = neural_f - self.neural_g = neural_g - - # set optimizer and networks - self.setup( - utils.default_prng_key(rng), - neural_f, - neural_g, - dim_data, - optimizer_f, - optimizer_g, - ) - - def setup( - self, - rng: jax.Array, - neural_f: BaseW2NeuralDual, - neural_g: BaseW2NeuralDual, - dim_data: int, - optimizer_f: optax.OptState, - optimizer_g: optax.OptState, - ) -> None: - """Setup all components required to train the network.""" - # split random number generator - rng, rng_f, rng_g = jax.random.split(rng, 3) - - # check setting of network architectures - warn_str = f"Setting of ICNN and the positive weights setting of the " \ - f"`W2NeuralDual` are not consistent. Proceeding with " \ - f"the `W2NeuralDual` setting, with positive weights " \ - f"being {self.pos_weights}." - if isinstance( - neural_f, models.ICNN - ) and neural_f.pos_weights is not self.pos_weights: - warnings.warn(warn_str, stacklevel=2) - neural_f.pos_weights = self.pos_weights - - if isinstance( - neural_g, models.ICNN - ) and neural_g.pos_weights is not self.pos_weights: - warnings.warn(warn_str, stacklevel=2) - neural_g.pos_weights = self.pos_weights - - self.state_f = neural_f.create_train_state( - rng_f, - optimizer_f, - (1, dim_data), # also include the batch dimension - ) - self.state_g = neural_g.create_train_state( - rng_g, - optimizer_g, - (1, dim_data), - ) - - # default to using back_and_forth with the non-convex models - if self.back_and_forth is None: - self.back_and_forth = isinstance(neural_f, models.MLP) - - if self.num_inner_iters == 1 and self.parallel_updates: - self.train_step_parallel = self.get_step_fn( - train=True, to_optimize="both" - ) - self.valid_step_parallel = self.get_step_fn( - train=False, to_optimize="both" - ) - self.train_fn = self.train_neuraldual_parallel - else: - if self.parallel_updates: - warnings.warn( - "parallel_updates set to True but disabling it " - "because num_inner_iters>1", - stacklevel=2 - ) - if self.back_and_forth: - raise NotImplementedError( - "back_and_forth not implemented without parallel updates" - ) - self.train_step_f = self.get_step_fn(train=True, to_optimize="f") - self.valid_step_f = self.get_step_fn(train=False, to_optimize="f") - self.train_step_g = self.get_step_fn(train=True, to_optimize="g") - self.valid_step_g = self.get_step_fn(train=False, to_optimize="g") - self.train_fn = self.train_neuraldual_alternating - - def __call__( # noqa: D102 - self, - trainloader_source: Iterator[jnp.ndarray], - trainloader_target: Iterator[jnp.ndarray], - validloader_source: Iterator[jnp.ndarray], - validloader_target: Iterator[jnp.ndarray], - callback: Optional[Callback_t] = None, - ) -> Union[potentials.DualPotentials, Tuple[potentials.DualPotentials, - Train_t]]: - logs = self.train_fn( - trainloader_source, - trainloader_target, - validloader_source, - validloader_target, - callback=callback, - ) - res = self.to_dual_potentials() - - return (res, logs) if self.logging else res - - def train_neuraldual_parallel( - self, - trainloader_source: Iterator[jnp.ndarray], - trainloader_target: Iterator[jnp.ndarray], - validloader_source: Iterator[jnp.ndarray], - validloader_target: Iterator[jnp.ndarray], - callback: Optional[Callback_t] = None, - ) -> Train_t: - """Training and validation with parallel updates.""" - try: - from tqdm.auto import tqdm - except ImportError: - tqdm = lambda _: _ - # define dict to contain source and target batch - train_batch, valid_batch = {}, {} - - # set logging dictionaries - train_logs = {"loss_f": [], "loss_g": [], "w_dist": [], "directions": []} - valid_logs = {"loss_f": [], "loss_g": [], "w_dist": []} - - for step in tqdm(range(self.num_train_iters)): - update_forward = not self.back_and_forth or step % 2 == 0 - if update_forward: - train_batch["source"] = jnp.asarray(next(trainloader_source)) - train_batch["target"] = jnp.asarray(next(trainloader_target)) - (self.state_f, self.state_g, loss, loss_f, loss_g, - w_dist) = self.train_step_parallel( - self.state_f, - self.state_g, - train_batch, - ) - else: - train_batch["target"] = jnp.asarray(next(trainloader_source)) - train_batch["source"] = jnp.asarray(next(trainloader_target)) - (self.state_g, self.state_f, loss, loss_f, loss_g, - w_dist) = self.train_step_parallel( - self.state_g, - self.state_f, - train_batch, - ) - - if self.logging and step % self.log_freq == 0: - self._update_logs(train_logs, loss_f, loss_g, w_dist) - train_logs["directions"].append( - "forward" if update_forward else "backward" - ) - - if callback is not None: - _ = callback(step, self.to_dual_potentials()) - - if not self.pos_weights: - # Only clip the weights of the f network - self.state_f = self.state_f.replace( - params=self._clip_weights_icnn(self.state_f.params) - ) - - # report the loss on an validation dataset periodically - if step != 0 and step % self.valid_freq == 0: - # get batch - valid_batch["source"] = jnp.asarray(next(validloader_source)) - valid_batch["target"] = jnp.asarray(next(validloader_target)) - - valid_loss_f, valid_loss_g, valid_w_dist = self.valid_step_parallel( - self.state_f, - self.state_g, - valid_batch, - ) - - if self.logging: - self._update_logs( - valid_logs, valid_loss_f, valid_loss_g, valid_w_dist - ) - - return {"train_logs": train_logs, "valid_logs": valid_logs} - - def train_neuraldual_alternating( - self, - trainloader_source: Iterator[jnp.ndarray], - trainloader_target: Iterator[jnp.ndarray], - validloader_source: Iterator[jnp.ndarray], - validloader_target: Iterator[jnp.ndarray], - callback: Optional[Callback_t] = None, - ) -> Train_t: - """Training and validation with alternating updates.""" - try: - from tqdm.auto import tqdm - except ImportError: - tqdm = lambda _: _ - # define dict to contain source and target batch - batch_g, batch_f, valid_batch = {}, {}, {} - - # set logging dictionaries - train_logs = {"loss_f": [], "loss_g": [], "w_dist": []} - valid_logs = {"loss_f": [], "loss_g": [], "w_dist": []} - - for step in tqdm(range(self.num_train_iters)): - # execute training steps - for _ in range(self.num_inner_iters): - # get train batch for potential g - batch_g["source"] = jnp.asarray(next(trainloader_source)) - batch_g["target"] = jnp.asarray(next(trainloader_target)) - - self.state_g, loss_g, _ = self.train_step_g( - self.state_f, self.state_g, batch_g - ) - - # get train batch for potential f - batch_f["source"] = jnp.asarray(next(trainloader_source)) - batch_f["target"] = jnp.asarray(next(trainloader_target)) - - self.state_f, loss_f, w_dist = self.train_step_f( - self.state_f, self.state_g, batch_f - ) - if not self.pos_weights: - # Only clip the weights of the f network - self.state_f = self.state_f.replace( - params=self._clip_weights_icnn(self.state_f.params) - ) - - if callback is not None: - callback(step, self.to_dual_potentials()) - - if self.logging and step % self.log_freq == 0: - self._update_logs(train_logs, loss_f, loss_g, w_dist) - - # report the loss on validation dataset periodically - if step != 0 and step % self.valid_freq == 0: - # get batch - valid_batch["source"] = jnp.asarray(next(validloader_source)) - valid_batch["target"] = jnp.asarray(next(validloader_target)) - - valid_loss_f, _ = self.valid_step_f( - self.state_f, self.state_g, valid_batch - ) - valid_loss_g, valid_w_dist = self.valid_step_g( - self.state_f, self.state_g, valid_batch - ) - - if self.logging: - self._update_logs( - valid_logs, valid_loss_f, valid_loss_g, valid_w_dist - ) - - return {"train_logs": train_logs, "valid_logs": valid_logs} - - def get_step_fn( - self, train: bool, to_optimize: Literal["f", "g", "parallel", "both"] - ): - """Create a parallel training and evaluation function.""" - - def loss_fn(params_f, params_g, f_value, g_value, g_gradient, batch): - """Loss function for both potentials.""" - # get two distributions - source, target = batch["source"], batch["target"] - - init_source_hat = g_gradient(params_g)(target) - - def g_value_partial(y: jnp.ndarray) -> jnp.ndarray: - """Lazy way of evaluating g if f's computation needs it.""" - return g_value(params_g)(y) - - f_value_partial = f_value(params_f, g_value_partial) - if self.conjugate_solver is not None: - finetune_source_hat = lambda y, x_init: self.conjugate_solver.solve( - f_value_partial, y, x_init=x_init - ).grad - finetune_source_hat = jax.vmap(finetune_source_hat) - source_hat_detach = jax.lax.stop_gradient( - finetune_source_hat(target, init_source_hat) - ) - else: - source_hat_detach = init_source_hat - - batch_dot = jax.vmap(jnp.dot) - - f_source = f_value_partial(source) - f_star_target = batch_dot(source_hat_detach, - target) - f_value_partial(source_hat_detach) - dual_source = f_source.mean() - dual_target = f_star_target.mean() - dual_loss = dual_source + dual_target - - if self.amortization_loss == "regression": - amor_loss = ((init_source_hat - source_hat_detach) ** 2).mean() - elif self.amortization_loss == "objective": - f_value_parameters_detached = f_value( - jax.lax.stop_gradient(params_f), g_value_partial - ) - amor_loss = ( - f_value_parameters_detached(init_source_hat) - - batch_dot(init_source_hat, target) - ).mean() - else: - raise ValueError("Amortization loss has been misspecified.") - - if to_optimize == "both": - loss = dual_loss + amor_loss - elif to_optimize == "f": - loss = dual_loss - elif to_optimize == "g": - loss = amor_loss - else: - raise ValueError( - f"Optimization target {to_optimize} has been misspecified." - ) - - if not self.pos_weights: - # Penalize the weights of both networks, even though one - # of them will be exactly clipped. - # Having both here is necessary in case this is being called with - # the potentials reversed with the back_and_forth. - loss += self.beta * self._penalize_weights_icnn(params_f) + \ - self.beta * self._penalize_weights_icnn(params_g) - - # compute Wasserstein-2 distance - C = jnp.mean(jnp.sum(source ** 2, axis=-1)) + \ - jnp.mean(jnp.sum(target ** 2, axis=-1)) - W2_dist = C - 2.0 * (f_source.mean() + f_star_target.mean()) - - return loss, (dual_loss, amor_loss, W2_dist) - - @jax.jit - def step_fn(state_f, state_g, batch): - """Step function of either training or validation.""" - grad_fn = jax.value_and_grad(loss_fn, argnums=[0, 1], has_aux=True) - if train: - # compute loss and gradients - (loss, (loss_f, loss_g, W2_dist)), (grads_f, grads_g) = grad_fn( - state_f.params, - state_g.params, - state_f.potential_value_fn, - state_g.potential_value_fn, - state_g.potential_gradient_fn, - batch, - ) - # update state - if to_optimize == "both": - return ( - state_f.apply_gradients(grads=grads_f), - state_g.apply_gradients(grads=grads_g), loss, loss_f, loss_g, - W2_dist - ) - if to_optimize == "f": - return state_f.apply_gradients(grads=grads_f), loss_f, W2_dist - if to_optimize == "g": - return state_g.apply_gradients(grads=grads_g), loss_g, W2_dist - raise ValueError("Optimization target has been misspecified.") - - # compute loss and gradients - (loss, (loss_f, loss_g, W2_dist)), _ = grad_fn( - state_f.params, - state_g.params, - state_f.potential_value_fn, - state_g.potential_value_fn, - state_g.potential_gradient_fn, - batch, - ) - - # do not update state - if to_optimize == "both": - return loss_f, loss_g, W2_dist - if to_optimize == "f": - return loss_f, W2_dist - if to_optimize == "g": - return loss_g, W2_dist - raise ValueError("Optimization target has been misspecified.") - - return step_fn - - def to_dual_potentials( - self, finetune_g: bool = True - ) -> potentials.DualPotentials: - r"""Return the Kantorovich dual potentials from the trained potentials. - - Args: - finetune_g: Run the conjugate solver to fine-tune the prediction. - - Returns: - A dual potential object - """ - f_value = self.state_f.potential_value_fn(self.state_f.params) - g_value_prediction = self.state_g.potential_value_fn( - self.state_g.params, f_value - ) - - def g_value_finetuned(y: jnp.ndarray) -> jnp.ndarray: - x_hat = jax.grad(g_value_prediction)(y) - grad_g_y = jax.lax.stop_gradient( - self.conjugate_solver.solve(f_value, y, x_init=x_hat).grad - ) - return -f_value(grad_g_y) + jnp.dot(grad_g_y, y) - - return potentials.DualPotentials( - f=f_value, - g=g_value_prediction if not finetune_g or self.conjugate_solver is None - else g_value_finetuned, - cost_fn=costs.SqEuclidean(), - corr=True - ) - - @staticmethod - def _clip_weights_icnn(params): - for k in params: - if k.startswith("w_z"): - params[k]["kernel"] = jnp.clip(params[k]["kernel"], a_min=0) - - return params - - @staticmethod - def _penalize_weights_icnn(params: Dict[str, jnp.ndarray]) -> float: - penalty = 0.0 - for k, param in params.items(): - if k.startswith("w_z"): - penalty += jnp.linalg.norm(jax.nn.relu(-param["kernel"])) - return penalty - - @staticmethod - def _update_logs( - logs: Dict[str, List[Union[float, str]]], - loss_f: jnp.ndarray, - loss_g: jnp.ndarray, - w_dist: jnp.ndarray, - ) -> None: - logs["loss_f"].append(float(loss_f)) - logs["loss_g"].append(float(loss_g)) - logs["w_dist"].append(float(w_dist)) diff --git a/ott/build/lib/ott/problems/__init__.py b/ott/build/lib/ott/problems/__init__.py deleted file mode 100644 index 5406247..0000000 --- a/ott/build/lib/ott/problems/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. -from . import linear, quadratic diff --git a/ott/build/lib/ott/problems/linear/__init__.py b/ott/build/lib/ott/problems/linear/__init__.py deleted file mode 100644 index 7d62b13..0000000 --- a/ott/build/lib/ott/problems/linear/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. -from . import barycenter_problem, linear_problem, potentials diff --git a/ott/build/lib/ott/problems/linear/barycenter_problem.py b/ott/build/lib/ott/problems/linear/barycenter_problem.py deleted file mode 100644 index ca5333a..0000000 --- a/ott/build/lib/ott/problems/linear/barycenter_problem.py +++ /dev/null @@ -1,200 +0,0 @@ -# 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. -from typing import Any, Dict, Optional, Sequence, Tuple - -import jax -import jax.numpy as jnp - -from ott.geometry import costs, geometry, segment - -__all__ = ["FreeBarycenterProblem", "FixedBarycenterProblem"] - - -@jax.tree_util.register_pytree_node_class -class FreeBarycenterProblem: - """Free Wasserstein barycenter problem :cite:`cuturi:14`. - - Args: - y: Array of shape ``[num_total_points, ndim]`` merging the points of all - measures. Alternatively, already segmented array of shape - ``[num_measures, max_measure_size, ndim]`` can be passed. - See also :func:`~ott.geometry.segment.segment_point_cloud`. - b: Array of shape ``[num_total_points,]`` containing the weights of all - the points within the measures that define the barycenter problem. - Same as ``y``, pre-segmented array of weights of shape - ``[num_measures, max_measure_size]`` can be passed. - If ``y`` is already pre-segmented, this array must be always specified. - weights: Array of shape ``[num_measures,]`` containing the weights of the - measures. - cost_fn: Cost function used. If `None`, - use the :class:`~ott.geometry.costs.SqEuclidean` cost. - epsilon: Epsilon regularization used to solve reg-OT problems. - kwargs: Keyword arguments :func:`~ott.geometry.segment.segment_point_cloud`. - Only used when ``y`` is not already segmented. When passing - ``segment_ids``, 2 arguments must be specified for jitting to work: - - - ``num_segments`` - the total number of measures. - - ``max_measure_size`` - maximum of support sizes of these measures. - """ - - def __init__( - self, - y: jnp.ndarray, - b: Optional[jnp.ndarray] = None, - weights: Optional[jnp.ndarray] = None, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, - **kwargs: Any, - ): - self._y = y - if y.ndim == 3 and b is None: - raise ValueError("Specify weights if `y` is already segmented.") - self._b = b - self._weights = weights - self.cost_fn = costs.SqEuclidean() if cost_fn is None else cost_fn - self.epsilon = epsilon - self._kwargs = kwargs - - if self._is_segmented: - # (num_measures, max_measure_size, ndim) - # (num_measures, max_measure_size) - assert self._y.shape[:2] == self._b.shape - else: - # (num_total_points, ndim) - # (num_total_points,) - assert self._b is None or self._y.shape[0] == self._b.shape[0] - - @property - def segmented_y_b(self) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Tuple of arrays containing the segmented measures and weights. - - - Segmented measures of shape ``[num_measures, max_measure_size, ndim]``. - - Segmented weights of shape ``[num_measures, max_measure_size]``. - """ - if self._is_segmented: - y, b = self._y, self._b - else: - y, b = segment.segment_point_cloud( - x=self._y, - a=self._b, - padding_vector=self.cost_fn._padder(self.ndim), - **self._kwargs - ) - return y, b - - @property - def flattened_y(self) -> jnp.ndarray: - """Array of shape ``[num_measures * (N_1 + N_2 + ...), ndim]``.""" - if self._is_segmented: - return self._y.reshape((-1, self._y.shape[-1])) - return self._y - - @property - def flattened_b(self) -> Optional[jnp.ndarray]: - """Array of shape ``[num_measures * (N_1 + N_2 + ...),]``.""" - return None if self._b is None else self._b.ravel() - - @property - def num_measures(self) -> int: - """Number of measures.""" - return self.segmented_y_b[0].shape[0] - - @property - def max_measure_size(self) -> int: - """Maximum number of points across all measures.""" - return self.segmented_y_b[0].shape[1] - - @property - def ndim(self) -> int: - """Number of dimensions of ``y``.""" - return self._y.shape[-1] - - @property - def weights(self) -> jnp.ndarray: - """Barycenter weights of shape ``[num_measures,]`` that sum to 1.""" - if self._weights is None: - return jnp.ones((self.num_measures,)) / self.num_measures - # Check that the number of measures coincides with the weights' size. - assert self._weights.shape[0] == self.num_measures - # By default, we assume that weights sum to 1, and enforce this if needed. - return self._weights / jnp.sum(self._weights) - - @property - def _is_segmented(self) -> bool: - return self._y.ndim == 3 - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return ([self._y, self._b, self._weights], { - "cost_fn": self.cost_fn, - "epsilon": self.epsilon, - **self._kwargs, - }) - - @classmethod - def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] - ) -> "FreeBarycenterProblem": - y, b, weights = children - return cls(y=y, b=b, weights=weights, **aux_data) - - -@jax.tree_util.register_pytree_node_class -class FixedBarycenterProblem: - """Fixed-support Wasserstein barycenter problem. - - Args: - geom: Geometry object. - a: batch of histograms of shape ``[batch, num_a]`` where ``num_a`` matches - the first value of the :attr:`~ott.geometry.Geometry.shape` attribute of - ``geom``. - weights: ``[batch,]`` positive weights summing to :math:`1`. Uniform by - default. - """ - - def __init__( - self, - geom: geometry.Geometry, - a: jnp.ndarray, - weights: Optional[jnp.ndarray] = None, - ): - self.geom = geom - self.a = a - self._weights = weights - - @property - def num_measures(self) -> int: - """Number of measures.""" - return self.a.shape[0] - - @property - def weights(self) -> jnp.ndarray: - """Barycenter weights of shape ``[num_measures,]`` that sum to :math`1`.""" - if self._weights is None: - return jnp.ones((self.num_measures,)) / self.num_measures - - # check that the number of measures coincides with the weights' size - assert self._weights.shape[0] == self.num_measures - # by default, we assume that weights sum to 1, and enforce this if needed - return self._weights / jnp.sum(self._weights) - - def tree_flatten(self): # noqa: D102 - return [self.geom, self.a, self._weights], None - - @classmethod - def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] - ) -> "FixedBarycenterProblem": - del aux_data - geom, a, weights = children - return cls(geom=geom, a=a, weights=weights) diff --git a/ott/build/lib/ott/problems/linear/linear_problem.py b/ott/build/lib/ott/problems/linear/linear_problem.py deleted file mode 100644 index 60d8cc8..0000000 --- a/ott/build/lib/ott/problems/linear/linear_problem.py +++ /dev/null @@ -1,132 +0,0 @@ -# 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. -from typing import Any, Callable, Dict, Optional, Sequence, Tuple - -import jax -import jax.numpy as jnp - -from ott.geometry import geometry - -__all__ = ["LinearProblem"] - -# TODO(michalk8): move to typing.py when refactoring the types -MarginalFunc = Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray] -TransportAppFunc = Callable[[jnp.ndarray, jnp.ndarray, jnp.ndarray, int], - jnp.ndarray] - - -@jax.tree_util.register_pytree_node_class -class LinearProblem: - r"""Linear OT problem. - - This class describes the main ingredients appearing in a linear OT problem. - Namely, a ``geom`` object (including cost structure/points) describing point - clouds or the support of measures, followed by probability masses ``a`` and - ``b``. Unbalancedness of the problem is also kept track of, through two - coefficients ``tau_a`` and ``tau_b``, which are both kept between 0 and 1 - (1 corresponding to a balanced OT problem). - - Args: - geom: The ground geometry cost of the linear problem. - a: The first marginal. If ``None``, it will be uniform. - b: The second marginal. If ``None``, it will be uniform. - tau_a: If :math:`<1`, defines how much unbalanced the problem is - on the first marginal. - tau_b: If :math:`< 1`, defines how much unbalanced the problem is - on the second marginal. - """ - - def __init__( - self, - geom: geometry.Geometry, - a: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, - tau_a: float = 1.0, - tau_b: float = 1.0, - labels_a: Optional[jnp.ndarray] = None, - labels_b: Optional[jnp.ndarray] = None, - ): - self.geom = geom - self._a = a - self._b = b - self.tau_a = tau_a - self.tau_b = tau_b - self.labels_a = labels_a - self.labels_b = labels_b - - @property - def a(self) -> jnp.ndarray: - """First marginal.""" - num_a = self.geom.shape[0] - return jnp.ones((num_a,)) / num_a if self._a is None else self._a - - @property - def b(self) -> jnp.ndarray: - """Second marginal.""" - num_b = self.geom.shape[1] - return jnp.ones((num_b,)) / num_b if self._b is None else self._b - - @property - def is_balanced(self) -> bool: - """Whether the problem is balanced.""" - return self.tau_a == 1.0 and self.tau_b == 1.0 - - @property - def is_uniform(self) -> bool: - """True if no weights ``a,b`` were passed, and have defaulted to uniform.""" - return self._a is None and self._b is None - - @property - def is_equal_size(self) -> bool: - """True if square shape, i.e. ``n == m``.""" - return self.geom.shape[0] == self.geom.shape[1] - - @property - def epsilon(self) -> float: - """Entropic regularization.""" - return self.geom.epsilon - - def get_transport_functions( - self, lse_mode: bool - ) -> Tuple[MarginalFunc, MarginalFunc, TransportAppFunc]: - """Instantiate useful functions for Sinkhorn depending on lse_mode.""" - geom = self.geom - if lse_mode: - marginal_a = lambda f, g: geom.marginal_from_potentials(f, g, 1) - marginal_b = lambda f, g: geom.marginal_from_potentials(f, g, 0) - app_transport = geom.apply_transport_from_potentials - else: - marginal_a = lambda f, g: geom.marginal_from_scalings( - geom.scaling_from_potential(f), geom.scaling_from_potential(g), 1 - ) - marginal_b = lambda f, g: geom.marginal_from_scalings( - geom.scaling_from_potential(f), geom.scaling_from_potential(g), 0 - ) - app_transport = lambda f, g, z, axis: geom.apply_transport_from_scalings( - geom.scaling_from_potential(f), geom.scaling_from_potential(g), z, - axis - ) - return marginal_a, marginal_b, app_transport - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return ([self.geom, self._a, self._b], { - "tau_a": self.tau_a, - "tau_b": self.tau_b - }) - - @classmethod - def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] - ) -> "LinearProblem": - return cls(*children, **aux_data) diff --git a/ott/build/lib/ott/problems/linear/potentials.py b/ott/build/lib/ott/problems/linear/potentials.py deleted file mode 100644 index 6885152..0000000 --- a/ott/build/lib/ott/problems/linear/potentials.py +++ /dev/null @@ -1,437 +0,0 @@ -# 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. -from typing import ( - Any, - Callable, - Dict, - Literal, - Optional, - Sequence, - Tuple, -) - -import jax -import jax.numpy as jnp -import jax.scipy as jsp -import jax.tree_util as jtu -import numpy as np - -from ott.geometry import costs -from ott.problems.linear import linear_problem - -try: - import matplotlib as mpl - import matplotlib.pyplot as plt -except ImportError: - mpl = plt = None - -__all__ = ["DualPotentials", "EntropicPotentials"] -Potential_t = Callable[[jnp.ndarray], float] - - -@jtu.register_pytree_node_class -class DualPotentials: - r"""The Kantorovich dual potential functions :math:`f` and :math:`g`. - - :math:`f` and :math:`g` are a pair of functions, candidates for the dual - OT Kantorovich problem, supposedly optimal for a given pair of measures. - - Args: - f: The first dual potential function. - g: The second dual potential function. - cost_fn: The cost function used to solve the OT problem. - corr: Whether the duals solve the problem in distance form, or correlation - form (as used for instance for ICNNs, see, e.g., top right of p.3 in - :cite:`makkuva:20`) - """ - - def __init__( - self, - f: Potential_t, - g: Potential_t, - *, - cost_fn: costs.CostFn, - corr: bool = False - ): - self._f = f - self._g = g - assert ( - not corr or type(cost_fn) == costs.SqEuclidean - ), "Duals in `corr` form can only be used with a squared-Euclidean cost." - self.cost_fn = cost_fn - self._corr = corr - - def transport(self, vec: jnp.ndarray, forward: bool = True) -> jnp.ndarray: - r"""Transport ``vec`` according to Brenier formula :cite:`brenier:91`. - - Uses Theorem 1.17 from :cite:`santambrogio:15` to compute an OT map when - given the Legendre transform of the dual potentials. - - That OT map can be recovered as :math:`x- (\nabla h^*)\circ \nabla f(x)`, - where :math:`h^*` is the Legendre transform of :math:`h`. For instance, - in the case :math:`h(\cdot) = \|\cdot\|^2, \nabla h(\cdot) = 2 \cdot\,`, - one has :math:`h^*(\cdot) = \|.\|^2 / 4`, and therefore - :math:`\nabla h^*(\cdot) = 0.5 \cdot\,`. - - When the dual potentials are solved in correlation form (only in the Sq. - Euclidean distance case), the maps are :math:`\nabla g` for forward, - :math:`\nabla f` for backward. - - Args: - vec: Points to transport, array of shape ``[n, d]``. - forward: Whether to transport the points from source to the target - distribution or vice-versa. - - Returns: - The transported points. - """ - from ott.geometry import costs - - vec = jnp.atleast_2d(vec) - if self._corr and isinstance(self.cost_fn, costs.SqEuclidean): - return self._grad_f(vec) if forward else self._grad_g(vec) - if forward: - return vec - self._grad_h_inv(self._grad_f(vec)) - return vec - self._grad_h_inv(self._grad_g(vec)) - - def distance(self, src: jnp.ndarray, tgt: jnp.ndarray) -> float: - r"""Evaluate Wasserstein distance between samples using dual potentials. - - This uses direct estimation of potentials against measures when dual - functions are provided in usual form. This expression is valid for any - cost function. - - When potentials are given in correlation form, as specified by the flag - ``corr``, the dual potentials solve the dual problem corresponding to the - minimization of the primal OT problem where the ground cost is - :math:`-2\langle x,y\rangle`. To recover the (squared) 2-Wasserstein - distance, terms are re-arranged and contributions from squared norms are - taken into account. - - Args: - src: Samples from the source distribution, array of shape ``[n, d]``. - tgt: Samples from the target distribution, array of shape ``[m, d]``. - - Returns: - Wasserstein distance using specified cost function. - """ - src, tgt = jnp.atleast_2d(src), jnp.atleast_2d(tgt) - f = jax.vmap(self.f) - g = jax.vmap(self.g) - out = jnp.mean(f(src)) + jnp.mean(g(tgt)) - if self._corr: - out = -2.0 * out + jnp.mean(jnp.sum(src ** 2, axis=-1)) - out += jnp.mean(jnp.sum(tgt ** 2, axis=-1)) - return out - - @property - def f(self) -> Potential_t: - """The first dual potential function.""" - return self._f - - @property - def g(self) -> Potential_t: - """The second dual potential function.""" - return self._g - - @property - def _grad_f(self) -> Callable[[jnp.ndarray], jnp.ndarray]: - """Vectorized gradient of the potential function :attr:`f`.""" - return jax.vmap(jax.grad(self.f, argnums=0)) - - @property - def _grad_g(self) -> Callable[[jnp.ndarray], jnp.ndarray]: - """Vectorized gradient of the potential function :attr:`g`.""" - return jax.vmap(jax.grad(self.g, argnums=0)) - - @property - def _grad_h_inv(self) -> Callable[[jnp.ndarray], jnp.ndarray]: - from ott.geometry import costs - - assert isinstance(self.cost_fn, costs.TICost), ( - "Cost must be a `TICost` and " - "provide access to Legendre transform of `h`." - ) - return jax.vmap(jax.grad(self.cost_fn.h_legendre)) - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return [], { - "f": self._f, - "g": self._g, - "cost_fn": self.cost_fn, - "corr": self._corr - } - - @classmethod - def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] - ) -> "DualPotentials": - return cls(*children, **aux_data) - - def plot_ot_map( - self, - source: jnp.ndarray, - target: jnp.ndarray, - samples: Optional[jnp.ndarray] = None, - forward: bool = True, - ax: Optional["plt.Axes"] = None, - scatter_kwargs: Optional[Dict[str, Any]] = None, - legend_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple["plt.Figure", "plt.Axes"]: - """Plot data and learned optimal transport map. - - Args: - source: samples from the source measure - target: samples from the target measure - samples: extra samples to transport, either ``source`` (if ``forward``) or - ``target`` (if not ``forward``) by default. - forward: use the forward map from the potentials if ``True``, - otherwise use the inverse map. - ax: axis to add the plot to - scatter_kwargs: additional kwargs passed into - :meth:`~matplotlib.axes.Axes.scatter` - legend_kwargs: additional kwargs passed into - :meth:`~matplotlib.axes.Axes.legend` - - Returns: - Figure and axes. - """ - if mpl is None: - raise RuntimeError("Please install `matplotlib` first.") - - if scatter_kwargs is None: - scatter_kwargs = {"alpha": 0.5} - if legend_kwargs is None: - legend_kwargs = { - "ncol": 3, - "loc": "upper center", - "bbox_to_anchor": (0.5, -0.05), - "edgecolor": "k" - } - - if ax is None: - fig = plt.figure(facecolor="white") - ax = fig.add_subplot(111) - else: - fig = ax.get_figure() - - # plot the source and target samples - if forward: - label_transport = r"$\nabla f(source)$" - source_color, target_color = "#1A254B", "#A7BED3" - else: - label_transport = r"$\nabla g(target)$" - source_color, target_color = "#A7BED3", "#1A254B" - - ax.scatter( - source[:, 0], - source[:, 1], - color=source_color, - label="source", - **scatter_kwargs, - ) - ax.scatter( - target[:, 0], - target[:, 1], - color=target_color, - label="target", - **scatter_kwargs, - ) - - # plot the transported samples - samples = (source if forward else target) if samples is None else samples - transported_samples = self.transport(samples, forward=forward) - ax.scatter( - transported_samples[:, 0], - transported_samples[:, 1], - color="#F2545B", - label=label_transport, - **scatter_kwargs, - ) - - for i in range(samples.shape[0]): - ax.arrow( - samples[i, 0], - samples[i, 1], - transported_samples[i, 0] - samples[i, 0], - transported_samples[i, 1] - samples[i, 1], - color=[0.5, 0.5, 1], - alpha=0.3, - ) - - ax.legend(**legend_kwargs) - return fig, ax - - def plot_potential( - self, - forward: bool = True, - quantile: float = 0.05, - kantorovich: bool = True, - ax: Optional["mpl.axes.Axes"] = None, - x_bounds: Tuple[float, float] = (-6, 6), - y_bounds: Tuple[float, float] = (-6, 6), - num_grid: int = 50, - contourf_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple["mpl.figure.Figure", "mpl.axes.Axes"]: - r"""Plot the potential. - - Args: - forward: use the forward map from the potentials - if ``True``, otherwise use the inverse map - quantile: quantile to filter the potentials with - kantorovich: whether to plot the Kantorovich potential - ax: axis to add the plot to - x_bounds: x-axis bounds of the plot - :math:`(x_{\text{min}}, x_{\text{max}})` - y_bounds: y-axis bounds of the plot - :math:`(y_{\text{min}}, y_{\text{max}})` - num_grid: number of points to discretize the domain into a grid - along each dimension - contourf_kwargs: additional kwargs passed into - :meth:`~matplotlib.axes.Axes.contourf` - - Returns: - Figure and axes. - """ - if contourf_kwargs is None: - contourf_kwargs = {} - - ax_specified = ax is not None - if not ax_specified: - fig, ax = plt.subplots(figsize=(6, 6), facecolor="white") - else: - fig = ax.get_figure() - - x1 = jnp.linspace(*x_bounds, num=num_grid) - x2 = jnp.linspace(*y_bounds, num=num_grid) - X1, X2 = jnp.meshgrid(x1, x2) - X12flat = jnp.hstack((X1.reshape(-1, 1), X2.reshape(-1, 1))) - Zflat = jax.vmap(self.f if forward else self.g)(X12flat) - if kantorovich: - Zflat = 0.5 * (jnp.linalg.norm(X12flat, axis=-1) ** 2) - Zflat - Zflat = np.asarray(Zflat) - vmin, vmax = np.quantile(Zflat, [quantile, 1.0 - quantile]) - Zflat = Zflat.clip(vmin, vmax) - Z = Zflat.reshape(X1.shape) - - CS = ax.contourf(X1, X2, Z, cmap="Blues", **contourf_kwargs) - ax.set_xlim(*x_bounds) - ax.set_ylim(*y_bounds) - fig.colorbar(CS, ax=ax) - if not ax_specified: - fig.tight_layout() - ax.set_title(r"$f$" if forward else r"$g$") - return fig, ax - - -@jtu.register_pytree_node_class -class EntropicPotentials(DualPotentials): - """Dual potential functions from finite samples :cite:`pooladian:21`. - - Args: - f_xy: The first dual potential vector of shape ``[n,]``. - g_xy: The second dual potential vector of shape ``[m,]``. - prob: Linear problem with :class:`~ott.geometry.pointcloud.PointCloud` - geometry that was used to compute the dual potentials using, e.g., - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. - f_xx: The first dual potential vector of shape ``[n,]`` used for debiasing - :cite:`pooladian:22`. - g_yy: The second dual potential vector of shape ``[m,]`` used for debiasing. - """ - - def __init__( - self, - f_xy: jnp.ndarray, - g_xy: jnp.ndarray, - prob: linear_problem.LinearProblem, - f_xx: Optional[jnp.ndarray] = None, - g_yy: Optional[jnp.ndarray] = None, - ): - # we pass directly the arrays and override the properties - # since only the properties need to be callable - super().__init__(f_xy, g_xy, cost_fn=prob.geom.cost_fn, corr=False) - self._prob = prob - self._f_xx = f_xx - self._g_yy = g_yy - - @property - def f(self) -> Potential_t: # noqa: D102 - return self._potential_fn(kind="f") - - @property - def g(self) -> Potential_t: # noqa: D102 - return self._potential_fn(kind="g") - - def _potential_fn(self, *, kind: Literal["f", "g"]) -> Potential_t: - from ott.geometry import pointcloud - - def callback( - x: jnp.ndarray, - *, - potential: jnp.ndarray, - y: jnp.ndarray, - weights: jnp.ndarray, - epsilon: float, - ) -> float: - x = jnp.atleast_2d(x) - assert x.shape[-1] == y.shape[-1], (x.shape, y.shape) - geom = pointcloud.PointCloud(x, y, cost_fn=self.cost_fn) - cost = geom.cost_matrix - z = (potential - cost) / epsilon - lse = -epsilon * jsp.special.logsumexp(z, b=weights, axis=-1) - return jnp.squeeze(lse) - - assert isinstance( - self._prob.geom, pointcloud.PointCloud - ), f"Expected point cloud geometry, found `{type(self._prob.geom)}`." - x, y = self._prob.geom.x, self._prob.geom.y - a, b = self._prob.a, self._prob.b - - if kind == "f": - # When seeking to evaluate 1st potential function, - # the 2nd set of potential values and support should be used, - # see proof of Prop. 2 in https://arxiv.org/pdf/2109.12004.pdf - potential, arr, weights = self._g, y, b - else: - potential, arr, weights = self._f, x, a - - potential_xy = jax.tree_util.Partial( - callback, - potential=potential, - y=arr, - weights=weights, - epsilon=self.epsilon, - ) - if not self.is_debiased: - return potential_xy - - ep = EntropicPotentials(self._f_xx, self._g_yy, prob=self._prob) - # switch the order because for `kind='f'` we require `f/x/a` in `other` - # which is accessed when `kind='g'` - potential_other = ep._potential_fn(kind="g" if kind == "f" else "f") - - return lambda x: (potential_xy(x) - potential_other(x)) - - @property - def is_debiased(self) -> bool: - """Whether the entropic map is debiased.""" - return self._f_xx is not None and self._g_yy is not None - - @property - def epsilon(self) -> float: - """Entropy regularizer.""" - return self._prob.geom.epsilon - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return [self._f, self._g, self._prob, self._f_xx, self._g_yy], {} diff --git a/ott/build/lib/ott/problems/quadratic/__init__.py b/ott/build/lib/ott/problems/quadratic/__init__.py deleted file mode 100644 index d2bb647..0000000 --- a/ott/build/lib/ott/problems/quadratic/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. -from . import gw_barycenter, quadratic_costs, quadratic_problem diff --git a/ott/build/lib/ott/problems/quadratic/gw_barycenter.py b/ott/build/lib/ott/problems/quadratic/gw_barycenter.py deleted file mode 100644 index 7acf962..0000000 --- a/ott/build/lib/ott/problems/quadratic/gw_barycenter.py +++ /dev/null @@ -1,325 +0,0 @@ -# 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 functools -from typing import Any, Dict, Literal, Optional, Sequence, Tuple, Union - -import jax -import jax.numpy as jnp - -from ott.geometry import costs, geometry, pointcloud, segment -from ott.math import utils as mu -from ott.problems.linear import barycenter_problem -from ott.problems.quadratic import quadratic_costs, quadratic_problem - -__all__ = ["GWBarycenterProblem"] - - -# TODO(michalk8): better abstraction (common superclass for Wasserstein bary) -@jax.tree_util.register_pytree_node_class -class GWBarycenterProblem(barycenter_problem.FreeBarycenterProblem): - """(Fused) Gromov-Wasserstein barycenter problem :cite:`peyre:16,vayer:19`. - - Args: - y: Array of shape ``[num_total_points, ndim]`` merging the points of all - measures. Alternatively, already segmented array of shape - ``[num_measures, max_measure_size, ndim]`` can be passed. - See also :func:`~ott.geometry.segment.segment_point_cloud`. - b: Array of shape ``[num_total_points,]`` containing the weights of all - the points within the measures that define the barycenter problem. - Same as ``y``, pre-segmented array of weights of shape - ``[num_measures, max_measure_size]`` can be passed. - If ``y`` is already pre-segmented, this array must be passed. - weights: Array of shape ``[num_measures,]`` containing the weights of the - barycenter problem. - costs: Alternative to ``y``, an array of shape - ``[num_measures, max_measure_size, max_measure_size]`` that defines padded - cost matrices for each measure. Used in the quadratic term. - Only one of ``y`` and ``cost`` can be specified. - y_fused: Array of shape ``[num_total_points, ndim_fused]`` containing - the data of the points of all measures used to define the linear term - in the fused case. Same as ``y``, it can be specified as a pre-segmented - array of shape ``[num_measures, max_measure_size, ndim_fused]``. - gw_loss: Gromov-Wasserstein loss. - fused_penalty: Multiplier of the linear term. Only used when - ``y_fused != None``. - scale_cost: Scaling of cost matrices passed to geometries. - kwargs: Keyword arguments for - :class:`~ott.problems.linear.barycenter_problem.BarycenterProblem`. - """ - - def __init__( - self, - y: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, - weights: Optional[jnp.ndarray] = None, - costs: Optional[jnp.ndarray] = None, - y_fused: Optional[jnp.ndarray] = None, - fused_penalty: float = 1.0, - gw_loss: Literal["sqeucl", "kl"] = "sqeucl", - scale_cost: Union[int, float, Literal["mean", "max_cost"]] = 1.0, - **kwargs: Any, - ): - assert y is None or costs is None, "Cannot specify both `y` and `costs`." - y = y if costs is None else costs - - super().__init__(y=y, b=b, weights=weights, **kwargs) - - self._y_fused = y_fused - self.fused_penalty = fused_penalty - self._loss_name = gw_loss - self.scale_cost = scale_cost - self._y_as_costs = costs is not None - - if self._y_as_costs: - # (num_measures, max_measure_size, max_measure_size) - _, n, m = self._y.shape - assert n == m, "Cost matrices must be square." - if self.is_fused: - seg_y = self._is_segmented - seg_fused = self._y_fused.ndim == 3 - if seg_y and seg_fused: - # (num_measures, max_measure_size, ndim_fused) - # (num_measures, max_measure_size, ndim) - assert self._y_fused.shape[:2] == self._y.shape[:2] - if not seg_y and not seg_fused: - # (num_total_points, ndim_fused), (num_total_points, ndim) - assert self._y_fused.shape[0] == self._y.shape[0] - # TODO(michalk8): in the future, consider checking the other 2 cases - # using `segmented_y` and `segmented_y_fused`? - - def update_barycenter( - self, transports: jnp.ndarray, a: jnp.ndarray - ) -> jnp.ndarray: - """Update the barycenter cost matrix. - - Uses the eq. 14 and 15 of :cite:`peyre:16`. - - Args: - transports: Transport maps of shape - ``[num_measures, bar_size, max_measure_size]``. - a: Barycenter weights of shape ``[bar_size,]``. - - Returns: - Update cost matrix of shape ``[bar_size, bar_size]``. - """ - - @functools.partial(jax.vmap, in_axes=[0, 0, 0, None]) - def project( - y: jnp.ndarray, - b: jnp.ndarray, - transport: jnp.ndarray, - fn: Optional[quadratic_costs.Loss], - ) -> jnp.ndarray: - geom = self._create_y_geometry(y, mask=b > 0.0) - fn, lin = (None, True) if fn is None else (fn.func, fn.is_linear) - - tmp = geom.apply_cost( - transport.T, - axis=0, - fn=fn, - is_linear=lin, - ) - return transport @ tmp - - fn = None if self._loss_name == "sqeucl" else self.gw_loss.h2 - y, b = self.segmented_y_b - weights = self.weights[:, None, None] - - barycenter = jnp.sum(weights * project(y, b, transports, fn), axis=0) - inv_a = jnp.where(a > 0, 1.0 / a, 1.0) - barycenter = (barycenter * inv_a[None, :]) * inv_a[:, None] - - # TODO(michalk8): in future, use `isinstanceof(self.gw_loss, ...)` - # once refactoring has been done - if self._loss_name == "kl": - return jnp.exp(barycenter) - return barycenter - - def update_features(self, transports: jnp.ndarray, - a: jnp.ndarray) -> Optional[jnp.ndarray]: - """Update the barycenter features in the fused case :cite:`vayer:19`. - - Uses :cite:`cuturi:14` eq. 8, and is implemented only - for the :class:`~ott.geometry.costs.SqEuclidean` cost. - - Args: - transports: Transport maps of shape - ``[num_measures, bar_size, max_measure_size]``. - a: Barycenter weights of shape ``[bar_size,]``. - - Returns: - Updated features of shape ``[bar_size, ndim_fused]``. - """ - if not self.is_fused: - raise RuntimeError( - "Updating features is available only in the fused case." - ) - - y_fused = self.segmented_y_fused - weights = self.weights[:, None, None] - inv_a = jnp.where(a > 0, 1.0 / a, 1.0) - transports = transports * inv_a[None, :, None] - - if self._loss_name == "sqeucl": - cost_fn = costs.SqEuclidean() - return jnp.sum( - weights * mu.barycentric_projection(transports, y_fused, cost_fn), - axis=0 - ) - raise NotImplementedError(self._loss_name) - - def _create_bary_geometry( - self, - cost_matrix: jnp.ndarray, - mask: Optional[jnp.ndarray] = None - ) -> geometry.Geometry: - return geometry.Geometry( - cost_matrix=cost_matrix, - src_mask=mask, - tgt_mask=mask, - epsilon=self.epsilon, - scale_cost=self.scale_cost - ) - - def _create_y_geometry( - self, - y: jnp.ndarray, - mask: Optional[jnp.ndarray] = None - ) -> geometry.Geometry: - if self._y_as_costs: - assert y.shape[0] == y.shape[1], y.shape - return geometry.Geometry( - y, - epsilon=self.epsilon, - scale_cost=self.scale_cost, - src_mask=mask, - tgt_mask=mask - ) - return pointcloud.PointCloud( - y, - epsilon=self.epsilon, - scale_cost=self.scale_cost, - cost_fn=self.cost_fn, - src_mask=mask, - tgt_mask=mask - ) - - def _create_fused_geometry( - self, - x: jnp.ndarray, - y: jnp.ndarray, - src_mask: Optional[jnp.ndarray] = None, - tgt_mask: Optional[jnp.ndarray] = None - ) -> pointcloud.PointCloud: - return pointcloud.PointCloud( - x, - y, - cost_fn=self.cost_fn, - epsilon=self.epsilon, - scale_cost=self.scale_cost, - src_mask=src_mask, - tgt_mask=tgt_mask - ) - - def _create_problem( - self, - state: "GWBarycenterState", # noqa: F821 - y: jnp.ndarray, - b: jnp.ndarray, - f: Optional[jnp.ndarray] = None - ) -> quadratic_problem.QuadraticProblem: - # TODO(michalk8): in future, mask in the problem for convenience? - bary_mask = state.a > 0.0 - y_mask = b > 0.0 - - geom_xx = self._create_bary_geometry(state.cost, mask=bary_mask) - geom_yy = self._create_y_geometry(y, mask=y_mask) - if self.is_fused: - assert f is not None - assert state.x.shape[1] == f.shape[1] - geom_xy = self._create_fused_geometry( - state.x, f, src_mask=bary_mask, tgt_mask=y_mask - ) - else: - geom_xy = None - - return quadratic_problem.QuadraticProblem( - geom_xx=geom_xx, - geom_yy=geom_yy, - geom_xy=geom_xy, - a=state.a, - b=b, - fused_penalty=self.fused_penalty, - ) - - @property - def is_fused(self) -> bool: - """Whether the problem is fused.""" - return self._y_fused is not None - - @property - def segmented_y_fused(self) -> Optional[jnp.ndarray]: - """Feature array of shape used in the fused case.""" - if not self.is_fused or self._y_fused.ndim == 3: - return self._y_fused - y_fused, _ = segment.segment_point_cloud( - x=self._y_fused, - padding_vector=self.cost_fn._padder(self.ndim_fused), - **self._kwargs - ) - return y_fused - - @property - def ndim(self) -> Optional[int]: # noqa: D102 - return None if self._y_as_costs else self._y.shape[-1] - - @property - def ndim_fused(self) -> Optional[int]: - """Number of dimensions of the fused term.""" - return self._y_fused.shape[-1] if self.is_fused else None - - @property - def gw_loss(self) -> quadratic_costs.GWLoss: - """Gromov-Wasserstein loss.""" - # TODO(michalk8): custom losses would require inverting some fns; - # `https://jax.readthedocs.io/en/latest/notebooks/ some fns; - # Writing_custom_interpreters_in_Jax.html#your-first-interpreter-invert` - # might be useful - if self._loss_name == "sqeucl": - return quadratic_costs.make_square_loss() - if self._loss_name == "kl": - return quadratic_costs.make_kl_loss() - raise NotImplementedError( - f"Loss `{self._loss_name}` is not yet implemented." - ) - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - (y, b, weights), aux = super().tree_flatten() - if self._y_as_costs: - children = [None, b, weights, y] - else: - children = [y, b, weights, None] - aux["fused_penalty"] = self.fused_penalty - aux["gw_loss"] = self._loss_name - aux["scale_cost"] = self.scale_cost - return children + [self._y_fused], aux - - @classmethod - def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] - ) -> "GWBarycenterProblem": - y, b, weights, costs, y_fused = children - return cls( - y=y, b=b, weights=weights, costs=costs, y_fused=y_fused, **aux_data - ) diff --git a/ott/build/lib/ott/problems/quadratic/quadratic_costs.py b/ott/build/lib/ott/problems/quadratic/quadratic_costs.py deleted file mode 100644 index 70f2bf5..0000000 --- a/ott/build/lib/ott/problems/quadratic/quadratic_costs.py +++ /dev/null @@ -1,78 +0,0 @@ -# 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. -from typing import Callable, NamedTuple - -import jax.numpy as jnp -import jax.scipy as jsp - -__all__ = ["make_square_loss", "make_kl_loss"] - - -class Loss(NamedTuple): # noqa: D101 - func: Callable[[jnp.ndarray], jnp.ndarray] - is_linear: bool - - -class GWLoss(NamedTuple): - r"""Efficient decomposition of the Gromov-Wasserstein loss function. - - The loss function :math:`L` is assumed to match the form given in eq. 5. of - :cite:`peyre:16`: - - .. math:: - L(x, y) = f_1(x) + f_2(y) - h_1(x) h_2(y) - - Args: - f1: First linear term. - f2: Second linear term. - h1: First quadratic term. - h2: Second quadratic term. - """ - f1: Loss - f2: Loss - h1: Loss - h2: Loss - - -def make_square_loss() -> GWLoss: - """Squared Euclidean loss for Gromov-Wasserstein. - - See Prop. 1 and Remark 1 of :cite:`peyre:16` for more information. - - Returns: - The squared Euclidean loss. - """ - f1 = Loss(lambda x: x ** 2, is_linear=False) - f2 = Loss(lambda y: y ** 2, is_linear=False) - h1 = Loss(lambda x: x, is_linear=True) - h2 = Loss(lambda y: 2.0 * y, is_linear=True) - return GWLoss(f1, f2, h1, h2) - - -def make_kl_loss(clipping_value: float = 1e-8) -> GWLoss: - r"""Kullback-Leibler loss for Gromov-Wasserstein. - - See Prop. 1 and Remark 1 of :cite:`peyre:16` for more information. - - Args: - clipping_value: Value used to avoid :math:`\log(0)`. - - Returns: - The KL loss. - """ - f1 = Loss(lambda x: -jsp.special.entr(x) - x, is_linear=False) - f2 = Loss(lambda y: y, is_linear=True) - h1 = Loss(lambda x: x, is_linear=True) - h2 = Loss(lambda y: jnp.log(jnp.clip(y, clipping_value)), is_linear=False) - return GWLoss(f1, f2, h1, h2) diff --git a/ott/build/lib/ott/problems/quadratic/quadratic_problem.py b/ott/build/lib/ott/problems/quadratic/quadratic_problem.py deleted file mode 100644 index 170fb3a..0000000 --- a/ott/build/lib/ott/problems/quadratic/quadratic_problem.py +++ /dev/null @@ -1,588 +0,0 @@ -# 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. -from typing import TYPE_CHECKING, Literal, Optional, Tuple, Union - -import jax -import jax.numpy as jnp -import jax.scipy as jsp - -from ott import utils -from ott.geometry import epsilon_scheduler, geometry, low_rank, pointcloud -from ott.problems.linear import linear_problem -from ott.problems.quadratic import quadratic_costs -from ott.types import Transport - -if TYPE_CHECKING: - from ott.solvers.linear import sinkhorn_lr - -__all__ = ["QuadraticProblem"] - - -@jax.tree_util.register_pytree_node_class -class QuadraticProblem: - r"""Quadratic OT problem. - - The quadratic loss of a single OT matrix is assumed to - have the form given in :cite:`peyre:16`, eq. 4. - - The two geometries below parameterize matrices :math:`C` and :math:`\bar{C}` - in that equation. The function :math:`L` (of two real values) in that equation - is assumed to match the form given in eq. 5., with our notations: - - .. math:: - L(x, y) = f_1(x) + f_2(y) - h_1(x) h_2(y) - - Args: - geom_xx: Ground geometry of the first space. - geom_yy: Ground geometry of the second space. - geom_xy: Geometry defining the linear penalty term for - fused Gromov-Wasserstein :cite:`vayer:19`. If :obj:`None`, the problem - reduces to a plain Gromov-Wasserstein problem :cite:`peyre:16`. - fused_penalty: Multiplier of the linear term in fused Gromov-Wasserstein, - i.e. ``problem = purely quadratic + fused_penalty * linear problem``. - scale_cost: option to rescale the cost matrices: - - - if :obj:`True`, use the default for each geometry. - - if :obj:`False`, keep the original scaling in geometries. - - if :class:`str`, use a specific method available in - :class:`~ott.geometry.geometry.Geometry` or - :class:`~ott.geometry.pointcloud.PointCloud`. - - if :obj:`None`, do not scale the cost matrices. - - a: The first marginal. If :obj:`None`, it will be uniform. - b: The second marginal. If :obj:`None`, it will be uniform. - loss: Gromov-Wasserstein loss function, see - :class:`~ott.problems.quadratic.quadratic_costs.GWLoss` for more - information. - tau_a: If :math:`< 1.0`, defines how much unbalanced the problem is on - the first marginal. - tau_b: If :math:`< 1.0`, defines how much unbalanced the problem is on - the second marginal. - gw_unbalanced_correction: Whether the unbalanced version of - :cite:`sejourne:21` is used. Otherwise, ``tau_a`` and ``tau_b`` - only affect the inner Sinkhorn loop. - ranks: Ranks of the cost matrices, see - :meth:`~ott.geometry.geometry.Geometry.to_LRCGeometry`. Used when - geometries are *not* :class:`~ott.geometry.pointcloud.PointCloud` with - `'sqeucl'` cost function. If `-1`, the geometries will not be converted - to low-rank. If :class:`tuple`, it specifies the ranks of ``geom_xx``, - ``geom_yy`` and ``geom_xy``, respectively. If :class:`int`, rank is shared - across all geometries. - tolerances: Tolerances used when converting geometries to low-rank. Used - when geometries are not :class:`~ott.geometry.pointcloud.PointCloud` with - `'sqeucl'` cost. If :class:`float`, it is shared across all geometries. - """ - - def __init__( - self, - geom_xx: geometry.Geometry, - geom_yy: geometry.Geometry, - geom_xy: Optional[geometry.Geometry] = None, - fused_penalty: float = 1.0, - scale_cost: Optional[Union[bool, float, str]] = False, - a: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, - loss: Union[Literal["sqeucl", "kl"], quadratic_costs.GWLoss] = "sqeucl", - tau_a: float = 1.0, - tau_b: float = 1.0, - gw_unbalanced_correction: bool = True, - ranks: Union[int, Tuple[int, ...]] = -1, - tolerances: Union[float, Tuple[float, ...]] = 1e-2, - labels_a: Optional[jnp.ndarray] = None, - labels_b: Optional[jnp.ndarray] = None, - n_labels: Optional[jnp.ndarray] = None, - block_diag_mat: Optional[jnp.ndarray] = None, - ): - self._geom_xx = geom_xx.set_scale_cost(scale_cost) - self._geom_yy = geom_yy.set_scale_cost(scale_cost) - self._geom_xy = ( - None if geom_xy is None else geom_xy.set_scale_cost(scale_cost) - ) - self.fused_penalty = fused_penalty - self.scale_cost = scale_cost - self._a = a - self._b = b - self.tau_a = tau_a - self.tau_b = tau_b - self.gw_unbalanced_correction = gw_unbalanced_correction - self.ranks = ranks - self.tolerances = tolerances - - self._loss_name = loss - if self._loss_name == "sqeucl": - self.loss = quadratic_costs.make_square_loss() - elif loss == "kl": - self.loss = quadratic_costs.make_kl_loss() - else: - self.loss = loss - self.labels_a = labels_a - self.labels_b = labels_b - self.n_labels = n_labels - self.block_diag_mat = block_diag_mat - - def marginal_dependent_cost_labeled( - self, - marginal_1: jnp.ndarray, - marginal_2: jnp.ndarray, - labels_1: jnp.ndarray, - labels_2: jnp.ndarray, - ) -> low_rank.LRCGeometry: - r"""Initialize cost term that depends on the marginals of the transport. - - Uses the first term in eq. 6, p. 1 of :cite:`peyre:16`. - - Let :math:`p` be the `[n,]` marginal of the transport matrix for samples - from :attr:`geom_xx` and :math:`q` the `[m,]` marginal of the - transport matrix for samples from :attr:`geom_yy`. - - When ``cost_xx`` (resp. ``cost_yy``) is the cost matrix of :attr:`geom_xx` - (resp. :attr:`geom_yy`), the cost term that depends on these marginals can - be written as: - - .. math:: - - \text{marginal_dep_term} = \text{lin1}(\text{cost_xx}) p \mathbb{1}_{m}^T - + \mathbb{1}_{n}(\text{lin2}(\text{cost_yy}) q)^T - - This helper function instantiates these two low-rank matrices and groups - them into a single low-rank cost geometry object. - - Args: - marginal_1: [n,], first marginal of transport matrix. - marginal_2: [m,], second marginal of transport matrix. - - Returns: - Low-rank geometry of rank 2, storing normalization constants. - """ - geom_xx, geom_yy = self.geom_xx, self.geom_yy - tmp1 = None - tmp2 = None - if self._loss_name == "sqeucl": # quadratic apply, efficient for LR - for l in jnp.unique(labels_1, size = self.n_labels): - tmp1_l = geom_xx.apply_square_cost(marginal_1, axis=1, sub_idx = labels_1 == l) - tmp2_l = geom_yy.apply_square_cost(marginal_2, axis=1, sub_idx = labels_2 == l) - if tmp1 is None: tmp1 = tmp1_l - else: tmp1 += tmp1_l - if tmp2 is None: tmp2 = tmp2_l - else: tmp2 += tmp2_l - else: - f1, f2 = self.linear_loss - for l in jnp.unique(labels_1, size = self.n_labels): - tmp1_l = apply_cost(geom_xx, marginal_1, axis=1, fn=f1, sub_idx = labels_1 == l) - tmp2_l = apply_cost(geom_yy, marginal_2, axis=1, fn=f2, sub_idx = labels_2 == l) - if tmp1 is None: tmp1 = tmp1_l - else: tmp1 += tmp1_l - if tmp2 is None: tmp2 = tmp2_l - else: tmp2 += tmp2_l - x_term = jnp.concatenate((tmp1, jnp.ones_like(tmp1)), axis=1) - y_term = jnp.concatenate((jnp.ones_like(tmp2), tmp2), axis=1) - return low_rank.LRCGeometry(cost_1=x_term, cost_2=y_term) - - def marginal_dependent_cost( - self, - marginal_1: jnp.ndarray, - marginal_2: jnp.ndarray, - ) -> low_rank.LRCGeometry: - r"""Initialize cost term that depends on the marginals of the transport. - - Uses the first term in eq. 6, p. 1 of :cite:`peyre:16`. - - Let :math:`p` be the `[n,]` marginal of the transport matrix for samples - from :attr:`geom_xx` and :math:`q` the `[m,]` marginal of the - transport matrix for samples from :attr:`geom_yy`. - - When ``cost_xx`` (resp. ``cost_yy``) is the cost matrix of :attr:`geom_xx` - (resp. :attr:`geom_yy`), the cost term that depends on these marginals can - be written as: - - .. math:: - - \text{marginal_dep_term} = \text{lin1}(\text{cost_xx}) p \mathbb{1}_{m}^T - + \mathbb{1}_{n}(\text{lin2}(\text{cost_yy}) q)^T - - This helper function instantiates these two low-rank matrices and groups - them into a single low-rank cost geometry object. - - Args: - marginal_1: [n,], first marginal of transport matrix. - marginal_2: [m,], second marginal of transport matrix. - - Returns: - Low-rank geometry of rank 2, storing normalization constants. - """ - geom_xx, geom_yy = self.geom_xx, self.geom_yy - if self._loss_name == "sqeucl": # quadratic apply, efficient for LR - tmp1 = geom_xx.apply_square_cost(marginal_1, axis=1) - tmp2 = geom_yy.apply_square_cost(marginal_2, axis=1) - else: - f1, f2 = self.linear_loss - tmp1 = apply_cost(geom_xx, marginal_1, axis=1, fn=f1) - tmp2 = apply_cost(geom_yy, marginal_2, axis=1, fn=f2) - x_term = jnp.concatenate((tmp1, jnp.ones_like(tmp1)), axis=1) - y_term = jnp.concatenate((jnp.ones_like(tmp2), tmp2), axis=1) - return low_rank.LRCGeometry(cost_1=x_term, cost_2=y_term) - - def cost_unbalanced_correction( - self, - transport_matrix: jnp.ndarray, - marginal_1: jnp.ndarray, - marginal_2: jnp.ndarray, - epsilon: epsilon_scheduler.Epsilon, - ) -> float: - r"""Calculate cost term from the quadratic divergence when unbalanced. - - In the unbalanced setting (``tau_a < 1.0 or tau_b < 1.0``), the - introduction of a quadratic divergence :cite:`sejourne:21` adds a term - to the GW local cost. - - Let :math:`a` [num_a,] be the target weights for samples - from geom_xx and :math:`b` [num_b,] be the target weights - for samples from `geom_yy`. Let :math:`P` [num_a, num_b] be the transport - matrix, :math:`P1` the first marginal and :math:`P^T1` the second marginal. - The term of the cost matrix coming from the quadratic KL in the - unbalanced case can be written as: - - `unbalanced_correction_term` = - :math:`tau_a / (1 - tau_a) * \sum(KL(P1|a))` - :math:`+ tau_b / (1 - tau_b) * \sum(KL(P^T1|b))` - :math:`+ epsilon * \sum(KL(P|ab'))` - - Args: - transport_matrix: jnp.ndarray[num_a, num_b], transport matrix. - marginal_1: jnp.ndarray[num_a,], marginal of the transport matrix - for samples from :attr:`geom_xx`. - marginal_2: jnp.ndarray[num_b,], marginal of the transport matrix - for samples from :attr:`geom_yy`. - epsilon: entropy regularizer. - - Returns: - The cost term. - """ - - def regularizer(tau: float) -> float: - return eps * tau / (1.0 - tau) - - eps = epsilon._target_init - marginal_1loga = jsp.special.xlogy(marginal_1, self.a).sum() - marginal_2logb = jsp.special.xlogy(marginal_2, self.b).sum() - - cost = eps * jsp.special.xlogy(transport_matrix, transport_matrix).sum() - if self.tau_a != 1.0: - cost += regularizer( - self.tau_a - ) * (-jsp.special.entr(marginal_1).sum() - marginal_1loga) - if self.tau_b != 1.0: - cost += regularizer( - self.tau_b - ) * (-jsp.special.entr(marginal_2).sum() - marginal_2logb) - return cost - - # TODO(michalk8): highly coupled to the pre-defined initializer, refactor - def init_transport_mass(self) -> float: - """Initialize the transport mass. - - Returns: - The sum of the elements of the normalized transport matrix. - """ - a = jax.lax.stop_gradient(self.a) - b = jax.lax.stop_gradient(self.b) - return a.sum() * b.sum() - - def update_lr_geom( - self, - lr_sink: "sinkhorn_lr.LRSinkhornOutput", - relative_epsilon: Optional[bool] = None, - ) -> geometry.Geometry: - """Recompute (possibly LRC) linearization using LR Sinkhorn output.""" - marginal_1 = lr_sink.marginal(1) - marginal_2 = lr_sink.marginal(0) - marginal_cost = self.marginal_dependent_cost(marginal_1, marginal_2) - - # Extract factors from LR Sinkhorn output - q, r, inv_sqg = lr_sink.q, lr_sink.r, 1.0 / jnp.sqrt(lr_sink.g) - # Distribute middle marginal evenly across both factors. - q, r = q * inv_sqg[None, :], r * inv_sqg[None, :] - - # Handle LRC Geometry case. - h1, h2 = self.quad_loss - geom_xx, geom_yy, geom_xy = self.geom_xx, self.geom_yy, self.geom_xy - tmp1 = apply_cost(geom_xx, q, axis=1, fn=h1) - tmp2 = apply_cost(geom_yy, r, axis=1, fn=h2) - if self.is_low_rank: - geom = low_rank.LRCGeometry( - cost_1=tmp1, cost_2=-tmp2, relative_epsilon=relative_epsilon - ) + marginal_cost - if self.is_fused: - geom = geom + geom_xy - else: - cost_matrix = marginal_cost.cost_matrix - jnp.dot(tmp1, tmp2.T) - cost_matrix += self.fused_penalty * self._fused_cost_matrix - geom = geometry.Geometry( - cost_matrix=cost_matrix, relative_epsilon=relative_epsilon - ) - return geom # noqa: RET504 - - def update_linearization( - self, - transport: Transport, - epsilon: Optional[Union[epsilon_scheduler.Epsilon, float]] = None, - old_transport_mass: float = 1.0, - relative_epsilon: Optional[bool] = None, - ) -> linear_problem.LinearProblem: - """Update linearization of GW problem by updating cost matrix. - - If the problem is balanced (``tau_a = 1.0 and tau_b = 1.0``), the equation - follows eq. 6, p. 1 of :cite:`peyre:16`. - - If the problem is unbalanced (``tau_a < 1.0 or tau_b < 1.0``), two cases are - possible, as explained in :meth:`init_linearization` above. - Finally, it is also possible to consider a Fused Gromov-Wasserstein problem. - Details about the resulting cost matrix are also given in - :meth:`init_linearization`. - - Args: - transport: Solution of the linearization of the quadratic problem. - epsilon: An epsilon scheduler or a float passed on to the linearization. - old_transport_mass: Sum of the elements of the transport matrix at the - previous iteration. - relative_epsilon: Whether to use relative epsilon in the linearized - geometry. - - Returns: - Updated linear OT problem, a new local linearization of GW problem. - """ - rescale_factor = 1.0 - unbalanced_correction = 0.0 - print("updating linearization") - if not self.is_balanced: - marginal_1 = transport.marginal(axis=1) - transport_mass = jax.lax.stop_gradient(marginal_1.sum()) - rescale_factor = jnp.sqrt(old_transport_mass / transport_mass) - - marginal_1 = transport.marginal(axis=1) * rescale_factor - marginal_2 = transport.marginal(axis=0) * rescale_factor - if self.labels_a is None: - marginal_cost = self.marginal_dependent_cost(marginal_1, marginal_2) - else: - marginal_cost = self.marginal_dependent_cost_labeled(marginal_1, marginal_2, - self.labels_a, self.labels_b) - - transport_matrix = transport.matrix * rescale_factor - if not self.is_balanced: - # Rescales transport for Unbalanced GW according to Sejourne et al. (2021) - transport_mass = jax.lax.stop_gradient(marginal_1.sum()) - epsilon = update_epsilon_unbalanced(epsilon, transport_mass) - unbalanced_correction = self.cost_unbalanced_correction( - transport_matrix, marginal_1, marginal_2, epsilon - ) - - h1, h2 = self.quad_loss - geom_xx, geom_yy = self.geom_xx, self.geom_yy - - tmp = apply_cost(geom_xx, transport_matrix, axis=1, fn=h1) - tmp = apply_cost(geom_yy, tmp.T, axis=1, fn=h2).T - - cost_matrix = marginal_cost.cost_matrix - tmp + unbalanced_correction - cost_matrix += self.fused_penalty * rescale_factor * self._fused_cost_matrix - if self.labels_a is not None: - print("Label considered for Sinkhorn run") - cost_matrix = cost_matrix + (1 / self.block_diag_mat -1) - geom = geometry.Geometry( - cost_matrix=cost_matrix, - epsilon=epsilon, - relative_epsilon=relative_epsilon, - ) - - return linear_problem.LinearProblem( - geom, self.a, self.b, tau_a=self.tau_a, tau_b=self.tau_b, - labels_a = self.labels_a, labels_b = self.labels_b - ) - - def update_lr_linearization( - self, - lr_sink: "sinkhorn_lr.LRSinkhornOutput", - *, - relative_epsilon: Optional[bool] = None, - ) -> linear_problem.LinearProblem: - """Update a Quad problem linearization using a LR Sinkhorn.""" - return linear_problem.LinearProblem( - self.update_lr_geom(lr_sink, relative_epsilon=relative_epsilon), - self.a, - self.b, - tau_a=self.tau_a, - tau_b=self.tau_b - ) - - @property - def _fused_cost_matrix(self) -> Union[float, jnp.ndarray]: - if not self.is_fused: - return 0.0 - geom_xy = self.geom_xy - - if isinstance(geom_xy, pointcloud.PointCloud) and geom_xy.is_online: - return geom_xy._compute_cost_matrix() * geom_xy.inv_scale_cost - return geom_xy.cost_matrix - - @property - def _is_low_rank_convertible(self) -> bool: - - def convertible(geom: geometry.Geometry) -> bool: - return isinstance(geom, low_rank.LRCGeometry) or ( - isinstance(geom, pointcloud.PointCloud) and geom.is_squared_euclidean - ) - - if self.is_low_rank: - return True - - geom_xx, geom_yy, geom_xy = self.geom_xx, self.geom_yy, self.geom_xy - # either explicitly via cost factorization or implicitly (e.g., a PC) - return self.ranks != -1 or ( - convertible(geom_xx) and convertible(geom_yy) and - (geom_xy is None or convertible(geom_xy)) - ) - - def to_low_rank( - self, - rng: Optional[jax.Array] = None, - ) -> "QuadraticProblem": - """Convert geometries to low-rank. - - Args: - rng: Random key for seeding. - - Returns: - Quadratic problem with low-rank geometries. - """ - - def convert( - vals: Union[int, float, Tuple[Union[int, float], ...]] - ) -> Tuple[Union[int, float], ...]: - size = 2 + self.is_fused - if isinstance(vals, (int, float)): - return (vals,) * 3 - assert len(vals) == size, vals - return vals + (None,) * (3 - size) - - if self.is_low_rank: - return self - - rng = utils.default_prng_key(rng) - rng1, rng2, rng3 = jax.random.split(rng, 3) - (geom_xx, geom_yy, geom_xy, *children), aux_data = self.tree_flatten() - (r1, r2, r3), (t1, t2, t3) = convert(self.ranks), convert(self.tolerances) - - geom_xx = geom_xx.to_LRCGeometry(rank=r1, tol=t1, rng=rng1) - geom_yy = geom_yy.to_LRCGeometry(rank=r2, tol=t2, rng=rng2) - if self.is_fused: - if isinstance( - geom_xy, pointcloud.PointCloud - ) and geom_xy.is_squared_euclidean: - geom_xy = geom_xy.to_LRCGeometry(scale=self.fused_penalty) - else: - geom_xy = geom_xy.to_LRCGeometry( - rank=r3, tol=t3, rng=rng3, scale=self.fused_penalty - ) - - return type(self).tree_unflatten( - aux_data, [geom_xx, geom_yy, geom_xy] + children - ) - - @property - def geom_xx(self) -> geometry.Geometry: - """Geometry of the first space.""" - return self._geom_xx - - @property - def geom_yy(self) -> geometry.Geometry: - """Geometry of the second space.""" - return self._geom_yy - - @property - def geom_xy(self) -> Optional[geometry.Geometry]: - """Geometry of the joint space.""" - return self._geom_xy - - @property - def a(self) -> jnp.ndarray: - """First marginal.""" - num_a = self.geom_xx.shape[0] - return jnp.ones((num_a,)) / num_a if self._a is None else self._a - - @property - def b(self) -> jnp.ndarray: - """Second marginal.""" - num_b = self.geom_yy.shape[0] - return jnp.ones((num_b,)) / num_b if self._b is None else self._b - - @property - def is_fused(self) -> bool: - """Whether the problem is fused.""" - return self.geom_xy is not None - - @property - def is_low_rank(self) -> bool: - """Whether all geometries are low-rank.""" - return ( - isinstance(self.geom_xx, low_rank.LRCGeometry) and - isinstance(self.geom_yy, low_rank.LRCGeometry) and - (not self.is_fused or isinstance(self.geom_xy, low_rank.LRCGeometry)) - ) - - @property - def linear_loss(self) -> Tuple[quadratic_costs.Loss, quadratic_costs.Loss]: - """Linear part of the Gromov-Wasserstein loss.""" - return self.loss.f1, self.loss.f2 - - @property - def quad_loss(self) -> Tuple[quadratic_costs.Loss, quadratic_costs.Loss]: - """Quadratic part of the Gromov-Wasserstein loss.""" - return self.loss.h1, self.loss.h2 - - @property - def is_balanced(self) -> bool: - """Whether the problem is balanced.""" - return ((not self.gw_unbalanced_correction) or - (self.tau_a == 1.0 and self.tau_b == 1.0)) - - def tree_flatten(self): # noqa: D102 - return ([self.geom_xx, self.geom_yy, self.geom_xy, self._a, self._b], { - "tau_a": self.tau_a, - "tau_b": self.tau_b, - "loss": self._loss_name, - "fused_penalty": self.fused_penalty, - "scale_cost": self.scale_cost, - "gw_unbalanced_correction": self.gw_unbalanced_correction, - "ranks": self.ranks, - "tolerances": self.tolerances - }) - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - geoms, (a, b) = children[:3], children[3:] - return cls(*geoms, a=a, b=b, **aux_data) - - -def update_epsilon_unbalanced( # noqa: D103 - epsilon: Union[float, epsilon_scheduler.Epsilon], transport_mass: float -) -> epsilon_scheduler.Epsilon: - if not isinstance(epsilon, epsilon_scheduler.Epsilon): - epsilon = epsilon_scheduler.Epsilon(epsilon, scale_epsilon=1.0) - return epsilon.set(scale_epsilon=epsilon._scale_epsilon * transport_mass) - - -def apply_cost( # noqa: D103 - geom: geometry.Geometry, arr: jnp.ndarray, *, axis: int, - fn: quadratic_costs.Loss -) -> jnp.ndarray: - return geom.apply_cost(arr, axis=axis, fn=fn.func, is_linear=fn.is_linear) diff --git a/ott/build/lib/ott/py.typed b/ott/build/lib/ott/py.typed deleted file mode 100755 index e69de29..0000000 diff --git a/ott/build/lib/ott/solvers/__init__.py b/ott/build/lib/ott/solvers/__init__.py deleted file mode 100644 index 1303312..0000000 --- a/ott/build/lib/ott/solvers/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. -from . import linear, quadratic, was_solver diff --git a/ott/build/lib/ott/solvers/linear/__init__.py b/ott/build/lib/ott/solvers/linear/__init__.py deleted file mode 100644 index 4b7c246..0000000 --- a/ott/build/lib/ott/solvers/linear/__init__.py +++ /dev/null @@ -1,30 +0,0 @@ -# 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. -from . import ( - acceleration, - continuous_barycenter, - discrete_barycenter, - implicit_differentiation, - lr_utils, - sinkhorn, - sinkhorn_lr, - univariate, -) -from ._solve import solve - -__all__ = [ - "acceleration", "continuous_barycenter", "discrete_barycenter", - "implicit_differentiation", "lr_utils", "sinkhorn", "sinkhorn_lr", "solve", - "univariate" -] diff --git a/ott/build/lib/ott/solvers/linear/_solve.py b/ott/build/lib/ott/solvers/linear/_solve.py deleted file mode 100644 index 2bca6a8..0000000 --- a/ott/build/lib/ott/solvers/linear/_solve.py +++ /dev/null @@ -1,60 +0,0 @@ -# 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. -from typing import Any, Optional, Union - -import jax.numpy as jnp - -from ott.geometry import geometry -from ott.problems.linear import linear_problem -from ott.solvers.linear import sinkhorn, sinkhorn_lr - -__all__ = ["solve"] - - -def solve( - geom: geometry.Geometry, - a: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, - tau_a: float = 1.0, - tau_b: float = 1.0, - rank: int = -1, - **kwargs: Any -) -> Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput]: - """Solve linear regularized OT problem using Sinkhorn iterations. - - Args: - geom: The ground geometry of the linear problem. - a: The first marginal. If :obj:`None`, it will be uniform. - b: The second marginal. If :obj:`None`, it will be uniform. - tau_a: If :math:`< 1`, defines how much unbalanced the problem is - on the first marginal. - tau_b: If :math:`< 1`, defines how much unbalanced the problem is - on the second marginal. - rank: - Rank constraint on the coupling to minimize the linear OT problem - :cite:`scetbon:21`. If :math:`-1`, no rank constraint is used. - kwargs: Keyword arguments for - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` or - :class:`~ott.solvers.linear.sinkhorn_lr.LRSinkhorn`, - depending on the ``rank``. - - Returns: - The Sinkhorn output. - """ - prob = linear_problem.LinearProblem(geom, a=a, b=b, tau_a=tau_a, tau_b=tau_b) - if rank > 0: - solver = sinkhorn_lr.LRSinkhorn(rank=rank, **kwargs) - else: - solver = sinkhorn.Sinkhorn(**kwargs) - return solver(prob) diff --git a/ott/build/lib/ott/solvers/linear/acceleration.py b/ott/build/lib/ott/solvers/linear/acceleration.py deleted file mode 100644 index 159322d..0000000 --- a/ott/build/lib/ott/solvers/linear/acceleration.py +++ /dev/null @@ -1,172 +0,0 @@ -# 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. -from typing import TYPE_CHECKING - -import jax -import jax.numpy as jnp - -from ott import utils - -if TYPE_CHECKING: - from ott.problems.linear import linear_problem - from ott.solvers.linear import sinkhorn - -__all__ = ["AndersonAcceleration", "Momentum"] - - -@utils.register_pytree_node -class AndersonAcceleration: - """Implements Anderson acceleration for Sinkhorn.""" - - # TODO(michalk8): use memory=0 as no Anderson acceleration? - memory: int = 2 # Number of iterates considered to form interpolation. - refresh_every: int = 1 # Recompute interpolation periodically. - ridge_identity: float = 1e-2 # Ridge used in the linear system. - - def extrapolation(self, xs: jnp.ndarray, fxs: jnp.ndarray) -> jnp.ndarray: - """Compute Anderson extrapolation from past observations.""" - # Remove -inf values to instantiate quadratic problem. All others - # remain since they might be caused by a valid issue. - fxs_clean = jnp.nan_to_num(fxs, nan=jnp.nan, posinf=jnp.inf, neginf=0.0) - xs_clean = jnp.nan_to_num(xs, nan=jnp.nan, posinf=jnp.inf, neginf=0.0) - residuals = fxs_clean - xs_clean - gram_matrix = jnp.matmul(residuals.T, residuals) - gram_matrix /= jnp.linalg.norm(gram_matrix) - - # Solve linear system to obtain weights - weights = jax.scipy.sparse.linalg.cg( - gram_matrix + self.ridge_identity * jnp.eye(xs.shape[1]), - jnp.ones(xs.shape[1]) - )[0] - weights /= jnp.sum(weights) - - # Recover linear combination and return it with NaN (caused - # by 0 weights leading to -jnp.inf potentials, mixed with weights - # coefficients of different signs), disambiguated to -inf. - combination = jnp.sum(fxs * weights[None, :], axis=1) - return jnp.where(jnp.isfinite(combination), combination, -jnp.inf) - - def update( - self, state: "sinkhorn.SinkhornState", iteration: int, - prob: "linear_problem.LinearProblem", lse_mode: bool - ) -> "sinkhorn.SinkhornState": - """Anderson acceleration update. - - When using Anderson acceleration, first update the dual variable f_u with - previous updates (if iteration count sufficiently large), then record - new iterations in array. - - Anderson acceleration always happens in potentials (not scalings) space, - regardless of the lse_mode setting. If the iteration count is large - enough the update below will output a potential variable. - - Args: - state: Sinkhorn state. - iteration: the current iteration. - prob: linear OT problem. - lse_mode: whether to compute in log-sum-exp or in scalings. - - Returns: - A potential variable. - """ - geom = prob.geom - trigger_update = jnp.logical_and( - iteration > self.memory, iteration % self.refresh_every == 0 - ) - fu = jnp.where( - trigger_update, self.extrapolation(state.old_fus, state.old_mapped_fus), - state.fu - ) - # If the interpolation was triggered, we store it in memory - # Otherwise we add the previous value (converting it to potential form if - # it was initially stored in scaling form). - old_fus = jnp.where( - trigger_update, - jnp.concatenate((state.old_fus[:, 1:], fu[:, None]), axis=1), - jnp.concatenate(( - state.old_fus[:, 1:], - (fu if lse_mode else geom.potential_from_scaling(fu))[:, None] - ), - axis=1) - ) - - # If update was triggered, ensure a scaling is returned, since the result - # from the extrapolation was outputted in potential form. - fu = jnp.where( - trigger_update, fu if lse_mode else geom.scaling_from_potential(fu), fu - ) - return state.set(potentials=(fu, state.gv), old_fus=old_fus) - - def init_maps( - self, pb, state: "sinkhorn.SinkhornState" - ) -> "sinkhorn.SinkhornState": - """Initialize log matrix used in Anderson acceleration with *NaN* values.""" - fus = jnp.ones((pb.geom.shape[0], self.memory)) * jnp.nan - return state.set(old_fus=fus, old_mapped_fus=fus) - - def update_history( - self, state: "sinkhorn.SinkhornState", pb, lse_mode: bool - ) -> "sinkhorn.SinkhornState": - """Update history of mapped dual variables.""" - f = state.fu if lse_mode else pb.geom.potential_from_scaling(state.fu) - mapped = jnp.concatenate((state.old_mapped_fus[:, 1:], f[:, None]), axis=1) - return state.set(old_mapped_fus=mapped) - - -@utils.register_pytree_node -class Momentum: - """Momentum for Sinkhorn updates. - - Can be either constant :cite:`thibault:21` or adaptive :cite:`lehmann:21`. - """ - - start: int = 0 - error_threshold: float = jnp.inf - value: float = 1.0 - inner_iterations: int = 1 - - def weight(self, state: "sinkhorn.SinkhornState", iteration: int) -> float: - """Compute momentum term if needed, using previously seen errors.""" - if self.start == 0: - return self.value - idx = self.start // self.inner_iterations - - return jax.lax.cond( - jnp.logical_and( - iteration >= self.start, state.errors[idx - 1, -1] - < self.error_threshold - ), lambda state: self.lehmann(state), lambda state: self.value, state - ) - - def lehmann(self, state: "sinkhorn.SinkhornState") -> float: - """Momentum formula :cite:`lehmann:21`, eq. 5.""" - idx = self.start // self.inner_iterations - error_ratio = jnp.minimum( - state.errors[idx - 1, -1] / state.errors[idx - 2, -1], 0.99 - ) - power = 1.0 / self.inner_iterations - return 2.0 / (1.0 + jnp.sqrt(1.0 - error_ratio ** power)) - - def __call__( # noqa: D102 - self, - weight: float, - value: jnp.ndarray, - new_value: jnp.ndarray, - lse_mode: bool = True - ) -> jnp.ndarray: - if lse_mode: - value = jnp.where(jnp.isfinite(value), value, 0.0) - return (1.0 - weight) * value + weight * new_value - value = jnp.where(value > 0.0, value, 1.0) - return value ** (1.0 - weight) * new_value ** weight diff --git a/ott/build/lib/ott/solvers/linear/continuous_barycenter.py b/ott/build/lib/ott/solvers/linear/continuous_barycenter.py deleted file mode 100644 index 8f60380..0000000 --- a/ott/build/lib/ott/solvers/linear/continuous_barycenter.py +++ /dev/null @@ -1,233 +0,0 @@ -# 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 functools -from typing import Any, NamedTuple, Optional, Tuple - -import jax -import jax.numpy as jnp - -from ott import utils -from ott.geometry import pointcloud -from ott.math import fixed_point_loop -from ott.math import utils as mu -from ott.problems.linear import barycenter_problem, linear_problem -from ott.solvers import was_solver - -__all__ = ["FreeBarycenterState", "FreeWassersteinBarycenter"] - - -class FreeBarycenterState(NamedTuple): - """Holds the state of the Wasserstein barycenter solver. - - Args: - costs: Holds the sequence of regularized GW costs seen through the outer - loop of the solver. - linear_convergence: Holds the sequence of bool convergence flags of the - inner Sinkhorn iterations. - errors: Holds sequence of vectors of errors of the Sinkhorn algorithm - at each iteration. - x: barycenter points. - a: barycenter weights. - """ - - costs: Optional[jnp.ndarray] = None - linear_convergence: Optional[jnp.ndarray] = None - errors: Optional[jnp.ndarray] = None - x: Optional[jnp.ndarray] = None - a: Optional[jnp.ndarray] = None - - def set(self, **kwargs: Any) -> "FreeBarycenterState": - """Return a copy of self, possibly with overwrites.""" - return self._replace(**kwargs) - - def update( - self, iteration: int, bar_prob: barycenter_problem.FreeBarycenterProblem, - linear_ot_solver: Any, store_errors: bool - ) -> "FreeBarycenterState": - """Update the state of the solver. - - Args: - iteration: the current iteration of the outer loop. - bar_prob: the barycenter problem. - linear_ot_solver: the linear OT solver to use. - store_errors: whether to store the errors of the inner loop. - - Returns: - The updated state. - """ - seg_y, seg_b = bar_prob.segmented_y_b - - @functools.partial(jax.vmap, in_axes=[None, None, 0, 0]) - def solve_linear_ot( - a: Optional[jnp.ndarray], x: jnp.ndarray, b: jnp.ndarray, y: jnp.ndarray - ): - out = linear_ot_solver( - linear_problem.LinearProblem( - pointcloud.PointCloud( - x, - y, - src_mask=a > 0.0, - tgt_mask=b > 0.0, - cost_fn=bar_prob.cost_fn, - epsilon=bar_prob.epsilon - ), a, b - ) - ) - return ( - out.reg_ot_cost, out.converged, out.matrix, - out.errors if store_errors else None - ) - - reg_ot_costs, convergeds, matrices, errors = solve_linear_ot( - self.a, self.x, seg_b, seg_y - ) - - cost = jnp.sum(reg_ot_costs * bar_prob.weights) - updated_costs = self.costs.at[iteration].set(cost) - converged = jnp.all(convergeds) - linear_convergence = self.linear_convergence.at[iteration].set(converged) - - if store_errors and self.errors is not None: - errors = self.errors.at[iteration, :, :].set(errors) - else: - errors = None - - # Approximation of barycenter as barycenter of barycenters per measure. - - barycenters_per_measure = mu.barycentric_projection( - matrices, seg_y, bar_prob.cost_fn - ) - - x_new = jax.vmap( - lambda w, y: bar_prob.cost_fn.barycenter(w, y)[0], in_axes=[None, 1] - )(bar_prob.weights, barycenters_per_measure) - - return self.set( - costs=updated_costs, - linear_convergence=linear_convergence, - errors=errors, - x=x_new - ) - - -@jax.tree_util.register_pytree_node_class -class FreeWassersteinBarycenter(was_solver.WassersteinSolver): - """Continuous Wasserstein barycenter solver :cite:`cuturi:14`.""" - - def __call__( # noqa: D102 - self, - bar_prob: barycenter_problem.FreeBarycenterProblem, - bar_size: int = 100, - x_init: Optional[jnp.ndarray] = None, - rng: Optional[jax.Array] = None, - ) -> FreeBarycenterState: - # TODO(michalk8): no reason for iterations to be outside this class - rng = utils.default_prng_key(rng) - return iterations(self, bar_size, bar_prob, x_init, rng) - - def init_state( - self, - bar_prob: barycenter_problem.FreeBarycenterProblem, - bar_size: int, - x_init: Optional[jnp.ndarray] = None, - rng: Optional[jax.Array] = None, - ) -> FreeBarycenterState: - """Initialize the state of the Wasserstein barycenter iterations. - - Args: - bar_prob: The barycenter problem. - bar_size: Size of the barycenter. - x_init: Initial barycenter estimate of shape ``[bar_size, ndim]``. - If `None`, ``bar_size`` points will be sampled from the input - measures according to their weights - :attr:`~ott.problems.linear.barycenter_problem.FreeBarycenterProblem.flattened_y`. - rng: Random key for seeding. - - Returns: - The initial barycenter state. - """ - if x_init is not None: - assert bar_size == x_init.shape[0] - x = x_init - else: - # sample randomly points in the support of the y measures - rng = utils.default_prng_key(rng) - indices_subset = jax.random.choice( - rng, - a=bar_prob.flattened_y.shape[0], - shape=(bar_size,), - replace=False, - p=bar_prob.flattened_b - ) - x = bar_prob.flattened_y[indices_subset, :] - - # TODO(cuturi) expand to non-uniform weights for barycenter. - a = jnp.ones((bar_size,)) / bar_size - num_iter = self.max_iterations - if self.store_inner_errors: - errors = -jnp.ones(( - num_iter, bar_prob.num_measures, - self.linear_ot_solver.outer_iterations - )) - else: - errors = None - return FreeBarycenterState( - -jnp.ones((num_iter,)), -jnp.ones((num_iter,)), errors, x, a - ) - - def output_from_state( # noqa: D102 - self, state: FreeBarycenterState - ) -> FreeBarycenterState: - # TODO(michalk8): create an output variable to match rest of the framework - return state - - -def iterations( - solver: FreeWassersteinBarycenter, bar_size: int, - bar_prob: barycenter_problem.FreeBarycenterProblem, x_init: jnp.ndarray, - rng: jax.Array -) -> FreeBarycenterState: - """Jittable Wasserstein barycenter outer loop.""" - - def cond_fn( - iteration: int, - constants: Tuple[FreeWassersteinBarycenter, - barycenter_problem.FreeBarycenterProblem], - state: FreeBarycenterState - ) -> bool: - solver, _ = constants - return solver._continue(state, iteration) - - def body_fn( - iteration, constants: Tuple[FreeWassersteinBarycenter, - barycenter_problem.FreeBarycenterProblem], - state: FreeBarycenterState, compute_error: bool - ) -> FreeBarycenterState: - del compute_error # Always assumed True - solver, bar_prob = constants - return state.update( - iteration, bar_prob, solver.linear_ot_solver, solver.store_inner_errors - ) - - state = fixed_point_loop.fixpoint_iter( - cond_fn=cond_fn, - body_fn=body_fn, - min_iterations=solver.min_iterations, - max_iterations=solver.max_iterations, - inner_iterations=1, - constants=(solver, bar_prob), - state=solver.init_state(bar_prob, bar_size, x_init, rng) - ) - - return solver.output_from_state(state) diff --git a/ott/build/lib/ott/solvers/linear/discrete_barycenter.py b/ott/build/lib/ott/solvers/linear/discrete_barycenter.py deleted file mode 100644 index dcfdc14..0000000 --- a/ott/build/lib/ott/solvers/linear/discrete_barycenter.py +++ /dev/null @@ -1,251 +0,0 @@ -# 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 functools -from typing import NamedTuple, Optional, Sequence - -import jax -import jax.numpy as jnp - -from ott.geometry import geometry -from ott.math import fixed_point_loop -from ott.problems.linear import barycenter_problem -from ott.solvers.linear import sinkhorn - -__all__ = ["SinkhornBarycenterOutput", "FixedBarycenter"] - - -class SinkhornBarycenterOutput(NamedTuple): # noqa: D101 - f: jnp.ndarray - g: jnp.ndarray - histogram: jnp.ndarray - errors: jnp.ndarray - - -@jax.tree_util.register_pytree_node_class -class FixedBarycenter: - """A Wasserstein barycenter solver for histograms on a common geometry. - - This solver uses a variant of the - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` algorithm proposed in - :cite:`janati:20a` to compute the barycenter of various measures supported on - the same (common to all) geometry. The geometry is assumed to be either - symmetric, or to describe costs between a set of points and another. In that - case all reference measures have support on the first measure, whereas the - barycenter is supported on the second. - - Args: - threshold: convergence threshold. The algorithm stops when the marginal - violations of all transport plans computed for that barycenter go below - that threshold. - norm_error: norm used to compute marginal deviation. - inner_iterations: number of iterations run before recomputing errors. - min_iterations: number of iterations run without checking whether - termination criterion is true. - max_iterations: maximal number of iterations. - lse_mode: sets computations in kernel (``False``) or log-sum-exp mode. - debiased: uses debiasing correction to avoid blur due to entropic - regularization. - """ - - def __init__( - self, - threshold: float = 1e-2, - norm_error: int = 1, - inner_iterations: float = 10, - min_iterations: int = 0, - max_iterations: int = 2000, - lse_mode: bool = True, - debiased: bool = False - ): - self.threshold = threshold - self.norm_error = norm_error - self.inner_iterations = inner_iterations - self.min_iterations = min_iterations - self.max_iterations = max_iterations - self.lse_mode = lse_mode - self.debiased = debiased - - def __call__( - self, - fixed_bp: barycenter_problem.FixedBarycenterProblem, - dual_initialization: Optional[jnp.ndarray] = None, - ) -> SinkhornBarycenterOutput: - """Solve barycenter problem, possibly using clever initialization. - - Args: - fixed_bp: Fixed barycenter problem. - dual_initialization: Initial value for the g_v potential/scalings, - one for each of the histograms described in ``fixed_bp``. If ``None``, - use initialization from :cite:`cuturi:15`, eq. 3.6. - - Returns: - The barycenter. - """ - geom = fixed_bp.geom - a = fixed_bp.a - num_a, num_b = geom.shape - - weights = fixed_bp.weights - - if dual_initialization is None: - # initialization strategy from :cite:`cuturi:15`, (3.6). - dual_initialization = geom.apply_cost(a.T, axis=0).T - dual_initialization -= jnp.average( - dual_initialization, weights=weights, axis=0 - )[jnp.newaxis, :] - - if self.debiased and not geom.is_symmetric: - raise ValueError("Geometry must be symmetric to use debiased option.") - norm_error = (self.norm_error,) - return _discrete_barycenter( - geom, a, weights, dual_initialization, self.threshold, norm_error, - self.inner_iterations, self.min_iterations, self.max_iterations, - self.lse_mode, self.debiased, num_a, num_b - ) - - def tree_flatten(self): # noqa: D102 - aux = vars(self).copy() - aux.pop("threshold") - return [ - self.threshold, - ], aux - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - return cls(**aux_data, threshold=children[0]) - - -@functools.partial(jax.jit, static_argnums=(5, 6, 7, 8, 9, 10, 11, 12)) -def _discrete_barycenter( - geom: geometry.Geometry, a: jnp.ndarray, weights: jnp.ndarray, - dual_initialization: jnp.ndarray, threshold: float, - norm_error: Sequence[int], inner_iterations: int, min_iterations: int, - max_iterations: int, lse_mode: bool, debiased: bool, num_a: int, num_b: int -) -> SinkhornBarycenterOutput: - """Jit'able function to compute discrete barycenters.""" - if lse_mode: - f_u = jnp.zeros_like(a) - g_v = dual_initialization - else: - f_u = jnp.ones_like(a) - g_v = geom.scaling_from_potential(dual_initialization) - # d below is as described in https://arxiv.org/abs/2006.02575. Note that - # d should be considered to be equal to eps log(d) with those notations - # if running in log-sum-exp mode. - d = jnp.zeros((num_b,)) if lse_mode else jnp.ones((num_b,)) - - if lse_mode: - parallel_update = jax.vmap( - lambda f, g, marginal, iter: geom. - update_potential(f, g, jnp.log(marginal), axis=1), - in_axes=[0, 0, 0, None] - ) - parallel_apply = jax.vmap( - lambda f_, g_, eps_: geom. - apply_lse_kernel(f_, g_, eps_, vec=None, axis=0)[0], - in_axes=[0, 0, None] - ) - else: - parallel_update = jax.vmap( - lambda f, g, marginal, iter: geom.update_scaling(g, marginal, axis=1), - in_axes=[0, 0, 0, None] - ) - parallel_apply = jax.vmap( - lambda f_, g_, eps_: geom.apply_kernel(f_, eps_, axis=0), - in_axes=[0, 0, None] - ) - - errors_fn = jax.vmap( - functools.partial( - sinkhorn.marginal_error, - geom=geom, - axis=1, - norm_error=norm_error, - lse_mode=lse_mode - ), - in_axes=[0, 0, 0] - ) - errors = -jnp.ones((max_iterations // inner_iterations + 1, len(norm_error))) - - const = (geom, a, weights) - - def cond_fn(iteration, const, state): # pylint: disable=unused-argument - errors = state[0] - return jnp.logical_or( - iteration == 0, errors[iteration // inner_iterations - 1, 0] > threshold - ) - - def body_fn(iteration, const, state, compute_error): - geom, a, weights = const - errors, d, f_u, g_v = state - - eps = geom._epsilon.at(iteration) # pylint: disable=protected-access - f_u = parallel_update(f_u, g_v, a, iteration) - # kernel_f_u stands for K times potential u if running in scaling mode, - # eps log K exp f / eps in lse mode. - kernel_f_u = parallel_apply(f_u, g_v, eps) - # b below is the running estimate for the barycenter if running in scaling - # mode, eps log b if running in lse mode. - if lse_mode: - b = jnp.average(kernel_f_u, weights=weights, axis=0) - else: - b = jnp.prod(kernel_f_u ** weights[:, jnp.newaxis], axis=0) - - if debiased: - if lse_mode: - b += d - d = 0.5 * ( - d + geom.update_potential( - jnp.zeros((num_a,)), d, b / eps, iteration=iteration, axis=0 - ) - ) - else: - b *= d - d = jnp.sqrt(d * geom.update_scaling(d, b, iteration=iteration, axis=0)) - if lse_mode: - g_v = b[jnp.newaxis, :] - kernel_f_u - else: - g_v = b[jnp.newaxis, :] / kernel_f_u - - # re-compute error if compute_error is True, else set to inf. - err = jnp.where( - jnp.logical_and(compute_error, iteration >= min_iterations), - jnp.mean(errors_fn(f_u, g_v, a)), jnp.inf - ) - - errors = errors.at[iteration // inner_iterations, :].set(err) - return errors, d, f_u, g_v - - state = (errors, d, f_u, g_v) - - state = fixed_point_loop.fixpoint_iter_backprop( - cond_fn, body_fn, min_iterations, max_iterations, inner_iterations, const, - state - ) - - errors, d, f_u, g_v = state - kernel_f_u = parallel_apply(f_u, g_v, geom.epsilon) - if lse_mode: - b = jnp.average(kernel_f_u, weights=weights, axis=0) - else: - b = jnp.prod(kernel_f_u ** weights[:, jnp.newaxis], axis=0) - - if debiased: - if lse_mode: - b += d - else: - b *= d - if lse_mode: - b = jnp.exp(b / geom.epsilon) - return SinkhornBarycenterOutput(f_u, g_v, b, errors) diff --git a/ott/build/lib/ott/solvers/linear/implicit_differentiation.py b/ott/build/lib/ott/solvers/linear/implicit_differentiation.py deleted file mode 100644 index fbf98ce..0000000 --- a/ott/build/lib/ott/solvers/linear/implicit_differentiation.py +++ /dev/null @@ -1,324 +0,0 @@ -# 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 dataclasses -from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple - -import jax -import jax.numpy as jnp - -from ott import utils -from ott.math import unbalanced_functions as uf - -if TYPE_CHECKING: - from ott.problems.linear import linear_problem - -LinOp_t = Callable[[jnp.ndarray], jnp.ndarray] -Solver_t = Callable[[LinOp_t, jnp.ndarray, Optional[LinOp_t], bool], - jnp.ndarray] - -__all__ = ["ImplicitDiff", "solve_jax_cg"] - - -@utils.register_pytree_node -class ImplicitDiff: - """Implicit differentiation of Sinkhorn algorithm. - - Args: - solver: Callable to compute the solution to a linear problem. The callable - expects a linear function, a vector, optionally another linear function - that implements the transpose of that function, and a boolean flag to - specify symmetry. This solver is by default one of :class:`lineax.CG` or - :class:`lineax.NormalCG` solvers, if the package can be imported, as - described in :func:`~ott.solvers.linear.lineax_implicit.solve_lineax`. - The :mod:`jax` alternative is described in - :func:`~ott.solvers.linear.implicit_differentiation.solve_jax_cg`. - Note that `lineax` solvers handle better poorly conditioned problems, - which arise typically when differentiating the solutions of balanced OT - problems (when ``tau_a==tau_b==1.0``). Relying on - :func:`~ott.solvers.linear.implicit_differentiation.solve_jax_cg` - for such cases might require hand-tuning ridge parameters, - in particular ``ridge_kernel`` and ``ridge_identity`` as described in its - doc. These parameters can be passed using ``solver_kwargs`` below. - solver_kwargs: keyword arguments passed on to the solver. - symmetric: flag used to figure out whether the linear system solved in the - implicit function theorem is symmetric or not. This happens when - ``tau_a==tau_b``, and when ``a == b``, or the precondition_fun - is the identity. The flag is False by default, and is also tested against - ``tau_a==tau_b``. It needs to be set manually by the user in the more - favorable case where the system is guaranteed to be symmetric. - precondition_fun: Function used to precondition, on both sides, the linear - system derived from first-order conditions of the regularized OT problem. - That linear system typically involves an equality between marginals (or - simple transform of these marginals when the problem is unbalanced) and - another function of the potentials. When that function is specified, that - function is applied on both sides of the equality, before being further - differentiated to provide the Jacobians needed for implicit function - theorem differentiation. - """ - - solver: Optional[Solver_t] = None - solver_kwargs: Optional[Dict[str, Any]] = None - symmetric: bool = False - precondition_fun: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None - - def solve( - self, - gr: Tuple[jnp.ndarray, jnp.ndarray], - ot_prob: "linear_problem.LinearProblem", - f: jnp.ndarray, - g: jnp.ndarray, - lse_mode: bool, - ) -> jnp.ndarray: - r"""Apply minus inverse of [hessian ``reg_ot_cost`` w.r.t. ``f``, ``g``]. - - This function is used to carry out implicit differentiation of ``sinkhorn`` - outputs, notably optimal potentials ``f`` and ``g``. That differentiation - requires solving a linear system, using (and inverting) the Jacobian of - (preconditioned) first-order conditions w.r.t. the reg-OT problem. - - Given a ``precondition_fun``, written here for short as :math:`h`, - the first order conditions for the dual energy - - .. math:: - - E(K, \epsilon, a, b, f, g) :=- + - \langle\exp^{f/\epsilon}, K \exp^{g/\epsilon}> - - form the basis of the Sinkhorn algorithm. To differentiate optimal solutions - to that problem, we exploit the fact that :math:`h(\nabla E = 0)` and - differentiate that identity to recover variations (Jacobians) of optimal - solutions :math:`f^\star, g^\star$` as a function of changes in the inputs. - The Jacobian of :math:`h(\nabla_{f,g} E = 0)` is a linear operator which, if - it were to be instantiated as a matrix, would be of size - :math:`(n+m) \times (n+m)`. When :math:`h` is the identity, that matrix is - the Hessian of :math:`E`, is symmetric and negative-definite - (:math:`E` is concave) and is structured as :math:`[A, B; B^T, D]`. More - generally, for other functions :math:`h`, the Jacobian of these - preconditioned - first order conditions is no longer symmetric (except if ``a==b``), and - has now a structure as :math:`[A, B; C, D]`. That system can - be still inverted more generic solvers. By default, :math:`h = \epsilon - \log`, as proposed in :cite:`cuturi:20a`. - - In both cases :math:`A` and :math:`D` are diagonal matrices, equal to the - row and - column marginals respectively, multiplied by the derivatives of :math:`h` - evaluated at those marginals, corrected (if handling the unbalanced case) - by the second derivative of the part of the objective that ties potentials - to the marginals (terms in ``phi_star``). When :math:`h` is the identity, - :math:`B` and :math:`B^T` are equal respectively to the OT matrix and its - transpose, i.e. :math:`n \times m` and :math:`m \times n` matrices. - When :math:`h` is not the identity, :math:`B` (resp. :math:`C`) is equal - to the OT matrix (resp. its transpose), rescaled on the left by the - application elementwise of :math:`h'` to the row (respectively column) - marginal sum of the transport. - - Note that we take great care in not instantiating these transport - matrices, to rely instead on calls to the ``app_transport`` method from the - ``Geometry`` object ``geom`` (which will either use potentials or scalings, - depending on ``lse_mode``) - - The Jacobian's diagonal + off-diagonal blocks structure allows to exploit - Schur complements. Depending on the sizes involved, it is better to - instantiate the Schur complement of the first or of the second diagonal - block. - - These linear systems are solved using the user-defined ``solver``, using - by default :mod:`lineax` solvers when available, or falling back on - :mod:`jax` when not. - - Args: - gr: 2-tuple, (vector of size ``n``, vector of size ``m``). - ot_prob: the instantiation of the regularized transport problem. - f: potential, w.r.t marginal a. - g: potential, w.r.t marginal b. - lse_mode: bool, log-sum-exp mode if True, kernel else. - - Returns: - A tuple of two vectors, of the same size as ``gr``. - """ - solver = _get_solver() if self.solver is None else self.solver - solver_kwargs = {} if self.solver_kwargs is None else self.solver_kwargs - geom = ot_prob.geom - marginal_a, marginal_b, app_transport = ( - ot_prob.get_transport_functions(lse_mode) - ) - if self.precondition_fun is None: - precond_fun = lambda x: geom.epsilon * jnp.log(x) - symmetric = False - else: - precond_fun = self.precondition_fun - symmetric = self.symmetric - - derivative = jax.vmap(jax.grad(precond_fun)) - - n, m = geom.shape - # pylint: disable=g-long-lambda - vjp_fg = lambda z: app_transport( - f, g, z * derivative(marginal_b(f, g)), axis=1 - ) / geom.epsilon - vjp_gf = lambda z: app_transport( - f, g, z * derivative(marginal_a(f, g)), axis=0 - ) / geom.epsilon - - if not symmetric: - vjp_fgt = lambda z: app_transport( - f, g, z, axis=0 - ) * derivative(marginal_b(f, g)) / geom.epsilon - vjp_gft = lambda z: app_transport( - f, g, z, axis=1 - ) * derivative(marginal_a(f, g)) / geom.epsilon - - diag_hess_a = ( - marginal_a(f, g) * derivative(marginal_a(f, g)) / geom.epsilon + - uf.diag_jacobian_of_marginal_fit( - ot_prob.a, f, ot_prob.tau_a, geom.epsilon, derivative - ) - ) - diag_hess_b = ( - marginal_b(f, g) * derivative(marginal_b(f, g)) / geom.epsilon + - uf.diag_jacobian_of_marginal_fit( - ot_prob.b, g, ot_prob.tau_b, geom.epsilon, derivative - ) - ) - n, m = geom.shape - # TODO(cuturi) consider materializing linear operator schur if size allows. - # Forks on using Schur complement of either A or D, depending on size. - if n > m: # if n is bigger, run m x m linear system. - inv_vjp_ff = lambda z: z / diag_hess_a - vjp_gg = lambda z: z * diag_hess_b - schur = lambda z: vjp_gg(z) - vjp_gf(inv_vjp_ff(vjp_fg(z))) - if not symmetric: - schur_t = lambda z: vjp_gg(z) - vjp_fgt(inv_vjp_ff(vjp_gft(z))) - else: - schur_t = None - res = gr[1] - vjp_gf(inv_vjp_ff(gr[0])) - sch = solver(schur, res, schur_t, symmetric, **solver_kwargs) - vjp_gr_f = inv_vjp_ff(gr[0] - vjp_fg(sch)) - vjp_gr_g = sch - else: - vjp_ff = lambda z: z * diag_hess_a - inv_vjp_gg = lambda z: z / diag_hess_b - schur = lambda z: vjp_ff(z) - vjp_fg(inv_vjp_gg(vjp_gf(z))) - - if not symmetric: - schur_t = lambda z: vjp_ff(z) - vjp_gft(inv_vjp_gg(vjp_fgt(z))) - else: - schur_t = None - res = gr[0] - vjp_fg(inv_vjp_gg(gr[1])) - sch = solver(schur, res, schur_t, symmetric, **solver_kwargs) - vjp_gr_g = inv_vjp_gg(gr[1] - vjp_gf(sch)) - vjp_gr_f = sch - - return jnp.concatenate((-vjp_gr_f, -vjp_gr_g)) - - def first_order_conditions( - self, prob, f: jnp.ndarray, g: jnp.ndarray, lse_mode: bool - ): - r"""Compute vector of first order conditions for the reg-OT problem. - - The output of this vector should be close to zero at optimality. - Upon completion of the Sinkhorn forward pass, its norm (using the norm - parameter defined using ``norm_error``) should be below the threshold - parameter. - - This error will be itself assumed to be close to zero when using implicit - differentiation. - - Args: - prob: definition of the linear optimal transport problem. - f: jnp.ndarray, first potential - g: jnp.ndarray, second potential - lse_mode: bool - - Returns: - a jnp.ndarray of size (size of ``n + m``) quantifying deviation to - optimality for variables ``f`` and ``g``. - """ - geom = prob.geom - marginal_a, marginal_b, _ = prob.get_transport_functions(lse_mode) - grad_a = uf.grad_of_marginal_fit(prob.a, f, prob.tau_a, geom.epsilon) - grad_b = uf.grad_of_marginal_fit(prob.b, g, prob.tau_b, geom.epsilon) - if self.precondition_fun is None: - precond_fun = lambda x: geom.epsilon * jnp.log(x) - else: - precond_fun = self.precondition_fun - - result_a = jnp.where( - prob.a > 0, - precond_fun(marginal_a(f, g)) - precond_fun(grad_a), 0.0 - ) - result_b = jnp.where( - prob.b > 0, - precond_fun(marginal_b(f, g)) - precond_fun(grad_b), 0.0 - ) - return jnp.concatenate((result_a, result_b)) - - def gradient( - self, prob: "linear_problem.LinearProblem", f: jnp.ndarray, - g: jnp.ndarray, lse_mode: bool, gr: Tuple[jnp.ndarray, jnp.ndarray] - ) -> "linear_problem.LinearProblem": - """Apply VJP to recover gradient in reverse mode differentiation.""" - # Applies first part of vjp to gr: inverse part of implicit function theorem - vjp_gr = self.solve(gr, prob, f, g, lse_mode) - - # Instantiates vjp of first order conditions of the objective, as a - # function of geom, a and b parameters (against which we differentiate) - foc_prob = lambda prob: self.first_order_conditions(prob, f, g, lse_mode) - - # Carries pullback onto original inputs, here geom, a and b. - _, pull_prob = jax.vjp(foc_prob, prob) - return pull_prob(vjp_gr) - - def replace(self, **kwargs: Any) -> "ImplicitDiff": # noqa: D102 - return dataclasses.replace(self, **kwargs) - - -def solve_jax_cg( - lin: LinOp_t, - b: jnp.ndarray, - lin_t: Optional[LinOp_t] = None, - symmetric: bool = False, - ridge_identity: float = 0.0, - ridge_kernel: float = 0.0, - **kwargs: Any -) -> jnp.ndarray: - """Wrapper around JAX native linear solvers. - - Args: - lin: Linear operator - b: vector. Returned `x` is such that `lin(x)=b` - lin_t: Linear operator, corresponding to transpose of `lin`. - symmetric: whether `lin` is symmetric. - ridge_kernel: promotes zero-sum solutions. Only use if `tau_a = tau_b = 1.0` - ridge_identity: handles rank deficient transport matrices (this happens - typically when rows/cols in cost/kernel matrices are collinear, or, - equivalently when two points from either measure are close). - kwargs: arguments passed to :func:`~jax.scipy.sparse.linalg.cg` - """ - op = lin if symmetric else lambda x: lin_t(lin(x)) - if ridge_kernel > 0.0 or ridge_identity > 0.0: - lin_reg = lambda x: op(x) + ridge_kernel * jnp.sum(x) + ridge_identity * x - else: - lin_reg = op - return jax.scipy.sparse.linalg.cg(lin_reg, b, **kwargs)[0] - - -def _get_solver() -> Solver_t: - """Get lineax solver when possible, default to jax.scipy else.""" - try: - from ott.solvers.linear import lineax_implicit - return lineax_implicit.solve_lineax - except ImportError: - return solve_jax_cg diff --git a/ott/build/lib/ott/solvers/linear/lineax_implicit.py b/ott/build/lib/ott/solvers/linear/lineax_implicit.py deleted file mode 100644 index 79b9e7c..0000000 --- a/ott/build/lib/ott/solvers/linear/lineax_implicit.py +++ /dev/null @@ -1,98 +0,0 @@ -# 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. -from typing import Any, Callable, Optional, TypeVar - -import equinox as eqx -import jax -import jax.numpy as jnp -import jax.tree_util as jtu -import lineax as lx -from jaxtyping import Array, Float, PyTree - -_T = TypeVar("_T") -_FlatPyTree = tuple[list[_T], jtu.PyTreeDef] - -__all__ = ["CustomTransposeLinearOperator", "solve_lineax"] - - -class CustomTransposeLinearOperator(lx.FunctionLinearOperator): - """Implement a linear operator that can specify its transpose directly.""" - fn: Callable[[PyTree[Float[Array, "..."]]], PyTree[Float[Array, "..."]]] - fn_t: Callable[[PyTree[Float[Array, "..."]]], PyTree[Float[Array, "..."]]] - input_structure: _FlatPyTree[jax.ShapeDtypeStruct] = eqx.static_field() - input_structure_t: _FlatPyTree[jax.ShapeDtypeStruct] = eqx.static_field() - tags: frozenset[object] - - def __init__(self, fn, fn_t, input_structure, input_structure_t, tags=()): - super().__init__(fn, input_structure, tags) - self.fn_t = eqx.filter_closure_convert(fn_t, input_structure_t) - self.input_structure_t = input_structure_t - - def transpose(self): - """Provide custom transposition operator from function.""" - return lx.FunctionLinearOperator(self.fn_t, self.input_structure_t) - - -def solve_lineax( - lin: Callable, - b: jnp.ndarray, - lin_t: Optional[Callable] = None, - symmetric: bool = False, - nonsym_solver: Optional[lx.AbstractLinearSolver] = None, - ridge_identity: float = 0.0, - ridge_kernel: float = 0.0, - **kwargs: Any -) -> jnp.ndarray: - """Wrapper around lineax solvers. - - Args: - lin: Linear operator - b: vector. Returned `x` is such that `lin(x)=b` - lin_t: Linear operator, corresponding to transpose of `lin`. - symmetric: whether `lin` is symmetric. - nonsym_solver: solver used when handling non-symmetric cases. Note that - :class:`~lineax.CG` is used by default in the symmetric case. - ridge_kernel: promotes zero-sum solutions. Only use if `tau_a = tau_b = 1.0` - ridge_identity: handles rank deficient transport matrices (this happens - typically when rows/cols in cost/kernel matrices are collinear, or, - equivalently when two points from either measure are close). - kwargs: arguments passed to :class:`~lineax.AbstractLinearSolver` linear - solver. - """ - input_structure = jax.eval_shape(lambda: b) - kwargs.setdefault("rtol", 1e-6) - kwargs.setdefault("atol", 1e-6) - - if ridge_kernel > 0.0 or ridge_identity > 0.0: - lin_reg = lambda x: lin(x) + ridge_kernel * jnp.sum(x) + ridge_identity * x - lin_t_reg = lambda x: lin_t(x) + ridge_kernel * jnp.sum( - x - ) + ridge_identity * x - else: - lin_reg, lin_t_reg = lin, lin_t - - if symmetric: - solver = lx.CG(**kwargs) - fn_operator = lx.FunctionLinearOperator( - lin_reg, input_structure, tags=lx.positive_semidefinite_tag - ) - return lx.linear_solve(fn_operator, b, solver).value - # In the non-symmetric case, use NormalCG by default, but consider - # user defined choice of alternative lx solver. - solver_type = lx.NormalCG if nonsym_solver is None else nonsym_solver - solver = solver_type(**kwargs) - fn_operator = CustomTransposeLinearOperator( - lin_reg, lin_t_reg, input_structure, input_structure - ) - return lx.linear_solve(fn_operator, b, solver).value diff --git a/ott/build/lib/ott/solvers/linear/lr_utils.py b/ott/build/lib/ott/solvers/linear/lr_utils.py deleted file mode 100644 index 8ade265..0000000 --- a/ott/build/lib/ott/solvers/linear/lr_utils.py +++ /dev/null @@ -1,371 +0,0 @@ -# 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. -from typing import NamedTuple, Optional, Tuple - -import jax -import jax.numpy as jnp -import jax.scipy as jsp - -from ott.math import fixed_point_loop -from ott.problems.linear import linear_problem - -__all__ = ["unbalanced_dykstra_lse", "unbalanced_dykstra_kernel"] - - -class State(NamedTuple): # noqa: D101 - v1: jnp.ndarray - v2: jnp.ndarray - u1: jnp.ndarray - u2: jnp.ndarray - g: jnp.ndarray - err: float - - -class Constants(NamedTuple): # noqa: D101 - a: jnp.ndarray - b: jnp.ndarray - rho_a: float - rho_b: float - supp_a: Optional[jnp.ndarray] = None - supp_b: Optional[jnp.ndarray] = None - - -def unbalanced_dykstra_lse( - c_q: jnp.ndarray, - c_r: jnp.ndarray, - c_g: jnp.ndarray, - gamma: float, - ot_prob: linear_problem.LinearProblem, - translation_invariant: bool = True, - tolerance: float = 1e-3, - min_iter: int = 0, - inner_iter: int = 10, - max_iter: int = 10000 -) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Dykstra's algorithm for the unbalanced - :class:`~ott.solvers.linear.sinkhorn_lr.LRSinkhorn` in LSE mode. - - Args: - c_q: Cost associated with :math:`Q`. - c_r: Cost associated with :math:`R`. - c_g: Cost associated with :math:`g`. - gamma: The (inverse of) the gradient step. - ot_prob: Unbalanced OT problem. - translation_invariant: Whether to use the translation invariant objective, - see :cite:`scetbon:23`, alg. 3. - tolerance: Convergence threshold. - min_iter: Minimum number of iterations. - inner_iter: Compute error every ``inner_iter``. - max_iter: Maximum number of iterations. - - Returns: - The :math:`Q`, :math:`R` and :math:`g` factors. - """ # noqa: D205 - - def _softm( - v: jnp.ndarray, - c: jnp.ndarray, - axis: int, - ) -> jnp.ndarray: - v = jnp.expand_dims(v, axis=1 - axis) - return jsp.special.logsumexp(v + c, axis=axis) - - def _error( - gamma: float, - new_state: State, - old_state: State, - ) -> float: - u1_err = jnp.linalg.norm(new_state.u1 - old_state.u1, ord=jnp.inf) - u2_err = jnp.linalg.norm(new_state.u2 - old_state.u2, ord=jnp.inf) - v1_err = jnp.linalg.norm(new_state.v1 - old_state.v1, ord=jnp.inf) - v2_err = jnp.linalg.norm(new_state.v2 - old_state.v2, ord=jnp.inf) - return (1.0 / gamma) * jnp.max(jnp.array([u1_err, u2_err, v1_err, v2_err])) - - def cond_fn( - iteration: int, - const: Constants, - state: State, - ) -> bool: - del iteration, const - return tolerance < state.err - - def body_fn( - iteration: int, const: Constants, state: State, compute_error: bool - ) -> State: - log_a, log_b = jnp.log(const.a), jnp.log(const.b) - rho_a, rho_b = const.rho_a, const.rho_b - - c_a = _get_ratio(const.rho_a, gamma) - c_b = _get_ratio(const.rho_b, gamma) - - if translation_invariant: - lam_a, lam_b = compute_lambdas(const, state, gamma, g=c_g, lse_mode=True) - - u1 = c_a * (log_a - _softm(state.v1, c_q, axis=1)) - u1 = u1 - lam_a / ((1.0 / gamma) + rho_a) - u2 = c_b * (log_b - _softm(state.v2, c_r, axis=1)) - u2 = u2 - lam_b / ((1.0 / gamma) + rho_b) - - state_lam = State( - v1=state.v1, v2=state.v2, u1=u1, u2=u2, g=state.g, err=state.err - ) - lam_a, lam_b = compute_lambdas( - const, state_lam, gamma, g=c_g, lse_mode=True - ) - - v1_trans = _softm(u1, c_q, axis=0) - v2_trans = _softm(u2, c_r, axis=0) - - g_trans = gamma * (lam_a + lam_b) + c_g - else: - u1 = c_a * (log_a - _softm(state.v1, c_q, axis=1)) - u2 = c_b * (log_b - _softm(state.v2, c_r, axis=1)) - - v1_trans = _softm(u1, c_q, axis=0) - v2_trans = _softm(u2, c_r, axis=0) - g_trans = c_g - - g = (1.0 / 3.0) * (g_trans + v1_trans + v2_trans) - v1 = g - v1_trans - v2 = g - v2_trans - - new_state = State(v1=v1, v2=v2, u1=u1, u2=u2, g=g, err=jnp.inf) - err = jax.lax.cond( - jnp.logical_and(compute_error, iteration >= min_iter), - _error, - lambda *_: state.err, - gamma, - new_state, - state, - ) - return State(v1=v1, v2=v2, u1=u1, u2=u2, g=g, err=err) - - n, m, r = c_q.shape[0], c_r.shape[0], c_g.shape[0] - constants = Constants( - a=ot_prob.a, - b=ot_prob.b, - rho_a=_rho(ot_prob.tau_a), - rho_b=_rho(ot_prob.tau_b), - supp_a=ot_prob.a > 0, - supp_b=ot_prob.b > 0, - ) - init_state = State( - v1=jnp.zeros(r), - v2=jnp.zeros(r), - u1=jnp.zeros(n), - u2=jnp.zeros(m), - g=c_g, - err=jnp.inf, - ) - - state: State = fixed_point_loop.fixpoint_iter_backprop( - cond_fn, body_fn, min_iter, max_iter, inner_iter, constants, init_state - ) - - q = jnp.exp(state.u1[:, None] + c_q + state.v1[None, :]) - r = jnp.exp(state.u2[:, None] + c_r + state.v2[None, :]) - g = jnp.exp(state.g) - - return q, r, g - - -def unbalanced_dykstra_kernel( - k_q: jnp.ndarray, - k_r: jnp.ndarray, - k_g: jnp.ndarray, - gamma: float, - ot_prob: linear_problem.LinearProblem, - translation_invariant: bool = True, - tolerance: float = 1e-3, - min_iter: int = 0, - inner_iter: int = 10, - max_iter: int = 10000 -) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Dykstra's algorithm for the unbalanced - :class:`~ott.solvers.linear.sinkhorn_lr.LRSinkhorn` in kernel mode. - - Args: - k_q: Kernel associated with :math:`Q`. - k_r: Kernel associated with :math:`R`. - k_g: Kernel associated with :math:`g`. - gamma: The (inverse of) the gradient step. - ot_prob: Unbalanced OT problem. - translation_invariant: Whether to use the translation invariant objective, - see :cite:`scetbon:23`, alg. 3. - tolerance: Convergence threshold. - min_iter: Minimum number of iterations. - inner_iter: Compute error every ``inner_iter``. - max_iter: Maximum number of iterations. - - Returns: - The :math:`Q`, :math:`R` and :math:`g` factors. - """ # noqa: D205 - - def _error( - gamma: float, - new_state: State, - old_state: State, - ) -> float: - u1_err = jnp.linalg.norm( - jnp.log(new_state.u1) - jnp.log(old_state.u1), ord=jnp.inf - ) - u2_err = jnp.linalg.norm( - jnp.log(new_state.u2) - jnp.log(old_state.u2), ord=jnp.inf - ) - v1_err = jnp.linalg.norm( - jnp.log(new_state.v1) - jnp.log(old_state.v1), ord=jnp.inf - ) - v2_err = jnp.linalg.norm( - jnp.log(new_state.v2) - jnp.log(old_state.v2), ord=jnp.inf - ) - return (1.0 / gamma) * jnp.max(jnp.array([u1_err, u2_err, v1_err, v2_err])) - - def cond_fn( - iteration: int, - const: Constants, - state: State, - ) -> bool: - del iteration, const - return tolerance < state.err - - def body_fn( - iteration: int, const: Constants, state: State, compute_error: bool - ) -> State: - c_a = _get_ratio(const.rho_a, gamma) - c_b = _get_ratio(const.rho_b, gamma) - - if translation_invariant: - lam_a, lam_b = compute_lambdas(const, state, gamma, g=k_g, lse_mode=False) - - u1 = jnp.where(const.supp_a, (const.a / (k_q @ state.v1)) ** c_a, 0.0) - u1 = u1 * jnp.exp(-lam_a / ((1.0 / gamma) + const.rho_a)) - u2 = jnp.where(const.supp_b, (const.b / (k_r @ state.v2)) ** c_b, 0.0) - u2 = u2 * jnp.exp(-lam_b / ((1.0 / gamma) + const.rho_b)) - - state_lam = State( - v1=state.v1, v2=state.v2, u1=u1, u2=u2, g=state.g, err=state.err - ) - lam_a, lam_b = compute_lambdas( - const, state_lam, gamma, g=k_g, lse_mode=False - ) - - v1_trans = k_q.T @ u1 - v2_trans = k_r.T @ u2 - - k_trans = jnp.exp(gamma * (lam_a + lam_b)) * k_g - g = (k_trans * v1_trans * v2_trans) ** (1.0 / 3.0) - else: - u1 = jnp.where(const.supp_a, (const.a / (k_q @ state.v1)) ** c_a, 0.0) - u2 = jnp.where(const.supp_b, (const.b / (k_r @ state.v2)) ** c_b, 0.0) - - v1_trans = k_q.T @ u1 - v2_trans = k_r.T @ u2 - - g = (k_g * v1_trans * v2_trans) ** (1.0 / 3.0) - - v1 = g / v1_trans - v2 = g / v2_trans - - new_state = State(v1=v1, v2=v2, u1=u1, u2=u2, g=g, err=jnp.inf) - err = jax.lax.cond( - jnp.logical_and(compute_error, iteration >= min_iter), - _error, - lambda *_: state.err, - gamma, - new_state, - state, - ) - return State(v1=v1, v2=v2, u1=u1, u2=u2, g=g, err=err) - - n, m, r = k_q.shape[0], k_r.shape[0], k_g.shape[0] - constants = Constants( - a=ot_prob.a, - b=ot_prob.b, - rho_a=_rho(ot_prob.tau_a), - rho_b=_rho(ot_prob.tau_b), - supp_a=ot_prob.a > 0.0, - supp_b=ot_prob.b > 0.0, - ) - init_state = State( - v1=jnp.ones(r), - v2=jnp.ones(r), - u1=jnp.ones(n), - u2=jnp.ones(m), - g=k_g, - err=jnp.inf - ) - - state: State = fixed_point_loop.fixpoint_iter_backprop( - cond_fn, body_fn, min_iter, max_iter, inner_iter, constants, init_state - ) - - q = state.u1[:, None] * k_q * state.v1[None, :] - r = state.u2[:, None] * k_r * state.v2[None, :] - - return q, r, state.g - - -def compute_lambdas( - const: Constants, state: State, gamma: float, g: jnp.ndarray, *, - lse_mode: bool -) -> Tuple[float, float]: - """TODO.""" - gamma_inv = 1.0 / gamma - rho_a = const.rho_a - rho_b = const.rho_b - - if lse_mode: - num_1 = jsp.special.logsumexp((-gamma_inv / rho_a) * state.u1, b=const.a) - num_2 = jsp.special.logsumexp((-gamma_inv / rho_b) * state.u2, b=const.b) - den = jsp.special.logsumexp(g - (state.v1 + state.v2)) - const_1 = num_1 - den - const_2 = num_2 - den - - ratio_1 = _get_ratio(rho_a, gamma) - ratio_2 = _get_ratio(rho_b, gamma) - harmonic = 1.0 / (1.0 - (ratio_1 * ratio_2)) - lam_1 = harmonic * gamma_inv * ratio_1 * (const_1 - ratio_2 * const_2) - lam_2 = harmonic * gamma_inv * ratio_2 * (const_2 - ratio_1 * const_1) - return lam_1, lam_2 - - num_1 = jnp.sum( - jnp.where( - const.supp_a, ((state.u1 ** (-gamma_inv / rho_a)) * const.a), 0.0 - ) - ) - num_2 = jnp.sum( - jnp.where( - const.supp_b, ((state.u2 ** (-gamma_inv / rho_b)) * const.b), 0.0 - ) - ) - den = jnp.sum(g / (state.v1 * state.v2)) - const_1 = jnp.log(num_1 / den) - const_2 = jnp.log(num_2 / den) - - ratio_1 = _get_ratio(rho_a, gamma) - ratio_2 = _get_ratio(rho_b, gamma) - harmonic = 1.0 / (1.0 - (ratio_1 * ratio_2)) - lam_1 = harmonic * gamma_inv * ratio_1 * (const_1 - ratio_2 * const_2) - lam_2 = harmonic * gamma_inv * ratio_2 * (const_2 - ratio_1 * const_1) - return lam_1, lam_2 - - -def _rho(tau: float) -> float: - tau = jnp.asarray(tau) # avoid division by 0 in Python, get NaN instead - return tau / (1.0 - tau) - - -def _get_ratio(rho: float, gamma: float) -> float: - gamma_inv = 1.0 / gamma - return jnp.where(jnp.isfinite(rho), rho / (rho + gamma_inv), 1.0) diff --git a/ott/build/lib/ott/solvers/linear/sinkhorn.py b/ott/build/lib/ott/solvers/linear/sinkhorn.py deleted file mode 100644 index beb3b61..0000000 --- a/ott/build/lib/ott/solvers/linear/sinkhorn.py +++ /dev/null @@ -1,1230 +0,0 @@ -# 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. -from typing import ( - Any, - Callable, - Literal, - Mapping, - NamedTuple, - Optional, - Sequence, - Tuple, - Union, -) - -import jax -import jax.experimental -import jax.numpy as jnp -import jax.scipy as jsp -import numpy as np - -from ott import utils -from ott.geometry import geometry -from ott.initializers.linear import initializers as init_lib -from ott.math import fixed_point_loop -from ott.math import unbalanced_functions as uf -from ott.math import utils as mu -from ott.problems.linear import linear_problem, potentials -from ott.solvers.linear import acceleration -from ott.solvers.linear import implicit_differentiation as implicit_lib - -__all__ = ["Sinkhorn", "SinkhornOutput"] - -ProgressCallbackFn_t = Callable[ - [Tuple[np.ndarray, np.ndarray, np.ndarray, "SinkhornState"]], None] - - -class SinkhornState(NamedTuple): - """Holds the state variables used to solve OT with Sinkhorn.""" - - potentials: Tuple[jnp.ndarray, ...] - errors: Optional[jnp.ndarray] = None - old_fus: Optional[jnp.ndarray] = None - old_mapped_fus: Optional[jnp.ndarray] = None - - def set(self, **kwargs: Any) -> "SinkhornState": - """Return a copy of self, with potential overwrites.""" - return self._replace(**kwargs) - - def solution_error( - self, - ot_prob: linear_problem.LinearProblem, - norm_error: Sequence[int], - *, - lse_mode: bool, - parallel_dual_updates: bool, - recenter: bool, - ) -> jnp.ndarray: - """State dependent function to return error.""" - fu, gv = self.fu, self.gv - if recenter and lse_mode: - fu, gv = self.recenter(fu, gv, ot_prob=ot_prob) - - return solution_error( - fu, - gv, - ot_prob, - norm_error=norm_error, - lse_mode=lse_mode, - parallel_dual_updates=parallel_dual_updates - ) - - def compute_kl_reg_cost( # noqa: D102 - self, ot_prob: linear_problem.LinearProblem, lse_mode: bool - ) -> float: - return compute_kl_reg_cost(self.fu, self.gv, ot_prob, lse_mode) - - def recenter( - self, - f: jnp.ndarray, - g: jnp.ndarray, - ot_prob: linear_problem.LinearProblem, - ) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Re-center dual potentials. - - If the ``ot_prob`` is balanced, the ``f`` potential is zero-centered. - Otherwise, use prop. 2 of :cite:`sejourne:22` re-center the potentials iff - ``tau_a < 1`` and ``tau_b < 1``. - - Args: - f: The first dual potential. - g: The second dual potential. - ot_prob: Linear OT problem. - - Returns: - The centered potentials. - """ - if ot_prob.is_balanced: - # center the potentials for numerical stability - is_finite = jnp.isfinite(f) - shift = jnp.sum(jnp.where(is_finite, f, 0.0)) / jnp.sum(is_finite) - return f - shift, g + shift - - if ot_prob.tau_a == 1.0 or ot_prob.tau_b == 1.0: - # re-centering wasn't done during the lse-step, ignore - return f, g - - rho_a = uf.rho(ot_prob.epsilon, ot_prob.tau_a) - rho_b = uf.rho(ot_prob.epsilon, ot_prob.tau_b) - tau = rho_a * rho_b / (rho_a + rho_b) - - shift = tau * ( - mu.logsumexp(-f / rho_a, b=ot_prob.a) - - mu.logsumexp(-g / rho_b, b=ot_prob.b) - ) - return f + shift, g - shift - - @property - def fu(self) -> jnp.ndarray: - """The first dual potential or scaling.""" - return self.potentials[0] - - @property - def gv(self) -> jnp.ndarray: - """The second dual potential or scaling.""" - return self.potentials[1] - - -def solution_error( - f_u: jnp.ndarray, - g_v: jnp.ndarray, - ot_prob: linear_problem.LinearProblem, - *, - norm_error: Sequence[int], - lse_mode: bool, - parallel_dual_updates: bool, -) -> jnp.ndarray: - """Given two potential/scaling solutions, computes deviation to optimality. - - When the ``ot_prob`` problem is balanced and the usual Sinkhorn updates are - used, this is simply deviation of the coupling's marginal to ``ot_prob.b``. - This is the case because the second (and last) update of the Sinkhorn - algorithm equalizes the row marginal of the coupling to ``ot_prob.a``. To - simplify the logic, this is parameterized by checking whether - `parallel_dual_updates = False`. - - When that flag is `True`, or when the problem is unbalanced, - additional quantities to qualify optimality must be taken into account. - - Args: - f_u: jnp.ndarray, potential or scaling - g_v: jnp.ndarray, potential or scaling - ot_prob: linear OT problem - norm_error: int, p-norm used to compute error. - lse_mode: True if log-sum-exp operations, False if kernel vector products. - parallel_dual_updates: Whether potentials/scalings were computed in - parallel. - - Returns: - a positive number quantifying how far from optimality current solution is. - """ - if ot_prob.is_balanced and not parallel_dual_updates: - return marginal_error( - f_u, g_v, ot_prob.b, ot_prob.geom, 0, norm_error, lse_mode - ) - - # In the unbalanced case, we compute the norm of the gradient. - # the gradient is equal to the marginal of the current plan minus - # the gradient of < z, rho_z(exp^(-h/rho_z) -1> where z is either a or b - # and h is either f or g. Note this is equal to z if rho_z → inf, which - # is the case when tau_z → 1.0 - if lse_mode: - grad_a = uf.grad_of_marginal_fit( - ot_prob.a, f_u, ot_prob.tau_a, ot_prob.epsilon - ) - grad_b = uf.grad_of_marginal_fit( - ot_prob.b, g_v, ot_prob.tau_b, ot_prob.epsilon - ) - else: - u = ot_prob.geom.potential_from_scaling(f_u) - v = ot_prob.geom.potential_from_scaling(g_v) - grad_a = uf.grad_of_marginal_fit( - ot_prob.a, u, ot_prob.tau_a, ot_prob.epsilon - ) - grad_b = uf.grad_of_marginal_fit( - ot_prob.b, v, ot_prob.tau_b, ot_prob.epsilon - ) - err = marginal_error(f_u, g_v, grad_a, ot_prob.geom, 1, norm_error, lse_mode) - err += marginal_error(f_u, g_v, grad_b, ot_prob.geom, 0, norm_error, lse_mode) - return err - - -def marginal_error( - f_u: jnp.ndarray, - g_v: jnp.ndarray, - target: jnp.ndarray, - geom: geometry.Geometry, - axis: int = 0, - norm_error: Sequence[int] = (1,), - lse_mode: bool = True -) -> jnp.asarray: - """Output how far Sinkhorn solution is w.r.t target. - - Args: - f_u: a vector of potentials or scalings for the first marginal. - g_v: a vector of potentials or scalings for the second marginal. - target: target marginal. - geom: Geometry object. - axis: axis (0 or 1) along which to compute marginal. - norm_error: (tuple of int) p's to compute p-norm between marginal/target - lse_mode: whether operating on scalings or potentials - - Returns: - Array of floats, quantifying difference between target / marginal. - """ - if lse_mode: - marginal = geom.marginal_from_potentials(f_u, g_v, axis=axis) - else: - marginal = geom.marginal_from_scalings(f_u, g_v, axis=axis) - norm_error = jnp.asarray(norm_error) - return jnp.sum( - jnp.abs(marginal - target) ** norm_error[:, jnp.newaxis], axis=1 - ) ** (1.0 / norm_error) - - -def compute_kl_reg_cost( - f: jnp.ndarray, g: jnp.ndarray, ot_prob: linear_problem.LinearProblem, - lse_mode: bool -) -> float: - r"""Compute objective of Sinkhorn for OT problem given dual solutions. - - The objective is evaluated for dual solution ``f`` and ``g``, using - information contained in ``ot_prob``. The objective is the regularized - optimal transport cost (i.e. the cost itself plus entropic and unbalanced - terms). Situations where marginals ``a`` or ``b`` in ot_prob have zero - coordinates are reflected in minus infinity entries in their corresponding - dual potentials. To avoid NaN that may result when multiplying 0's by infinity - values, ``jnp.where`` is used to cancel these contributions. - - Args: - f: jnp.ndarray, potential - g: jnp.ndarray, potential - ot_prob: linear optimal transport problem. - lse_mode: bool, whether to compute total mass in lse or kernel mode. - - Returns: - The regularized transport cost. - """ - supp_a = ot_prob.a > 0 - supp_b = ot_prob.b > 0 - fa = ot_prob.geom.potential_from_scaling(ot_prob.a) - if ot_prob.tau_a == 1.0: - div_a = jnp.sum(jnp.where(supp_a, ot_prob.a * (f - fa), 0.0)) - else: - rho_a = uf.rho(ot_prob.epsilon, ot_prob.tau_a) - div_a = -jnp.sum( - jnp.where(supp_a, ot_prob.a * uf.phi_star(-(f - fa), rho_a), 0.0) - ) - - gb = ot_prob.geom.potential_from_scaling(ot_prob.b) - if ot_prob.tau_b == 1.0: - div_b = jnp.sum(jnp.where(supp_b, ot_prob.b * (g - gb), 0.0)) - else: - rho_b = uf.rho(ot_prob.epsilon, ot_prob.tau_b) - div_b = -jnp.sum( - jnp.where(supp_b, ot_prob.b * uf.phi_star(-(g - gb), rho_b), 0.0) - ) - - # Using https://arxiv.org/pdf/1910.12958.pdf (24) - if lse_mode: - total_sum = jnp.sum(ot_prob.geom.marginal_from_potentials(f, g)) - else: - u = ot_prob.geom.scaling_from_potential(f) - v = ot_prob.geom.scaling_from_potential(g) - total_sum = jnp.sum(ot_prob.geom.marginal_from_scalings(u, v)) - return div_a + div_b + ot_prob.epsilon * ( - jnp.sum(ot_prob.a) * jnp.sum(ot_prob.b) - total_sum - ) - - -class SinkhornOutput(NamedTuple): - """Holds the output of a Sinkhorn solver applied to a problem. - - Objects of this class contain both solutions and problem definition of a - regularized OT problem, along several methods that can be used to access its - content, to, for instance, materialize an OT matrix or apply it to a vector - (without having to materialize it when not needed). - - Args: - f: dual variables vector of size ``ot.prob.shape[0]`` returned by Sinkhorn - g: dual variables vector of size ``ot.prob.shape[1]`` returned by Sinkhorn - errors: vector or errors, along iterations. This vector is of size - ``max_iterations // inner_iterations`` where those were the parameters - passed on to the :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. - For each entry indexed at ``i``, ``errors[i]`` can be either a real - non-negative value (meaning the algorithm recorded that error at the - ``i * inner_iterations`` iteration), a ``jnp.inf`` value (meaning the - algorithm computed that iteration but did not compute its error, because, - for instance, ``i < min_iterations // inner_iterations``), or a ``-1``, - meaning that execution was terminated before that iteration, because the - criterion was found to be smaller than ``threshold``. - reg_ot_cost: the regularized optimal transport cost. By default this is - the linear contribution + KL term. See - :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.ent_reg_cost`, - :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.primal_cost` and - :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.dual_cost` for other - objective values. - ot_prob: stores the definition of the OT problem, including geometry, - marginals, unbalanced regularizers, etc. - threshold: convergence threshold used to control the termination of the - algorithm. - converged: whether the output corresponds to a solution whose error is - below the convergence threshold. - inner_iterations: number of iterations that were run between two - computations of errors. - """ - - potentials: Tuple[jnp.ndarray, ...] - errors: Optional[jnp.ndarray] = None - reg_ot_cost: Optional[float] = None - ot_prob: Optional[linear_problem.LinearProblem] = None - threshold: Optional[jnp.ndarray] = None - converged: Optional[bool] = None - inner_iterations: Optional[int] = None - - def set(self, **kwargs: Any) -> "SinkhornOutput": - """Return a copy of self, with potential overwrites.""" - return self._replace(**kwargs) - - def set_cost( # noqa: D102 - self, ot_prob: linear_problem.LinearProblem, lse_mode: bool, - use_danskin: bool - ) -> "SinkhornOutput": - f = jax.lax.stop_gradient(self.f) if use_danskin else self.f - g = jax.lax.stop_gradient(self.g) if use_danskin else self.g - return self.set(reg_ot_cost=compute_kl_reg_cost(f, g, ot_prob, lse_mode)) - - @property - def dual_cost(self) -> jnp.ndarray: - """Return dual transport cost, without considering regularizer.""" - a, b = self.ot_prob.a, self.ot_prob.b - dual_cost = jnp.sum(jnp.where(a > 0.0, a * self.f, 0)) - dual_cost += jnp.sum(jnp.where(b > 0.0, b * self.g, 0)) - return dual_cost - - @property - def primal_cost(self) -> float: - """Return transport cost of current transport solution at geometry.""" - return self.transport_cost_at_geom(other_geom=self.geom) - - @property - def ent_reg_cost(self) -> float: - r"""Entropy regularized cost. - - This outputs - - .. math:: - - \langle P^{\star},C\rangle - \varepsilon H(P^{\star}) + - \rho_a\text{KL}(P^{\star} 1|a) + \rho_b\text{KL}(1^T P^{\star}|b), - - where :math:`P^{\star}, a, b` is the coupling returned by the - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` and the two marginal weight - vectors; :math:`\rho_a=\varepsilon \tau_a / (1-\tau_a)` and - :m ath:`\rho_b=\varepsilon \tau_b / (1-\tau_b)` are obtained when the problem - is unbalanced from parameters ``tau_a`` and ``tau_b``. Note that the last - two terms vanish in the balanced case, when ``tau_a==tau_b==1``. - """ - ent_a = jnp.sum(jsp.special.entr(self.ot_prob.a)) - ent_b = jnp.sum(jsp.special.entr(self.ot_prob.b)) - return self.reg_ot_cost - self.geom.epsilon * (ent_a + ent_b) - - @property - def kl_reg_cost(self) -> float: - r"""KL regularized OT transport cost. - - This outputs - - .. math:: - - \langle P^{\star}, C \rangle + \varepsilon KL(P^{\star},ab^T) + - \rho_a\text{KL}(P^{\star} 1|a) + \rho_b\text{KL}(1^T P^{\star}|b), - - where :math:`P^{\star}, a, b` are the coupling returned by the - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` algorithm and the two - marginal weight vectors, respectively, and - :math:`\rho_a=\varepsilon \tau_a / (1-\tau_a)` and - :math:`\rho_b=\varepsilon \tau_b / (1-\tau_b)` are obtained when the problem - is unbalanced from parameters ``tau_a`` and ``tau_b``. Note that the last - two terms vanish in the balanced case, when ``tau_a==tau_b==1``. This - quantity coincides with :attr:`reg_ot_cost`, which is computed using - dual variables. - """ - return self.reg_ot_cost - - def transport_cost_at_geom( - self, other_geom: geometry.Geometry - ) -> jnp.ndarray: - r"""Return bare transport cost of current solution at any geometry. - - In order to compute cost, we check first if the geometry can be converted - to a low-rank cost geometry in order to speed up computations, without - having to materialize the full cost matrix. If this is not possible, - we resort to instantiating both transport matrix and cost matrix. - - Args: - other_geom: geometry whose cost matrix is used to evaluate the transport - cost. - - Returns: - the transportation cost at :math:`C`, i.e. :math:`\langle P, C \rangle`. - """ - # TODO(cuturi): handle online mode for non Euclidean pointcloud geometries. - # TODO(michalk8): handle SqEucl point cloud is not converted to LRCGeom - if other_geom.can_LRC: - geom = other_geom.to_LRCGeometry() - return jnp.sum(self.apply(geom.cost_1.T) * geom.cost_2.T) - return jnp.sum(self.matrix * other_geom.cost_matrix) - - @property - def geom(self) -> geometry.Geometry: # noqa: D102 - return self.ot_prob.geom - - @property - def a(self) -> jnp.ndarray: # noqa: D102 - return self.ot_prob.a - - @property - def b(self) -> jnp.ndarray: # noqa: D102 - return self.ot_prob.b - - @property - def n_iters(self) -> int: # noqa: D102 - """Returns the total number of iterations that were needed to terminate.""" - return jnp.sum(self.errors != -1) * self.inner_iterations - - @property - def scalings(self) -> Tuple[jnp.ndarray, jnp.ndarray]: # noqa: D102 - u = self.ot_prob.geom.scaling_from_potential(self.f) - v = self.ot_prob.geom.scaling_from_potential(self.g) - return u, v - - @property - def matrix(self) -> jnp.ndarray: - """Transport matrix if it can be instantiated.""" - try: - return self.ot_prob.geom.transport_from_potentials(self.f, self.g) - except ValueError: - return self.ot_prob.geom.transport_from_scalings(*self.scalings) - - @property - def transport_mass(self) -> float: - """Sum of transport matrix.""" - return self.marginal(0).sum() - - def apply( - self, - inputs: jnp.ndarray, - axis: int = 0, - lse_mode: bool = True - ) -> jnp.ndarray: - """Apply the transport to a ndarray; axis=1 for its transpose.""" - geom = self.ot_prob.geom - if lse_mode: - return geom.apply_transport_from_potentials( - self.f, self.g, inputs, axis=axis - ) - u = geom.scaling_from_potential(self.f) - v = geom.scaling_from_potential(self.g) - return geom.apply_transport_from_scalings(u, v, inputs, axis=axis) - - def marginal(self, axis: int) -> jnp.ndarray: # noqa: D102 - return self.ot_prob.geom.marginal_from_potentials(self.f, self.g, axis=axis) - - def cost_at_geom(self, other_geom: geometry.Geometry) -> float: - """Return reg-OT cost for matrix, evaluated at other cost matrix.""" - return ( - jnp.sum(self.matrix * other_geom.cost_matrix) - - self.geom.epsilon * jnp.sum(jax.scipy.special.entr(self.matrix)) - ) - - def to_dual_potentials(self) -> potentials.EntropicPotentials: - """Return the entropic map estimator.""" - return potentials.EntropicPotentials(self.f, self.g, self.ot_prob) - - @property - def f(self) -> jnp.ndarray: - """The first dual potential.""" - return self.potentials[0] - - @property - def g(self) -> jnp.ndarray: - """The second dual potential.""" - return self.potentials[1] - - -@jax.tree_util.register_pytree_node_class -class Sinkhorn: - r"""Sinkhorn solver. - - The Sinkhorn algorithm is a fixed point iteration that solves a regularized - optimal transport (reg-OT) problem between two measures. - The optimization variables are a pair of vectors (called potentials, or - scalings when parameterized as exponential of the former). Calling this - function returns therefore a pair of optimal vectors. In addition to these, - it also returns the objective value achieved by these optimal vectors; - a vector of size ``max_iterations/inner_iterations`` that records the vector - of values recorded to monitor convergence, throughout the execution of the - algorithm (padded with `-1` if convergence happens before), as well as a - boolean to signify whether the algorithm has converged within the number of - iterations specified by the user. - - The reg-OT problem is specified by two measures, of respective sizes ``n`` and - ``m``. From the viewpoint of the ``sinkhorn`` function, these two measures are - only seen through a triplet (``geom``, ``a``, ``b``), where ``geom`` is a - ``Geometry`` object, and ``a`` and ``b`` are weight vectors of respective - sizes ``n`` and ``m``. Starting from two initial values for those potentials - or scalings (both can be defined by the user by passing value in - ``init_dual_a`` or ``init_dual_b``), the Sinkhorn algorithm will use - elementary operations that are carried out by the ``geom`` object. - - Math: - Given a geometry ``geom``, which provides a cost matrix :math:`C` with its - regularization parameter :math:`\varepsilon`, (or a kernel matrix :math:`K`) - the reg-OT problem consists in finding two vectors `f`, `g` of size ``n``, - ``m`` that maximize the following criterion. - - .. math:: - - \arg\max_{f, g}{- \langle a, \phi_a^{*}(-f) \rangle - \langle b, - \phi_b^{*}(-g) \rangle - \varepsilon \langle e^{f/\varepsilon}, - e^{-C/\varepsilon} e^{-g/\varepsilon}} \rangle - - where :math:`\phi_a(z) = \rho_a z(\log z - 1)` is a scaled entropy, and - :math:`\phi_a^{*}(z) = \rho_a e^{z/\varepsilon}`, its Legendre transform. - - That problem can also be written, instead, using positive scaling vectors - `u`, `v` of size ``n``, ``m``, handled with the kernel - :math:`K := e^{-C/\varepsilon}`, - - .. math:: - - \arg\max_{u, v >0} - \langle a,\phi_a^{*}(-\varepsilon\log u) \rangle + - \langle b, \phi_b^{*}(-\varepsilon\log v) \rangle - \langle u, K v \rangle - - Both of these problems corresponds, in their *primal* formulation, to - solving the unbalanced optimal transport problem with a variable matrix - :math:`P` of size ``n`` x ``m``: - - .. math:: - - \arg\min_{P>0} \langle P,C \rangle -\varepsilon \text{KL}(P | ab^T) - + \rho_a \text{KL}(P\mathbf{1}_m | a) + \rho_b \text{KL}(P^T \mathbf{1}_n - | b) - - where :math:`KL` is the generalized Kullback-Leibler divergence. - - The very same primal problem can also be written using a kernel :math:`K` - instead of a cost :math:`C` as well: - - .. math:: - - \arg\min_{P} \varepsilon \text{KL}(P|K) - + \rho_a \text{KL}(P\mathbf{1}_m | a) + - \rho_b \text{KL}(P^T \mathbf{1}_n | b) - - The *original* OT problem taught in linear programming courses is recovered - by using the formulation above relying on the cost :math:`C`, and letting - :math:`\varepsilon \rightarrow 0`, and :math:`\rho_a, \rho_b \rightarrow - \infty`. - In that case the entropy disappears, whereas the :math:`KL` regularization - above become constraints on the marginals of :math:`P`: This results in a - standard min cost flow problem. This problem is not handled for now in this - toolbox, which focuses exclusively on the case :math:`\varepsilon > 0`. - - The *balanced* regularized OT problem is recovered for finite - :math:`\varepsilon > 0` but letting :math:`\rho_a, \rho_b \rightarrow - \infty`. This problem can be shown to be equivalent to a matrix scaling - problem, which can be solved using the Sinkhorn fixed-point algorithm. - To handle the case :math:`\rho_a, \rho_b \rightarrow \infty`, the - ``sinkhorn`` function uses parameters ``tau_a`` and ``tau_b`` equal - respectively to :math:`\rho_a /(\varepsilon + \rho_a)` and - :math:`\rho_b / (\varepsilon + \rho_b)` instead. Setting either of these - parameters to 1 corresponds to setting the corresponding - :math:`\rho_a, \rho_b` to :math:`\infty`. - - The Sinkhorn algorithm solves the reg-OT problem by seeking optimal - :math:`f`, :math:`g` potentials (or alternatively their parameterization - as positive scaling vectors :math:`u`, :math:`v`), rather than solving the - primal problem in :math:`P`. This is mostly for efficiency (potentials and - scalings have a ``n + m`` memory footprint, rather than ``n m`` required - to store `P`). This is also because both problems are, in fact, equivalent, - since the optimal transport :math:`P^{\star}` can be recovered from - optimal potentials :math:`f^{\star}`, :math:`g^{\star}` or scaling - :math:`u^{\star}`, :math:`v^{\star}`, using the geometry's cost or kernel - matrix respectively: - - .. math:: - - P^{\star} = \exp\left(\frac{f^{\star}\mathbf{1}_m^T + \mathbf{1}_n g^{*T}- - C}{\varepsilon}\right) \text{ or } P^{\star} = \text{diag}(u^{\star}) K - \text{diag}(v^{\star}) - - By default, the Sinkhorn algorithm solves this dual problem in :math:`f, g` - or :math:`u, v` using block coordinate ascent, i.e. devising an update for - each :math:`f` and :math:`g` (resp. :math:`u` and :math:`v`) that cancels - their respective gradients, one at a time. These two iterations are repeated - ``inner_iterations`` times, after which the norm of these gradients will be - evaluated and compared with the ``threshold`` value. The iterations are then - repeated as long as that error exceeds ``threshold``. - - Note on Sinkhorn updates: - The boolean flag ``lse_mode`` sets whether the algorithm is run in either: - - - log-sum-exp mode (``lse_mode=True``), in which case it is directly - defined in terms of updates to `f` and `g`, using log-sum-exp - computations. This requires access to the cost matrix :math:`C`, as it is - stored, or possibly computed on the fly by ``geom``. - - - kernel mode (``lse_mode=False``), in which case it will require access - to a matrix vector multiplication operator :math:`z \rightarrow K z`, - where :math:`K` is either instantiated from :math:`C` as - :math:`\exp(-C/\varepsilon)`, or provided directly. In that case, rather - than optimizing on :math:`f` and :math:`g`, it is more convenient to - optimize on their so called scaling formulations, - :math:`u := \exp(f / \varepsilon)` and :math:`v := \exp(g / \varepsilon)`. - While faster (applying matrices is faster than applying ``lse`` repeatedly - over lines), this mode is also less stable numerically, notably for - smaller :math:`\varepsilon`. - - In the source code, the variables ``f_u`` or ``g_v`` can be either regarded - as potentials (real) or scalings (positive) vectors, depending on the choice - of ``lse_mode`` by the user. Once optimization is carried out, we only - return dual variables in potential form, i.e. ``f`` and ``g``. - - In addition to standard Sinkhorn updates, the user can also use heavy-ball - type updates using a ``momentum`` parameter in ]0,2[. We also implement a - strategy that tries to set that parameter adaptively at - ``chg_momentum_from`` iterations, as a function of progress in the error, - as proposed in the literature. - - Another upgrade to the standard Sinkhorn updates provided to the users lies - in using Anderson acceleration. This can be parameterized by setting the - otherwise null ``anderson`` to a positive integer. When selected,the - algorithm will recompute, every ``refresh_anderson_frequency`` (set by - default to 1) an extrapolation of the most recently computed ``anderson`` - iterates. When using that option, notice that differentiation (if required) - can only be carried out using implicit differentiation, and that all - momentum related parameters are ignored. - - The ``parallel_dual_updates`` flag is set to ``False`` by default. In that - setting, ``g_v`` is first updated using the latest values for ``f_u`` and - ``g_v``, before proceeding to update ``f_u`` using that new value for - ``g_v``. When the flag is set to ``True``, both ``f_u`` and ``g_v`` are - updated simultaneously. Note that setting that choice to ``True`` requires - using some form of averaging (e.g. ``momentum=0.5``). Without this, and on - its own ``parallel_dual_updates`` won't work. - - Differentiation: - The optimal solutions ``f`` and ``g`` and the optimal objective - (``reg_ot_cost``) outputted by the Sinkhorn algorithm can be differentiated - w.r.t. relevant inputs ``geom``, ``a`` and ``b``. In the default setting, - implicit differentiation of the optimality conditions (``implicit_diff`` - not equal to ``None``), this has two consequences, treating ``f`` and ``g`` - differently from ``reg_ot_cost``. - - - The termination criterion used to stop Sinkhorn (cancellation of - gradient of objective w.r.t. ``f_u`` and ``g_v``) is used to differentiate - ``f`` and ``g``, given a change in the inputs. These changes are computed - by solving a linear system. The arguments starting with - ``implicit_solver_*`` allow to define the linear solver that is used, and - to control for two types or regularization (we have observed that, - depending on the architecture, linear solves may require higher ridge - parameters to remain stable). The optimality conditions in Sinkhorn can be - analyzed as satisfying a ``z=z'`` condition, which are then - differentiated. It might be beneficial (e.g., as in :cite:`cuturi:20a`) - to use a preconditioning function ``precondition_fun`` to differentiate - instead ``h(z) = h(z')``. - - - The objective ``reg_ot_cost`` returned by Sinkhorn uses the so-called - envelope (or Danskin's) theorem. In that case, because it is assumed that - the gradients of the dual variables ``f_u`` and ``g_v`` w.r.t. dual - objective are zero (reflecting the fact that they are optimal), small - variations in ``f_u`` and ``g_v`` due to changes in inputs (such as - ``geom``, ``a`` and ``b``) are considered negligible. As a result, - ``stop_gradient`` is applied on dual variables ``f_u`` and ``g_v`` when - evaluating the ``reg_ot_cost`` objective. Note that this approach is - `invalid` when computing higher order derivatives. In that case the - ``use_danskin`` flag must be set to ``False``. - - An alternative yet more costly way to differentiate the outputs of the - Sinkhorn iterations is to use unrolling, i.e. reverse mode differentiation - of the Sinkhorn loop. This is possible because Sinkhorn iterations are - wrapped in a custom fixed point iteration loop, defined in - ``fixed_point_loop``, rather than a standard while loop. This is to ensure - the end result of this fixed point loop can also be differentiated, if - needed, using standard JAX operations. To ensure differentiability, - the ``fixed_point_loop.fixpoint_iter_backprop`` loop does checkpointing of - state variables (here ``f_u`` and ``g_v``) every ``inner_iterations``, and - backpropagates automatically, block by block, through blocks of - ``inner_iterations`` at a time. - - Note: - * The Sinkhorn algorithm may not converge within the maximum number of - iterations for possibly several reasons: - - 1. the regularizer (defined as ``epsilon`` in the geometry ``geom`` - object) is too small. Consider either switching to ``lse_mode=True`` - (at the price of a slower execution), increasing ``epsilon``, or, - alternatively, if you are unable or unwilling to increase ``epsilon``, - either increase ``max_iterations`` or ``threshold``. - 2. the probability weights ``a`` and ``b`` do not have the same total - mass, while using a balanced (``tau_a=tau_b=1.0``) setup. - Consider either normalizing ``a`` and ``b``, or set either ``tau_a`` - and/or ``tau_b<1.0``. - 3. OOMs issues may arise when storing either cost or kernel matrices that - are too large in ``geom``. In the case where, the ``geom`` geometry is - a ``PointCloud``, some of these issues might be solved by setting the - ``online`` flag to ``True``. This will trigger a re-computation on the - fly of the cost/kernel matrix. - - * The weight vectors ``a`` and ``b`` can be passed on with coordinates that - have zero weight. This is then handled by relying on simple arithmetic for - ``inf`` values that will likely arise (due to :math:`\log 0` when - ``lse_mode`` is ``True``, or divisions by zero when ``lse_mode`` is - ``False``). Whenever that arithmetic is likely to produce ``NaN`` values - (due to ``-inf * 0``, or ``-inf - -inf``) in the forward pass, we use - ``jnp.where`` conditional statements to carry ``inf`` rather than ``NaN`` - values. In the reverse mode differentiation, the inputs corresponding to - these 0 weights (a location `x`, or a row in the corresponding cost/kernel - matrix), and the weight itself will have ``NaN`` gradient values. This is - reflects that these gradients are undefined, since these points were not - considered in the optimization and have therefore no impact on the output. - - Args: - lse_mode: ``True`` for log-sum-exp computations, ``False`` for kernel - multiplication. - threshold: tolerance used to stop the Sinkhorn iterations. This is - typically the deviation between a target marginal and the marginal of the - current primal solution when either or both tau_a and tau_b are 1.0 - (balanced or semi-balanced problem), or the relative change between two - successive solutions in the unbalanced case. - norm_error: power used to define p-norm of error for marginal/target. - inner_iterations: the Sinkhorn error is not recomputed at each - iteration but every ``inner_iterations`` instead. - min_iterations: the minimum number of Sinkhorn iterations carried - out before the error is computed and monitored. - max_iterations: the maximum number of Sinkhorn iterations. If - ``max_iterations`` is equal to ``min_iterations``, Sinkhorn iterations are - run by default using a :func:`jax.lax.scan` loop rather than a custom, - unroll-able :func:`jax.lax.while_loop` that monitors convergence. - In that case the error is not monitored and the ``converged`` - flag will return ``False`` as a consequence. - momentum: Momentum instance. - anderson: AndersonAcceleration instance. - implicit_diff: instance used to solve implicit differentiation. Unrolls - iterations if None. - parallel_dual_updates: updates potentials or scalings in parallel if True, - sequentially (in Gauss-Seidel fashion) if False. - recenter_potentials: Whether to re-center the dual potentials. - If the problem is balanced, the ``f`` potential is zero-centered for - numerical stability. Otherwise, use the approach of :cite:`sejourne:22` - to achieve faster convergence. Only used when ``lse_mode = True`` and - ``tau_a < 1`` and ``tau_b < 1``. - use_danskin: when ``True``, it is assumed the entropy regularized cost - is evaluated using optimal potentials that are frozen, i.e. whose - gradients have been stopped. This is useful when carrying out first order - differentiation, and is only valid (as with ``implicit_differentiation``) - when the algorithm has converged with a low tolerance. - initializer: how to compute the initial potentials/scalings. This refers to - a few possible classes implemented following the template in - :class:`~ott.initializers.linear.SinkhornInitializer`. - progress_fn: callback function which gets called during the Sinkhorn - iterations, so the user can display the error at each iteration, - e.g., using a progress bar. See :func:`~ott.utils.default_progress_fn` - for a basic implementation. - kwargs_init: keyword arguments when creating the initializer. - """ - - def __init__( - self, - lse_mode: bool = True, - threshold: float = 1e-3, - norm_error: int = 1, - inner_iterations: int = 10, - min_iterations: int = 0, - max_iterations: int = 2000, - momentum: Optional[acceleration.Momentum] = None, - anderson: Optional[acceleration.AndersonAcceleration] = None, - parallel_dual_updates: bool = False, - recenter_potentials: bool = False, - use_danskin: Optional[bool] = None, - implicit_diff: Optional[implicit_lib.ImplicitDiff - ] = implicit_lib.ImplicitDiff(), # noqa: B008 - initializer: Union[Literal["default", "gaussian", "sorting", "subsample"], - init_lib.SinkhornInitializer] = "default", - progress_fn: Optional[ProgressCallbackFn_t] = None, - kwargs_init: Optional[Mapping[str, Any]] = None, - ): - self.lse_mode = lse_mode - self.threshold = threshold - self.inner_iterations = inner_iterations - self.min_iterations = min_iterations - self.max_iterations = max_iterations - self._norm_error = norm_error - self.anderson = anderson - self.implicit_diff = implicit_diff - - if momentum is not None: - self.momentum = acceleration.Momentum( - momentum.start, momentum.error_threshold, momentum.value, - self.inner_iterations - ) - else: - # Use no momentum if using Anderson or unrolling. - if self.anderson is not None or self.implicit_diff is None: - self.momentum = acceleration.Momentum( - inner_iterations=self.inner_iterations - ) - else: - # no momentum - self.momentum = acceleration.Momentum() - - self.parallel_dual_updates = parallel_dual_updates - self.recenter_potentials = recenter_potentials - self.initializer = initializer - self.progress_fn = progress_fn - self.kwargs_init = {} if kwargs_init is None else kwargs_init - - # Force implicit_differentiation to True when using Anderson acceleration, - # Reset all momentum parameters to default (i.e. no momentum) - if anderson: - self.implicit_diff = ( - implicit_lib.ImplicitDiff() - if self.implicit_diff is None else self.implicit_diff - ) - self.momentum = acceleration.Momentum( - inner_iterations=self.inner_iterations - ) - - # By default, use Danskin theorem to differentiate - # the objective when using implicit_lib. - self.use_danskin = ((self.implicit_diff is not None) - if use_danskin is None else use_danskin) - - def __call__( - self, - ot_prob: linear_problem.LinearProblem, - init: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray]] = (None, None), - rng: Optional[jax.Array] = None, - ) -> SinkhornOutput: - """Run Sinkhorn algorithm. - - Args: - ot_prob: Linear OT problem. - init: Initial dual potentials/scalings f_u and g_v, respectively. - Any `None` values will be initialized using the initializer. - rng: Random number generator key for stochastic initialization. - - Returns: - The Sinkhorn output. - """ - rng = utils.default_prng_key(rng) - initializer = self.create_initializer() - init_dual_a, init_dual_b = initializer( - ot_prob, *init, lse_mode=self.lse_mode, rng=rng - ) - return run(ot_prob, self, (init_dual_a, init_dual_b)) - - def lse_step( - self, ot_prob: linear_problem.LinearProblem, state: SinkhornState, - iteration: int - ) -> SinkhornState: - """Sinkhorn LSE update.""" - - def k(tau_i: float, tau_j: float) -> float: - num = -tau_j * (tau_a - 1) * (tau_b - 1) * (tau_i - 1) - denom = (tau_j - 1) * (tau_a * (tau_b - 1) + tau_b * (tau_a - 1)) - return num / denom - - def xi(tau_i: float, tau_j: float) -> float: - k_ij = k(tau_i, tau_j) - return k_ij / (1.0 - k_ij) - - def smin( - potential: jnp.ndarray, marginal: jnp.ndarray, tau: float - ) -> float: - rho = uf.rho(ot_prob.epsilon, tau) - return -rho * mu.logsumexp(-potential / rho, b=marginal) - - # only for an unbalanced problems with `tau_{a,b} < 1` - recenter = ( - self.recenter_potentials and ot_prob.tau_a < 1.0 and ot_prob.tau_b < 1.0 - ) - w = self.momentum.weight(state, iteration) - tau_a, tau_b = ot_prob.tau_a, ot_prob.tau_b - old_fu, old_gv = state.fu, state.gv - - if recenter: - k11, k22 = k(tau_a, tau_a), k(tau_b, tau_b) - xi12, xi21 = xi(tau_a, tau_b), xi(tau_b, tau_a) - - # update g potential - new_gv = tau_b * ot_prob.geom.update_potential( - old_fu, old_gv, jnp.log(ot_prob.b), iteration, axis=0 - ) - if recenter: - new_gv -= k22 * smin(old_fu, ot_prob.a, tau_a) - new_gv += xi21 * smin(new_gv, ot_prob.b, tau_b) - gv = self.momentum(w, old_gv, new_gv, self.lse_mode) - - if not self.parallel_dual_updates: - old_gv = gv - - # update f potential - new_fu = tau_a * ot_prob.geom.update_potential( - old_fu, old_gv, jnp.log(ot_prob.a), iteration, axis=1 - ) - if recenter: - new_fu -= k11 * smin(old_gv, ot_prob.b, tau_b) - new_fu += xi12 * smin(new_fu, ot_prob.a, tau_a) - fu = self.momentum(w, old_fu, new_fu, self.lse_mode) - - return state.set(potentials=(fu, gv)) - - def kernel_step( - self, ot_prob: linear_problem.LinearProblem, state: SinkhornState, - iteration: int - ) -> SinkhornState: - """Sinkhorn multiplicative update.""" - w = self.momentum.weight(state, iteration) - old_gv = state.gv - new_gv = ot_prob.geom.update_scaling( - state.fu, ot_prob.b, iteration, axis=0 - ) ** ot_prob.tau_b - gv = self.momentum(w, state.gv, new_gv, self.lse_mode) - new_fu = ot_prob.geom.update_scaling( - old_gv if self.parallel_dual_updates else gv, - ot_prob.a, - iteration, - axis=1 - ) ** ot_prob.tau_a - fu = self.momentum(w, state.fu, new_fu, self.lse_mode) - return state.set(potentials=(fu, gv)) - - def one_iteration( - self, ot_prob: linear_problem.LinearProblem, state: SinkhornState, - iteration: int, compute_error: bool - ) -> SinkhornState: - """Carries out one Sinkhorn iteration. - - Depending on lse_mode, these iterations can be either in: - - - log-space for numerical stability. - - scaling space, using standard kernel-vector multiply operations. - - Args: - ot_prob: the transport problem definition - state: SinkhornState named tuple. - iteration: the current iteration of the Sinkhorn loop. - compute_error: flag to indicate this iteration computes/stores an error - - Returns: - The updated state. - """ - # When running updates in parallel (Gauss-Seidel mode), old_g_v will be - # used to update f_u, rather than the latest g_v computed in this loop. - # Unused otherwise. - if self.anderson: - state = self.anderson.update(state, iteration, ot_prob, self.lse_mode) - - if self.lse_mode: # In lse_mode, run additive updates. - print("lse step") - state = self.lse_step(ot_prob, state, iteration) - else: - state = self.kernel_step(ot_prob, state, iteration) - - if self.anderson: - state = self.anderson.update_history(state, ot_prob, self.lse_mode) - - # re-computes error if compute_error is True, else set it to inf. - err = jax.lax.cond( - jnp.logical_or( - iteration == self.max_iterations - 1, - jnp.logical_and(compute_error, iteration >= self.min_iterations) - ), - lambda state, prob: state.solution_error( - prob, - self.norm_error, - lse_mode=self.lse_mode, - parallel_dual_updates=self.parallel_dual_updates, - recenter=self.recenter_potentials - )[0], - lambda *_: jnp.inf, - state, - ot_prob, - ) - errors = state.errors.at[iteration // self.inner_iterations, :].set(err) - state = state.set(errors=errors) - - if self.progress_fn is not None: - jax.experimental.io_callback( - self.progress_fn, None, - (iteration, self.inner_iterations, self.max_iterations, state) - ) - return state - - def _converged(self, state: SinkhornState, iteration: int) -> bool: - err = state.errors[iteration // self.inner_iterations - 1, 0] - return jnp.logical_and(iteration > 0, err < self.threshold) - - def _diverged(self, state: SinkhornState, iteration: int) -> bool: - err = state.errors[iteration // self.inner_iterations - 1, 0] - return jnp.logical_not(jnp.isfinite(err)) - - def _continue(self, state: SinkhornState, iteration: int) -> bool: - """Continue while not(converged) and not(diverged).""" - return jnp.logical_and( - jnp.logical_not(self._diverged(state, iteration)), - jnp.logical_not(self._converged(state, iteration)) - ) - - @property - def outer_iterations(self) -> int: - """Upper bound on number of times inner_iterations are carried out. - - This integer can be used to set constant array sizes to track the algorithm - progress, notably errors. - """ - return np.ceil(self.max_iterations / self.inner_iterations).astype(int) - - def init_state( - self, ot_prob: linear_problem.LinearProblem, init: Tuple[jnp.ndarray, - jnp.ndarray] - ) -> SinkhornState: - """Return the initial state of the loop.""" - errors = -jnp.ones((self.outer_iterations, len(self.norm_error))) - state = SinkhornState(init, errors=errors) - return self.anderson.init_maps(ot_prob, state) if self.anderson else state - - def output_from_state( - self, ot_prob: linear_problem.LinearProblem, state: SinkhornState - ) -> SinkhornOutput: - """Create an output from a loop state. - - Note: - When differentiating the regularized OT cost, and assuming Sinkhorn has - run to convergence, Danskin's (or the envelope) - `theorem `_ - :cite:`danskin:67,bertsekas:71` - states that the resulting OT cost as a function of the inputs - (``geometry``, ``a``, ``b``) behaves locally as if the dual optimal - potentials were frozen and did not vary with those inputs. - - Notice this is only valid, as when using ``implicit_differentiation`` - mode, if the Sinkhorn algorithm outputs potentials that are near optimal. - namely when the threshold value is set to a small tolerance. - - The flag ``use_danskin`` controls whether that assumption is made. By - default, that flag is set to the value of ``implicit_differentiation`` if - not specified. If you wish to compute derivatives of order 2 and above, - set ``use_danskin`` to ``False``. - - Args: - ot_prob: the transport problem. - state: a SinkhornState. - - Returns: - A SinkhornOutput. - """ - geom = ot_prob.geom - - f = state.fu if self.lse_mode else geom.potential_from_scaling(state.fu) - g = state.gv if self.lse_mode else geom.potential_from_scaling(state.gv) - if self.recenter_potentials: - f, g = state.recenter(f, g, ot_prob=ot_prob) - - # By convention, the algorithm is said to have converged if the algorithm - # has not nan'ed during iterations (notice some errors might be infinite, - # this convention is used when the error is not recomputed), and if the - # last recorded error is lower than the threshold. Note that this will be - # the case if either the algorithm terminated earlier (in which case the - # last state.errors[-1] = -1 by convention) or if the algorithm carried out - # the maximal number of iterations and its last recorded error (at -1 - # position) is lower than the threshold. - - converged = jnp.logical_and( - jnp.logical_not(jnp.any(jnp.isnan(state.errors))), state.errors[-1] - < self.threshold - )[0] - - return SinkhornOutput((f, g), - errors=state.errors[:, 0], - threshold=jnp.array(self.threshold), - converged=converged, - inner_iterations=self.inner_iterations) - - @property - def norm_error(self) -> Tuple[int, ...]: - """Powers used to compute the p-norm between marginal/target.""" - # To change momentum adaptively, one needs errors in ||.||_1 norm. - # In that case, we add this exponent to the list of errors to compute, - # notably if that was not the error requested by the user. - if self.momentum and self.momentum.start > 0 and self._norm_error != 1: - return self._norm_error, 1 - return self._norm_error, - - # TODO(michalk8): in the future, enforce this (+ in GW) via abstract method - def create_initializer(self) -> init_lib.SinkhornInitializer: # noqa: D102 - if isinstance(self.initializer, init_lib.SinkhornInitializer): - return self.initializer - if self.initializer == "default": - return init_lib.DefaultInitializer() - if self.initializer == "gaussian": - return init_lib.GaussianInitializer() - if self.initializer == "sorting": - return init_lib.SortingInitializer(**self.kwargs_init) - if self.initializer == "subsample": - return init_lib.SubsampleInitializer(**self.kwargs_init) - raise NotImplementedError( - f"Initializer `{self.initializer}` is not yet implemented." - ) - - def tree_flatten(self): # noqa: D102 - aux = vars(self).copy() - aux["norm_error"] = aux.pop("_norm_error") - aux.pop("threshold") - return [self.threshold], aux - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - return cls(**aux_data, threshold=children[0]) - - -def run( - ot_prob: linear_problem.LinearProblem, solver: Sinkhorn, - init: Tuple[jnp.ndarray, ...] -) -> SinkhornOutput: - """Run loop of the solver, outputting a state upgraded to an output.""" - #iter_fun = _iterations_implicit if solver.implicit_diff else iterations - iter_fun = iterations - out = iter_fun(ot_prob, solver, init) - # Be careful here, the geom and the cost are injected at the end, where it - # does not interfere with the implicit differentiation. - out = out.set_cost(ot_prob, solver.lse_mode, solver.use_danskin) - return out.set(ot_prob=ot_prob) - - -def iterations( - ot_prob: linear_problem.LinearProblem, solver: Sinkhorn, - init: Tuple[jnp.ndarray, ...] -) -> SinkhornOutput: - """Jittable Sinkhorn loop. args contain initialization variables.""" - - def cond_fn( - iteration: int, const: Tuple[linear_problem.LinearProblem, Sinkhorn], - state: SinkhornState - ) -> bool: - _, solver = const - return solver._continue(state, iteration) - - def body_fn( - iteration: int, const: Tuple[linear_problem.LinearProblem, Sinkhorn], - state: SinkhornState, compute_error: bool - ) -> SinkhornState: - ot_prob, solver = const - return solver.one_iteration(ot_prob, state, iteration, compute_error) - - # Run the Sinkhorn loop. Choose either a standard fixpoint_iter loop if - # differentiation is implicit, otherwise switch to the backprop friendly - # version of that loop if unrolling to differentiate. - if solver.implicit_diff: - fix_point = fixed_point_loop.fixpoint_iter - else: - fix_point = fixed_point_loop.fixpoint_iter_backprop - - const = ot_prob, solver - state = solver.init_state(ot_prob, init) - state = fix_point( - cond_fn, body_fn, solver.min_iterations, solver.max_iterations, - solver.inner_iterations, const, state - ) - return solver.output_from_state(ot_prob, state) - - -def _iterations_taped( - ot_prob: linear_problem.LinearProblem, solver: Sinkhorn, - init: Tuple[jnp.ndarray, ...] -) -> Tuple[SinkhornOutput, Tuple[jnp.ndarray, jnp.ndarray, - linear_problem.LinearProblem, Sinkhorn]]: - """Run forward pass of the Sinkhorn algorithm storing side information.""" - state = iterations(ot_prob, solver, init) - return state, (state.f, state.g, ot_prob, solver) - - -def _iterations_implicit_bwd(res, gr: SinkhornOutput): - """Run Sinkhorn in backward mode, using implicit differentiation. - - Args: - res: residual data sent from fwd pass, used for computations below. In this - case consists in the output itself, as well as inputs against which we - wish to differentiate. - gr: gradients w.r.t outputs of fwd pass, here w.r.t size f, g, errors. Note - that differentiability w.r.t. errors is not handled, and only f, g is - considered. - - Returns: - a tuple of gradients: PyTree for geom, one jnp.ndarray for each of a and b. - """ - f, g, ot_prob, solver = res - out = solver.implicit_diff.gradient( - ot_prob, f, g, solver.lse_mode, gr.potentials - ) - return *out, None, None - - -# sets threshold, norm_errors, geom, a and b to be differentiable, as those are -# non-static. Only differentiability w.r.t. geom, a and b will be used. -_iterations_implicit = jax.custom_vjp(iterations) -_iterations_implicit.defvjp(_iterations_taped, _iterations_implicit_bwd) diff --git a/ott/build/lib/ott/solvers/linear/sinkhorn_lr.py b/ott/build/lib/ott/solvers/linear/sinkhorn_lr.py deleted file mode 100644 index da949da..0000000 --- a/ott/build/lib/ott/solvers/linear/sinkhorn_lr.py +++ /dev/null @@ -1,822 +0,0 @@ -# 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. -from typing import ( - Any, - Callable, - Literal, - Mapping, - NamedTuple, - Optional, - Tuple, - Union, -) - -import jax -import jax.experimental -import jax.numpy as jnp -import jax.scipy as jsp -import numpy as np - -from ott.geometry import geometry -from ott.initializers.linear import initializers_lr -from ott.math import fixed_point_loop -from ott.math import utils as mu -from ott.problems.linear import linear_problem -from ott.solvers.linear import lr_utils, sinkhorn - -__all__ = ["LRSinkhorn", "LRSinkhornOutput"] - -ProgressCallbackFn_t = Callable[ - [Tuple[np.ndarray, np.ndarray, np.ndarray, "LRSinkhornState"]], None] - - -class LRSinkhornState(NamedTuple): - """State of the Low Rank Sinkhorn algorithm.""" - q: jnp.ndarray - r: jnp.ndarray - g: jnp.ndarray - gamma: float - costs: jnp.ndarray - errors: jnp.ndarray - crossed_threshold: bool - - def compute_error( # noqa: D102 - self, previous_state: "LRSinkhornState" - ) -> float: - err_q = mu.gen_js(self.q, previous_state.q, c=1.0) - err_r = mu.gen_js(self.r, previous_state.r, c=1.0) - err_g = mu.gen_js(self.g, previous_state.g, c=1.0) - - return ((1.0 / self.gamma) ** 2) * (err_q + err_r + err_g) - - def reg_ot_cost( # noqa: D102 - self, - ot_prob: linear_problem.LinearProblem, - *, - epsilon: float, - use_danskin: bool = False - ) -> float: - """For LR Sinkhorn, this defaults to the primal cost of LR solution.""" - return compute_reg_ot_cost( - self.q, - self.r, - self.g, - ot_prob, - epsilon=epsilon, - use_danskin=use_danskin - ) - - def solution_error( # noqa: D102 - self, ot_prob: linear_problem.LinearProblem, norm_error: Tuple[int, ...] - ) -> jnp.ndarray: - return solution_error(self.q, self.r, ot_prob, norm_error) - - def set(self, **kwargs: Any) -> "LRSinkhornState": - """Return a copy of self, with potential overwrites.""" - return self._replace(**kwargs) - - -def compute_reg_ot_cost( - q: jnp.ndarray, - r: jnp.ndarray, - g: jnp.ndarray, - ot_prob: linear_problem.LinearProblem, - epsilon: float, - use_danskin: bool = False -) -> float: - """Compute the regularized OT cost, here the primal cost of the LR solution. - - Args: - q: first factor of solution - r: second factor of solution - g: weights of solution - ot_prob: linear problem - epsilon: Entropic regularization. - use_danskin: if True, use Danskin's theorem :cite:`danskin:67,bertsekas:71` - to avoid computing the gradient of the cost function. - - Returns: - regularized OT cost, the (primal) transport cost of the low-rank solution. - """ - - def ent(x: jnp.ndarray) -> float: - # generalized entropy - return jnp.sum(jsp.special.entr(x) + x) - - tau_a, tau_b = ot_prob.tau_a, ot_prob.tau_b - - q = jax.lax.stop_gradient(q) if use_danskin else q - r = jax.lax.stop_gradient(r) if use_danskin else r - g = jax.lax.stop_gradient(g) if use_danskin else g - - cost = jnp.sum(ot_prob.geom.apply_cost(r, axis=1) * q * (1.0 / g)[None, :]) - cost -= epsilon * (ent(q) + ent(r) + ent(g)) - if tau_a != 1.0: - cost += tau_a / (1.0 - tau_a) * mu.gen_kl(jnp.sum(q, axis=1), ot_prob.a) - if tau_b != 1.0: - cost += tau_b / (1.0 - tau_b) * mu.gen_kl(jnp.sum(r, axis=1), ot_prob.b) - - return cost - - -def solution_error( - q: jnp.ndarray, r: jnp.ndarray, ot_prob: linear_problem.LinearProblem, - norm_error: Tuple[int, ...] -) -> jnp.ndarray: - """Compute solution error. - - Since only balanced case is available for LR, this is marginal deviation. - - Args: - q: first factor of solution. - r: second factor of solution. - ot_prob: linear problem. - norm_error: int, p-norm used to compute error. - - Returns: - one or possibly many numbers quantifying deviation to true marginals. - """ - norm_error = jnp.array(norm_error) - # Update the error - err = jnp.sum( - jnp.abs(jnp.sum(q, axis=1) - ot_prob.a) ** norm_error[:, None], axis=1 - ) ** (1.0 / norm_error) - err += jnp.sum( - jnp.abs(jnp.sum(r, axis=1) - ot_prob.b) ** norm_error[:, None], axis=1 - ) ** (1.0 / norm_error) - err += jnp.sum( - jnp.abs(jnp.sum(q, axis=0) - jnp.sum(r, axis=0)) ** norm_error[:, None], - axis=1 - ) ** (1.0 / norm_error) - - return err - - -class LRSinkhornOutput(NamedTuple): - """Transport interface for a low-rank Sinkhorn solution.""" - - q: jnp.ndarray - r: jnp.ndarray - g: jnp.ndarray - costs: jnp.ndarray - # TODO(michalk8): must be called `errors`, because of `store_inner_errors` - # in future, enforce via class hierarchy - errors: jnp.ndarray - ot_prob: linear_problem.LinearProblem - epsilon: float - inner_iterations: int - # TODO(michalk8): Optional is an artifact of the current impl., refactor - reg_ot_cost: Optional[float] = None - - def set(self, **kwargs: Any) -> "LRSinkhornOutput": - """Return a copy of self, with potential overwrites.""" - return self._replace(**kwargs) - - def set_cost( # noqa: D102 - self, - ot_prob: linear_problem.LinearProblem, - lse_mode: bool, - use_danskin: bool = False - ) -> "LRSinkhornOutput": - del lse_mode - return self.set(reg_ot_cost=self.compute_reg_ot_cost(ot_prob, use_danskin)) - - def compute_reg_ot_cost( # noqa: D102 - self, - ot_prob: linear_problem.LinearProblem, - use_danskin: bool = False, - ) -> float: - return compute_reg_ot_cost( - self.q, - self.r, - self.g, - ot_prob, - epsilon=self.epsilon, - use_danskin=use_danskin - ) - - @property - def geom(self) -> geometry.Geometry: # noqa: D102 - return self.ot_prob.geom - - @property - def a(self) -> jnp.ndarray: # noqa: D102 - return self.ot_prob.a - - @property - def b(self) -> jnp.ndarray: # noqa: D102 - return self.ot_prob.b - - @property - def n_iters(self) -> int: # noqa: D102 - return jnp.sum(self.errors != -1) * self.inner_iterations - - @property - def converged(self) -> bool: # noqa: D102 - return jnp.logical_and( - jnp.any(self.costs == -1), jnp.all(jnp.isfinite(self.costs)) - ) - - @property - def matrix(self) -> jnp.ndarray: - """Transport matrix if it can be instantiated.""" - return (self.q * self._inv_g) @ self.r.T - - def apply(self, inputs: jnp.ndarray, axis: int = 0) -> jnp.ndarray: - """Apply the transport to a array; axis=1 for its transpose.""" - q, r = (self.q, self.r) if axis == 1 else (self.r, self.q) - # for `axis=0`: (batch, m), (m, r), (r,), (r, n) - return ((inputs @ r) * self._inv_g) @ q.T - - def marginal(self, axis: int) -> jnp.ndarray: # noqa: D102 - length = self.q.shape[0] if axis == 0 else self.r.shape[0] - return self.apply(jnp.ones(length,), axis=axis) - - def cost_at_geom(self, other_geom: geometry.Geometry) -> float: - """Return OT cost for current solution, evaluated at any cost matrix.""" - return jnp.sum(self.q * other_geom.apply_cost(self.r, axis=1) * self._inv_g) - - def transport_cost_at_geom(self, other_geom: geometry.Geometry) -> float: - """Return (by recomputing it) bare transport cost of current solution.""" - return self.cost_at_geom(other_geom) - - @property - def primal_cost(self) -> float: - """Return (by recomputing it) transport cost of current solution.""" - return self.transport_cost_at_geom(other_geom=self.geom) - - @property - def transport_mass(self) -> float: - """Sum of transport matrix.""" - return self.marginal(0).sum() - - @property - def _inv_g(self) -> jnp.ndarray: - return 1.0 / self.g - - -@jax.tree_util.register_pytree_node_class -class LRSinkhorn(sinkhorn.Sinkhorn): - r"""Low-Rank Sinkhorn solver for linear reg-OT problems. - - The algorithm is described in :cite:`scetbon:21` and the implementation - contained here is adapted from `LOT `_. - - The algorithm minimizes a non-convex problem. It therefore requires special - care to initialization and convergence. Convergence is evaluated on successive - evaluations of the objective. - - Args: - rank: Rank constraint on the coupling to minimize the linear OT problem - gamma: The (inverse of) gradient step size used by mirror descent. - gamma_rescale: Whether to rescale :math:`\gamma` every iteration as - described in :cite:`scetbon:22b`. - epsilon: Entropic regularization added on top of low-rank problem. - initializer: How to initialize the :math:`Q`, :math:`R` and :math:`g` - factors. - lse_mode: Whether to run computations in LSE or kernel mode. - inner_iterations: Number of inner iterations used by the algorithm before - re-evaluating progress. - use_danskin: Use Danskin theorem to evaluate gradient of objective w.r.t. - input parameters. Only `True` handled at this moment. - progress_fn: callback function which gets called during the Sinkhorn - iterations, so the user can display the error at each iteration, - e.g., using a progress bar. See :func:`~ott.utils.default_progress_fn` - for a basic implementation. - kwargs_dys: Keyword arguments passed to :meth:`dykstra_update_lse`, - :meth:`dykstra_update_kernel` or one of the functions defined in - :mod:`ott.solvers.linear`, depending on whether the problem - is balanced and on the ``lse_mode``. - kwargs_init: Keyword arguments for - :class:`~ott.initializers.linear.initializers_lr.LRInitializer`. - kwargs: Keyword arguments for - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. - """ - - def __init__( - self, - rank: int, - gamma: float = 10.0, - gamma_rescale: bool = True, - epsilon: float = 0.0, - initializer: Union[Literal["random", "rank2", "k-means", - "generalized-k-means"], - initializers_lr.LRInitializer] = "random", - lse_mode: bool = True, - inner_iterations: int = 10, - use_danskin: bool = True, - kwargs_dys: Optional[Mapping[str, Any]] = None, - kwargs_init: Optional[Mapping[str, Any]] = None, - progress_fn: Optional[ProgressCallbackFn_t] = None, - **kwargs: Any, - ): - kwargs["implicit_diff"] = None # not yet implemented - super().__init__( - lse_mode=lse_mode, - inner_iterations=inner_iterations, - use_danskin=use_danskin, - **kwargs - ) - self.rank = rank - self.gamma = gamma - self.gamma_rescale = gamma_rescale - self.epsilon = epsilon - self.initializer = initializer - self.progress_fn = progress_fn - # can be `None` - self.kwargs_dys = {} if kwargs_dys is None else kwargs_dys - self.kwargs_init = {} if kwargs_init is None else kwargs_init - - def __call__( - self, - ot_prob: linear_problem.LinearProblem, - init: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray], - Optional[jnp.ndarray]] = (None, None, None), - rng: Optional[jax.Array] = None, - **kwargs: Any, - ) -> LRSinkhornOutput: - """Run low-rank Sinkhorn. - - Args: - ot_prob: Linear OT problem. - init: Initial values for the low-rank factors: - - - :attr:`~ott.solvers.linear.sinkhorn_lr.LRSinkhornOutput.q`. - - :attr:`~ott.solvers.linear.sinkhorn_lr.LRSinkhornOutput.r`. - - :attr:`~ott.solvers.linear.sinkhorn_lr.LRSinkhornOutput.g`. - - Any `None` values will be initialized using the initializer. - rng: Random key for seeding. - kwargs: Additional arguments when calling the initializer. - - Returns: - The low-rank Sinkhorn output. - """ - initializer = self.create_initializer(ot_prob) - init = initializer(ot_prob, *init, rng=rng, **kwargs) - return run(ot_prob, self, init) - - def _get_costs( - self, - ot_prob: linear_problem.LinearProblem, - state: LRSinkhornState, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, float]: - log_q, log_r, log_g = ( - mu.safe_log(state.q), mu.safe_log(state.r), mu.safe_log(state.g) - ) - - inv_g = 1.0 / state.g[None, :] - tmp = ot_prob.geom.apply_cost(state.r, axis=1) - - grad_q = tmp * inv_g - grad_r = ot_prob.geom.apply_cost(state.q) * inv_g - grad_g = -jnp.sum(state.q * tmp, axis=0) / (state.g ** 2) - - grad_q += self.epsilon * log_q - grad_r += self.epsilon * log_r - grad_g += self.epsilon * log_g - - if self.gamma_rescale: - norm_q = jnp.max(jnp.abs(grad_q)) ** 2 - norm_r = jnp.max(jnp.abs(grad_r)) ** 2 - norm_g = jnp.max(jnp.abs(grad_g)) ** 2 - gamma = self.gamma / jnp.max(jnp.array([norm_q, norm_r, norm_g])) - else: - gamma = self.gamma - - eps_factor = 1.0 / (self.epsilon * gamma + 1.0) - gamma *= eps_factor - - c_q = -gamma * grad_q + eps_factor * log_q - c_r = -gamma * grad_r + eps_factor * log_r - c_g = -gamma * grad_g + eps_factor * log_g - - return c_q, c_r, c_g, gamma - - # TODO(michalk8): move to `lr_utils` when refactoring this - def dykstra_update_lse( - self, - c_q: jnp.ndarray, - c_r: jnp.ndarray, - h: jnp.ndarray, - gamma: float, - ot_prob: linear_problem.LinearProblem, - min_entry_value: float = 1e-6, - tolerance: float = 1e-3, - min_iter: int = 0, - inner_iter: int = 10, - max_iter: int = 10000 - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Run Dykstra's algorithm.""" - # shortcuts for problem's definition. - r = self.rank - n, m = ot_prob.geom.shape - loga, logb = jnp.log(ot_prob.a), jnp.log(ot_prob.b) - - h_old = h - g1_old, g2_old = jnp.zeros(r), jnp.zeros(r) - f1, f2 = jnp.zeros(n), jnp.zeros(m) - - w_gi, w_gp = jnp.zeros(r), jnp.zeros(r) - w_q, w_r = jnp.zeros(r), jnp.zeros(r) - err = jnp.inf - state_inner = f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err - constants = c_q, c_r, loga, logb - - def cond_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...] - ) -> bool: - del iteration, constants - *_, err = state_inner - return err > tolerance - - def _softm( - f: jnp.ndarray, g: jnp.ndarray, c: jnp.ndarray, axis: int - ) -> jnp.ndarray: - return jsp.special.logsumexp( - gamma * (f[:, None] + g[None, :] - c), axis=axis - ) - - def body_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...], compute_error: bool - ) -> Tuple[jnp.ndarray, ...]: - # TODO(michalk8): in the future, use `NamedTuple` - f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err = state_inner - c_q, c_r, loga, logb = constants - - # First Projection - f1 = jnp.where( - jnp.isfinite(loga), - (loga - _softm(f1, g1_old, c_q, axis=1)) / gamma + f1, loga - ) - f2 = jnp.where( - jnp.isfinite(logb), - (logb - _softm(f2, g2_old, c_r, axis=1)) / gamma + f2, logb - ) - - h = h_old + w_gi - h = jnp.maximum(jnp.log(min_entry_value) / gamma, h) - w_gi += h_old - h - h_old = h - - # Update couplings - g_q = _softm(f1, g1_old, c_q, axis=0) - g_r = _softm(f2, g2_old, c_r, axis=0) - - # Second Projection - h = (1.0 / 3.0) * (h_old + w_gp + w_q + w_r) - h += g_q / (3.0 * gamma) - h += g_r / (3.0 * gamma) - g1 = h + g1_old - g_q / gamma - g2 = h + g2_old - g_r / gamma - - w_q = w_q + g1_old - g1 - w_r = w_r + g2_old - g2 - w_gp = h_old + w_gp - h - - q, r, _ = recompute_couplings(f1, g1, c_q, f2, g2, c_r, h, gamma) - - g1_old = g1 - g2_old = g2 - h_old = h - - err = jax.lax.cond( - jnp.logical_and(compute_error, iteration >= min_iter), - lambda: solution_error(q, r, ot_prob, self.norm_error)[0], lambda: err - ) - - return f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err - - def recompute_couplings( - f1: jnp.ndarray, - g1: jnp.ndarray, - c_q: jnp.ndarray, - f2: jnp.ndarray, - g2: jnp.ndarray, - c_r: jnp.ndarray, - h: jnp.ndarray, - gamma: float, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - q = jnp.exp(gamma * (f1[:, None] + g1[None, :] - c_q)) - r = jnp.exp(gamma * (f2[:, None] + g2[None, :] - c_r)) - g = jnp.exp(gamma * h) - return q, r, g - - state_inner = fixed_point_loop.fixpoint_iter_backprop( - cond_fn, body_fn, min_iter, max_iter, inner_iter, constants, state_inner - ) - - f1, f2, g1_old, g2_old, h_old, _, _, _, _, _ = state_inner - return recompute_couplings(f1, g1_old, c_q, f2, g2_old, c_r, h_old, gamma) - - def dykstra_update_kernel( - self, - k_q: jnp.ndarray, - k_r: jnp.ndarray, - k_g: jnp.ndarray, - gamma: float, - ot_prob: linear_problem.LinearProblem, - min_entry_value: float = 1e-6, - tolerance: float = 1e-3, - min_iter: int = 0, - inner_iter: int = 10, - max_iter: int = 10000 - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Run Dykstra's algorithm.""" - # shortcuts for problem's definition. - rank = self.rank - n, m = ot_prob.geom.shape - a, b = ot_prob.a, ot_prob.b - supp_a, supp_b = a > 0, b > 0 - - g_old = k_g - v1_old, v2_old = jnp.ones(rank), jnp.ones(rank) - u1, u2 = jnp.ones(n), jnp.ones(m) - - q_gi, q_gp = jnp.ones(rank), jnp.ones(rank) - q_q, q_r = jnp.ones(rank), jnp.ones(rank) - err = jnp.inf - state_inner = u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err - constants = k_q, k_r, k_g, a, b - - def cond_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...] - ) -> bool: - del iteration, constants - *_, err = state_inner - return err > tolerance - - def body_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...], compute_error: bool - ) -> Tuple[jnp.ndarray, ...]: - # TODO(michalk8): in the future, use `NamedTuple` - u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err = state_inner - k_q, k_r, k_g, a, b = constants - - # First Projection - u1 = jnp.where(supp_a, a / jnp.dot(k_q, v1_old), 0.0) - u2 = jnp.where(supp_b, b / jnp.dot(k_r, v2_old), 0.0) - g = jnp.maximum(min_entry_value, g_old * q_gi) - q_gi = (g_old * q_gi) / g - g_old = g - - # Second Projection - v1_trans = jnp.dot(k_q.T, u1) - v2_trans = jnp.dot(k_r.T, u2) - g = (g_old * q_gp * v1_old * q_q * v1_trans * v2_old * q_r * - v2_trans) ** (1 / 3) - v1 = g / v1_trans - v2 = g / v2_trans - q_gp = (g_old * q_gp) / g - q_q = (q_q * v1_old) / v1 - q_r = (q_r * v2_old) / v2 - v1_old = v1 - v2_old = v2 - g_old = g - - # Compute Couplings - q, r, _ = recompute_couplings(u1, v1, k_q, u2, v2, k_r, g) - - err = jax.lax.cond( - jnp.logical_and(compute_error, iteration >= min_iter), - lambda: solution_error(q, r, ot_prob, self.norm_error)[0], lambda: err - ) - - return u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err - - def recompute_couplings( - u1: jnp.ndarray, - v1: jnp.ndarray, - k_q: jnp.ndarray, - u2: jnp.ndarray, - v2: jnp.ndarray, - k_r: jnp.ndarray, - g: jnp.ndarray, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - q = u1.reshape((-1, 1)) * k_q * v1.reshape((1, -1)) - r = u2.reshape((-1, 1)) * k_r * v2.reshape((1, -1)) - return q, r, g - - state_inner = fixed_point_loop.fixpoint_iter_backprop( - cond_fn, body_fn, min_iter, max_iter, inner_iter, constants, state_inner - ) - - u1, u2, v1_old, v2_old, g_old, _, _, _, _, _ = state_inner - return recompute_couplings(u1, v1_old, k_q, u2, v2_old, k_r, g_old) - - def lse_step( - self, ot_prob: linear_problem.LinearProblem, state: LRSinkhornState, - iteration: int - ) -> LRSinkhornState: - """LR Sinkhorn LSE update.""" - c_q, c_r, c_g, gamma = self._get_costs(ot_prob, state) - - if ot_prob.is_balanced: - c_q, c_r, h = c_q / -gamma, c_r / -gamma, c_g / gamma - q, r, g = self.dykstra_update_lse( - c_q, c_r, h, gamma, ot_prob, **self.kwargs_dys - ) - else: - q, r, g = lr_utils.unbalanced_dykstra_lse( - c_q, c_r, c_g, gamma, ot_prob, **self.kwargs_dys - ) - return state.set(q=q, g=g, r=r, gamma=gamma) - - def kernel_step( - self, ot_prob: linear_problem.LinearProblem, state: LRSinkhornState, - iteration: int - ) -> LRSinkhornState: - """LR Sinkhorn Kernel update.""" - c_q, c_r, c_g, gamma = self._get_costs(ot_prob, state) - c_q, c_r, c_g = jnp.exp(c_q), jnp.exp(c_r), jnp.exp(c_g) - - if ot_prob.is_balanced: - q, r, g = self.dykstra_update_kernel( - c_q, c_r, c_g, gamma, ot_prob, **self.kwargs_dys - ) - else: - q, r, g = lr_utils.unbalanced_dykstra_kernel( - c_q, c_r, c_g, gamma, ot_prob, **self.kwargs_dys - ) - return state.set(q=q, g=g, r=r, gamma=gamma) - - def one_iteration( - self, ot_prob: linear_problem.LinearProblem, state: LRSinkhornState, - iteration: int, compute_error: bool - ) -> LRSinkhornState: - """Carries out one low-rank Sinkhorn iteration. - - Depending on lse_mode, these iterations can be either in: - - - log-space for numerical stability. - - scaling space, using standard kernel-vector multiply operations. - - Args: - ot_prob: the transport problem definition - state: LRSinkhornState named tuple. - iteration: the current iteration of the Sinkhorn outer loop. - compute_error: flag to indicate this iteration computes/stores an error - - Returns: - The updated state. - """ - previous_state = state - it = iteration // self.inner_iterations - if self.lse_mode: # In lse_mode, run additive updates. - state = self.lse_step(ot_prob, state, iteration) - else: - state = self.kernel_step(ot_prob, state, iteration) - - # re-computes error if compute_error is True, else set it to inf. - cost = jax.lax.cond( - jnp.logical_and(compute_error, iteration >= self.min_iterations), - lambda: state.reg_ot_cost(ot_prob, epsilon=self.epsilon), - lambda: jnp.inf - ) - error = state.compute_error(previous_state) - crossed_threshold = jnp.logical_or( - state.crossed_threshold, - jnp.logical_and( - state.errors[it - 1] >= self.threshold, error < self.threshold - ) - ) - - state = state.set( - costs=state.costs.at[it].set(cost), - errors=state.errors.at[it].set(error), - crossed_threshold=crossed_threshold, - ) - - if self.progress_fn is not None: - jax.experimental.io_callback( - self.progress_fn, None, - (iteration, self.inner_iterations, self.max_iterations, state) - ) - - return state - - @property - def norm_error(self) -> Tuple[int]: # noqa: D102 - return self._norm_error, - - def create_initializer( - self, prob: linear_problem.LinearProblem - ) -> initializers_lr.LRInitializer: - """Create a low-rank Sinkhorn initializer. - - Args: - prob: Linear OT problem used to determine the initializer. - - Returns: - Low-rank initializer. - """ - if isinstance(self.initializer, initializers_lr.LRInitializer): - assert self.initializer.rank == self.rank, \ - f"Expected initializer's rank to be `{self.rank}`," \ - f"found `{self.initializer.rank}`." - return self.initializer - return initializers_lr.LRInitializer.from_solver( - self, kind=self.initializer, **self.kwargs_init - ) - - def init_state( - self, ot_prob: linear_problem.LinearProblem, - init: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] - ) -> LRSinkhornState: - """Return the initial state of the loop.""" - q, r, g = init - return LRSinkhornState( - q=q, - r=r, - g=g, - gamma=self.gamma, - costs=-jnp.ones(self.outer_iterations), - errors=-jnp.ones(self.outer_iterations), - crossed_threshold=False, - ) - - def output_from_state( - self, ot_prob: linear_problem.LinearProblem, state: LRSinkhornState - ) -> LRSinkhornOutput: - """Create an output from a loop state. - - Args: - ot_prob: the transport problem. - state: a LRSinkhornState. - - Returns: - A LRSinkhornOutput. - """ - return LRSinkhornOutput( - q=state.q, - r=state.r, - g=state.g, - ot_prob=ot_prob, - costs=state.costs, - errors=state.errors, - epsilon=self.epsilon, - inner_iterations=self.inner_iterations, - ) - - def _converged(self, state: LRSinkhornState, iteration: int) -> bool: - - def conv_crossed(prev_err: float, curr_err: float) -> bool: - return jnp.logical_and( - prev_err < self.threshold, curr_err < self.threshold - ) - - def conv_not_crossed(prev_err: float, curr_err: float) -> bool: - return jnp.logical_and(curr_err < prev_err, curr_err < self.threshold) - - # for convergence error, we consider 2 possibilities: - # 1. we either crossed the convergence threshold; in this case we require - # that the previous error was also below the threshold - # 2. we haven't crossed the threshold; in this case, we can be below or - # above the threshold: - # if we're above, we wait until we reach the convergence threshold and - # then, the above condition applies - # if we're below and we improved w.r.t. the previous iteration, - # we have converged; otherwise we continue, since we may be stuck - # in a local minimum (e.g., during the initial iterations) - - it = iteration // self.inner_iterations - return jax.lax.cond( - state.crossed_threshold, conv_crossed, conv_not_crossed, - state.errors[it - 2], state.errors[it - 1] - ) - - def _diverged(self, state: LRSinkhornState, iteration: int) -> bool: - it = iteration // self.inner_iterations - return jnp.logical_and( - jnp.logical_not(jnp.isfinite(state.errors[it - 1])), - jnp.logical_not(jnp.isfinite(state.costs[it - 1])) - ) - - -def run( - ot_prob: linear_problem.LinearProblem, - solver: LRSinkhorn, - init: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray], - Optional[jnp.ndarray]], -) -> LRSinkhornOutput: - """Run loop of the solver, outputting a state upgraded to an output.""" - out = sinkhorn.iterations(ot_prob, solver, init) - out = out.set_cost( - ot_prob, lse_mode=solver.lse_mode, use_danskin=solver.use_danskin - ) - return out.set(ot_prob=ot_prob) diff --git a/ott/build/lib/ott/solvers/linear/univariate.py b/ott/build/lib/ott/solvers/linear/univariate.py deleted file mode 100644 index 1bb7b13..0000000 --- a/ott/build/lib/ott/solvers/linear/univariate.py +++ /dev/null @@ -1,346 +0,0 @@ -# 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. -from typing import NamedTuple, Optional, Tuple, Union - -import jax -import jax.numpy as jnp - -from ott import utils -from ott.geometry import costs, pointcloud -from ott.math import utils as mu -from ott.problems.linear import linear_problem - -__all__ = [ - "UnivariateOutput", "UnivariateSolver", "uniform_distance", - "quantile_distance" -] - -Distance_t = Tuple[float, Optional[jnp.ndarray], Optional[jnp.ndarray]] - - -class UnivariateOutput(NamedTuple): # noqa: D101 - """Output of the :class:`~ott.solvers.linear.UnivariateSolver`. - - Objects of this class contain both solutions and problem definition of a - univariate OT problem. - - Args: - prob: OT problem between 2 weighted ``[n, d]`` and ``[m, d]`` point clouds. - ot_costs: ``[d,]`` optimal transport cost values, computed independently - along each of the ``d`` slices. - paired_indices: ``None`` if no transport was computed / recorded (e.g. when - using quantiles or subsampling approximations). Otherwise, output a tensor - of shape ``[d, 2, m+n]``, of ``m+n`` pairs of indices, for which the - optimal transport assigns mass, on each slice of the ``d`` slices - described in the dataset. Namely, for each index ``0<=k jnp.ndarray: - """Outputs a ``[d, n, m]`` tensor of all ``[n, m]`` transport matrices. - - This tensor will be extremely sparse, since it will have at most ``d(n+m)`` - non-zero values, out of ``dnm`` total entries. - """ - assert self.paired_indices is not None, \ - "[d, n, m] tensor of transports cannot be computed, likely because an" \ - " approximate method was used (using either subsampling or quantiles)." - - n, m = self.prob.geom.shape - if self.prob.is_equal_size and self.prob.is_uniform: - transport_matrices_from_indices = jax.vmap( - lambda idx, idy: jnp.eye(n)[idx, :][:, idy].T, in_axes=[0, 0] - ) - return transport_matrices_from_indices( - self.paired_indices[:, 0, :], self.paired_indices[:, 1, :] - ) - - # raveled indexing of entries. - indices = self.paired_indices[:, 0] * m + self.paired_indices[:, 1] - # segment sum is needed to collect several contributions - return jax.vmap( - lambda idx, mass: jax.ops.segment_sum( - mass, idx, indices_are_sorted=True, num_segments=n * m - ).reshape(n, m), - in_axes=[0, 0] - )(indices, self.mass_paired_indices) - - @property - def mean_transport_matrix(self) -> jnp.ndarray: - """Return the mean transport matrix, averaged over slices.""" - return jnp.mean(self.transport_matrices, axis=0) - - -@jax.tree_util.register_pytree_node_class -class UnivariateSolver: - r"""Univariate solver to compute 1D OT distance over slices of data. - - Computes 1-Dimensional optimal transport distance between two $d$-dimensional - point clouds. The total distance is the sum of univariate Wasserstein - distances on the $d$ slices of data: given two weighted point-clouds, stored - as ``[n, d]`` and ``[m, d]`` in a - :class:`~ott.problems.linear.linear_problem.LinearProblem` object, with - respective weights ``a`` and ``b``, the solver - computes ``d`` OT distances between each of these ``[n, 1]`` and ``[m, 1]`` - slices. The distance is computed using the analytical formula by default, - which involves sorting each of the slices independently. The optimal transport - matrices are also outputted when possible (described in sparse form, i.e. - pairs of indices and mass transferred between those indices). - - When weights ``a`` and ``b`` are uniform, and ``n=m``, the computation only - involves comparing sorted entries per slice, and ``d`` assignments are given. - - The user may also supply a ``num_subsamples`` parameter to extract as many - points from the original point cloud, sampled with probability masses ``a`` - and ``b``. This then simply applied the method above to the subsamples, to - output ``d`` costs, but assignments are not provided. - - When the problem is not uniform or not of equal size, the method defaults to - an inversion of the CDF, and outputs both costs and transport matrix in sparse - form. - - When a ``quantiles`` argument is passed, either specifying explicit quantiles - or a grid of quantiles, the distance is evaluated by comparing the quantiles - of the two point clouds on each slice. The OT costs are returned but - assignments are not provided. - - Args: - num_subsamples: Option to reduce the size of inputs by doing random - subsampling, taken into account marginal probabilities. - quantiles: When a vector or a number of quantiles is passed, the distance - is computed by evaluating the cost function on the sectional (one for each - dimension) quantiles of the two point cloud distributions described in the - problem. - """ - - def __init__( - self, - num_subsamples: Optional[int] = None, - quantiles: Optional[Union[int, jnp.ndarray]] = None, - ): - self._quantiles = quantiles - self.num_subsamples = num_subsamples - - @property - def quantiles(self) -> Optional[jnp.ndarray]: - """Quantiles' values used to evaluate OT cost.""" - if self._quantiles is None: - return None - if isinstance(self._quantiles, int): - return jnp.linspace(0.0, 1.0, self._quantiles) - return self._quantiles - - @property - def num_quantiles(self) -> int: - """Number of quantiles used to evaluate OT cost.""" - return 0 if self.quantiles is None else self.quantiles.shape[0] - - def __call__( - self, - prob: linear_problem.LinearProblem, - return_transport: bool = True, - rng: Optional[jax.Array] = None, - ) -> UnivariateOutput: - """Computes Univariate Distance between the ``d`` dimensional slices. - - Args: - prob: Problem with a :attr:`~ott.problems.linear.LinearProblem.geom` - attribute, the two point clouds ``x`` and ``y`` - (of respective sizes ``[n, d]`` and ``[m, d]``) and a ground - `TI cost ` between two scalars. - The ``[n,]`` and ``[m,]`` size probability weights vectors are stored - in attributes `:attr:`~ott.problems.linear.LinearProblem.a` and - :attr:`~ott.problems.linear.LinearProblem.b`. - return_transport: Whether to also return pairs of matched indices used to - compute optimal transport matrices. - rng: Used for random downsampling, if specified in the solver. - - Returns: - An output object, that computes ``d`` OT costs, in addition to, possibly, - paired lists of indices and their corresponding masses, on each of the - ``d`` dimensional slices of the input. - """ - geom = prob.geom - assert isinstance(geom, pointcloud.PointCloud), \ - "Geometry object in problem must be a PointCloud." - assert isinstance(geom.cost_fn, costs.TICost), \ - "Geometry's cost must be translation invariant." - - rng = utils.default_prng_key(rng) - - if self.num_subsamples: - x, y = self._subsample(prob, rng) - is_uniform_same_size = True - else: - # check if problem has the property uniform / same number of points - x, y = geom.x, geom.y - is_uniform_same_size = prob.is_uniform and prob.is_equal_size - - if self.quantiles is not None: - assert prob.is_uniform, \ - "The 'quantiles' method can only be used with uniform marginals." - out = _quant_dist(x, y, geom.cost_fn, self.quantiles, self.num_quantiles) - elif is_uniform_same_size: - return_transport = return_transport and not self.num_subsamples - out = uniform_distance(x, y, geom.cost_fn, return_transport) - else: - fn = jax.vmap(quantile_distance, in_axes=[1, 1, None, None, None, None]) - out = fn(x, y, geom.cost_fn, prob.a, prob.b, return_transport) - - return UnivariateOutput(prob, *out) - - def _subsample(self, prob: linear_problem.LinearProblem, - rng: jax.Array) -> Tuple[jnp.ndarray, jnp.ndarray]: - n, m = prob.geom.shape - x, y = prob.geom.x, prob.geom.y - - if prob.is_uniform: - x = x[jnp.linspace(0, n, num=self.num_subsamples).astype(int), :] - y = y[jnp.linspace(0, m, num=self.num_subsamples).astype(int), :] - return x, y - - rng1, rng2 = jax.random.split(rng, 2) - x = jax.random.choice(rng1, x, (self.num_subsamples,), p=prob.a, axis=0) - y = jax.random.choice(rng2, y, (self.num_subsamples,), p=prob.b, axis=0) - return x, y - - def tree_flatten(self): # noqa: D102 - return None, (self.num_subsamples, self._quantiles) - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - del children - return cls(*aux_data) - - -def uniform_distance( - x: jnp.ndarray, - y: jnp.ndarray, - cost_fn: costs.TICost, - return_transport: bool = True -) -> Distance_t: - """Distance between two equal-size families of uniformly weighted values x/y. - - Args: - x: Vector ``[n,]`` of real values. - y: Vector ``[n,]`` of real values. - cost_fn: Translation invariant cost function, i.e. ``c(x, y) = h(x - y)``. - return_transport: whether to return mapped pairs. - - Returns: - optimal transport cost, a list of ``n+m`` paired indices, and their - corresponding transport mass. Note that said mass can be null in some - entries, but sums to 1.0 - """ - n = x.shape[0] - i_x, i_y = jnp.argsort(x, axis=0), jnp.argsort(y, axis=0) - x = jnp.take_along_axis(x, i_x, axis=0) - y = jnp.take_along_axis(y, i_y, axis=0) - ot_costs = jax.vmap(cost_fn.h, in_axes=[0])(x.T - y.T) / n - - if return_transport: - paired_indices = jnp.stack([i_x, i_y]).transpose([2, 0, 1]) - mass_paired_indices = jnp.ones((n,)) / n - return ot_costs, paired_indices, mass_paired_indices - - return ot_costs, None, None - - -def quantile_distance( - x: jnp.ndarray, - y: jnp.ndarray, - cost_fn: costs.TICost, - a: jnp.ndarray, - b: jnp.ndarray, - return_transport: bool = True, -) -> Distance_t: - """Computes distance between quantile functions of distributions (a,x)/(b,y). - - Args: - x: Vector ``[n,]`` of real values. - y: Vector ``[m,]`` of real values. - cost_fn: Translation invariant cost function, i.e. ``c(x, y) = h(x - y)``. - a: Vector ``[n,]`` of non-negative weights summing to 1. - b: Vector ``[m,]`` of non-negative weights summing to 1. - return_transport: whether to return mapped pairs. - - Returns: - optimal transport cost, a list of ``n + m`` paired indices, and their - corresponding transport mass. Note that said mass can be null in some - entries, but sums to 1.0 - - Notes: - Inspired by :func:`~scipy.stats.wasserstein_distance`, - but can be used with other costs, not just :math:`c(x, y) = |x - y|`. - """ - x, i_x = mu.sort_and_argsort(x, argsort=True) - y, i_y = mu.sort_and_argsort(y, argsort=True) - - all_values = jnp.concatenate([x, y]) - all_values_sorted, all_values_sorter = mu.sort_and_argsort( - all_values, argsort=True - ) - - x_pdf = jnp.concatenate([a[i_x], jnp.zeros_like(b)])[all_values_sorter] - y_pdf = jnp.concatenate([jnp.zeros_like(a), b[i_y]])[all_values_sorter] - - x_cdf = jnp.cumsum(x_pdf) - y_cdf = jnp.cumsum(y_pdf) - - x_y_cdfs = jnp.concatenate([x_cdf, y_cdf]) - quantile_levels, _ = mu.sort_and_argsort(x_y_cdfs, argsort=False) - - i_x_cdf_inv = jnp.searchsorted(x_cdf, quantile_levels) - x_cdf_inv = all_values_sorted[i_x_cdf_inv] - i_y_cdf_inv = jnp.searchsorted(y_cdf, quantile_levels) - y_cdf_inv = all_values_sorted[i_y_cdf_inv] - - diff_q = jnp.diff(quantile_levels) - cost = jnp.sum( - jax.vmap(cost_fn.h)(y_cdf_inv[1:, None] - x_cdf_inv[1:, None]) * diff_q - ) - if not return_transport: - return cost, None, None - - n = x.shape[0] - - i_in_sorted_x_of_quantile = all_values_sorter[i_x_cdf_inv] % n - i_in_sorted_y_of_quantile = all_values_sorter[i_y_cdf_inv] - n - - orig_i = i_x[i_in_sorted_x_of_quantile][1:] - orig_j = i_y[i_in_sorted_y_of_quantile][1:] - - return cost, jnp.stack([orig_i, orig_j]), diff_q - - -def _quant_dist( - x: jnp.ndarray, y: jnp.ndarray, cost_fn: costs.TICost, q: jnp.ndarray, - n_q: int -) -> Tuple[jnp.ndarray, None, None]: - x_q = jnp.quantile(x, q, axis=0) - y_q = jnp.quantile(y, q, axis=0) - ot_costs = jax.vmap(cost_fn.pairwise, in_axes=[1, 1])(x_q, y_q) - - return ot_costs / n_q, None, None diff --git a/ott/build/lib/ott/solvers/quadratic/__init__.py b/ott/build/lib/ott/solvers/quadratic/__init__.py deleted file mode 100644 index 560ac3d..0000000 --- a/ott/build/lib/ott/solvers/quadratic/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -# 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. -from . import ( - gromov_wasserstein, - gromov_wasserstein_lr, - gw_barycenter, - lower_bound, -) -from ._solve import solve diff --git a/ott/build/lib/ott/solvers/quadratic/_solve.py b/ott/build/lib/ott/solvers/quadratic/_solve.py deleted file mode 100644 index 9cdefec..0000000 --- a/ott/build/lib/ott/solvers/quadratic/_solve.py +++ /dev/null @@ -1,91 +0,0 @@ -# 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. -from typing import Any, Literal, Optional, Union - -import jax.numpy as jnp - -from ott.geometry import geometry -from ott.problems.quadratic import quadratic_costs, quadratic_problem -from ott.solvers.quadratic import gromov_wasserstein as gw -from ott.solvers.quadratic import gromov_wasserstein_lr as lrgw - -__all__ = ["solve"] - - -def solve( - geom_xx: geometry.Geometry, - geom_yy: geometry.Geometry, - geom_xy: Optional[geometry.Geometry] = None, - fused_penalty: float = 1.0, - a: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, - tau_a: float = 1.0, - tau_b: float = 1.0, - loss: Union[Literal["sqeucl", "kl"], quadratic_costs.GWLoss] = "sqeucl", - gw_unbalanced_correction: bool = True, - rank: int = -1, - **kwargs: Any, -) -> Union[gw.GWOutput, lrgw.LRGWOutput]: - """Solve quadratic regularized OT problem using a Gromov-Wasserstein solver. - - Args: - geom_xx: Ground geometry of the first space. - geom_yy: Ground geometry of the second space. - geom_xy: Geometry defining the linear penalty term for - fused Gromov-Wasserstein :cite:`vayer:19`. If :obj:`None`, the problem - reduces to a plain Gromov-Wasserstein problem :cite:`peyre:16`. - fused_penalty: Multiplier of the linear term in fused Gromov-Wasserstein, - i.e. ``problem = purely quadratic + fused_penalty * linear problem``. - a: The first marginal. If :obj:`None`, it will be uniform. - b: The second marginal. If :obj:`None`, it will be uniform. - tau_a: If :math:`< 1`, defines how much unbalanced the problem is - on the first marginal. - tau_b: If :math:`< 1`, defines how much unbalanced the problem is - on the second marginal. - loss: Gromov-Wasserstein loss function, see - :class:`~ott.problems.quadratic.quadratic_costs.GWLoss` for more - information. If ``rank > 0``, ``'sqeucl'`` is always used. - gw_unbalanced_correction: Whether the unbalanced version of - :cite:`sejourne:21` is used. Otherwise, ``tau_a`` and ``tau_b`` - only affect the resolution of the linearization of the GW problem - in the inner loop. Only used when ``rank = -1``. - rank: Rank constraint on the coupling to minimize the quadratic OT problem - :cite:`scetbon:22`. If :math:`-1`, no rank constraint is used. - kwargs: Keyword arguments for - :class:`~ott.solvers.quadratic.gromov_wasserstein.GromovWasserstein` or - :class:`~ott.solvers.quadratic.gromov_wasserstein_lr.LRGromovWasserstein`, - depending on the ``rank`` - - Returns: - The Gromov-Wasserstein output. - """ - prob = quadratic_problem.QuadraticProblem( - geom_xx=geom_xx, - geom_yy=geom_yy, - geom_xy=geom_xy, - fused_penalty=fused_penalty, - a=a, - b=b, - tau_a=tau_a, - tau_b=tau_b, - loss=loss, - gw_unbalanced_correction=gw_unbalanced_correction - ) - - if rank > 0: - solver = lrgw.LRGromovWasserstein(rank=rank, **kwargs) - else: - solver = gw.GromovWasserstein(rank=rank, **kwargs) - - return solver(prob) diff --git a/ott/build/lib/ott/solvers/quadratic/gromov_wasserstein.py b/ott/build/lib/ott/solvers/quadratic/gromov_wasserstein.py deleted file mode 100644 index 7272bd1..0000000 --- a/ott/build/lib/ott/solvers/quadratic/gromov_wasserstein.py +++ /dev/null @@ -1,461 +0,0 @@ -# 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. -from typing import ( - Any, - Callable, - Dict, - Literal, - Mapping, - NamedTuple, - Optional, - Sequence, - Tuple, - Union, -) - -import jax -import jax.experimental -import jax.numpy as jnp -import numpy as np - -from ott import utils -from ott.geometry import geometry -from ott.initializers.quadratic import initializers as quad_initializers -from ott.math import fixed_point_loop -from ott.problems.linear import linear_problem -from ott.problems.quadratic import quadratic_problem -from ott.solvers import was_solver -from ott.solvers.linear import sinkhorn, sinkhorn_lr - -__all__ = ["GromovWasserstein", "GWOutput"] - -LinearOutput = Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput] - -ProgressCallbackFn_t = Callable[ - [Tuple[np.ndarray, np.ndarray, np.ndarray, "GWState"]], None -] - - -class GWOutput(NamedTuple): - """Holds the output of the Gromov-Wasserstein solver. - - Args: - costs: Holds the sequence of regularized GW costs seen through the outer - loop of the solver. - linear_convergence: Holds the sequence of bool convergence flags of the - inner Sinkhorn iterations. - converged: Convergence flag for the outer GW iterations. - errors: Holds sequence of vectors of errors of the Sinkhorn algorithm - at each iteration. - linear_state: State used to solve and store solutions to the local - linearization of GW. - geom: The geometry underlying the local linearization. - old_transport_mass: Holds total mass of transport at previous iteration. - """ - - costs: Optional[jnp.ndarray] = None - linear_convergence: Optional[jnp.ndarray] = None - converged: bool = False - errors: Optional[jnp.ndarray] = None - linear_state: Optional[LinearOutput] = None - geom: Optional[geometry.Geometry] = None - # Intermediate values. - old_transport_mass: float = 1.0 - - def set(self, **kwargs: Any) -> "GWOutput": - """Return a copy of self, possibly with overwrites.""" - return self._replace(**kwargs) - - @property - def matrix(self) -> jnp.ndarray: - """Transport matrix.""" - return self._rescale_factor * self.linear_state.matrix - - def apply(self, inputs: jnp.ndarray, axis: int = 0) -> jnp.ndarray: - """Apply the transport to an array; axis=1 for its transpose.""" - return self._rescale_factor * self.linear_state.apply(inputs, axis=axis) - - @property - def reg_gw_cost(self) -> float: - """Regularized optimal transport cost of the linearization.""" - return self.linear_state.reg_ot_cost - - @property - def _rescale_factor(self) -> float: - return jnp.sqrt( - self.old_transport_mass / self.linear_state.transport_mass - ) - - @property - def primal_cost(self) -> float: - """Return transport cost of current linear OT solution at geometry.""" - return self.linear_state.transport_cost_at_geom(other_geom=self.geom) - - @property - def n_iters(self) -> int: # noqa: D102 - if self.errors is None: - return -1 - return jnp.sum(self.errors[:, 0] != -1) - - -class GWState(NamedTuple): - """State of the Gromov-Wasserstein solver. - - Attributes: - costs: Holds the sequence of regularized GW costs seen through the outer - loop of the solver. - linear_convergence: Holds the sequence of bool convergence flags of the - inner Sinkhorn iterations. - linear_state: State used to solve and store solutions to the local - linearization of GW. - linear_pb: Local linearization of the quadratic GW problem. - old_transport_mass: Intermediary value of the mass of the transport matrix. - rngs: Random keys passed to low-rank initializers at every GW iteration - when not using warm start. - errors: Holds sequence of vectors of errors of the Sinkhorn algorithm - at each iteration. - """ - - costs: jnp.ndarray - linear_convergence: jnp.ndarray - linear_state: LinearOutput - linear_pb: linear_problem.LinearProblem - old_transport_mass: float - rngs: Optional[jax.Array] = None - errors: Optional[jnp.ndarray] = None - - def set(self, **kwargs: Any) -> "GWState": - """Return a copy of self, possibly with overwrites.""" - return self._replace(**kwargs) - - def update( # noqa: D102 - self, - iteration: int, - linear_sol: LinearOutput, - linear_pb: linear_problem.LinearProblem, - store_errors: bool, - old_transport_mass: float, - ) -> "GWState": - costs = self.costs.at[iteration].set(linear_sol.reg_ot_cost) - errors = None - if store_errors and self.errors is not None: - errors = self.errors.at[iteration, :].set(linear_sol.errors) - linear_convergence = self.linear_convergence.at[iteration].set( - linear_sol.converged - ) - - return self.set( - linear_state=linear_sol, - linear_pb=linear_pb, - costs=costs, - linear_convergence=linear_convergence, - errors=errors, - old_transport_mass=old_transport_mass, - ) - - -@jax.tree_util.register_pytree_node_class -class GromovWasserstein(was_solver.WassersteinSolver): - """Gromov-Wasserstein solver :cite:`peyre:16`. - - .. seealso:: - Low-rank Gromov-Wasserstein :cite:`scetbon:23` is implemented in - :class:`~ott.solvers.quadratic.gromov_wasserstein_lr.LRGromovWasserstein`. - - Args: - args: Positional arguments for - :class:`~ott.solvers.was_solver.WassersteinSolver`. - warm_start: Whether to initialize Sinkhorn calls using values - from the previous iteration. If :obj:`None`, warm starts are not used for - standard Sinkhorn. - relative_epsilon: Whether to use relative epsilon in the linearized - geometry. - quad_initializer: Quadratic initializer. If the solver is entropic, - :class:`~ott.initializers.quadratic.initializers.QuadraticInitializer` - is always used. - progress_fn: callback function which gets called during the - Gromov-Wasserstein iterations, so the user can display the error at each - iteration, e.g., using a progress bar. - See :func:`~ott.utils.default_progress_fn` for a basic implementation. - kwargs_init: Keyword arguments when creating the initializer. - kwargs: Keyword arguments for - :class:`~ott.solvers.was_solver.WassersteinSolver`. - """ - - def __init__( - self, - *args: Any, - warm_start: Optional[bool] = None, - relative_epsilon: Optional[bool] = None, - quad_initializer: Optional[ - Union[ - Literal["random", "rank2", "k-means", "generalized-k-means"], - quad_initializers.BaseQuadraticInitializer, - ] - ] = None, - progress_fn: Optional[ProgressCallbackFn_t] = None, - kwargs_init: Optional[Mapping[str, Any]] = None, - kwargs_sinkhorn: Any = {}, - **kwargs: Any, - ): - super().__init__(kwargs=kwargs_sinkhorn, *args, **kwargs) - assert not self.is_low_rank, ( - "For low-rank GW, use " - "`ott.solvers.quadratic.gromov_wasserstein_lr.LRGromovWasserstein`." - ) - self._warm_start = warm_start - self.relative_epsilon = relative_epsilon - self.quad_initializer = quad_initializer - self.progress_fn = progress_fn - self.kwargs_init = {} if kwargs_init is None else kwargs_init - - def __call__( - self, - prob: quadratic_problem.QuadraticProblem, - init: Optional[linear_problem.LinearProblem] = None, - rng: Optional[jax.Array] = None, - **kwargs: Any, - ) -> GWOutput: - """Run the Gromov-Wasserstein solver. - - Args: - prob: Quadratic OT problem. - init: Initial linearization of the quadratic problem. If `None`, it will - be computed using the initializer. - rng: Random number key. - kwargs: Keyword arguments used when calling the initializer. - - Returns: - The Gromov-Wasserstein output. - """ - rng = utils.default_prng_key(rng) - rng1, rng2 = jax.random.split(rng, 2) - - # if prob._is_low_rank_convertible: - # prob = prob.to_low_rank() - - if init is None: - initializer = self.create_initializer(prob) - init = initializer( - prob, - epsilon=self.epsilon, - rng=rng1, - relative_epsilon=self.relative_epsilon, - **kwargs, - ) - print("GW called") - out = iterations(self, prob, init, rng2) - # TODO(lpapaxanthoos): remove stop_gradient when using backprop - if self.is_low_rank: - linearization = prob.update_lr_linearization( - jax.lax.stop_gradient(out.linear_state), - relative_epsilon=self.relative_epsilon, - ) - else: - linearization = prob.update_linearization( - jax.lax.stop_gradient(out.linear_state), - epsilon=self.epsilon, - old_transport_mass=jax.lax.stop_gradient( - out.old_transport_mass - ), - relative_epsilon=self.relative_epsilon, - ) - - linear_state = out.linear_state.set_cost(linearization, True, True) - iteration = jnp.sum(out.costs != -1) - converged = jnp.logical_and( - iteration < self.max_iterations, jnp.all(out.linear_convergence) - ) - return out.set( - linear_state=linear_state, - geom=linearization.geom, - converged=converged, - ) - - def init_state( - self, - prob: quadratic_problem.QuadraticProblem, - init: linear_problem.LinearProblem, - rng: jax.Array, - ) -> GWState: - """Initialize the state of the Gromov-Wasserstein iterations. - - Args: - prob: Quadratic OT problem. - init: Initial linearization of the quadratic problem. - rng: Random key for low-rank initializers. Only used when - :attr:`warm_start` is `False`. - - Returns: - The initial Gromov-Wasserstein state. - """ - linear_state = self.linear_ot_solver(init) - num_iter = self.max_iterations - transport_mass = prob.init_transport_mass() - if self.store_inner_errors: - errors = -jnp.ones( - (num_iter, self.linear_ot_solver.outer_iterations) - ) - else: - errors = None - - return GWState( - costs=-jnp.ones((num_iter,)), - linear_convergence=-jnp.ones((num_iter,)), - linear_state=linear_state, - linear_pb=init, - old_transport_mass=transport_mass, - rngs=jax.random.split(rng, num_iter), - errors=errors, - ) - - def output_from_state( - self, - state: GWState, - ) -> GWOutput: - """Create an output from a loop state. - - Arguments: - state: A GWState. - - Returns: - A GWOutput. - """ - return GWOutput( - costs=state.costs, - linear_convergence=state.linear_convergence, - errors=state.errors, - linear_state=state.linear_state, - geom=state.linear_pb.geom, - old_transport_mass=state.old_transport_mass, - ) - - def create_initializer( - self, prob: quadratic_problem.QuadraticProblem - ) -> quad_initializers.BaseQuadraticInitializer: - """Create quadratic, possibly low-rank initializer. - - Args: - prob: Quadratic OT problem used to determine the initializer. - - Returns: - The initializer. - """ - if isinstance( - self.quad_initializer, quad_initializers.BaseQuadraticInitializer - ): - return self.quad_initializer - # no other options implemented, use the default - return quad_initializers.QuadraticInitializer(**self.kwargs_init) - - @property - def warm_start(self) -> bool: - """Whether to initialize Sinkhorn using previous solutions.""" - return ( - self.is_low_rank if self._warm_start is None else self._warm_start - ) - - def tree_flatten( - self, - ) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - children, aux_data = super().tree_flatten() - aux_data["warm_start"] = self._warm_start - aux_data["progress_fn"] = self.progress_fn - aux_data["relative_epsilon"] = self.relative_epsilon - aux_data["quad_initializer"] = self.quad_initializer - aux_data["kwargs_init"] = self.kwargs_init - return children, aux_data - - -def iterations( - solver: GromovWasserstein, - prob: quadratic_problem.QuadraticProblem, - init: linear_problem.LinearProblem, - rng: jax.Array, -) -> GWOutput: - """Jittable Gromov-Wasserstein outer loop.""" - - def cond_fn( - iteration: int, solver: GromovWasserstein, state: GWState - ) -> bool: - return solver._continue(state, iteration) - - def body_fn( - iteration: int, - solver: GromovWasserstein, - state: GWState, - compute_error: bool, - ) -> GWState: - del compute_error # always assumed true for the outer loop of GW - - lin_state = state.linear_state - if solver.is_low_rank: - rng = state.rngs[iteration] - init = ( - (lin_state.q, lin_state.r, lin_state.g) - if solver.warm_start - else (None, None, None) - ) - linear_pb = prob.update_lr_linearization( - state.linear_state, relative_epsilon=solver.relative_epsilon - ) - out = solver.linear_ot_solver(linear_pb, init=init, rng=rng) - else: - init = ( - (lin_state.f, lin_state.g) - if solver.warm_start - else (None, None) - ) - linear_pb = prob.update_linearization( - lin_state, - solver.epsilon, - state.old_transport_mass, - relative_epsilon=solver.relative_epsilon, - ) - prob.linear_problem = linear_pb - out = solver.linear_ot_solver(linear_pb, init=init) - - old_transport_mass = jax.lax.stop_gradient( - state.linear_state.transport_mass - ) - new_state = state.update( - iteration, - out, - linear_pb, - solver.store_inner_errors, - old_transport_mass, - ) - - # Inner iterations is currently fixed to 1. - inner_iterations = 1 - if solver.progress_fn is not None: - jax.experimental.io_callback( - solver.progress_fn, - None, - (iteration, inner_iterations, solver.max_iterations, state), - ) - - return new_state - - state = fixed_point_loop.fixpoint_iter( - cond_fn=cond_fn, - body_fn=body_fn, - min_iterations=solver.min_iterations, - max_iterations=solver.max_iterations, - inner_iterations=1, - constants=solver, - state=solver.init_state(prob, init, rng=rng), - ) - - return solver.output_from_state(state) diff --git a/ott/build/lib/ott/solvers/quadratic/gromov_wasserstein_lr.py b/ott/build/lib/ott/solvers/quadratic/gromov_wasserstein_lr.py deleted file mode 100644 index cb12911..0000000 --- a/ott/build/lib/ott/solvers/quadratic/gromov_wasserstein_lr.py +++ /dev/null @@ -1,897 +0,0 @@ -# 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. -"""A Jax implementation of the unbalanced low-rank GW algorithm.""" -from typing import ( - Any, - Callable, - Literal, - Mapping, - NamedTuple, - Optional, - Tuple, - Union, -) - -import jax -import jax.experimental -import jax.numpy as jnp -import jax.scipy as jsp -import numpy as np - -from ott import utils -from ott.geometry import geometry, low_rank -from ott.initializers.linear import initializers_lr -from ott.math import fixed_point_loop -from ott.math import unbalanced_functions as uf -from ott.math import utils as mu -from ott.problems.quadratic import quadratic_problem -from ott.solvers.linear import lr_utils, sinkhorn - -__all__ = ["LRGromovWasserstein", "LRGWOutput"] - -ProgressCallbackFn_t = Callable[ - [Tuple[np.ndarray, np.ndarray, np.ndarray, "LRGWState"]], None] - - -class LRGWState(NamedTuple): - """State of the low-rank GW algorithm.""" - q: jnp.ndarray - r: jnp.ndarray - g: jnp.ndarray - gamma: float - costs: jnp.ndarray - errors: jnp.ndarray - crossed_threshold: bool - - def compute_error( # noqa: D102 - self, previous_state: "LRGWState" - ) -> float: - err_q = mu.gen_js(self.q, previous_state.q, c=1.0) - err_r = mu.gen_js(self.r, previous_state.r, c=1.0) - err_g = mu.gen_js(self.g, previous_state.g, c=1.0) - - return ((1.0 / self.gamma) ** 2) * (err_q + err_r + err_g) - - def reg_gw_cost( # noqa: D102 - self, - ot_prob: quadratic_problem.QuadraticProblem, - *, - epsilon: float, - use_danskin: bool = False - ) -> float: - return compute_reg_gw_cost( - self.q, - self.r, - self.g, - ot_prob, - epsilon=epsilon, - use_danskin=use_danskin - ) - - def set(self, **kwargs: Any) -> "LRGWState": - """Return a copy of self, with potential overwrites.""" - return self._replace(**kwargs) - - -def compute_reg_gw_cost( - q: jnp.ndarray, - r: jnp.ndarray, - g: jnp.ndarray, - ot_prob: quadratic_problem.QuadraticProblem, - epsilon: float, - use_danskin: bool = False -) -> float: - """Compute the regularized OT cost, here the primal cost of the LR solution. - - Args: - q: first factor of solution - r: second factor of solution - g: weights of solution - ot_prob: linear problem - epsilon: Entropic regularization. - use_danskin: if True, use Danskin's theorem :cite:`danskin:67,bertsekas:71` - to avoid computing the gradient of the cost function. - - Returns: - regularized OT cost, the (primal) transport cost of the low-rank solution. - """ - - def ent(x: jnp.ndarray) -> float: - # generalized entropy - return jnp.sum(jsp.special.entr(x) + x) - - q = jax.lax.stop_gradient(q) if use_danskin else q - r = jax.lax.stop_gradient(r) if use_danskin else r - g = jax.lax.stop_gradient(g) if use_danskin else g - - out = LRGWOutput( - q=q, - r=r, - g=g, - ot_prob=ot_prob, - costs=None, - errors=None, - epsilon=None, - inner_iterations=None, - ) - - cost = out.primal_cost - epsilon * (ent(q) + ent(r) + ent(g)) - if ot_prob.tau_a != 1.0: - rho_a = uf.rho(1.0, ot_prob.tau_a) - cost += rho_a * mu.gen_kl(jnp.sum(q, axis=1), ot_prob.a) - if ot_prob.tau_b != 1.0: - rho_b = uf.rho(1.0, ot_prob.tau_b) - cost += rho_b * mu.gen_kl(jnp.sum(r, axis=1), ot_prob.b) - - return cost - - -class LRGWOutput(NamedTuple): - """Transport interface for a low-rank GW solution.""" - q: jnp.ndarray - r: jnp.ndarray - g: jnp.ndarray - costs: jnp.ndarray - # TODO(michalk8): must be called `errors`, because of `store_inner_errors` - # in future, enforce via class hierarchy - errors: jnp.ndarray - ot_prob: quadratic_problem.QuadraticProblem - epsilon: float - inner_iterations: int - reg_gw_cost: Optional[float] = None - - def set(self, **kwargs: Any) -> "LRGWOutput": - """Return a copy of self, with potential overwrites.""" - return self._replace(**kwargs) - - def set_cost( # noqa: D102 - self, - ot_prob: quadratic_problem.QuadraticProblem, - lse_mode: bool, - use_danskin: bool = False - ) -> "LRGWOutput": - del lse_mode - return self.set(reg_gw_cost=self.compute_reg_gw_cost(ot_prob, use_danskin)) - - def compute_reg_gw_cost( # noqa: D102 - self, - ot_prob: quadratic_problem.QuadraticProblem, - use_danskin: bool = False, - ) -> float: - return compute_reg_gw_cost( - self.q, - self.r, - self.g, - ot_prob, - epsilon=self.epsilon, - use_danskin=use_danskin - ) - - @property - def geom(self) -> geometry.Geometry: # noqa: D102 - """Linearized geometry.""" - return _linearized_geometry(self.ot_prob, q=self.q, r=self.r, g=self.g) - - @property - def a(self) -> jnp.ndarray: # noqa: D102 - return self.ot_prob.a - - @property - def b(self) -> jnp.ndarray: # noqa: D102 - return self.ot_prob.b - - @property - def n_iters(self) -> int: # noqa: D102 - return jnp.sum(self.errors != -1) * self.inner_iterations - - @property - def converged(self) -> bool: # noqa: D102 - return jnp.logical_and( - jnp.any(self.costs == -1), jnp.all(jnp.isfinite(self.costs)) - ) - - @property - def matrix(self) -> jnp.ndarray: - """Transport matrix if it can be instantiated.""" - return (self.q * self._inv_g) @ self.r.T - - def apply(self, inputs: jnp.ndarray, axis: int = 0) -> jnp.ndarray: - """Apply the transport to a array; axis=1 for its transpose.""" - q, r = (self.q, self.r) if axis == 1 else (self.r, self.q) - # for `axis=0`: (batch, m), (m, r), (r,), (r, n) - return ((inputs @ r) * self._inv_g) @ q.T - - def marginal(self, axis: int) -> jnp.ndarray: # noqa: D102 - length = self.q.shape[0] if axis == 0 else self.r.shape[0] - return self.apply(jnp.ones(length,), axis=axis) - - def cost_at_geom(self, other_geom: geometry.Geometry) -> float: - """Return OT cost for current solution, evaluated at any cost matrix.""" - return jnp.sum(self.q * other_geom.apply_cost(self.r, axis=1) * self._inv_g) - - def transport_cost_at_geom(self, other_geom: geometry.Geometry) -> float: - """Return (by recomputing it) bare transport cost of current solution.""" - return self.cost_at_geom(other_geom) - - @property - def primal_cost(self) -> float: - """Return (by recomputing it) transport cost of current solution.""" - geom_xx, geom_yy = self.ot_prob.geom_xx, self.ot_prob.geom_yy - marginal_a = self.ot_prob.a if self.ot_prob.tau_a == 1.0 else self.q.sum(1) - marginal_b = self.ot_prob.b if self.ot_prob.tau_b == 1.0 else self.r.sum(1) - - quad_cost = 0.5 * self.transport_cost_at_geom(other_geom=self.geom) - quad_cost += jnp.vdot(geom_xx.apply_square_cost(marginal_a), marginal_a) - quad_cost += jnp.vdot(geom_yy.apply_square_cost(marginal_b), marginal_b) - - if not self.ot_prob.is_fused: - return quad_cost - - alpha = self.ot_prob.fused_penalty / (self.ot_prob.fused_penalty + 1.0) - norm_g = jnp.linalg.norm(self.g, ord=1) - - lin_cost = self.cost_at_geom(self.ot_prob.geom_xy) - return alpha * norm_g * lin_cost + (1.0 - alpha) * quad_cost - - @property - def transport_mass(self) -> float: - """Sum of transport matrix.""" - return self.marginal(0).sum() - - @property - def _inv_g(self) -> jnp.ndarray: - return 1.0 / self.g - - -@jax.tree_util.register_pytree_node_class -class LRGromovWasserstein(sinkhorn.Sinkhorn): - r"""Low-rank Gromov-Wasserstein solver :cite:`scetbon:23`. - - The algorithm minimizes a non-convex problem. It therefore requires special - care to initialization and convergence. Convergence is evaluated on successive - evaluations of the objective. - - .. warning:: - This solver only for the **unbalanced** case. Balanced case is implemented - in :class:`~ott.solvers.quadratic.gromov_wasserstein.GromovWasserstein` - and will be unified here in the future release. - - Args: - rank: Rank constraint on the coupling to minimize the linear OT problem - gamma: The (inverse of) gradient step size used by mirror descent. - gamma_rescale: Whether to rescale :math:`\gamma` every iteration as - described in :cite:`scetbon:22b`. - epsilon: Entropic regularization added on top of low-rank problem. - initializer: How to initialize the :math:`Q`, :math:`R` and :math:`g` - factors. - lse_mode: Whether to run computations in LSE or kernel mode. - inner_iterations: Number of inner iterations used by the algorithm before - re-evaluating progress. - use_danskin: Use Danskin theorem to evaluate gradient of objective w.r.t. - input parameters. Only `True` handled at this moment. - implicit_diff: Whether to use implicit differentiation. Currently, only - ``implicit_diff = False`` is implemented. - progress_fn: callback function which gets called during the GW - iterations, so the user can display the error at each iteration, - e.g., using a progress bar. See :func:`~ott.utils.default_progress_fn` - for a basic implementation. - kwargs_dys: Keyword arguments passed to :meth:`dykstra_update_lse`, - :meth:`dykstra_update_kernel` or one of the functions defined in - :mod:`ott.solvers.linear`, depending on the ``lse_mode``. - kwargs_init: Keyword arguments for - :class:`~ott.initializers.linear.initializers_lr.LRInitializer`. - kwargs: Keyword arguments for - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. - """ - - def __init__( - self, - rank: int, - gamma: float = 10.0, - gamma_rescale: bool = True, - epsilon: float = 0.0, - initializer: Union[Literal["random", "rank2", "k-means", - "generalized-k-means"], - initializers_lr.LRInitializer] = "random", - lse_mode: bool = True, - inner_iterations: int = 10, - use_danskin: bool = True, - implicit_diff: bool = False, - kwargs_dys: Optional[Mapping[str, Any]] = None, - kwargs_init: Optional[Mapping[str, Any]] = None, - progress_fn: Optional[ProgressCallbackFn_t] = None, - **kwargs: Any, - ): - assert not implicit_diff, "Implicit diff. not yet implemented." - super().__init__( - lse_mode=lse_mode, - inner_iterations=inner_iterations, - use_danskin=use_danskin, - implicit_diff=implicit_diff, - **kwargs - ) - self.rank = rank - self.gamma = gamma - self.gamma_rescale = gamma_rescale - self.epsilon = epsilon - self.initializer = initializer - self.progress_fn = progress_fn - # can be `None` - self.kwargs_dys = {} if kwargs_dys is None else kwargs_dys - self.kwargs_init = {} if kwargs_init is None else kwargs_init - - def __call__( - self, - ot_prob: quadratic_problem.QuadraticProblem, - init: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray], - Optional[jnp.ndarray]] = (None, None, None), - rng: Optional[jax.Array] = None, - **kwargs: Any, - ) -> LRGWOutput: - """Run low-rank Gromov-Wasserstein solver. - - Args: - ot_prob: Linear OT problem. - init: Initial values for the low-rank factors: - - - :attr:`~ott.solvers.linear.sinkhorn_lr.LRGWOutput.q`. - - :attr:`~ott.solvers.linear.sinkhorn_lr.LRGWOutput.r`. - - :attr:`~ott.solvers.linear.sinkhorn_lr.LRGWOutput.g`. - - Any `None` values will be initialized using the initializer. - rng: Random key for seeding. - kwargs: Additional arguments when calling the initializer. - - Returns: - The low-rank GW output. - """ - rng = utils.default_prng_key(rng) - rng_lrc, rng_init = jax.random.split(rng) - - if ot_prob._is_low_rank_convertible: - ot_prob = ot_prob.to_low_rank(rng=rng_lrc) - - initializer = self.create_initializer(ot_prob) - init = initializer(ot_prob, *init, rng=rng_init, **kwargs) - return run(ot_prob, self, init) - - def _get_costs( - self, - ot_prob: quadratic_problem.QuadraticProblem, - state: LRGWState, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, float]: - q, r, g = state.q, state.r, state.g - log_q, log_r, log_g = mu.safe_log(q), mu.safe_log(r), mu.safe_log(g) - inv_g = 1.0 / g[None, :] - lin_geom = _linearized_geometry(ot_prob, q=q, r=r, g=g) - - tmp = lin_geom.apply_cost(r, axis=1) - grad_q = tmp * inv_g - if ot_prob.tau_a != 1.0: # unbalanced grad - grad_q += 2.0 * ot_prob.geom_xx.apply_square_cost(q.sum(1), axis=1) - - grad_r = lin_geom.apply_cost(q, axis=0) * inv_g - if ot_prob.tau_b != 1.0: # unbalanced grad - grad_r += 2.0 * ot_prob.geom_yy.apply_square_cost(r.sum(1), axis=1) - - omega_quad = jnp.sum(q * tmp, axis=0) - grad_g = -omega_quad / (g ** 2) - - if ot_prob.is_fused: - alpha = ot_prob.fused_penalty / (ot_prob.fused_penalty + 1.0) - norm_g = jnp.linalg.norm(g, ord=1) - - tmp = ot_prob.geom_xy.apply_cost(r, axis=1) - lin_grad_q = tmp * inv_g * norm_g - lin_grad_r = ot_prob.geom_xy.apply_cost(q) * inv_g * norm_g - - omega_lin = jnp.sum(q * tmp, axis=0) - lin_grad_g = -omega_lin / (g ** 2) * norm_g + jnp.sum(q * tmp * inv_g) - - grad_q = alpha * lin_grad_q + (1.0 - alpha) * grad_q - grad_r = alpha * lin_grad_r + (1.0 - alpha) * grad_r - grad_g = alpha * lin_grad_g + (1.0 - alpha) * grad_g - - grad_q += self.epsilon * log_q - grad_r += self.epsilon * log_r - grad_g += self.epsilon * log_g - - if self.gamma_rescale: - norm_q = jnp.max(jnp.abs(grad_q)) ** 2 - norm_r = jnp.max(jnp.abs(grad_r)) ** 2 - norm_g = jnp.max(jnp.abs(grad_g)) ** 2 - gamma = self.gamma / jnp.max(jnp.array([norm_q, norm_r, norm_g])) - else: - gamma = self.gamma - - eps_factor = 1.0 / (self.epsilon * gamma + 1.0) - gamma *= eps_factor - - c_q = -gamma * grad_q + eps_factor * log_q - c_r = -gamma * grad_r + eps_factor * log_r - c_g = -gamma * grad_g + eps_factor * log_g - - return c_q, c_r, c_g, gamma - - # TODO(michalk8): move to `lr_utils` when refactoring this the future - def dykstra_update_lse( - self, - c_q: jnp.ndarray, - c_r: jnp.ndarray, - h: jnp.ndarray, - gamma: float, - ot_prob: quadratic_problem.QuadraticProblem, - min_entry_value: float = 1e-6, - tolerance: float = 1e-3, - min_iter: int = 0, - inner_iter: int = 10, - max_iter: int = 10000 - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Run Dykstra's algorithm.""" - # shortcuts for problem's definition. - r = self.rank - n, m = ot_prob.geom_xx.shape[0], ot_prob.geom_yy.shape[0] - loga, logb = jnp.log(ot_prob.a), jnp.log(ot_prob.b) - - h_old = h - g1_old, g2_old = jnp.zeros(r), jnp.zeros(r) - f1, f2 = jnp.zeros(n), jnp.zeros(m) - - w_gi, w_gp = jnp.zeros(r), jnp.zeros(r) - w_q, w_r = jnp.zeros(r), jnp.zeros(r) - err = jnp.inf - state_inner = f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err - constants = c_q, c_r, loga, logb - - def cond_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...] - ) -> bool: - del iteration, constants - *_, err = state_inner - return err > tolerance - - def _softm( - f: jnp.ndarray, g: jnp.ndarray, c: jnp.ndarray, axis: int - ) -> jnp.ndarray: - return jsp.special.logsumexp( - gamma * (f[:, None] + g[None, :] - c), axis=axis - ) - - def body_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...], compute_error: bool - ) -> Tuple[jnp.ndarray, ...]: - # TODO(michalk8): in the future, use `NamedTuple` - f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err = state_inner - c_q, c_r, loga, logb = constants - - # First Projection - f1 = jnp.where( - jnp.isfinite(loga), - (loga - _softm(f1, g1_old, c_q, axis=1)) / gamma + f1, loga - ) - f2 = jnp.where( - jnp.isfinite(logb), - (logb - _softm(f2, g2_old, c_r, axis=1)) / gamma + f2, logb - ) - - h = h_old + w_gi - h = jnp.maximum(jnp.log(min_entry_value) / gamma, h) - w_gi += h_old - h - h_old = h - - # Update couplings - g_q = _softm(f1, g1_old, c_q, axis=0) - g_r = _softm(f2, g2_old, c_r, axis=0) - - # Second Projection - h = (1.0 / 3.0) * (h_old + w_gp + w_q + w_r) - h += g_q / (3.0 * gamma) - h += g_r / (3.0 * gamma) - g1 = h + g1_old - g_q / gamma - g2 = h + g2_old - g_r / gamma - - w_q = w_q + g1_old - g1 - w_r = w_r + g2_old - g2 - w_gp = h_old + w_gp - h - - q, r, _ = recompute_couplings(f1, g1, c_q, f2, g2, c_r, h, gamma) - - g1_old = g1 - g2_old = g2 - h_old = h - - err = jax.lax.cond( - jnp.logical_and(compute_error, iteration >= min_iter), - lambda: dykstra_solution_error(q, r, ot_prob, self.norm_error)[0], - lambda: err - ) - - return f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err - - def recompute_couplings( - f1: jnp.ndarray, - g1: jnp.ndarray, - c_q: jnp.ndarray, - f2: jnp.ndarray, - g2: jnp.ndarray, - c_r: jnp.ndarray, - h: jnp.ndarray, - gamma: float, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - q = jnp.exp(gamma * (f1[:, None] + g1[None, :] - c_q)) - r = jnp.exp(gamma * (f2[:, None] + g2[None, :] - c_r)) - g = jnp.exp(gamma * h) - return q, r, g - - state_inner = fixed_point_loop.fixpoint_iter_backprop( - cond_fn, body_fn, min_iter, max_iter, inner_iter, constants, state_inner - ) - - f1, f2, g1_old, g2_old, h_old, _, _, _, _, _ = state_inner - return recompute_couplings(f1, g1_old, c_q, f2, g2_old, c_r, h_old, gamma) - - def dykstra_update_kernel( - self, - k_q: jnp.ndarray, - k_r: jnp.ndarray, - k_g: jnp.ndarray, - gamma: float, - ot_prob: quadratic_problem.QuadraticProblem, - min_entry_value: float = 1e-6, - tolerance: float = 1e-3, - min_iter: int = 0, - inner_iter: int = 10, - max_iter: int = 10000 - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Run Dykstra's algorithm.""" - # shortcuts for problem's definition. - del gamma - rank = self.rank - n, m = ot_prob.geom_xx.shape[0], ot_prob.geom_yy.shape[0] - a, b = ot_prob.a, ot_prob.b - supp_a, supp_b = a > 0, b > 0 - - g_old = k_g - v1_old, v2_old = jnp.ones(rank), jnp.ones(rank) - u1, u2 = jnp.ones(n), jnp.ones(m) - - q_gi, q_gp = jnp.ones(rank), jnp.ones(rank) - q_q, q_r = jnp.ones(rank), jnp.ones(rank) - err = jnp.inf - state_inner = u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err - constants = k_q, k_r, k_g, a, b - - def cond_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...] - ) -> bool: - del iteration, constants - *_, err = state_inner - return err > tolerance - - def body_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...], compute_error: bool - ) -> Tuple[jnp.ndarray, ...]: - # TODO(michalk8): in the future, use `NamedTuple` - u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err = state_inner - k_q, k_r, k_g, a, b = constants - - # First Projection - u1 = jnp.where(supp_a, a / jnp.dot(k_q, v1_old), 0.0) - u2 = jnp.where(supp_b, b / jnp.dot(k_r, v2_old), 0.0) - g = jnp.maximum(min_entry_value, g_old * q_gi) - q_gi = (g_old * q_gi) / g - g_old = g - - # Second Projection - v1_trans = jnp.dot(k_q.T, u1) - v2_trans = jnp.dot(k_r.T, u2) - g = (g_old * q_gp * v1_old * q_q * v1_trans * v2_old * q_r * - v2_trans) ** (1 / 3) - v1 = g / v1_trans - v2 = g / v2_trans - q_gp = (g_old * q_gp) / g - q_q = (q_q * v1_old) / v1 - q_r = (q_r * v2_old) / v2 - v1_old = v1 - v2_old = v2 - g_old = g - - # Compute Couplings - q, r, _ = recompute_couplings(u1, v1, k_q, u2, v2, k_r, g) - - err = jax.lax.cond( - jnp.logical_and(compute_error, iteration >= min_iter), - lambda: dykstra_solution_error(q, r, ot_prob, self.norm_error)[0], - lambda: err - ) - - return u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err - - def recompute_couplings( - u1: jnp.ndarray, - v1: jnp.ndarray, - k_q: jnp.ndarray, - u2: jnp.ndarray, - v2: jnp.ndarray, - k_r: jnp.ndarray, - g: jnp.ndarray, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - q = u1.reshape((-1, 1)) * k_q * v1.reshape((1, -1)) - r = u2.reshape((-1, 1)) * k_r * v2.reshape((1, -1)) - return q, r, g - - state_inner = fixed_point_loop.fixpoint_iter_backprop( - cond_fn, body_fn, min_iter, max_iter, inner_iter, constants, state_inner - ) - - u1, u2, v1_old, v2_old, g_old, _, _, _, _, _ = state_inner - return recompute_couplings(u1, v1_old, k_q, u2, v2_old, k_r, g_old) - - def lse_step( - self, ot_prob: quadratic_problem.QuadraticProblem, state: LRGWState, - iteration: int - ) -> LRGWState: - """Low-rank GW LSE update.""" - c_q, c_r, c_g, gamma = self._get_costs(ot_prob, state) - - if ot_prob.is_balanced: - c_q, c_r, h = c_q / -gamma, c_r / -gamma, c_g / gamma - q, r, g = self.dykstra_update_lse( - c_q, c_r, h, gamma, ot_prob, **self.kwargs_dys - ) - else: - q, r, g = lr_utils.unbalanced_dykstra_lse( - c_q, c_r, c_g, gamma, ot_prob, **self.kwargs_dys - ) - return state.set(q=q, g=g, r=r, gamma=gamma) #, (c_q, c_r, c_g) - - def kernel_step( - self, ot_prob: quadratic_problem.QuadraticProblem, state: LRGWState, - iteration: int - ) -> LRGWState: - """Low-rank GW kernel update.""" - c_q, c_r, c_g, gamma = self._get_costs(ot_prob, state) - c_q, c_r, c_g = jnp.exp(c_q), jnp.exp(c_r), jnp.exp(c_g) - - if ot_prob.is_balanced: - q, r, g = self.dykstra_update_kernel( - c_q, c_r, c_g, gamma, ot_prob, **self.kwargs_dys - ) - else: - q, r, g = lr_utils.unbalanced_dykstra_kernel( - c_q, c_r, c_g, gamma, ot_prob, **self.kwargs_dys - ) - return state.set(q=q, g=g, r=r, gamma=gamma) - - def one_iteration( - self, ot_prob: quadratic_problem.QuadraticProblem, state: LRGWState, - iteration: int, compute_error: bool - ) -> LRGWState: - """Carries out one low-rank GW iteration. - - Depending on lse_mode, these iterations can be either in: - - - log-space for numerical stability. - - scaling space, using standard kernel-vector multiply operations. - - Args: - ot_prob: the transport problem definition - state: the current state. - iteration: the current iteration of the GW outer loop. - compute_error: flag to indicate this iteration computes/stores an error - - Returns: - The updated state. - """ - previous_state = state - - it = iteration // self.inner_iterations - if self.lse_mode: # In lse_mode, run additive updates. - state = self.lse_step(ot_prob, state, iteration) - else: - state = self.kernel_step(ot_prob, state, iteration) - - # re-computes error if compute_error is True, else set it to inf. - cost = jax.lax.cond( - jnp.logical_and(compute_error, iteration >= self.min_iterations), - lambda: state.reg_gw_cost(ot_prob, epsilon=self.epsilon), - lambda: jnp.inf - ) - error = state.compute_error(previous_state) - crossed_threshold = jnp.logical_or( - state.crossed_threshold, - jnp.logical_and( - state.errors[it - 1] >= self.threshold, error < self.threshold - ) - ) - - state = state.set( - costs=state.costs.at[it].set(cost), - errors=state.errors.at[it].set(error), - crossed_threshold=crossed_threshold, - ) - - if self.progress_fn is not None: - jax.experimental.io_callback( - self.progress_fn, None, - (iteration, self.inner_iterations, self.max_iterations, state) - ) - - return state - - @property - def norm_error(self) -> Tuple[int]: # noqa: D102 - return self._norm_error, - - def create_initializer( - self, - prob: quadratic_problem.QuadraticProblem, - ) -> initializers_lr.LRInitializer: - """Create a low-rank GW initializer. - - Args: - prob: Quadratic OT problem used to determine the initializer. - - Returns: - Low-rank initializer. - """ - if isinstance(self.initializer, initializers_lr.LRInitializer): - assert self.initializer.rank == self.rank, \ - f"Expected initializer's rank to be `{self.rank}`," \ - f"found `{self.initializer.rank}`." - return self.initializer - - return initializers_lr.LRInitializer.from_solver( - self, kind=self.initializer, **self.kwargs_init - ) - - def init_state( - self, ot_prob: quadratic_problem.QuadraticProblem, - init: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] - ) -> LRGWState: - """Return the initial state of the loop.""" - q, r, g = init - return LRGWState( - q=q, - r=r, - g=g, - gamma=self.gamma, - costs=-jnp.ones(self.outer_iterations), - errors=-jnp.ones(self.outer_iterations), - crossed_threshold=False, - ) - - def output_from_state( - self, ot_prob: quadratic_problem.QuadraticProblem, state: LRGWState - ) -> LRGWOutput: - """Create an output from a loop state. - - Args: - ot_prob: the transport problem. - state: GW state. - - Returns: - A LRGWOutput. - """ - return LRGWOutput( - q=state.q, - r=state.r, - g=state.g, - ot_prob=ot_prob, - costs=state.costs, - errors=state.errors, - epsilon=self.epsilon, - inner_iterations=self.inner_iterations, - ) - - def _converged(self, state: LRGWState, iteration: int) -> bool: - - def conv_crossed(prev_err: float, curr_err: float) -> bool: - return jnp.logical_and( - prev_err < self.threshold, curr_err < self.threshold - ) - - def conv_not_crossed(prev_err: float, curr_err: float) -> bool: - return jnp.logical_and(curr_err < prev_err, curr_err < self.threshold) - - # for convergence error, we consider 2 possibilities: - # 1. we either crossed the convergence threshold; in this case we require - # that the previous error was also below the threshold - # 2. we haven't crossed the threshold; in this case, we can be below or - # above the threshold: - # if we're above, we wait until we reach the convergence threshold and - # then, the above condition applies - # if we're below and we improved w.r.t. the previous iteration, - # we have converged; otherwise we continue, since we may be stuck - # in a local minimum (e.g., during the initial iterations) - - it = iteration // self.inner_iterations - return jax.lax.cond( - state.crossed_threshold, conv_crossed, conv_not_crossed, - state.errors[it - 2], state.errors[it - 1] - ) - - def _diverged(self, state: LRGWState, iteration: int) -> bool: - it = iteration // self.inner_iterations - return jnp.logical_and( - jnp.logical_not(jnp.isfinite(state.errors[it - 1])), - jnp.logical_not(jnp.isfinite(state.costs[it - 1])) - ) - - -def run( - ot_prob: quadratic_problem.QuadraticProblem, - solver: LRGromovWasserstein, - init: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray], - Optional[jnp.ndarray]], -) -> LRGWOutput: - """Run loop of the solver, outputting a state upgraded to an output.""" - out = sinkhorn.iterations(ot_prob, solver, init) - out = out.set_cost( - ot_prob, lse_mode=solver.lse_mode, use_danskin=solver.use_danskin - ) - return out.set(ot_prob=ot_prob) - - -def dykstra_solution_error( - q: jnp.ndarray, r: jnp.ndarray, ot_prob: quadratic_problem.QuadraticProblem, - norm_error: Tuple[int, ...] -) -> jnp.ndarray: - """Compute solution error. - - Since only balanced case is available for LR, this is marginal deviation. - - Args: - q: first factor of solution. - r: second factor of solution. - ot_prob: linear problem. - norm_error: int, p-norm used to compute error. - - Returns: - one or possibly many numbers quantifying deviation to true marginals. - """ - norm_error = jnp.array(norm_error) - # Update the error - err = jnp.sum( - jnp.abs(jnp.sum(q, axis=1) - ot_prob.a) ** norm_error[:, None], axis=1 - ) ** (1.0 / norm_error) - err += jnp.sum( - jnp.abs(jnp.sum(r, axis=1) - ot_prob.b) ** norm_error[:, None], axis=1 - ) ** (1.0 / norm_error) - err += jnp.sum( - jnp.abs(jnp.sum(q, axis=0) - jnp.sum(r, axis=0)) ** norm_error[:, None], - axis=1 - ) ** (1.0 / norm_error) - - return err - - -def _linearized_geometry( - prob: quadratic_problem.QuadraticProblem, - *, - q: jnp.ndarray, - r: jnp.ndarray, - g: jnp.ndarray, -) -> low_rank.LRCGeometry: - inv_sqrt_g = 1.0 / jnp.sqrt(g[None, :]) - - # TODO(michalk8): below is for squared loss, handle KL loss in the future; - # will need to be updated in many other places as well - tmp1 = -4.0 * prob.geom_xx.apply_cost(q, axis=1) * inv_sqrt_g - tmp2 = prob.geom_yy.apply_cost(r, axis=1) * inv_sqrt_g - return low_rank.LRCGeometry(tmp1, tmp2) diff --git a/ott/build/lib/ott/solvers/quadratic/gw_barycenter.py b/ott/build/lib/ott/solvers/quadratic/gw_barycenter.py deleted file mode 100644 index f0d350b..0000000 --- a/ott/build/lib/ott/solvers/quadratic/gw_barycenter.py +++ /dev/null @@ -1,339 +0,0 @@ -# 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. -from functools import partial -from typing import Any, Dict, NamedTuple, Optional, Sequence, Tuple, Union - -import jax -import jax.numpy as jnp - -from ott import utils -from ott.geometry import pointcloud -from ott.math import fixed_point_loop -from ott.problems.linear import linear_problem -from ott.problems.quadratic import gw_barycenter -from ott.solvers import was_solver -from ott.solvers.quadratic import gromov_wasserstein - -__all__ = ["GWBarycenterState", "GromovWassersteinBarycenter"] - - -class GWBarycenterState(NamedTuple): - """State of the GW barycenter problem. - - Args: - cost: Barycenter cost matrix of shape ``[bar_size, bar_size]``. - x: Barycenter features of shape ``[bar_size, ndim_fused]``. - Only used in the fused case. - a: Weights of the barycenter of shape ``[bar_size,]``. - errors: Array of shape - ``[max_iter, num_measures, quad_max_iter, lin_outer_iter]`` containing - the GW errors at each iteration. - costs: Array of shape ``[max_iter,]`` containing the cost at each iteration. - costs_bary: Array of shape ``[max_iter, num_measures]`` containing the - cost between the individual measures and the barycenter at each iteration. - gw_convergence: Array of shape ``[max_iter,]`` containing the convergence - of all GW problems at each iteration. - """ - cost: Optional[jnp.ndarray] = None - x: Optional[jnp.ndarray] = None - a: Optional[jnp.ndarray] = None - errors: Optional[jnp.ndarray] = None - costs: Optional[jnp.ndarray] = None - costs_bary: Optional[jnp.ndarray] = None - gw_convergence: Optional[jnp.ndarray] = None - - def set(self, **kwargs: Any) -> "GWBarycenterState": - """Return a copy of self, possibly with overwrites.""" - return self._replace(**kwargs) - - @property - def n_iters(self) -> int: - """Number of iterations.""" - if self.gw_convergence is None: - return -1 - return jnp.sum(self.gw_convergence != -1) - - -@jax.tree_util.register_pytree_node_class -class GromovWassersteinBarycenter(was_solver.WassersteinSolver): - """Gromov-Wasserstein barycenter solver. - - Args: - epsilon: Entropy regularizer. - min_iterations: Minimum number of iterations. - max_iterations: Maximum number of outermost iterations. - threshold: Convergence threshold. - store_inner_errors: Whether to store the errors of the GW solver, as well - as its linear solver, at each iteration for each measure. - quad_solver: The GW solver. - kwargs: Keyword argument for - :class:`~ott.solvers.quadratic.gromov_wasserstein.GromovWasserstein`. - Only used when ``quad_solver = None``. - """ - - def __init__( - self, - epsilon: Optional[float] = None, - min_iterations: int = 5, - max_iterations: int = 50, - threshold: float = 1e-3, - store_inner_errors: bool = False, - quad_solver: Optional[gromov_wasserstein.GromovWasserstein] = None, - # TODO(michalk8): maintain the API compatibility with `was_solver` - # but makes passing kwargs with the same name to `quad_solver` impossible - # will be fixed when refactoring the solvers - # note that `was_solver` also suffers from this - **kwargs: Any, - ): - super().__init__( - epsilon=epsilon, - min_iterations=min_iterations, - max_iterations=max_iterations, - threshold=threshold, - store_inner_errors=store_inner_errors, - ) - if quad_solver is None: - kwargs["epsilon"] = epsilon - # TODO(michalk8): store only GW errors? - kwargs["store_inner_errors"] = store_inner_errors - self._quad_solver = gromov_wasserstein.GromovWasserstein(**kwargs) - else: - self._quad_solver = quad_solver - - def __call__( - self, problem: gw_barycenter.GWBarycenterProblem, bar_size: int, - **kwargs: Any - ) -> GWBarycenterState: - """Solver the (fused) GW barycenter problem. - - Args: - problem: The GW barycenter problem. - bar_size: Size of the barycenter. - kwargs: Keyword arguments for :meth:`init_state`. - - Returns: - The solution. - """ - state = self.init_state(problem, bar_size, **kwargs) - state = iterations(self, problem, state) - return self.output_from_state(state) - - def init_state( - self, - problem: gw_barycenter.GWBarycenterProblem, - bar_size: int, - bar_init: Optional[Union[jnp.ndarray, Tuple[jnp.ndarray, - jnp.ndarray]]] = None, - a: Optional[jnp.ndarray] = None, - rng: Optional[jax.Array] = None, - ) -> GWBarycenterState: - """Initialize the (fused) Gromov-Wasserstein barycenter state. - - Args: - problem: The barycenter problem. - bar_size: Size of the barycenter. - bar_init: Initial barycenter value. Can be one of the following: - - - ``None`` - randomly initialize the barycenter. - - :class:`jax.numpy.ndarray` - barycenter cost matrix of shape - ``[bar_size, bar_size]``. - Only used in the non-fused case. - - :class:`tuple` of :class:`jax.numpy.ndarray` - the first array - corresponds to a cost matrix of shape ``[bar_size, bar_size]``, - the second array is a ``[bar_size, ndim_fused]`` feature matrix used - in the fused case. - - a: An array of shape ``[bar_size,]`` containing the barycenter weights. - rng: Random key for seeding used when ``bar_init = None``. - - Returns: - The initial barycenter state. - """ - if a is None: - a = jnp.ones((bar_size,)) / bar_size - else: - assert a.shape == (bar_size,) - - if bar_init is None: - rng = utils.default_prng_key(rng) - _, b = problem.segmented_y_b - rngs = jax.random.split(rng, problem.num_measures) - linear_solver = self._quad_solver.linear_ot_solver - - transports = init_transports(linear_solver, rngs, a, b, problem.epsilon) - x = problem.update_features(transports, a) if problem.is_fused else None - cost = problem.update_barycenter(transports, a) - else: - cost, x = bar_init if isinstance(bar_init, tuple) else (bar_init, None) - assert cost.shape == (bar_size, bar_size) - if problem.is_fused: - assert x is not None, "Barycenter features are not initialized." - assert x.shape == (bar_size, problem.ndim_fused) - - num_iter = self.max_iterations - if self.store_inner_errors: - # TODO(michalk8): in the future, think about how to do this in general - errors = -jnp.ones(( - num_iter, problem.num_measures, self._quad_solver.max_iterations, - self._quad_solver.linear_ot_solver.outer_iterations - )) - else: - errors = None - - costs = -jnp.ones((num_iter,)) - costs_bary = -jnp.ones((num_iter, problem.num_measures)) - gw_convergence = -jnp.ones((num_iter,)) - return GWBarycenterState( - cost=cost, - x=x, - a=a, - errors=errors, - costs=costs, - costs_bary=costs_bary, - gw_convergence=gw_convergence - ) - - def update_state( - self, - state: GWBarycenterState, - iteration: int, - problem: gw_barycenter.GWBarycenterProblem, - store_errors: bool = True, - ) -> Tuple[float, bool, jnp.ndarray, Optional[jnp.ndarray]]: - """Solve the (fused) Gromov-Wasserstein barycenter problem.""" - - def solve_gw( - state: GWBarycenterState, b: jnp.ndarray, y: jnp.ndarray, - f: Optional[jnp.ndarray] - ) -> Tuple[float, bool, jnp.ndarray, Optional[jnp.ndarray]]: - quad_problem = problem._create_problem(state, y=y, b=b, f=f) - out = self._quad_solver(quad_problem) - return ( - out.reg_gw_cost, out.converged, out.matrix, - out.errors if store_errors else None - ) - - in_axes = [None, 0, 0] - in_axes += [0] if problem.is_fused else [None] - solve_fn = jax.vmap(solve_gw, in_axes=in_axes) - - y, b = problem.segmented_y_b - y_f = problem.segmented_y_fused - costs, convergeds, transports, errors = solve_fn(state, b, y, y_f) - - cost = jnp.sum(costs * problem.weights) - costs_bary = state.costs_bary.at[iteration].set(costs) - costs = state.costs.at[iteration].set(cost) - - converged = jnp.all(convergeds) - gw_convergence = state.gw_convergence.at[iteration].set(converged) - - if self.store_inner_errors: - errors = state.errors.at[iteration, ...].set(errors) - else: - errors = None - - x = problem.update_features( - transports, state.a - ) if problem.is_fused else state.x - cost = problem.update_barycenter(transports, state.a) - return state.set( - cost=cost, - x=x, - costs=costs, - costs_bary=costs_bary, - errors=errors, - gw_convergence=gw_convergence - ) - - def output_from_state(self, state: GWBarycenterState) -> GWBarycenterState: - """No-op.""" - # TODO(michalk8): just for consistency with continuous barycenter - # will be refactored in the future to create an output - return state - - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - children, aux = super().tree_flatten() - return children + [self._quad_solver], aux - - @classmethod - def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] - ) -> "GromovWassersteinBarycenter": - epsilon, _, threshold, quad_solver = children - return cls( - epsilon=epsilon, - threshold=threshold, - quad_solver=quad_solver, - **aux_data, - ) - - -@partial(jax.vmap, in_axes=[None, 0, None, 0, None]) -def init_transports( - solver, rng: jax.Array, a: jnp.ndarray, b: jnp.ndarray, - epsilon: Optional[float] -) -> jnp.ndarray: - """Initialize random 2D point cloud and solve the linear OT problem. - - Args: - solver: Linear OT solver. - rng: Random key for seeding. - a: Source marginals (e.g., for barycenter) of shape ``[bar_size,]``. - b: Target marginals of shape ``[max_measure_size,]``. - epsilon: Entropy regularization. - - Returns: - Transport map of shape ``[bar_size, max_measure_size]``. - """ - rng1, rng2 = jax.random.split(rng, 2) - x = jax.random.normal(rng1, shape=(len(a), 2)) - y = jax.random.normal(rng2, shape=(len(b), 2)) - geom = pointcloud.PointCloud( - x, y, epsilon=epsilon, src_mask=a > 0, tgt_mask=b > 0 - ) - problem = linear_problem.LinearProblem(geom, a=a, b=b) - return solver(problem).matrix - - -def iterations( # noqa: D103 - solver: GromovWassersteinBarycenter, - problem: gw_barycenter.GWBarycenterProblem, init_state: GWBarycenterState -) -> GWBarycenterState: - - def cond_fn( - iteration: int, constants: GromovWassersteinBarycenter, - state: GWBarycenterState - ) -> bool: - solver, _ = constants - return solver._continue(state, iteration) - - def body_fn( - iteration, constants: Tuple[GromovWassersteinBarycenter, - gw_barycenter.GWBarycenterProblem], - state: GWBarycenterState, compute_error: bool - ) -> GWBarycenterState: - del compute_error # always assumed true - solver, problem = constants - return solver.update_state(state, iteration, problem) - - return fixed_point_loop.fixpoint_iter( - cond_fn=cond_fn, - body_fn=body_fn, - min_iterations=solver.min_iterations, - max_iterations=solver.max_iterations, - inner_iterations=1, - constants=(solver, problem), - state=init_state, - ) diff --git a/ott/build/lib/ott/solvers/quadratic/lower_bound.py b/ott/build/lib/ott/solvers/quadratic/lower_bound.py deleted file mode 100644 index f0868ad..0000000 --- a/ott/build/lib/ott/solvers/quadratic/lower_bound.py +++ /dev/null @@ -1,96 +0,0 @@ -# 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. -from typing import TYPE_CHECKING, Any, Optional - -import jax -import jax.tree_util as jtu - -from ott.geometry import pointcloud -from ott.problems.quadratic import quadratic_problem -from ott.solvers import linear -from ott.solvers.linear import sinkhorn - -if TYPE_CHECKING: - from ott.geometry import distrib_costs - -__all__ = ["LowerBoundSolver"] - - -@jtu.register_pytree_node_class -class LowerBoundSolver: - """Lower bound OT solver. - - Computes the third lower bound distance from :cite:`memoli:11`, def. 6.3. - - Args: - epsilon: Entropy regularization for the resulting linear problem. - distrib_cost: Univariate Wasserstein cost, used to compare two point clouds - in different spaces, where each point is seen as its distribution of costs - to other points in its point cloud. - """ - - def __init__( - self, - epsilon: Optional[float] = None, - distrib_cost: Optional["distrib_costs.UnivariateWasserstein"] = None, - ): - from ott.geometry import distrib_costs - - self.epsilon = epsilon - self.distrib_cost = ( - distrib_costs.UnivariateWasserstein() - if distrib_cost is None else distrib_cost - ) - - def __call__( - self, - prob: quadratic_problem.QuadraticProblem, - epsilon: Optional[float] = None, - rng: Optional[jax.Array] = None, - **kwargs: Any - ) -> sinkhorn.SinkhornOutput: - """Compute a lower-bound for the GW problem using a simple linearization. - - This solver handles a quadratic problem by computing a proxy ``[n, m]`` - cost-matrix, injecting it into a linear OT solver to output a first an OT - matrix that can be used either to linearize/initialize the resolution - ot the GW problem, or more simply as a simple GW solution. - - Args: - prob: Quadratic OT problem. - epsilon: Entropic regularization passed on to solve the linearization of - the quadratic problem using 1D costs. - rng: Random key, possibly used when computing 1D costs when using - subsampling. - kwargs: Keyword arguments for :func:`~ott.solvers.linear.solve`. - - Returns: - A linear OT output, an approximation of the OT coupling obtained using - the lower bound provided by :cite:`memoli:11`. - """ - dists_xx = prob.geom_xx.cost_matrix - dists_yy = prob.geom_yy.cost_matrix - - geom_xy = pointcloud.PointCloud( - dists_xx, dists_yy, cost_fn=self.distrib_cost, epsilon=self.epsilon - ) - return linear.solve(geom_xy, **kwargs) - - def tree_flatten(self): # noqa: D102 - return (self.epsilon, self.distrib_cost), None - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - del aux_data - return cls(*children) diff --git a/ott/build/lib/ott/solvers/was_solver.py b/ott/build/lib/ott/solvers/was_solver.py deleted file mode 100644 index 6b25718..0000000 --- a/ott/build/lib/ott/solvers/was_solver.py +++ /dev/null @@ -1,131 +0,0 @@ -# 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. -from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple, Union - -import jax -import jax.numpy as jnp - -if TYPE_CHECKING: - from ott.solvers.linear import continuous_barycenter, sinkhorn, sinkhorn_lr - -__all__ = ["WassersteinSolver"] - -State = Union[ - "sinkhorn.SinkhornState", - "sinkhorn_lr.LRSinkhornState", - "continuous_barycenter.FreeBarycenterState", -] - - -# TODO(michalk8): refactor to have generic nested solver API -@jax.tree_util.register_pytree_node_class -class WassersteinSolver: - """A generic solver for problems that use a linear problem in inner loop.""" - - def __init__( - self, - epsilon: Optional[float] = None, - rank: int = -1, - linear_ot_solver: Optional[ - Union["sinkhorn.Sinkhorn", "sinkhorn_lr.LRSinkhorn"] - ] = None, - min_iterations: int = 5, - max_iterations: int = 50, - threshold: float = 1e-3, - store_inner_errors: bool = False, - kwargs: Any = {}, - ): - from ott.solvers.linear import sinkhorn, sinkhorn_lr - - default_epsilon = 1.0 - # Set epsilon value to default if needed, but keep track of whether None was - # passed to handle the case where a linear_ot_solver is passed or not. - self.epsilon = epsilon if epsilon is not None else default_epsilon - self.rank = rank - self.linear_ot_solver = linear_ot_solver - if self.linear_ot_solver is None: - # Detect if user requests low-rank solver. In that case the - # default_epsilon makes little sense, since it was designed for GW. - if self.is_low_rank: - if epsilon is None: - # Use default entropic regularization in LRSinkhorn if None was passed - self.linear_ot_solver = sinkhorn_lr.LRSinkhorn( - rank=self.rank, **kwargs - ) - else: - # If epsilon is passed, use it to replace the default LRSinkhorn value - self.linear_ot_solver = sinkhorn_lr.LRSinkhorn( - rank=self.rank, epsilon=self.epsilon, **kwargs - ) - else: - # When using Entropic GW, epsilon is not handled inside Sinkhorn, - # but rather added back to the Geometry object re-instantiated - # when linearizing the problem. Therefore, no need to pass it to solver. - self.linear_ot_solver = sinkhorn.Sinkhorn(**kwargs) - - self.min_iterations = min_iterations - self.max_iterations = max_iterations - self.threshold = threshold - self.store_inner_errors = store_inner_errors - self._kwargs = kwargs - - @property - def is_low_rank(self) -> bool: - """Whether the solver is low-rank.""" - return self.rank > 0 - - def tree_flatten( - self, - ) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 - return ( - [self.epsilon, self.linear_ot_solver, self.threshold], - { - "min_iterations": self.min_iterations, - "max_iterations": self.max_iterations, - "rank": self.rank, - "store_inner_errors": self.store_inner_errors, - **self._kwargs, - }, - ) - - @classmethod - def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] - ) -> "WassersteinSolver": - epsilon, linear_ot_solver, threshold = children - return cls( - epsilon=epsilon, - linear_ot_solver=linear_ot_solver, - threshold=threshold, - **aux_data, - ) - - def _converged(self, state: State, iteration: int) -> bool: - costs, i, tol = state.costs, iteration, self.threshold - return jnp.logical_and( - i >= 2, jnp.isclose(costs[i - 2], costs[i - 1], rtol=tol) - ) - - def _diverged(self, state: State, iteration: int) -> bool: - return jnp.logical_not(jnp.isfinite(state.costs[iteration - 1])) - - def _continue(self, state: State, iteration: int) -> bool: - """Continue while not(converged) and not(diverged).""" - return jnp.logical_or( - iteration <= 2, - jnp.logical_and( - jnp.logical_not(self._diverged(state, iteration)), - jnp.logical_not(self._converged(state, iteration)), - ), - ) diff --git a/ott/build/lib/ott/tools/__init__.py b/ott/build/lib/ott/tools/__init__.py deleted file mode 100644 index 75dbfae..0000000 --- a/ott/build/lib/ott/tools/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# 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. -from . import ( - gaussian_mixture, - k_means, - plot, - segment_sinkhorn, - sinkhorn_divergence, - soft_sort, -) diff --git a/ott/build/lib/ott/tools/gaussian_mixture/__init__.py b/ott/build/lib/ott/tools/gaussian_mixture/__init__.py deleted file mode 100644 index 3195d77..0000000 --- a/ott/build/lib/ott/tools/gaussian_mixture/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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. -from . import fit_gmm_pair, gaussian, gaussian_mixture, gaussian_mixture_pair diff --git a/ott/build/lib/ott/tools/gaussian_mixture/fit_gmm.py b/ott/build/lib/ott/tools/gaussian_mixture/fit_gmm.py deleted file mode 100644 index fe869ab..0000000 --- a/ott/build/lib/ott/tools/gaussian_mixture/fit_gmm.py +++ /dev/null @@ -1,303 +0,0 @@ -# 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. -r"""Fit a Gaussian mixture model. - -Sample usage: - -# initialize GMM with K-means++ -gmm_init = fit_gmm.initialize( - rng=rng, - points=my_points, - point_weights=None, - n_components=COMPONENTS) - -# refine GMM parameters using EM -gmm = fit_gmm.fit_model_em( - gmm=gmm_init, - points=my_points, - point_weights=None, - steps=10, - verbose=True) - - -We fit the model using EM. Below we'll use notation following -https://en.wikipedia.org/wiki/Expectation%E2%80%93maximization_algorithm - -Our data X is generated by a Gaussian mixture with unknown parameters \Theta. -We denote the (unobserved) component that gave rise to each point as Z. - -In EM we start with an initial estimate of $\Theta$, $\Theta^{(0)}$, and we -then iteratively update it via - -$$\Theta^{(t+1)} = \argmax_{\Theta} Q(\Theta|\Theta^{(t)})$$ - -where - -$$ -Q(\Theta|\Theta(t)) = E_{Z|X,\Theta^{(t)}} \left[ \log L(\Theta; X, Z) \right] -$$ -""" - -from typing import Optional - -import jax -import jax.numpy as jnp - -from ott.tools.gaussian_mixture import gaussian_mixture - -__all__ = ["initialize", "fit_model_em"] - -# EM algorithm for parameter estimation - - -def get_assignment_probs( - gmm: gaussian_mixture.GaussianMixture, points: jnp.ndarray -) -> jnp.ndarray: - r"""Get component assignment probabilities used in the E step of EM. - - Here we compute the component assignment probabilities p(Z|X, \Theta^{(t)}) - that we need to compute the expectation used for Q(\Theta|\Theta^{(t)}). - - Args: - gmm: GMM model - points: set of samples being fitted, shape (n, n_dimensions) - - Returns: - An array of assignment probabilities with shape (n, n_components) - """ - return jnp.exp(gmm.get_log_component_posterior(points)) - - -def get_q( - gmm: gaussian_mixture.GaussianMixture, - assignment_probs: jnp.ndarray, - points: jnp.ndarray, - point_weights: Optional[jnp.ndarray] = None, -) -> float: - r"""Get Q(\Theta|\Theta^{(t)}). - - Args: - gmm: GaussianMixture with parameters \Theta - assignment_probs: p(Z|X, \Theta^{(t)}) as computed by get_assignment_probs - points: observations X - point_weights: optional set of weights for the samples. If None, use - a weight of 1/n where n is the number of points. - - Returns: - Q(\Theta|\Theta^{(t)}) - """ - # log P(X, Z| \Theta) = log P(X|Z, \Theta) + log P(Z|\Theta) - loglik = (gmm.conditional_log_prob(points) + gmm.log_component_weights()) - if point_weights is None: - point_weights = jnp.ones(points.shape[0]) - return ( - jnp.sum(point_weights * jnp.sum(assignment_probs * loglik, axis=-1)) / - jnp.sum(point_weights) - ) - - -def log_prob_loss( - gmm: gaussian_mixture.GaussianMixture, - points: jnp.ndarray, - point_weights: Optional[jnp.ndarray] = None, -) -> float: - """Loss function: weighted mean of (-log prob of observations). - - Args: - gmm: GMM model - points: set of samples being fitted - point_weights: optional set of weights for the samples. If None, use - a weight of 1/n where n is the number of points. - - Returns: - The GMM loss for the points. - """ - if point_weights is None: - return -jnp.mean(gmm.log_prob(points)) - return -jnp.sum(point_weights * gmm.log_prob(points)) / jnp.sum(point_weights) - - -def fit_model_em( - gmm: gaussian_mixture.GaussianMixture, - points: jnp.ndarray, - point_weights: Optional[jnp.ndarray], - steps: int, - jit: bool = True, - verbose: bool = False, -) -> gaussian_mixture.GaussianMixture: - """Fit a GMM using the EM algorithm. - - Args: - gmm: initial GMM model - points: set of samples to fit, shape (n, n_dimensions) - point_weights: optional set of weights for points, shape (n,). If None, - uses equal weights for all points. - steps: number of steps of EM to perform - jit: if True, compile functions - verbose: if True, print the loss at each step - - Returns: - A GMM with updated parameters. - """ - if point_weights is None: - point_weights = jnp.ones(points.shape[:-1]) - loss_fn = log_prob_loss - get_q_fn = get_q - e_step_fn = get_assignment_probs - m_step_fn = gaussian_mixture.GaussianMixture.from_points_and_assignment_probs - if jit: - loss_fn = jax.jit(loss_fn) - get_q_fn = jax.jit(get_q_fn) - e_step_fn = jax.jit(e_step_fn) - m_step_fn = jax.jit(m_step_fn) - - for i in range(steps): - assignment_probs = e_step_fn(gmm, points) - gmm_new = m_step_fn(points, point_weights, assignment_probs) - if gmm_new.has_nans(): - raise ValueError("NaNs in fit.") - if verbose: - loss = loss_fn(gmm_new, points, point_weights) - q = get_q_fn( - gmm=gmm_new, - assignment_probs=assignment_probs, - points=points, - point_weights=point_weights - ) - print(f"{i} q={q} -log prob={loss}") # noqa: T201 - gmm = gmm_new - return gmm - - -# KMeans++ for initialization -# See https://en.wikipedia.org/wiki/K-means%2B%2B for details - - -def _get_dist_sq(points: jnp.ndarray, loc: jnp.ndarray) -> jnp.ndarray: - """Get the squared distance from each point to each loc.""" - - def _dist_sq_one_loc(points: jnp.ndarray, loc: jnp.ndarray) -> jnp.ndarray: - return jnp.sum((points - loc[None]) ** 2, axis=-1) - - dist_sq_fn = jax.vmap(_dist_sq_one_loc, in_axes=(None, 0), out_axes=1) - return dist_sq_fn(points, loc) - - -def _get_locs( - rng: jax.Array, points: jnp.ndarray, n_components: int -) -> jnp.ndarray: - """Get the initial component means. - - Args: - rng: jax.random key - points: (n, n_dimensions) array of observations - n_components: desired number of components - - Returns: - (n_components, n_dimensions) array of means. - """ - points = points.copy() - n_points = points.shape[0] - weights = jnp.ones(n_points) / n_points - rng, subrng = jax.random.split(rng) - index = jax.random.choice(subrng, a=points.shape[0], p=weights) - loc = points[index] - points = jnp.concatenate([points[:index], points[index + 1:]], axis=0) - - locs = loc[None] - for _ in range(n_components - 1): - dist_sq = _get_dist_sq(points, locs) - min_dist_sq = jnp.min(dist_sq, axis=-1) - weights = min_dist_sq / jnp.sum(min_dist_sq) - rng, subrng = jax.random.split(rng) - index = jax.random.choice(subrng, a=points.shape[0], p=weights) - loc = points[index] - points = jnp.concatenate([points[:index], points[index + 1:]], axis=0) - locs = jnp.concatenate([locs, loc[None]], axis=0) - return locs - - -def from_kmeans_plusplus( - rng: jax.Array, - points: jnp.ndarray, - point_weights: Optional[jnp.ndarray], - n_components: int, -) -> gaussian_mixture.GaussianMixture: - """Initialize a GMM via a single pass of K-means++. - - Args: - rng: jax.random key - points: (n, n_dimensions) array of observations - point_weights: (n,) array of weights for points - n_components: desired number of components - - Returns: - An initial Gaussian mixture model. - - Raises: - ValueError if any fitted parameters are non-finite. - """ - rng, subrng = jax.random.split(rng) - locs = _get_locs(rng=subrng, points=points, n_components=n_components) - dist_sq = _get_dist_sq(points, locs) - assignment_prob = (dist_sq == jnp.min(dist_sq, - axis=-1)[:, None]).astype(points.dtype) - del dist_sq - - if point_weights is None: - point_weights = jnp.ones_like(points[..., 0]) - return gaussian_mixture.GaussianMixture.from_points_and_assignment_probs( - points=points, - point_weights=point_weights, - assignment_probs=assignment_prob - ) - - -def initialize( - rng: jax.Array, - points: jnp.ndarray, - point_weights: Optional[jnp.ndarray], - n_components: int, - n_attempts: int = 50, - verbose: bool = False -) -> gaussian_mixture.GaussianMixture: - """Initialize a GMM via K-means++ with retries on failure. - - Args: - rng: jax.random key - points: (n, n_dimensions) array of observations - point_weights: (n,) array of weights for points - n_components: desired number of components - n_attempts: number of attempts to initialize before failing - verbose: if True, print status information - - Returns: - An initial Gaussian mixture model. - - Raises: - ValueError if initialization was unsuccessful after n_attempts attempts. - """ - for attempt in range(n_attempts): - rng, subrng = jax.random.split(rng) - try: - return from_kmeans_plusplus( - rng=subrng, - points=points, - point_weights=point_weights, - n_components=n_components - ) - except ValueError: - if verbose: - print(f"Failed to initialize, attempt {attempt}.") # noqa: T201 - raise ValueError("Failed to initialize.") diff --git a/ott/build/lib/ott/tools/gaussian_mixture/fit_gmm_pair.py b/ott/build/lib/ott/tools/gaussian_mixture/fit_gmm_pair.py deleted file mode 100644 index 7ecde26..0000000 --- a/ott/build/lib/ott/tools/gaussian_mixture/fit_gmm_pair.py +++ /dev/null @@ -1,393 +0,0 @@ -# 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. -r"""Fit 2 GMMs to 2 point clouds using likelihood and (approx) W2 distance. - -Suppose we have two large point clouds and want to estimate a coupling and a -W2 distance between them. :cite:`delon:20` propose fitting a GMM to each -point cloud while simultaneously minimizing a Wasserstein-like distance -called MW2 between the fitted GMMs. MW2 is an upper bound on W2, -the Wasserstein distance between the GMMs. Here we implement -their algorithm as well as a generalization that allows for reweightings using -generalized, penalized expectation-maximization -(see section 6.2 of :cite:`delon:20`). - -As in `fit_gmm.py`, we assume that the observations $X_0$ and $X_1$ from -batches 0 and 1 are generated by GMMs with parameters $\Theta_0$ and $\Theta_1$, -respectively. We will use $\Theta$ to denote the combined parameters -for the two GMMs. We denote the (unobserved) components that gave rise to the -observations $X_i$ as $Z_i$. - -Our goal is to maximize a weighted sum of the likelihood of the observations $X$ -under the fitted GMMs and a measure of distance, $MW_2$, between the fitted -GMMs. The problem would be a straightforward maximization exercise if we knew -the components $Z$ that generated each observation $X$. Because the $Z$ are -unobserved, however, we use EM: - -We start with an initial estimate of $\Theta$, $\Theta^{(t)}$. - -* The E-step: We use the current $\Theta^{(t)}$ to estimate the likelihood of - all possible cluster attributions for each observation $X$. - -* The M-step: We form the function $Q(\Theta|\Theta^{(t)})$, - the log likelihood of our observations averaged over all possible - assignments. We then obtain an updated parameter estimate, $\Theta^{(t+1)}$, - by numerically maximizing the sum of $Q$ and our GMM distance penalty. - -It can be shown that if we maximize the penalized $Q$ above, this procedure will -increase or leave unchanged the penalized log likelihood for $\Theta$. We -iterate over these two steps until convergence. Note that the resulting -estimate for $\Theta$ may only be a *local* maximum of the penalized -likelihood function. - - -Sample usage: - -# (Note that we usually initialize a pair to a single GMM that we fit to a -# pooled set, then the two GMMs separate as we optimize the pair.) -pair_init = gaussian_mixture_pair.GaussianMixturePair( - gmm0=gmm0, - gmm1=gmm1, - epsilon=1.e-2, - tau=1.) -fit_model_em_fn = fit_gmm_pair.get_fit_model_em_fn( - weight_transport=0.1, - weight_splitting=1., - epsilon=pair_init.epsilon, - jit=True) -pair, loss = fit_model_em_fn( - pair=pair_init, - points0=samples_gmm0, - points1=samples_gmm1, - point_weights0=None, - point_weights1=None, - em_steps=30, - m_steps=20, - verbose=True) -""" -# TODO(geoffd): look into refactoring so we jit higher level functions - -import functools -import math -from typing import Callable, NamedTuple, Optional, Tuple - -import jax -import jax.numpy as jnp - -from ott.tools.gaussian_mixture import ( - fit_gmm, - gaussian_mixture, - gaussian_mixture_pair, -) - -__all__ = ["get_fit_model_em_fn"] - -LOG2 = math.log(2) - - -class Observations(NamedTuple): - """Weighted observations and their E-step assignment probabilities.""" - - points: jnp.ndarray - point_weights: jnp.ndarray - assignment_probs: jnp.ndarray - - -# Model fit - - -def get_q( - gmm: gaussian_mixture.GaussianMixture, obs: Observations -) -> jnp.ndarray: - r"""Get Q(\Theta|\Theta^{(t)}). - - Here Q is the log likelihood for our observations based on the current - parameter estimates for \Theta and averaged over the current component - assignment probabilities. See the overview of EM above for more details. - - Args: - gmm: GMM model parameterized by Theta - obs: weighted observations with component assignments computed in the E step - for \Theta^{(t)} - - Returns: - Q(\Theta|\Theta^{(t)}) - """ - # Q = E_Z log p(X, Z| Theta) - # = \sum_Z P(Z|X, Theta^(t)) [log p(X, Z | Theta)] - # Here P(Z|X, theta^(t)) is the set of assignment probabilities - # we computed in the E step. - # log p(X, Z| theta) is given by - log_p_x_z = ( - gmm.conditional_log_prob(obs.points) + # p(X | Z, theta) - gmm.log_component_weights() - ) # p(Z | theta) - return ( - jnp.sum( - obs.point_weights * - jnp.sum(log_p_x_z * obs.assignment_probs, axis=-1), - axis=0 - ) / jnp.sum(obs.point_weights, axis=0) - ) - - -# Objective function - - -@functools.lru_cache -def get_objective_fn(weight_transport: float): - """Get the total loss function with static parameters in a closure. - - Args: - weight_transport: weight for the transport penalty - - Returns: - A function that returns the objective for a GaussianMixturePair. - """ - - def _objective_fn( - pair: gaussian_mixture_pair.GaussianMixturePair, - obs0: Observations, - obs1: Observations, - ) -> jnp.ndarray: - """Compute the objective function for a pair of GMMs. - - Args: - pair: pair of GMMs + coupling for which to evaluate the objective - obs0: first set of observations - obs1: second set of observations - - Returns: - The objective to be minimized in the M-step. - """ - q0 = get_q(gmm=pair.gmm0, obs=obs0) - q1 = get_q(gmm=pair.gmm1, obs=obs1) - cost_matrix = pair.get_cost_matrix() - sinkhorn_output = pair.get_sinkhorn(cost_matrix=cost_matrix) - transport_penalty = sinkhorn_output.reg_ot_cost - return q0 + q1 - weight_transport * transport_penalty - - return _objective_fn - - -def print_losses( - iteration: int, weight_transport: float, - pair: gaussian_mixture_pair.GaussianMixturePair, obs0: Observations, - obs1: Observations -): - """Print the loss components for diagnostic purposes.""" - q0 = get_q(gmm=pair.gmm0, obs=obs0) - q1 = get_q(gmm=pair.gmm1, obs=obs1) - cost_matrix = pair.get_cost_matrix() - sinkhorn_output = pair.get_sinkhorn(cost_matrix=cost_matrix) - transport_penalty = sinkhorn_output.reg_ot_cost - objective = q0 + q1 - weight_transport * transport_penalty - - print( # noqa: T201 - f"{iteration:3d} {q0:.3f} {q1:.3f} " - f"transport:{transport_penalty:.3f} " - f"objective:{objective:.3f}" - ) - - -# The E-step for a single GMM - - -def do_e_step( # noqa: D103 - e_step_fn: Callable[[gaussian_mixture.GaussianMixture, jnp.ndarray], - jnp.ndarray], - gmm: gaussian_mixture.GaussianMixture, - points: jnp.ndarray, - point_weights: jnp.ndarray, -) -> Observations: - assignment_probs = e_step_fn(gmm, points) - return Observations( - points=points, - point_weights=point_weights, - assignment_probs=assignment_probs - ) - - -# The M-step - - -def get_m_step_fn(learning_rate: float, objective_fn, jit: bool): - """Get a function that performs the M-step of the EM algorithm. - - We precompile and precompute a few quantities that we put into a closure. - - Args: - learning_rate: learning rate to use for the Adam optimizer - objective_fn: the objective function to maximize - jit: if True, precompile key methods - - Returns: - A function that performs the M-step of EM. - """ - import optax - - def _m_step_fn( - pair: gaussian_mixture_pair.GaussianMixturePair, - obs0: Observations, - obs1: Observations, - steps: int, - ) -> gaussian_mixture_pair.GaussianMixturePair: - """Perform the M-step on a pair of Gaussian mixtures. - - Args: - pair: GMM parameters to optimize - obs0: first set of observations - obs1: second set of observations - steps: number of optimization steps to use when maximizing the objective - - Returns: - A GaussianMixturePair with updated parameters. - """ - state = opt_init((pair,)) - - for _ in range(steps): - grad_objective = grad_objective_fn(pair, obs0, obs1) - updates, state = opt_update(grad_objective, state, (pair,)) - (pair,) = optax.apply_updates((pair,), updates) - for j, gmm in enumerate((pair.gmm0, pair.gmm1)): - if gmm.has_nans(): - raise ValueError(f"NaN in gmm{j}") - return pair - - grad_objective_fn = jax.grad(objective_fn, argnums=(0,)) - if jit: - grad_objective_fn = jax.jit(grad_objective_fn) - - opt_init, opt_update = optax.chain( - # Set the parameters of Adam. Note the learning_rate is not here. - optax.scale_by_adam(b1=0.9, b2=0.999, eps=1e-8), - optax.scale(learning_rate) - ) - - return _m_step_fn - - -def get_fit_model_em_fn( - weight_transport: float, - learning_rate: float = 0.001, - jit: bool = True, -): - """Get a function that performs penalized EM. - - We precompile and precompute a few quantities that we put into a closure. - - Args: - weight_transport: weight for the transportation loss in the total loss - learning_rate: learning rate to use for the Adam optimizer - jit: if True, precompile key methods - - Returns: - A function that performs generalized, penalized EM. - """ - objective_fn = get_objective_fn(weight_transport=weight_transport) - e_step_fn = fit_gmm.get_assignment_probs - - if jit: - objective_fn = jax.jit(objective_fn) - e_step_fn = jax.jit(e_step_fn) - - m_step_fn = get_m_step_fn( - learning_rate=learning_rate, objective_fn=objective_fn, jit=jit - ) - - def _fit_model_em( - pair: gaussian_mixture_pair.GaussianMixturePair, - points0: jnp.ndarray, - points1: jnp.ndarray, - point_weights0: Optional[jnp.ndarray], - point_weights1: Optional[jnp.ndarray], - em_steps: int, - m_steps: int = 50, - verbose: bool = False, - ) -> Tuple[gaussian_mixture_pair.GaussianMixturePair, float]: - """Optimize a GaussianMixturePair using penalized EM. - - Args: - pair: GaussianMixturePair to optimize - points0: observations associated with pair.gmm0 - points1: observations associated with pair.gmm1 - point_weights0: weights for points0 - point_weights1: weights for points1 - em_steps: number of EM steps to perform - m_steps: number of gradient descent steps to perform in the M-step - verbose: if True, print status messages - - Returns: - An updated GaussianMixturePair and the final loss. - """ - if point_weights0 is None: - point_weights0 = jnp.ones(points0.shape[0]) - if point_weights1 is None: - point_weights1 = jnp.ones(points1.shape[0]) - - if pair.lock_gmm1: - obs1 = do_e_step( - e_step_fn=e_step_fn, - gmm=pair.gmm1, - points=points1, - point_weights=point_weights1 - ) - - for i in range(em_steps): - # E-step - obs0 = do_e_step( - e_step_fn=e_step_fn, - gmm=pair.gmm0, - points=points0, - point_weights=point_weights0 - ) - if not pair.lock_gmm1: - obs1 = do_e_step( - e_step_fn=e_step_fn, - gmm=pair.gmm1, - points=points1, - point_weights=point_weights1 - ) - - # print current losses - if verbose: - print_losses( - iteration=i, - weight_transport=weight_transport, - pair=pair, - obs0=obs0, - obs1=obs1 - ) - - # the M-step - pair = m_step_fn(pair=pair, obs0=obs0, obs1=obs1, steps=m_steps) - - # final E-step before computing the loss - obs0 = do_e_step( - e_step_fn=e_step_fn, - gmm=pair.gmm0, - points=points0, - point_weights=point_weights0 - ) - if not pair.lock_gmm1: - obs1 = do_e_step( - e_step_fn=e_step_fn, - gmm=pair.gmm1, - points=points1, - point_weights=point_weights1 - ) - - loss = objective_fn(pair=pair, obs0=obs0, obs1=obs1) - return pair, loss - - return _fit_model_em diff --git a/ott/build/lib/ott/tools/gaussian_mixture/gaussian.py b/ott/build/lib/ott/tools/gaussian_mixture/gaussian.py deleted file mode 100644 index 722f465..0000000 --- a/ott/build/lib/ott/tools/gaussian_mixture/gaussian.py +++ /dev/null @@ -1,217 +0,0 @@ -# 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 Optional, Union - -import jax -import jax.numpy as jnp - -from ott.tools.gaussian_mixture import scale_tril - -__all__ = ["Gaussian"] - -LOG2PI = math.log(2.0 * math.pi) - - -@jax.tree_util.register_pytree_node_class -class Gaussian: - """Normal distribution.""" - - def __init__(self, loc: jnp.ndarray, scale: scale_tril.ScaleTriL): - self._loc = loc - self._scale = scale - - @classmethod - def from_samples( - cls, - points: jnp.ndarray, - weights: Optional[jnp.ndarray] = None - ) -> "Gaussian": - """Construct a Gaussian from weighted samples. - - Unbiased, weighted covariance formula from `GSL - `_. - - Args: - points: [n x d] array of samples - weights: [n] array of weights - - Returns: - Gaussian. - """ - n = points.shape[0] - if weights is None: - weights = jnp.ones(n) / n - - mean = weights.dot(points) - centered_x = (points - mean) - scaled_centered_x = centered_x * weights.reshape(-1, 1) - cov = scaled_centered_x.T.dot(centered_x) / (1 - weights.dot(weights)) - return cls.from_mean_and_cov(mean=mean, cov=cov) - - @classmethod - def from_random( - cls, - rng: jax.Array, - n_dimensions: int, - stdev_mean: float = 0.1, - stdev_cov: float = 0.1, - ridge: Union[float, jnp.ndarray] = 0, - ) -> "Gaussian": - """Construct a random Gaussian. - - Args: - rng: jax.random key - n_dimensions: desired covariance dimensions - stdev_mean: standard deviation of location and log eigenvalues - (means for both are 0) - stdev_cov: standard deviated of the covariance - ridge: Offset for means. - - Returns: - A random Gaussian. - """ - rng, subrng0, subrng1 = jax.random.split(rng, num=3) - loc = jax.random.normal(subrng0, shape=(n_dimensions,)) * stdev_mean + ridge - scale = scale_tril.ScaleTriL.from_random( - subrng1, n_dimensions=n_dimensions, stdev=stdev_cov - ) - return cls(loc=loc, scale=scale) - - @classmethod - def from_mean_and_cov(cls, mean: jnp.ndarray, cov: jnp.ndarray) -> "Gaussian": - """Construct a Gaussian from a mean and covariance.""" - scale = scale_tril.ScaleTriL.from_covariance(cov) - return cls(loc=mean, scale=scale) - - @property - def loc(self) -> jnp.ndarray: - """Mean of the Gaussian.""" - return self._loc - - @property - def scale(self) -> scale_tril.ScaleTriL: - """Scale of the Gaussian.""" - return self._scale - - @property - def n_dimensions(self) -> int: - """Dimensionality of the Gaussian.""" - return self.loc.shape[-1] - - def covariance(self) -> jnp.ndarray: - """Covariance of the Gaussian.""" - return self.scale.covariance() - - def to_z(self, x: jnp.ndarray) -> jnp.ndarray: - r"""Transform :math:`x` to :math:`z = \frac{x - loc}{scale}`.""" - return self.scale.centered_to_z(x_centered=x - self.loc) - - def from_z(self, z: jnp.ndarray) -> jnp.ndarray: - r"""Transform :math:`z` to :math:`x = loc + scale \cdot z`.""" - return self.scale.z_to_centered(z=z) + self.loc - - def log_prob( - self, - x: jnp.ndarray, # (?, d) - ) -> jnp.ndarray: # (?, d) - """Log probability for a Gaussian with a diagonal covariance.""" - d = x.shape[-1] - z = self.to_z(x) - log_det = self.scale.log_det_covariance() - return ( - -0.5 * (d * LOG2PI + log_det[None] + jnp.sum(z ** 2, axis=-1)) - ) # (?, k) - - def sample(self, rng: jax.Array, size: int) -> jnp.ndarray: - """Generate samples from the distribution.""" - std_samples_t = jax.random.normal(rng, shape=(self.n_dimensions, size)) - return self.loc[None] + ( - jnp.swapaxes( - jnp.matmul(self.scale.cholesky(), std_samples_t), - axis1=-2, - axis2=-1 - ) - ) - - def w2_dist(self, other: "Gaussian") -> jnp.ndarray: - r"""Wasserstein distance :math:`W_2^2` to another Gaussian. - - .. math:: - - W_2^2 = ||\mu_0-\mu_1||^2 + - \text{trace} ( (\Lambda_0^\frac{1}{2} - \Lambda_1^\frac{1}{2})^2 ) - - Args: - other: other Gaussian - - Returns: - The :math:`W_2^2` distance between self and other - """ - delta_mean = jnp.sum((self.loc - other.loc) ** 2, axis=-1) - delta_sigma = self.scale.w2_dist(other.scale) - return delta_mean + delta_sigma - - def f_potential(self, dest: "Gaussian", points: jnp.ndarray) -> jnp.ndarray: - """Optimal potential for W2 distance between Gaussians. Evaluated on points. - - Args: - dest: Gaussian object - points: samples - - Returns: - Dual potential, f - """ - scale_matrix = self.scale.gaussian_map(dest_scale=dest.scale) - centered_x = points - self.loc - scaled_x = (scale_matrix @ centered_x.T) - - @jax.vmap - def batch_inner_product(x, y): - return x.dot(y) - - return ( - 0.5 * batch_inner_product(points, points) - - 0.5 * batch_inner_product(centered_x, scaled_x.T) - - points.dot(dest.loc) - ) - - def transport(self, dest: "Gaussian", points: jnp.ndarray) -> jnp.ndarray: - """Transport points according to map between two Gaussian measures. - - Args: - dest: Gaussian object - points: samples - - Returns: - Transported samples - """ - return self.scale.transport( - dest_scale=dest.scale, points=points - self.loc[None] - ) + dest.loc[None] - - def tree_flatten(self): # noqa: D102 - children = (self.loc, self.scale) - aux_data = {} - return children, aux_data - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - return cls(*children, **aux_data) - - def __hash__(self): - return jax.tree_util.tree_flatten(self).__hash__() - - def __eq__(self, other): - return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture.py b/ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture.py deleted file mode 100644 index 27a5689..0000000 --- a/ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture.py +++ /dev/null @@ -1,329 +0,0 @@ -# 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. -from typing import List, Tuple, Union - -import jax -import jax.numpy as jnp - -from ott.tools.gaussian_mixture import ( - gaussian, - linalg, - probabilities, - scale_tril, -) - -__all__ = ["GaussianMixture"] - - -def get_summary_stats_from_points_and_assignment_probs( - points: jnp.ndarray, point_weights: jnp.ndarray, - assignment_probs: jnp.ndarray -) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Get component summary stats from points and component probabilities. - - Args: - points: array of points, shape (n, n_dim) - point_weights: array of weights for the points, shape (n,) - assignment_probs: array of component assignment probabilities for the - points, shape (n, n_components) - - Returns: - Tuple containing for each component, - * the sample mean for each component, shape (n_components, n_dim) - * the sample covariance for each component, - shape (n_components, n_dim, n_dim) - * the weight for each component, - shape (n_components,) - """ - - def component_from_points(points, point_weights, assignment_probs): - component_weight = ( - jnp.sum(point_weights * assignment_probs) / jnp.sum(point_weights) - ) - component_mean, component_cov = linalg.get_mean_and_cov( - points=points, weights=point_weights * assignment_probs - ) - return component_mean, component_cov, component_weight - - components_from_points_fn = jax.vmap( - component_from_points, in_axes=(None, None, 1), out_axes=0 - ) - - return components_from_points_fn(points, point_weights, assignment_probs) - - -@jax.tree_util.register_pytree_node_class -class GaussianMixture: - """Gaussian Mixture model.""" - - def __init__( - self, loc: jnp.ndarray, scale_params: jnp.ndarray, - component_weight_ob: probabilities.Probabilities - ): - self._loc = loc - self._scale_params = scale_params - self._component_weight_ob = component_weight_ob - - @classmethod - def from_random( - cls, - rng: jax.Array, - n_components: int, - n_dimensions: int, - stdev_mean: float = 0.1, - stdev_cov: float = 0.1, - stdev_weights: float = 0.1, - ridge: Union[float, jnp.array] = 0, - ) -> "GaussianMixture": - """Construct a random GMM.""" - loc = [] - scale_params = [] - for _ in range(n_components): - rng, subrng = jax.random.split(rng) - component = gaussian.Gaussian.from_random( - subrng, - n_dimensions=n_dimensions, - stdev_mean=stdev_mean, - stdev_cov=stdev_cov, - ridge=ridge, - ) - loc.append(component.loc) - scale_params.append(component.scale.params) - loc = jnp.stack(loc, axis=0) - scale_params = jnp.stack(scale_params, axis=0) - weight_ob = probabilities.Probabilities.from_random( - subrng, - n_dimensions=n_components, - stdev=stdev_weights, - ) - return cls( - loc=loc, scale_params=scale_params, component_weight_ob=weight_ob - ) - - @classmethod - def from_mean_cov_component_weights( - cls, mean: jnp.ndarray, cov: jnp.ndarray, component_weights: jnp.ndarray - ): - """Construct a GMM from means, covariances, and component weights.""" - scale_params = [] - for i in range(cov.shape[0]): - scale_params.append(scale_tril.ScaleTriL.from_covariance(cov[i]).params) - scale_params = jnp.stack(scale_params, axis=0) - weight_ob = probabilities.Probabilities.from_probs(component_weights) - return cls( - loc=mean, scale_params=scale_params, component_weight_ob=weight_ob - ) - - @classmethod - def from_points_and_assignment_probs( - cls, - points: jnp.ndarray, - point_weights: jnp.ndarray, - assignment_probs: jnp.ndarray, - ) -> "GaussianMixture": - """Estimate a GMM from points and a set of component probabilities.""" - mean, cov, wts = get_summary_stats_from_points_and_assignment_probs( - points=points, - point_weights=point_weights, - assignment_probs=assignment_probs - ) - return cls.from_mean_cov_component_weights( - mean=mean, cov=cov, component_weights=wts - ) - - @property - def dtype(self): - """Dtype of the GMM parameters.""" - return self.loc.dtype - - @property - def n_dimensions(self): - """Number of dimensions of the GMM parameters.""" - return self._loc.shape[-1] - - @property - def n_components(self): - """Number of components of the GMM parameters.""" - return self._loc.shape[-2] - - @property - def loc(self) -> jnp.ndarray: - """Location parameters of the GMM.""" - return self._loc - - @property - def scale_params(self) -> jnp.ndarray: - """Scale parameters of the GMM.""" - return self._scale_params - - @property - def cholesky(self) -> jnp.ndarray: - """Cholesky decomposition of the GMM covariance matrices.""" - size = self.n_dimensions - - def _get_cholesky(scale_params): - return scale_tril.ScaleTriL(params=scale_params, size=size).cholesky() - - return jax.vmap(_get_cholesky, in_axes=0, out_axes=0)(self.scale_params) - - @property - def covariance(self) -> jnp.ndarray: - """Covariance matrices of the GMM.""" - size = self.n_dimensions - - def _get_covariance(scale_params): - return scale_tril.ScaleTriL(params=scale_params, size=size).covariance() - - return jax.vmap(_get_covariance, in_axes=0, out_axes=0)(self.scale_params) - - @property - def component_weight_ob(self) -> probabilities.Probabilities: - """Component weight object.""" - return self._component_weight_ob - - @property - def component_weights(self) -> jnp.ndarray: - """Component weights probabilities.""" - return self._component_weight_ob.probs() - - def log_component_weights(self) -> jnp.ndarray: - """Log component weights probabilities.""" - return self._component_weight_ob.log_probs() - - def _get_normal( - self, loc: jnp.ndarray, scale_params: jnp.ndarray - ) -> gaussian.Gaussian: - size = loc.shape[-1] - return gaussian.Gaussian( - loc=loc, scale=scale_tril.ScaleTriL(params=scale_params, size=size) - ) - - def get_component(self, index: int) -> gaussian.Gaussian: - """Specified GMM component.""" - return self._get_normal( - loc=self.loc[index], scale_params=self.scale_params[index] - ) - - def components(self) -> List[gaussian.Gaussian]: - """List of all GMM components.""" - return [self.get_component(i) for i in range(self.n_components)] - - def sample(self, rng: jax.Array, size: int) -> jnp.ndarray: - """Generate samples from the distribution.""" - subrng0, subrng1 = jax.random.split(rng) - component = self.component_weight_ob.sample(rng=subrng0, size=size) - std_samples = jax.random.normal(subrng1, shape=(size, self.n_dimensions)) - - def _transform_single_component(k, scale, loc): - - def _transform_single_value(single_component, single_x): - return jax.lax.cond( - single_component == k, - lambda x: jnp.matmul(scale, x[:, None])[:, 0] + loc, jnp.zeros_like, - single_x - ) - - return jax.vmap(_transform_single_value)(component, std_samples) - - return jnp.sum( - jax.vmap(_transform_single_component) - (jnp.arange(self.n_components), self.cholesky, self.loc), - axis=0 - ) - - def conditional_log_prob(self, x: jnp.ndarray) -> jnp.ndarray: - """Compute the component-conditional log probability of x. - - Args: - x: (n, n_dimensions) array of points - - Returns: - (n, n_components) array of the log probability of x conditioned on it - having come from each component. - """ - - def _log_prob_single_component( - loc: jnp.ndarray, scale_params: jnp.ndarray, x: jnp.ndarray - ): - norm = self._get_normal(loc=loc, scale_params=scale_params) - return norm.log_prob(x) - - conditional_log_prob_fn = jax.vmap( - _log_prob_single_component, in_axes=(0, 0, None), out_axes=1 - ) - return conditional_log_prob_fn(self._loc, self._scale_params, x) - - def log_prob(self, x: jnp.ndarray) -> jnp.ndarray: - """Compute the log probability of the observations x. - - Args: - x: (n, n_dimensions) array of points - - Returns: - (n,) array of log probabilities. - """ - # p(x) = \sum_i p(x|c_i) p(c_i) - log_prob_conditional = self.conditional_log_prob(x) - log_component_weight = self.log_component_weights() - return jax.scipy.special.logsumexp( - log_prob_conditional + log_component_weight[None, :], axis=-1 - ) - - def get_log_component_posterior(self, x: jnp.ndarray) -> jnp.ndarray: - """Compute the posterior probability that x came from each component. - - Args: - x: (n, n_dimensions) array of points - - Returns: - (n, n_components) array of poster component log probabilities. - """ - # p(x | c_i) = p(x, c_i) / p(c_i) => p(x, c_i) = p(x | c_i) p(c_i) - # p(c_i | x) = p(x, c_i) / p(x) - # = p(x | c_i) p(c_i) / sum_j(p(x | c_j)p(c_j)) - log_prob_conditional = self.conditional_log_prob(x) - log_component_weight = self.log_component_weights() - log_prob_unnorm = log_prob_conditional + log_component_weight[None, :] - return log_prob_unnorm - jax.scipy.special.logsumexp( - log_prob_unnorm, axis=-1, keepdims=True - ) - - def has_nans(self) -> bool: # noqa: D102 - for leaf in jax.tree_util.tree_leaves(self): - if jnp.any(~jnp.isfinite(leaf)): - return True - return False - - def tree_flatten(self): # noqa: D102 - children = (self.loc, self.scale_params, self.component_weight_ob) - aux_data = {} - return children, aux_data - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - return cls(*children, **aux_data) - - def __repr__(self): - class_name = type(self).__name__ - children, aux = self.tree_flatten() - return "{}({})".format( - class_name, ", ".join([repr(c) for c in children] + - [f"{k}: {repr(v)}" for k, v in aux.items()]) - ) - - def __hash__(self): - return jax.tree_util.tree_flatten(self).__hash__() - - def __eq__(self, other): - return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture_pair.py b/ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture_pair.py deleted file mode 100644 index 62e04c0..0000000 --- a/ott/build/lib/ott/tools/gaussian_mixture/gaussian_mixture_pair.py +++ /dev/null @@ -1,222 +0,0 @@ -# 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. -from typing import Any - -import jax -import jax.numpy as jnp - -from ott.geometry import costs, geometry, pointcloud -from ott.problems.linear import linear_problem -from ott.solvers.linear import sinkhorn -from ott.tools.gaussian_mixture import gaussian_mixture - -__all__ = ["GaussianMixturePair"] - - -@jax.tree_util.register_pytree_node_class -class GaussianMixturePair: - """Coupled pair of Gaussian mixture models. - - Includes methods used in estimating an optimal pairing between GMM components - using the Wasserstein-like method described in :cite:`delon:20`, - as well as generalization that allows for the reweighting of components. - - :cite:`delon:20` propose fitting a pair of GMMs to a pair - of point clouds in such a way that the sum of the log likelihood of the - points minus a weighted penalty involving a Wasserstein-like distance between - the GMMs. Their proposed algorithm involves using EM in which a balanced - Sinkhorn algorithm is used to estimate a coupling between the GMMs at each - step of EM. - - Our generalization of this algorithm allows for a mismatch between the - marginals of the coupling and the GMM component weights. This mismatch can be - interpreted as components being reweighted rather than being transported. - We penalize reweighting with a generalized KL-divergence penalty, and we give - the option to use the unbalanced Sinkhorn algorithm rather than the balanced - to compute the divergence between GMMs. - """ - - def __init__( - self, - gmm0: gaussian_mixture.GaussianMixture, - gmm1: gaussian_mixture.GaussianMixture, - epsilon: float = 1e-2, - tau: float = 1.0, - lock_gmm1: bool = False, - ): - """Constructor. - - When fitting a pair of coupled GMMs with *no* reweighting of components - using the algorithm in :cite:`delon:20`, set tau = 1. The coupling between - components will be determined via the balanced Sinkhorn algorithm. - - When fitting a pair of coupled GMMs in which reweighting of components is - allowed, set tau to a value in (0, 1). The resulting coupling will penalize - the generalized KL divergence between the coupling's marginals and the GMM - component weights with a weight of rho = epsilon tau / (1 - tau). - - Args: - gmm0: first GMM in the pair - gmm1: second GMM in the pair - epsilon: regularization weight to use for the Sinkhorn algorithm - tau: encodes the weight, rho, to use for the generalized KL divergence - between the coupling's marginals and GMM component weights as - rho = epsilon tau / (1 - tau) - lock_gmm1: indicates whether the parameters of gmm1 should be modified - during optimization - """ # noqa: D401 - self._gmm0 = gmm0 - self._gmm1 = gmm1 - self._epsilon = epsilon - self._tau = tau - self._lock_gmm1 = lock_gmm1 - - @property - def dtype(self): # noqa: D102 - return self.gmm0.dtype - - @property - def gmm0(self): # noqa: D102 - return self._gmm0 - - @property - def gmm1(self): # noqa: D102 - return self._gmm1 - - @property - def epsilon(self): # noqa: D102 - return self._epsilon - - @property - def tau(self): # noqa: D102 - return self._tau - - @property - def rho(self): # noqa: D102 - return self.epsilon * self.tau / (1.0 - self.tau) - - @property - def lock_gmm1(self): # noqa: D102 - return self._lock_gmm1 - - def get_bures_geometry(self) -> pointcloud.PointCloud: - """Get a Bures Geometry for the two GMMs.""" - mean0 = self.gmm0.loc - dimension = mean0.shape[-1] - cov0 = self.gmm0.covariance - cov0 = cov0.reshape(cov0.shape[:-2] + (dimension * dimension,)) - x = jnp.concatenate([mean0, cov0], axis=-1) - mean1 = self.gmm1.loc - cov1 = self.gmm1.covariance - cov1 = cov1.reshape(cov1.shape[:-2] + (dimension * dimension,)) - y = jnp.concatenate([mean1, cov1], axis=-1) - return pointcloud.PointCloud( - x=x, - y=y, - cost_fn=costs.Bures(dimension=dimension), - epsilon=self.epsilon - ) - - def get_cost_matrix(self) -> jnp.ndarray: - """Get matrix of :math:`W_2^2` costs between all pairs of components.""" - return self.get_bures_geometry().cost_matrix - - def get_sinkhorn( - self, cost_matrix: jnp.ndarray, **kwargs: Any - ) -> sinkhorn.SinkhornOutput: - """Get the output of Sinkhorn's method for a given cost matrix.""" - # We use a Geometry here rather than the PointCloud created in - # get_bures_geometry to avoid recomputing the cost matrix, since - # the cost matrix is quite expensive - geom = geometry.Geometry(cost_matrix=cost_matrix, epsilon=self.epsilon) - prob = linear_problem.LinearProblem( - geom, - a=self.gmm0.component_weights, - b=self.gmm1.component_weights, - tau_a=self.tau, - tau_b=self.tau - ) - return sinkhorn.Sinkhorn(**kwargs)(prob) - - def get_normalized_sinkhorn_coupling( - self, - sinkhorn_output: sinkhorn.SinkhornOutput, - ) -> jnp.ndarray: - """Get the normalized coupling matrix for the specified Sinkhorn output. - - Args: - sinkhorn_output: Sinkhorn algorithm output as returned by - :meth:`get_sinkhorn`. - - Returns: - A coupling matrix that tells how much of the mass of each component of - :attr:`gmm0` is mapped to each component of :attr:`gmm1`. - """ - return sinkhorn_output.matrix / jnp.sum(sinkhorn_output.matrix) - - def tree_flatten(self): - """Method used by jax.tree_util to flatten a GaussianMixturePair. - - We control the subset of parameters that we will optimize in fit_gmm_pair - by selectively placing them in either children (the parameters to optimize) - or aux_data (the parameters to leave alone). - - Returns: - A tuple of child pytrees and a dict of auxiliary data. - """ # noqa: D401 - children = [self.gmm0] - aux_data = { - "epsilon": self.epsilon, - "tau": self.tau, - "lock_gmm1": self.lock_gmm1 - } - if self.lock_gmm1: - aux_data["gmm1"] = self.gmm1 - else: - children.append(self.gmm1) - return tuple(children), aux_data - - @classmethod - def tree_unflatten(cls, aux_data, children): - """Method used by jax.tree_util to unflatten a GaussianMixturePair. - - tree_flatten controls which parameters get optimized by placing them in - either children or aux_data; here we invert the process. - - Args: - aux_data: auxiliary data that is passed to the constructor as kwargs - children: child pytrees passed to the constructor as args - - Returns: - A GaussianMixturePair. - """ # noqa: D401 - children = list(children) - if "gmm1" in aux_data: - gmm1 = aux_data.pop("gmm1") - children.insert(1, gmm1) - return cls(*children, **aux_data) - - def __repr__(self): - class_name = type(self).__name__ - children, aux = self.tree_flatten() - return "{}({})".format( - class_name, ", ".join([repr(c) for c in children] + - [f"{k}: {repr(v)}" for k, v in aux.items()]) - ) - - def __hash__(self): - return jax.tree_util.tree_flatten(self).__hash__() - - def __eq__(self, other): - return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/ott/build/lib/ott/tools/gaussian_mixture/linalg.py b/ott/build/lib/ott/tools/gaussian_mixture/linalg.py deleted file mode 100644 index dfc7d4a..0000000 --- a/ott/build/lib/ott/tools/gaussian_mixture/linalg.py +++ /dev/null @@ -1,138 +0,0 @@ -# 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. -from typing import Callable, Iterable, List, Tuple - -import jax -import jax.numpy as jnp - - -def get_mean_and_var( - points: jnp.ndarray, # (n, d) - weights: jnp.ndarray, # (n,) -) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Get the mean and variance of a weighted set of points.""" - weights_sum = jnp.sum(weights, axis=-1) # (1,) - mean = ( - # matmul((1, n), (n, d)) -> (1, d) - jnp.matmul(weights, points) / weights_sum - ) - # center points - centered = points - mean[None, :] # (n, d) - (1, d) - var = ( - # matmul((1, n), (n, d)) -> (1, d) - jnp.matmul(weights, centered ** 2) / weights_sum - ) - return mean, var - - -def get_mean_and_cov( - points: jnp.ndarray, # (n, d) - weights: jnp.ndarray, # (n,) -) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Get the mean and covariance of a weighted set of points.""" - weights_sum = jnp.sum(weights, axis=-1, keepdims=True) # (1,) - mean = ( - # matmul((1, n), (n, d)) -> (1, d) - jnp.matmul(weights, points) / weights_sum - ) - # center points - centered = points - mean[None, :] # (n, d) - (1, d) - cov = ( - jnp.matmul( - # (1, n) (d, n) - weights[None, :] * jnp.swapaxes(centered, axis1=-2, axis2=-1), - # (n, d) - centered - ) / weights_sum - ) - return mean, cov - - -def flat_to_tril(x: jnp.ndarray, size: int) -> jnp.ndarray: - """Map flat values to lower triangular matrices. - - Args: - x: flat values - size: size of lower triangular matrices. x should have shape - (..., size(size+1)/2), and the final matrices should have shape - (..., size, size). - - Returns: - Lower triangular matrices. - """ - m = jnp.zeros(x.shape[:-1] + (size, size)) - tril = jnp.tril_indices(size) - return m.at[..., tril[0], tril[1]].set(x) - - -def tril_to_flat(m: jnp.ndarray) -> jnp.ndarray: - """Flatten lower triangular matrices. - - Args: - m: lower triangular matrices of shape (..., size, size) - - Returns: - A vector of shape (..., size (size+1) // 2) - """ - size = m.shape[-1] - tril = jnp.tril_indices(size) - return m[..., tril[0], tril[1]] - - -def apply_to_diag( - m: jnp.ndarray, fn: Callable[[jnp.ndarray], jnp.ndarray] -) -> jnp.ndarray: - """Apply a function to the diagonal of a matrix.""" - size = m.shape[-1] - diag = jnp.diagonal(m, axis1=-2, axis2=-1) - ind = jnp.arange(size) - return m.at[..., ind, ind].set(fn(diag)) - - -def matrix_powers( - m: jnp.ndarray, - powers: Iterable[float], -) -> List[jnp.ndarray]: - """Raise a real, symmetric matrix to multiple powers.""" - eigs, q = jnp.linalg.eigh(m) - qt = jnp.swapaxes(q, axis1=-2, axis2=-1) - ret = [] - for power in powers: - ret.append(jnp.matmul(jnp.expand_dims(eigs ** power, -2) * q, qt)) - return ret - - -def invmatvectril( - m: jnp.ndarray, x: jnp.ndarray, lower: bool = True -) -> jnp.ndarray: - """Multiply x by the inverse of a triangular matrix. - - Args: - m: triangular matrix, shape (d, d) - x: array of points, shape (n, d) - lower: if True, m is lower triangular; otherwise m is upper triangular - - Returns: - m^{-1} x - """ - return jnp.transpose( - jax.scipy.linalg.solve_triangular(m, jnp.transpose(x), lower=lower) - ) - - -def get_random_orthogonal(rng: jax.Array, dim: int) -> jnp.ndarray: - """Get a random orthogonal matrix with the specified dimension.""" - m = jax.random.normal(rng, shape=[dim, dim]) - q, _ = jnp.linalg.qr(m) - return q diff --git a/ott/build/lib/ott/tools/gaussian_mixture/probabilities.py b/ott/build/lib/ott/tools/gaussian_mixture/probabilities.py deleted file mode 100644 index c0ae2b2..0000000 --- a/ott/build/lib/ott/tools/gaussian_mixture/probabilities.py +++ /dev/null @@ -1,100 +0,0 @@ -# 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. -from typing import Optional - -import jax -import jax.numpy as jnp - -__all__ = ["Probabilities"] - - -@jax.tree_util.register_pytree_node_class -class Probabilities: - """Parameterized array of probabilities of length n. - - The internal representation is a length n-1 unconstrained array. We convert - to a length n simplex by appending a 0 and taking a softmax. - """ - - _params: jnp.ndarray - - def __init__(self, params): - self._params = params - - @classmethod - def from_random( - cls, - rng: jax.Array, - n_dimensions: int, - stdev: Optional[float] = 0.1, - ) -> "Probabilities": - """Construct a random Probabilities.""" - return cls(params=jax.random.normal(rng, shape=(n_dimensions - 1,)) * stdev) - - @classmethod - def from_probs(cls, probs: jnp.ndarray) -> "Probabilities": - """Construct Probabilities from a vector of probabilities.""" - log_probs = jnp.log(probs) - log_probs_normalized, norm = log_probs[:-1], log_probs[-1] - log_probs_normalized -= norm - return cls(params=log_probs_normalized) - - @property - def params(self): # noqa: D102 - return self._params - - @property - def dtype(self): # noqa: D102 - return self._params.dtype - - def unnormalized_log_probs(self) -> jnp.ndarray: - """Get the unnormalized log probabilities.""" - return jnp.concatenate([self._params, jnp.zeros((1,))], axis=-1) - - def log_probs(self) -> jnp.ndarray: - """Get the log probabilities.""" - return jax.nn.log_softmax(self.unnormalized_log_probs()) - - def probs(self) -> jnp.ndarray: - """Get the probabilities.""" - return jax.nn.softmax(self.unnormalized_log_probs()) - - def sample(self, rng: jax.Array, size: int) -> jnp.ndarray: - """Sample from the distribution.""" - return jax.random.categorical( - rng, logits=self.unnormalized_log_probs(), shape=(size,) - ) - - def tree_flatten(self): # noqa: D102 - children = (self.params,) - aux_data = {} - return children, aux_data - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - return cls(*children, **aux_data) - - def __repr__(self): - class_name = type(self).__name__ - children, aux = self.tree_flatten() - return "{}({})".format( - class_name, ", ".join([repr(c) for c in children] + - [f"{k}: {repr(v)}" for k, v in aux.items()]) - ) - - def __hash__(self): - return jax.tree_util.tree_flatten(self).__hash__() - - def __eq__(self, other): - return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/ott/build/lib/ott/tools/gaussian_mixture/scale_tril.py b/ott/build/lib/ott/tools/gaussian_mixture/scale_tril.py deleted file mode 100644 index b26dded..0000000 --- a/ott/build/lib/ott/tools/gaussian_mixture/scale_tril.py +++ /dev/null @@ -1,214 +0,0 @@ -# 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. -from typing import Optional, Tuple - -import jax -import jax.numpy as jnp - -from ott.geometry import costs -from ott.math import matrix_square_root -from ott.tools.gaussian_mixture import linalg - -__all__ = ["ScaleTriL"] - - -@jax.tree_util.register_pytree_node_class -class ScaleTriL: - """Pytree for a lower triangular Cholesky-factored covariance matrix.""" - - def __init__(self, params: jnp.ndarray, size: int): - self._params = params - self._size = size - - @classmethod - def from_points_and_weights( - cls, - points: jnp.ndarray, - weights: jnp.ndarray, - ) -> Tuple[jnp.ndarray, "ScaleTriL"]: - """Get a mean and a ScaleTriL from a set of points and weights.""" - mean, cov = linalg.get_mean_and_cov(points=points, weights=weights) - return mean, cls.from_covariance(cov) - - @classmethod - def from_random( - cls, - rng: jax.Array, - n_dimensions: int, - stdev: Optional[float] = 0.1, - ) -> "ScaleTriL": - """Construct a random ScaleTriL. - - Args: - rng: pseudo-random number generator key - n_dimensions: number of dimensions - stdev: desired standard deviation (around 0) for the log eigenvalues - - Returns: - A ScaleTriL. - """ - # generate a random orthogonal matrix - rng, subrng = jax.random.split(rng) - q = linalg.get_random_orthogonal(subrng, dim=n_dimensions) - - # generate random eigenvalues - eigs = stdev * jnp.exp(jax.random.normal(rng, shape=(n_dimensions,))) - - # random positive definite matrix - sigma = q * jnp.expand_dims(eigs, -2) @ q.T - - # cholesky factorization - chol = jnp.linalg.cholesky(sigma) - # flatten - m = linalg.apply_to_diag(chol, jnp.log) - flat = linalg.tril_to_flat(m) - return cls(params=flat, size=n_dimensions) - - @classmethod - def from_cholesky(cls, cholesky: jnp.ndarray) -> "ScaleTriL": - """Construct ScaleTriL from a Cholesky factor of a covariance matrix.""" - m = linalg.apply_to_diag(cholesky, jnp.log) - flat = linalg.tril_to_flat(m) - return cls(params=flat, size=cholesky.shape[-1]) - - @classmethod - def from_covariance( - cls, - covariance: jnp.ndarray, - ) -> "ScaleTriL": - """Construct ScaleTriL from a covariance matrix.""" - cholesky = jnp.linalg.cholesky(covariance) - return cls.from_cholesky(cholesky) - - @property - def params(self) -> jnp.ndarray: - """Internal representation.""" - return self._params - - @property - def size(self) -> int: - """Size of the covariance matrix.""" - return self._size - - @property - def dtype(self): - """Data type of the covariance matrix.""" - return self._params.dtype - - def cholesky(self) -> jnp.ndarray: - """Get a lower triangular Cholesky factor for the covariance matrix.""" - m = linalg.flat_to_tril(self._params, size=self._size) - return linalg.apply_to_diag(m, jnp.exp) - - def covariance(self) -> jnp.ndarray: - """Get the covariance matrix.""" - cholesky = self.cholesky() - return cholesky @ cholesky.T - - def covariance_sqrt(self) -> jnp.ndarray: - """Get the square root of the covariance matrix.""" - return linalg.matrix_powers(self.covariance(), (0.5,))[0] - - def log_det_covariance(self) -> jnp.ndarray: - """Get the log of the determinant of the covariance matrix.""" - diag = jnp.diagonal(self.cholesky(), axis1=-2, axis2=-1) - return 2.0 * jnp.sum(jnp.log(diag), axis=-1) - - def centered_to_z(self, x_centered: jnp.ndarray) -> jnp.ndarray: - """Map centered points to standardized centered points (i.e. cov(z) = I).""" - return linalg.invmatvectril(m=self.cholesky(), x=x_centered, lower=True) - - def z_to_centered(self, z: jnp.ndarray) -> jnp.ndarray: - """Scale standardized points to points with the specified covariance.""" - return (self.cholesky() @ z.T).T - - def w2_dist(self, other: "ScaleTriL") -> jnp.ndarray: - r"""Wasserstein distance W_2^2 to another Gaussian with same mean. - - Args: - other: Scale for the other Gaussian - - Returns: - The W_2^2 distance - """ - dimension = self.size - - def _flatten_cov(cov: jnp.ndarray) -> jnp.ndarray: - cov = cov.reshape(cov.shape[:-2] + (dimension * dimension,)) - return jnp.concatenate([jnp.zeros(dimension), cov], axis=-1) - - x0 = _flatten_cov(self.covariance()) - x1 = _flatten_cov(other.covariance()) - cost_fn = costs.Bures(dimension=dimension) - return (cost_fn.norm(x0) + cost_fn.norm(x1) + cost_fn.pairwise(x0, x1))[ - ..., - ] - - def gaussian_map(self, dest_scale: "ScaleTriL") -> jnp.ndarray: - """Scaling matrix used in transport between 0-mean Gaussians. - - Sigma_mu^{-1/2} @ - [Sigma_mu ^{1/2} Sigma_nu Sigma_mu ^{1/2}]^{1/2} - @ Sigma_mu ^{-1/2} - - Args: - dest_scale: destination Scale - - Returns: - Gaussian scaling matrix, same dimension as self.covaraince - """ - sqrt0, sqrt0_inv = linalg.matrix_powers(self.covariance(), (0.5, -0.5)) - sigma1 = dest_scale.covariance() - m = matrix_square_root.sqrtm_only( - jnp.matmul(sqrt0, jnp.matmul(sigma1, sqrt0)) - ) - return jnp.matmul(sqrt0_inv, jnp.matmul(m, sqrt0_inv)) - - def transport( - self, dest_scale: "ScaleTriL", points: jnp.ndarray - ) -> jnp.ndarray: - """Apply Monge map, computed between two 0-mean Gaussians, to points. - - Args: - dest_scale: destination Scale - points: points to transport - - Returns: - Points transported to a Gaussian with the new scale. - """ - m = self.gaussian_map(dest_scale) - return (m @ points.T).T - - def tree_flatten(self): # noqa: D102 - children = (self.params,) - aux_data = {"size": self.size} - return children, aux_data - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - return cls(*children, **aux_data) - - def __repr__(self): - class_name = type(self).__name__ - children, aux = self.tree_flatten() - return "{}({})".format( - class_name, ", ".join([repr(c) for c in children] + - [f"{k}: {repr(v)}" for k, v in aux.items()]) - ) - - def __hash__(self): - return jax.tree_util.tree_flatten(self).__hash__() - - def __eq__(self, other): - return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/ott/build/lib/ott/tools/k_means.py b/ott/build/lib/ott/tools/k_means.py deleted file mode 100644 index 2a12f6e..0000000 --- a/ott/build/lib/ott/tools/k_means.py +++ /dev/null @@ -1,412 +0,0 @@ -# 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 functools -import math -from typing import Callable, Literal, NamedTuple, Optional, Tuple, Union - -import jax -import jax.numpy as jnp - -from ott import utils -from ott.geometry import costs, pointcloud -from ott.math import fixed_point_loop - -__all__ = ["k_means", "KMeansOutput"] - -Init_t = Union[Literal["k-means++", "random"], - Callable[[pointcloud.PointCloud, int, jnp.ndarray], jnp.ndarray]] - - -class KPPState(NamedTuple): # noqa: D101 - rng: jax.Array - centroids: jnp.ndarray - centroid_dists: jnp.ndarray - - -class KMeansState(NamedTuple): # noqa: D101 - centroids: jnp.ndarray - prev_assignment: jnp.ndarray - assignment: jnp.ndarray - errors: jnp.ndarray - center_shift: float - - -class KMeansConst(NamedTuple): # noqa: D101 - geom: pointcloud.PointCloud - x_weights: jnp.ndarray - - @property - def x(self) -> jnp.ndarray: - """Array of shape ``[n, ndim]`` containing the unweighted point cloud.""" - return self.geom.x - - @property - def weighted_x(self): - """Array of shape ``[n, ndim]`` containing the weighted point cloud.""" - return self.x_weights[:, :-1] - - @property - def weights(self) -> jnp.ndarray: - """Array of shape ``[n, 1]`` containing weights for each point.""" - return self.x_weights[:, -1:] - - -class KMeansOutput(NamedTuple): - """Output of the :func:`~ott.tools.k_means.k_means` algorithm. - - Args: - centroids: Array of shape ``[k, ndim]`` containing the centroids. - assignment: Array of shape ``[n,]`` containing the labels. - converged: Whether the algorithm has converged. - iteration: The number of iterations run. - error: (Weighted) sum of squared distances from each point to its closest - center. - inner_errors: Array of shape ``[max_iterations,]`` containing the ``error`` - at every iteration. - """ - centroids: jnp.ndarray - assignment: jnp.ndarray - converged: bool - iteration: int - error: float - inner_errors: Optional[jnp.ndarray] - - @classmethod - def _from_state( - cls, - state: KMeansState, - *, - tol: float, - store_inner_errors: bool = False - ) -> "KMeansOutput": - errs = state.errors - mask = errs == -1 - error = jnp.nanmin(jnp.where(mask, jnp.nan, errs)) - - assignment_same = jnp.all(state.prev_assignment == state.assignment) - tol_satisfied = jnp.logical_or(jnp.any(mask), (errs[-2] - errs[-1]) <= tol) - converged = jnp.logical_or(assignment_same, tol_satisfied) - - return cls( - centroids=state.centroids, - assignment=state.assignment.astype(int), - converged=converged, - iteration=jnp.sum(~mask), - error=error, - inner_errors=errs if store_inner_errors else None, - ) - - -def _random_init( - geom: pointcloud.PointCloud, k: int, rng: jax.Array -) -> jnp.ndarray: - ixs = jnp.arange(geom.shape[0]) - ixs = jax.random.choice(rng, ixs, shape=(k,), replace=False) - return geom.subset(ixs, None).x - - -def _k_means_plus_plus( - geom: pointcloud.PointCloud, - k: int, - rng: jax.Array, - n_local_trials: Optional[int] = None, -) -> jnp.ndarray: - - def init_fn(geom: pointcloud.PointCloud, rng: jax.Array) -> KPPState: - rng, next_rng = jax.random.split(rng, 2) - ix = jax.random.choice(rng, jnp.arange(geom.shape[0]), shape=()) - centroids = jnp.full((k, geom.cost_rank), jnp.inf).at[0].set(geom.x[ix]) - dists = geom.subset([ix], None).cost_matrix[0] - return KPPState(rng=next_rng, centroids=centroids, centroid_dists=dists) - - def body_fn( - iteration: int, const: Tuple[pointcloud.PointCloud, jnp.ndarray], - state: KPPState, compute_error: bool - ) -> KPPState: - del compute_error - rng, next_rng = jax.random.split(state.rng, 2) - geom, ixs = const - - # no need to normalize when `replace=True` - probs = state.centroid_dists - ixs = jax.random.choice( - rng, ixs, shape=(n_local_trials,), p=probs, replace=True - ) - geom = geom.subset(ixs, None) - - candidate_dists = jnp.minimum(geom.cost_matrix, state.centroid_dists) - best_ix = jnp.argmin(candidate_dists.sum(1)) - - centroids = state.centroids.at[iteration + 1].set(geom.x[best_ix]) - centroid_dists = candidate_dists[best_ix] - - return KPPState( - rng=next_rng, centroids=centroids, centroid_dists=centroid_dists - ) - - if n_local_trials is None: - n_local_trials = 2 + int(math.log(k)) - assert n_local_trials > 0, n_local_trials - - state = init_fn(geom, rng) - constants = (geom, jnp.arange(geom.shape[0])) - state = fixed_point_loop.fixpoint_iter( - lambda *_, **__: True, - body_fn, - min_iterations=k - 1, - max_iterations=k - 1, - inner_iterations=1, - constants=constants, - state=state - ) - - return state.centroids - - -@functools.partial(jax.vmap, in_axes=[None, 0, 0, 0], out_axes=0) -def _reallocate_centroids( - const: KMeansConst, - ix: jnp.ndarray, - centroid: jnp.ndarray, - weight: jnp.ndarray, -) -> Tuple[jnp.ndarray, jnp.ndarray]: - is_empty = weight <= 0.0 - new_centroid = (1 - is_empty) * centroid + is_empty * const.x[ix] # (ndim,) - centroid_to_remove = is_empty * const.weighted_x[ix] # (ndim,) - weight_to_remove = is_empty * const.weights[ix] # (1,) - return new_centroid, jnp.concatenate([centroid_to_remove, weight_to_remove]) - - -def _update_assignment( - const: KMeansConst, - centroids: jnp.ndarray, -) -> Tuple[jnp.ndarray, jnp.ndarray]: - (x, _, *args), aux_data = const.geom.tree_flatten() - cost_matrix = type( - const.geom - ).tree_unflatten(aux_data, [x, centroids] + args).cost_matrix - - assignment = jnp.argmin(cost_matrix, axis=1) - dist_to_centers = cost_matrix[jnp.arange(len(assignment)), assignment] - return assignment, dist_to_centers - - -def _update_centroids( - const: KMeansConst, k: int, assignment: jnp.ndarray, - dist_to_centers: jnp.ndarray -) -> jnp.ndarray: - # TODO(michalk8): - # cannot put `k` into `const`, see https://github.com/ott-jax/ott/issues/129 - x_weights = jax.ops.segment_sum(const.x_weights, assignment, num_segments=k) - centroids, ws = x_weights[:, :-1], x_weights[:, -1:] - - far_ixs = jnp.argsort(dist_to_centers)[-k:][::-1] - centroids, to_remove = _reallocate_centroids(const, far_ixs, centroids, ws) - to_remove = jax.ops.segment_sum( - to_remove, assignment[far_ixs], num_segments=k - ) - centroids -= to_remove[:, :-1] - ws -= to_remove[:, -1:] - - return centroids * jnp.where(ws > 0.0, 1.0 / ws, 1.0) - - -@functools.partial(jax.vmap, in_axes=[0] + [None] * 9) -def _k_means( - rng: jax.Array, - geom: pointcloud.PointCloud, - k: int, - weights: Optional[jnp.ndarray] = None, - init: Init_t = "k-means++", - n_local_trials: Optional[int] = None, - tol: float = 1e-4, - min_iterations: int = 0, - max_iterations: int = 300, - store_inner_errors: bool = False, -) -> KMeansOutput: - - def init_fn(init: Init_t) -> KMeansState: - if init == "k-means++": - init = functools.partial( - _k_means_plus_plus, n_local_trials=n_local_trials - ) - elif init == "random": - init = _random_init - if not callable(init): - raise TypeError( - f"Expected `init` to be 'k-means++', 'random' " - f"or a callable, found `{init_fn!r}`." - ) - - centroids = init(geom, k, rng) - if centroids.shape != (k, geom.cost_rank): - raise ValueError( - f"Expected initial centroids to have shape " - f"`{k, geom.cost_rank}`, found `{centroids.shape}`." - ) - n = geom.shape[0] - # TODO(michalk8): find a better solution for the below error - # not using floats for the assignment when `fixpoint_iter_backprop` is used: - - # .../jax/_src/dtypes.py:370: - # .0 = - # > CUB = set.intersection(*(UB[n] for n in N)) - # E jax._src.traceback_util.UnfilteredStackTrace: - # KeyError: dtype([('float0', 'V')]) - # E The stack trace below excludes JAX-internal frames. - # E The preceding is the original exception that occurred, unmodified. - prev_assignment = jnp.full((n,), -2.0) - assignment = jnp.full((n,), -1.0) - errors = jnp.full((max_iterations,), -1.0) - - return KMeansState( - centroids=centroids, - prev_assignment=prev_assignment, - assignment=assignment, - center_shift=jnp.inf, - errors=errors, - ) - - def cond_fn(iteration: int, const: KMeansConst, state: KMeansState) -> bool: - del iteration, const - assignment_not_same = jnp.any(state.prev_assignment != state.assignment) - tol_not_satisfied = state.center_shift > tol - return jnp.logical_and(assignment_not_same, tol_not_satisfied) - - def body_fn( - iteration: int, const: KMeansConst, state: KMeansState, - compute_error: bool - ) -> KMeansState: - del compute_error - - assignment, dist_to_centers = _update_assignment(const, state.centroids) - centroids = _update_centroids(const, k, assignment, dist_to_centers) - err = jnp.sum(const.weights[:, 0] * dist_to_centers) - center_shift = jnp.linalg.norm(state.centroids - centroids, ord="fro") ** 2 - - return KMeansState( - centroids=centroids, - prev_assignment=state.assignment, - assignment=assignment.astype(float), - center_shift=center_shift, - errors=state.errors.at[iteration].set(err) - ) - - def finalize_fn(const: KMeansConst, state: KMeansState) -> KMeansState: - last_iter = jnp.sum(state.errors != -1) - 1 - - assignment, dist_to_centers = _update_assignment(const, state.centroids) - err = jnp.sum(const.weights[:, 0] * dist_to_centers) - - return state._replace( - assignment=assignment.astype(float), - errors=state.errors.at[last_iter].set(err) - ) - - force_scan = min_iterations == max_iterations - fixpoint_fn = ( # prefer auto-diff if possible - fixed_point_loop.fixpoint_iter if force_scan else - fixed_point_loop.fixpoint_iter_backprop - ) - x_weights = jnp.hstack([weights[:, None] * geom.x, weights[:, None]]) - const = KMeansConst(geom, x_weights) - - state = fixpoint_fn( - cond_fn, - body_fn, - min_iterations=min_iterations, - max_iterations=max_iterations, - inner_iterations=1, - constants=const, - state=init_fn(init) - ) - state = jax.lax.cond( - jnp.all(state.prev_assignment == state.assignment), (lambda _, s: s), - finalize_fn, const, state - ) - - return KMeansOutput._from_state( - state, tol=tol, store_inner_errors=store_inner_errors - ) - - -def k_means( - geom: Union[jnp.ndarray, pointcloud.PointCloud], - k: int, - weights: Optional[jnp.ndarray] = None, - init: Init_t = "k-means++", - n_init: int = 10, - n_local_trials: Optional[int] = None, - tol: float = 1e-4, - min_iterations: int = 0, - max_iterations: int = 300, - store_inner_errors: bool = False, - rng: Optional[jax.Array] = None, -) -> KMeansOutput: - r"""K-means clustering using Lloyd's algorithm :cite:`lloyd:82`. - - Args: - geom: Point cloud of shape ``[n, ndim]`` to cluster. If passed as an array, - :class:`~ott.geometry.costs.SqEuclidean` cost is assumed. - k: The number of clusters. - weights: The weights of input points. These weights are considered when - computing the centroids and inertia. If ``None``, use uniform weights. - init: Initialization method. Can be one of the following: - - - **'k-means++'** - select initial centroids that are - :math:`\mathcal{O}(\log k)`-optimal :cite:`arthur:07`. - - **'random'** - randomly select ``k`` points from the ``geom``. - - :func:`callable` - a function which takes the point cloud, the number of - clusters and a random key and returns the centroids as an array of shape - ``[k, ndim]``. - - n_init: Number of times k-means will run with different initial seeds. - n_local_trials: Number of local trials when ``init = 'k-means++'``. - If ``None``, :math:`2 + \lfloor log(k) \rfloor` is used. - tol: Relative tolerance with respect to the Frobenius norm of the centroids' - shift between two consecutive iterations. - min_iterations: Minimum number of iterations. - max_iterations: Maximum number of iterations. - store_inner_errors: Whether to store the errors (inertia) at each iteration. - rng: Random key for seeding the initializations. - - Returns: - The k-means clustering. - """ - assert geom.shape[ - 0] >= k, f"Cannot cluster `{geom.shape[0]}` points into `{k}` clusters." - if isinstance(geom, jnp.ndarray): - geom = pointcloud.PointCloud(geom) - if isinstance(geom.cost_fn, costs.Cosine): - geom = geom._cosine_to_sqeucl() - assert geom.is_squared_euclidean - rng = utils.default_prng_key(rng) - - if geom.is_online: - # to allow materializing the cost matrix - children, aux_data = geom.tree_flatten() - aux_data["batch_size"] = None - geom = type(geom).tree_unflatten(aux_data, children) - - if weights is None: - weights = jnp.ones(geom.shape[0]) - assert weights.shape == (geom.shape[0],) - - rngs = jax.random.split(rng, n_init) - out = _k_means( - rngs, geom, k, weights, init, n_local_trials, tol, min_iterations, - max_iterations, store_inner_errors - ) - best_ix = jnp.argmin(out.error) - return jax.tree_util.tree_map(lambda arr: arr[best_ix], out) diff --git a/ott/build/lib/ott/tools/plot.py b/ott/build/lib/ott/tools/plot.py deleted file mode 100644 index bd1f42e..0000000 --- a/ott/build/lib/ott/tools/plot.py +++ /dev/null @@ -1,247 +0,0 @@ -# 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. -from typing import List, Optional, Sequence, Tuple, Union - -import jax.numpy as jnp -import numpy as np -import scipy - -from ott.geometry import pointcloud -from ott.solvers.linear import sinkhorn, sinkhorn_lr -from ott.solvers.quadratic import gromov_wasserstein - -try: - import matplotlib.pyplot as plt - from matplotlib import animation -except ImportError: - plt = animation = None - -# TODO(michalk8): make sure all outputs conform to a unified transport interface -Transport = Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput, - gromov_wasserstein.GWOutput] - - -def bidimensional(x: jnp.ndarray, - y: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]: - """Apply PCA to reduce to bi-dimensional data.""" - if x.shape[1] < 3: - return x, y - - u, s, _ = scipy.sparse.linalg.svds( - np.array(jnp.concatenate([x, y], axis=0)), k=2 - ) - proj = u * s - k = x.shape[0] - return proj[:k], proj[k:] - - -class Plot: - """Plot an optimal transport map between two point clouds. - - This object can either plot or update a plot, to create animations as a - :class:`~matplotlib.animation.FuncAnimation`, which can in turned be saved to - disk at will. There are two design principles here: - - #. we do not rely on saving to/loading from disk to create animations - #. we try as much as possible to disentangle the transport problem from - its visualization. - - We use 2D scatter plots by default, relying on PCA visualization for d>3 data. - This step requires a conversion to a numpy array, in order to compute leading - singular values. This tool is therefore not designed having performance in - mind. - - Args: - fig: Specify figure object. Created by default - ax: Specify axes objects. Created by default - threshold: value below which links in transportation matrix won't be - plotted. This value should be negative when using animations. - scale: scale used for marker plots. - show_lines: whether to show OT lines, as described in ``ot.matrix`` argument - cmap: color map used to plot line colors. - scale_alpha_by_coupling: use or not the coupling's value as proxy for alpha - alpha: default alpha value for lines. - title: title of the plot. - """ - - def __init__( - self, - fig: Optional["plt.Figure"] = None, - ax: Optional["plt.Axes"] = None, - threshold: float = -1.0, - scale: int = 200, - show_lines: bool = True, - cmap: str = "cool", - scale_alpha_by_coupling: bool = False, - alpha: float = 0.7, - title: Optional[str] = None - ): - if plt is None: - raise RuntimeError("Please install `matplotlib` first.") - - if ax is None and fig is None: - fig, ax = plt.subplots() - elif fig is None: - fig = plt.gcf() - elif ax is None: - ax = plt.gca() - self.fig = fig - self.ax = ax - self._show_lines = show_lines - self._lines = [] - self._points_x = None - self._points_y = None - self._threshold = threshold - self._scale = scale - self._cmap = cmap - self._scale_alpha_by_coupling = scale_alpha_by_coupling - self._alpha = alpha - self._title = title - - def _scatter(self, ot: Transport): - """Compute the position and scales of the points on a 2D plot.""" - if not isinstance(ot.geom, pointcloud.PointCloud): - raise ValueError("So far we only plot PointCloud geometry.") - - x, y = ot.geom.x, ot.geom.y - a, b = ot.a, ot.b - x, y = bidimensional(x, y) - scales_x = a * self._scale * a.shape[0] - scales_y = b * self._scale * b.shape[0] - return x, y, scales_x, scales_y - - def _mapping(self, x: jnp.ndarray, y: jnp.ndarray, matrix: jnp.ndarray): - """Compute the lines representing the mapping between the 2 point clouds.""" - # Only plot the lines with a cost above the threshold. - u, v = jnp.where(matrix > self._threshold) - c = matrix[jnp.where(matrix > self._threshold)] - xy = jnp.concatenate([x[u], y[v]], axis=-1) - - # Check if we want to adjust transparency. - scale_alpha_by_coupling = self._scale_alpha_by_coupling - - # We can only adjust transparency if max(c) != min(c). - if scale_alpha_by_coupling: - min_matrix, max_matrix = jnp.min(c), jnp.max(c) - scale_alpha_by_coupling = max_matrix != min_matrix - - result = [] - for i in range(xy.shape[0]): - strength = jnp.max(jnp.array(matrix.shape)) * c[i] - if scale_alpha_by_coupling: - normalized_strength = (c[i] - min_matrix) / (max_matrix - min_matrix) - alpha = self._alpha * float(normalized_strength) - else: - alpha = self._alpha - - # Matplotlib's transparency is sensitive to numerical errors. - alpha = np.clip(alpha, 0.0, 1.0) - - start, end = xy[i, [0, 2]], xy[i, [1, 3]] - result.append((start, end, strength, alpha)) - - return result - - def __call__(self, ot: Transport) -> List["plt.Artist"]: - """Plot couplings in 2-D, using PCA if data is higher dimensional.""" - x, y, sx, sy = self._scatter(ot) - self._points_x = self.ax.scatter( - *x.T, s=sx, edgecolors="k", marker="o", label="x" - ) - self._points_y = self.ax.scatter( - *y.T, s=sy, edgecolors="k", marker="X", label="y" - ) - self.ax.legend(fontsize=15) - if not self._show_lines: - return [] - - lines = self._mapping(x, y, ot.matrix) - cmap = plt.get_cmap(self._cmap) - self._lines = [] - for start, end, strength, alpha in lines: - line, = self.ax.plot( - start, - end, - linewidth=0.5 + 4 * strength, - color=cmap(strength), - zorder=0, - alpha=alpha - ) - self._lines.append(line) - if self._title is not None: - self.ax.set_title(self._title) - return [self._points_x, self._points_y] + self._lines - - def update(self, - ot: Transport, - title: Optional[str] = None) -> List["plt.Artist"]: - """Update a plot with a transport instance.""" - x, y, _, _ = self._scatter(ot) - self._points_x.set_offsets(x) - self._points_y.set_offsets(y) - if not self._show_lines: - return [] - - new_lines = self._mapping(x, y, ot.matrix) - cmap = plt.get_cmap(self._cmap) - for line, new_line in zip(self._lines, new_lines): - start, end, strength, alpha = new_line - - line.set_data(start, end) - line.set_linewidth(0.5 + 4 * strength) - line.set_color(cmap(strength)) - line.set_alpha(alpha) - - # Maybe add new lines to the plot. - num_lines = len(self._lines) - num_to_plot = len(new_lines) if self._show_lines else 0 - for i in range(num_lines, num_to_plot): - start, end, strength, alpha = new_lines[i] - - line, = self.ax.plot( - start, - end, - linewidth=0.5 + 4 * strength, - color=cmap(strength), - zorder=0, - alpha=alpha - ) - self._lines.append(line) - - self._lines = self._lines[:num_to_plot] # Maybe remove some - if title is not None: - self.ax.set_title(title) - return [self._points_x, self._points_y] + self._lines - - def animate( - self, - transports: Sequence[Transport], - titles: Optional[Sequence[str]] = None, - frame_rate: float = 10.0 - ) -> "animation.FuncAnimation": - """Make an animation from several transports.""" - _ = self(transports[0]) - if titles is None: - titles = [None for _ in np.arange(0, len(transports))] - assert len(titles) == len(transports), ( - f"titles/transports lengths differ `{len(titles)}`/`{len(transports)}`." - ) - return animation.FuncAnimation( - self.fig, - lambda i: self.update(transports[i], titles[i]), - np.arange(0, len(transports)), - init_func=lambda: self.update(transports[0], titles[0]), - interval=1000 / frame_rate, - blit=True - ) diff --git a/ott/build/lib/ott/tools/segment_sinkhorn.py b/ott/build/lib/ott/tools/segment_sinkhorn.py deleted file mode 100644 index 3e6b29d..0000000 --- a/ott/build/lib/ott/tools/segment_sinkhorn.py +++ /dev/null @@ -1,143 +0,0 @@ -# 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. -from types import MappingProxyType -from typing import Any, Mapping, Optional, Tuple - -import jax.numpy as jnp - -from ott.geometry import costs, pointcloud, segment -from ott.problems.linear import linear_problem -from ott.solvers.linear import sinkhorn - - -def segment_sinkhorn( - x: jnp.ndarray, - y: jnp.ndarray, - num_segments: Optional[int] = None, - max_measure_size: Optional[int] = None, - cost_fn: Optional[costs.CostFn] = None, - segment_ids_x: Optional[jnp.ndarray] = None, - segment_ids_y: Optional[jnp.ndarray] = None, - indices_are_sorted: bool = False, - num_per_segment_x: Optional[Tuple[int, ...]] = None, - num_per_segment_y: Optional[Tuple[int, ...]] = None, - weights_x: Optional[jnp.ndarray] = None, - weights_y: Optional[jnp.ndarray] = None, - sinkhorn_kwargs: Mapping[str, Any] = MappingProxyType({}), - **kwargs: Any -) -> jnp.ndarray: - """Compute regularized OT cost between subsets of vectors in `x` and `y`. - - Helper function designed to compute Sinkhorn regularized OT cost between - several point clouds of varying size, in parallel, using padding. - In practice, The inputs `x` and `y` (and their weight vectors `weights_x` and - `weights_y`) are assumed to be large weighted point clouds, that describe - points taken from multiple measures. To extract several subsets of points, we - provide two interfaces. The first interface assumes that a vector of id's is - passed, describing for each point of `x` (resp. `y`) to which measure the - point belongs to. The second interface assumes that `x` and `y` were simply - formed by concatenating several measures contiguously, and that only indices - that segment these groups are needed to recover them. - - For both interfaces, both `x` and `y` should contain the same total number of - segments. Each segment will be padded as necessary, all segments rearranged as - a tensor, and :func:`jax.vmap` used to evaluate Sinkhorn divergences in - parallel. - - Args: - x: Array of input points, of shape `[num_x, feature]`. - Multiple segments are held in this single array. - y: Array of target points, of shape `[num_y, feature]`. - num_segments: Number of segments contained in `x` and `y`. - Providing this is required for JIT compilation to work, - see also :func:`~ott.geometry.segment.segment_point_cloud`. - max_measure_size: Total size of measures after padding. Should ideally be - set to an upper bound on points clouds processed with the segment - interface. Providing this is required for JIT compilation to work. - cost_fn: Cost function, defaults to - :class:`~ott.geometry.costs.SqEuclidean`. - segment_ids_x: **1st interface** The segment ID for which each row of `x` - belongs. This is a similar interface to `jax.ops.segment_sum`. - segment_ids_y: **1st interface** The segment ID for which each row of `y` - belongs. - indices_are_sorted: **1st interface** Whether `segment_ids_x` and - `segment_ids_y` are sorted. - num_per_segment_x: **2nd interface** Number of points in each segment in - `x`. For example, [100, 20, 30] would imply that `x` is segmented into - three arrays of length `[100]`, `[20]`, and `[30]` respectively. - num_per_segment_y: **2nd interface** Number of points in each segment in - `y`. - weights_x: Weights of each input points, arranged in the same segmented - order as `x`. - weights_y: Weights of each input points, arranged in the same segmented - order as `y`. - sinkhorn_kwargs: Optionally a dict containing the keywords arguments for - calls for the :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver, - called three times to evaluate for each segment the Sinkhorn regularized - OT cost between `x`/`y`, `x`/`x`, and `y`/`y` (except when `static_b` is - `True`, in which case `y`/`y` is not evaluated). - kwargs: keywords arguments passed to form - :class:`~ott.geometry.pointcloud.PointCloud` geometry objects from the - subsets of points and masses selected in `x` and `y`, possibly a - :class:`~ott.geometry.costs.CostFn` or an entropy regularizer. - - Returns: - An array of Sinkhorn regularized OT costs for each segment. - """ - # instantiate padding vector - dim = x.shape[1] - if cost_fn is None: - # default padder - padding_vector = costs.CostFn._padder(dim=dim) - else: - padding_vector = cost_fn._padder(dim=dim) - - def eval_fn( - padded_x: jnp.ndarray, - padded_y: jnp.ndarray, - padded_weight_x: jnp.ndarray, - padded_weight_y: jnp.ndarray, - ) -> float: - mask_x = padded_weight_x > 0.0 - mask_y = padded_weight_y > 0.0 - - geom = pointcloud.PointCloud( - padded_x, - padded_y, - cost_fn=cost_fn, - src_mask=mask_x, - tgt_mask=mask_y, - **kwargs, - ) - prob = linear_problem.LinearProblem( - geom, a=padded_weight_x, b=padded_weight_y - ) - solver = sinkhorn.Sinkhorn(**sinkhorn_kwargs) - return solver(prob).reg_ot_cost - - return segment._segment_interface( - x, - y, - eval_fn, - num_segments=num_segments, - max_measure_size=max_measure_size, - segment_ids_x=segment_ids_x, - segment_ids_y=segment_ids_y, - indices_are_sorted=indices_are_sorted, - num_per_segment_x=num_per_segment_x, - num_per_segment_y=num_per_segment_y, - weights_x=weights_x, - weights_y=weights_y, - padding_vector=padding_vector - ) diff --git a/ott/build/lib/ott/tools/sinkhorn_divergence.py b/ott/build/lib/ott/tools/sinkhorn_divergence.py deleted file mode 100644 index 572075e..0000000 --- a/ott/build/lib/ott/tools/sinkhorn_divergence.py +++ /dev/null @@ -1,354 +0,0 @@ -# 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. -from types import MappingProxyType -from typing import Any, Mapping, Optional, Tuple, Type - -import jax.numpy as jnp - -from ott import utils -from ott.geometry import costs, geometry, pointcloud, segment -from ott.problems.linear import linear_problem, potentials -from ott.solvers import linear -from ott.solvers.linear import acceleration, sinkhorn - -__all__ = [ - "sinkhorn_divergence", "segment_sinkhorn_divergence", - "SinkhornDivergenceOutput" -] - -Potentials_t = Tuple[jnp.ndarray, jnp.ndarray] - - -@utils.register_pytree_node -class SinkhornDivergenceOutput: # noqa: D101 - divergence: float - potentials: Tuple[Potentials_t, Potentials_t, Potentials_t] - geoms: Tuple[geometry.Geometry, geometry.Geometry, geometry.Geometry] - errors: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray], - Optional[jnp.ndarray]] - converged: Tuple[bool, bool, bool] - a: jnp.ndarray - b: jnp.ndarray - n_iters: Tuple[int, int, int] - - def to_dual_potentials(self) -> "potentials.EntropicPotentials": - """Return dual estimators :cite:`pooladian:22`, eq. 8.""" - geom_xy, *_ = self.geoms - prob_xy = linear_problem.LinearProblem(geom_xy, a=self.a, b=self.b) - (f_xy, g_xy), (f_x, _), (_, g_y) = self.potentials - return potentials.EntropicPotentials( - f_xy, g_xy, prob_xy, f_xx=f_x, g_yy=g_y - ) - - def tree_flatten(self): # noqa: D102 - return [ - self.divergence, - self.potentials, - self.geoms, - self.a, - self.b, - ], { - "n_iters": self.n_iters, - "converged": self.converged, - "errors": self.errors - } - - @classmethod - def tree_unflatten(cls, aux_data, children): # noqa: D102 - div, pots, geoms, a, b = children - return cls(div, pots, geoms, a=a, b=b, **aux_data) - - -def sinkhorn_divergence( - geom: Type[geometry.Geometry], - *args: Any, - a: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, - sinkhorn_kwargs: Mapping[str, Any] = MappingProxyType({}), - static_b: bool = False, - share_epsilon: bool = True, - symmetric_sinkhorn: bool = True, - **kwargs: Any, -) -> SinkhornDivergenceOutput: - """Compute Sinkhorn divergence defined by a geometry, weights, parameters. - - Args: - geom: Type of the geometry. - args: Positional arguments to - :meth:`~ott.geometry.geometry.Geometry.prepare_divergences` that are - specific to each geometry. - a: the weight of each input point. The sum of all elements of `a` must - match that of `b` to converge. - b: the weight of each target point. The sum of all elements of `b` must - match that of `a` to converge. - sinkhorn_kwargs: keywords arguments for - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` that is called twice - if ``static_b = True`` else 3 times. - static_b: if True, divergence of measure `b` against itself is **not** - computed. - share_epsilon: if True, enforces that the same epsilon regularizer is shared - for all 2 or 3 terms of the Sinkhorn divergence. In that case, the epsilon - will be by default that used when comparing x to y (contained in the first - geometry). This flag is set to True by default, because in the default - setting, the epsilon regularization is a function of the mean of the cost - matrix. - symmetric_sinkhorn: Use Sinkhorn updates in Eq. 25 of :cite:`feydy:19` for - symmetric terms comparing x/x and y/y. - kwargs: keywords arguments to the generic class. This is specific to each - geometry. - - Returns: - Sinkhorn divergence value, three pairs of potentials, three costs. - """ - geoms = geom.prepare_divergences(*args, static_b=static_b, **kwargs) - geom_xy, geom_x, geom_y, *_ = geoms + (None,) * 3 - num_a, num_b = geom_xy.shape - - if share_epsilon: - if isinstance(geom_x, geometry.Geometry): - geom_x = geom_x.copy_epsilon(geom_xy) - if isinstance(geom_y, geometry.Geometry): - geom_y = geom_y.copy_epsilon(geom_xy) - - a = jnp.ones(num_a) / num_a if a is None else a - b = jnp.ones(num_b) / num_b if b is None else b - return _sinkhorn_divergence( - geom_xy, - geom_x, - geom_y, - a=a, - b=b, - symmetric_sinkhorn=symmetric_sinkhorn, - **sinkhorn_kwargs - ) - - -def _sinkhorn_divergence( - geometry_xy: geometry.Geometry, - geometry_xx: geometry.Geometry, - geometry_yy: Optional[geometry.Geometry], - a: jnp.ndarray, - b: jnp.ndarray, - symmetric_sinkhorn: bool, - **kwargs: Any, -) -> SinkhornDivergenceOutput: - """Compute the (unbalanced) Sinkhorn divergence for the wrapper function. - - This definition includes a correction depending on the total masses of each - measure, as defined in :sejourne:19:, eq. 15. - - Args: - geometry_xy: a Cost object able to apply kernels with a certain epsilon, - between the views X and Y. - geometry_xx: a Cost object able to apply kernels with a certain epsilon, - between elements of the view X. - geometry_yy: a Cost object able to apply kernels with a certain epsilon, - between elements of the view Y. - a: jnp.ndarray[n]: the weight of each input point. The sum of - all elements of ``b`` must match that of ``a`` to converge. - b: jnp.ndarray[m]: the weight of each target point. The sum of - all elements of ``b`` must match that of ``a`` to converge. - symmetric_sinkhorn: Use Sinkhorn updates in Eq. 25 of :cite:`feydy:19` for - symmetric terms comparing x/x and y/y. - kwargs: Keyword arguments to :func:`~ott.solvers.linear.sinkhorn.Sinkhorn`. - - Returns: - SinkhornDivergenceOutput named tuple. - """ - # When computing a Sinkhorn divergence, the (x,y) terms and (x,x) / (y,y) - # terms are computed independently. The user might want to pass some - # sinkhorn_kwargs to parameterize Sinkhorn's behavior, but those should - # only apply to the (x,y) part. For the (x,x) / (y,y) part we fall back - # on a simpler choice (parallel_dual_updates + momentum 0.5) that is known - # to work well in such settings. In the future we might want to give some - # freedom on setting parameters for the (x,x)/(y,y) part. - # Since symmetric terms are computed assuming a = b, the linear systems - # arising in implicit differentiation (if used) of the potentials computed for - # the symmetric parts should be marked as symmetric. - kwargs_symmetric = kwargs.copy() - if symmetric_sinkhorn: - kwargs_symmetric.update( - parallel_dual_updates=True, - momentum=acceleration.Momentum(start=0, value=0.5), - anderson=None, - ) - implicit_diff = kwargs.get("implicit_diff", None) - if implicit_diff is not None: - kwargs_symmetric["implicit_diff"] = implicit_diff.replace(symmetric=True) - - out_xy = linear.solve(geometry_xy, a, b, **kwargs) - out_xx = linear.solve(geometry_xx, a, a, **kwargs_symmetric) - if geometry_yy is None: - # Create dummy output, corresponds to scenario where static_b is True. - # This choice ensures that `converged`` of this dummy output is True. - out_yy = sinkhorn.SinkhornOutput( - (None, None), - errors=jnp.array([-jnp.inf]), - reg_ot_cost=0.0, - threshold=0.0, - inner_iterations=0, - ) - else: - out_yy = linear.solve(geometry_yy, b, b, **kwargs_symmetric) - - div = ( - out_xy.reg_ot_cost - 0.5 * (out_xx.reg_ot_cost + out_yy.reg_ot_cost) + - 0.5 * geometry_xy.epsilon * (jnp.sum(a) - jnp.sum(b)) ** 2 - ) - return SinkhornDivergenceOutput( - divergence=div, - potentials=((out_xy.f, out_xy.g), (out_xx.f, out_xx.g), - (out_yy.f, out_yy.g)), - geoms=(geometry_xy, geometry_xx, geometry_yy), - errors=(out_xy.errors, out_xx.errors, out_yy.errors), - converged=(out_xy.converged, out_xx.converged, out_yy.converged), - a=a, - b=b, - n_iters=(out_xy.n_iters, out_xx.n_iters, out_yy.n_iters), - ) - - -def segment_sinkhorn_divergence( - x: jnp.ndarray, - y: jnp.ndarray, - num_segments: Optional[int] = None, - max_measure_size: Optional[int] = None, - cost_fn: Optional[costs.CostFn] = None, - segment_ids_x: Optional[jnp.ndarray] = None, - segment_ids_y: Optional[jnp.ndarray] = None, - indices_are_sorted: bool = False, - num_per_segment_x: Optional[Tuple[int, ...]] = None, - num_per_segment_y: Optional[Tuple[int, ...]] = None, - weights_x: Optional[jnp.ndarray] = None, - weights_y: Optional[jnp.ndarray] = None, - sinkhorn_kwargs: Mapping[str, Any] = MappingProxyType({}), - static_b: bool = False, - share_epsilon: bool = True, - symmetric_sinkhorn: bool = False, - **kwargs: Any -) -> jnp.ndarray: - """Compute Sinkhorn divergence between subsets of vectors in `x` and `y`. - - Helper function designed to compute Sinkhorn divergences between several point - clouds of varying size, in parallel, using padding for efficiency. - In practice, The inputs `x` and `y` (and their weight vectors `weights_x` and - `weights_y`) are assumed to be large weighted point clouds, that describe - points taken from multiple measures. To extract several subsets of points, we - provide two interfaces. The first interface assumes that a vector of id's is - passed, describing for each point of `x` (resp. `y`) to which measure the - point belongs to. The second interface assumes that `x` and `y` were simply - formed by concatenating several measures contiguously, and that only indices - that segment these groups are needed to recover them. - - For both interfaces, both `x` and `y` should contain the same total number of - segments. Each segment will be padded as necessary, all segments rearranged as - a tensor, and `vmap` used to evaluate Sinkhorn divergences in parallel. - - Args: - x: Array of input points, of shape `[num_x, feature]`. - Multiple segments are held in this single array. - y: Array of target points, of shape `[num_y, feature]`. - num_segments: Number of segments contained in `x` and `y`. - Providing this is required for JIT compilation to work, - see also :func:`~ott.geometry.segment.segment_point_cloud`. - max_measure_size: Total size of measures after padding. Should ideally be - set to an upper bound on points clouds processed with the segment - interface. Should also be smaller than total length of `x` or `y`. - Providing this is required for JIT compilation to work. - cost_fn: Cost function, - defaults to :class:`~ott.geometry.costs.SqEuclidean`. - segment_ids_x: **1st interface** The segment ID for which each row of `x` - belongs. This is a similar interface to :func:`jax.ops.segment_sum`. - segment_ids_y: **1st interface** The segment ID for which each row of `y` - belongs. - indices_are_sorted: **1st interface** Whether `segment_ids_x` and - `segment_ids_y` are sorted. - num_per_segment_x: **2nd interface** Number of points in each segment in - `x`. For example, [100, 20, 30] would imply that `x` is segmented into - three arrays of length `[100]`, `[20]`, and `[30]` respectively. - num_per_segment_y: **2nd interface** Number of points in each segment in - `y`. - weights_x: Weights of each input points, arranged in the same segmented - order as `x`. - weights_y: Weights of each input points, arranged in the same segmented - order as `y`. - sinkhorn_kwargs: Optionally a dict containing the keywords arguments for - calls to the `sinkhorn` function, called three times to evaluate for each - segment the Sinkhorn regularized OT cost between `x`/`y`, `x`/`x`, and - `y`/`y` (except when `static_b` is `True`, in which case `y`/`y` is not - evaluated) - static_b: if True, divergence of measure b against itself is NOT computed - share_epsilon: if True, enforces that the same epsilon regularizer is shared - for all 2 or 3 terms of the Sinkhorn divergence. In that case, the epsilon - will be by default that used when comparing x to y (contained in the first - geometry). This flag is set to True by default, because in the default - setting, the epsilon regularization is a function of the mean of the cost - matrix. - symmetric_sinkhorn: Use Sinkhorn updates in Eq. 25 of :cite:`feydy:19` for - symmetric terms comparing x/x and y/y. - kwargs: keywords arguments passed to form - :class:`~ott.geometry.pointcloud.PointCloud` geometry objects from the - subsets of points and masses selected in `x` and `y`, this could be for - instance entropy regularization float, scheduler or normalization. - - Returns: - An array of Sinkhorn divergences for each segment. - """ - # instantiate padding vector - dim = x.shape[1] - if cost_fn is None: - # default padder - padding_vector = costs.CostFn._padder(dim=dim) - else: - padding_vector = cost_fn._padder(dim=dim) - - def eval_fn( - padded_x: jnp.ndarray, - padded_y: jnp.ndarray, - padded_weight_x: jnp.ndarray, - padded_weight_y: jnp.ndarray, - ) -> float: - mask_x = padded_weight_x > 0.0 - mask_y = padded_weight_y > 0.0 - return sinkhorn_divergence( - pointcloud.PointCloud, - padded_x, - padded_y, - a=padded_weight_x, - b=padded_weight_y, - sinkhorn_kwargs=sinkhorn_kwargs, - static_b=static_b, - share_epsilon=share_epsilon, - symmetric_sinkhorn=symmetric_sinkhorn, - cost_fn=cost_fn, - src_mask=mask_x, - tgt_mask=mask_y, - **kwargs - ).divergence - - return segment._segment_interface( - x, - y, - eval_fn, - num_segments=num_segments, - max_measure_size=max_measure_size, - segment_ids_x=segment_ids_x, - segment_ids_y=segment_ids_y, - indices_are_sorted=indices_are_sorted, - num_per_segment_x=num_per_segment_x, - num_per_segment_y=num_per_segment_y, - weights_x=weights_x, - weights_y=weights_y, - padding_vector=padding_vector - ) diff --git a/ott/build/lib/ott/tools/soft_sort.py b/ott/build/lib/ott/tools/soft_sort.py deleted file mode 100644 index faf52f2..0000000 --- a/ott/build/lib/ott/tools/soft_sort.py +++ /dev/null @@ -1,706 +0,0 @@ -# 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 functools -from typing import Any, Callable, Optional, Tuple, Union - -import jax -import jax.numpy as jnp -import jax.tree_util as jtu -import numpy as np - -from ott import utils -from ott.geometry import costs, pointcloud -from ott.problems.linear import linear_problem -from ott.solvers import linear -from ott.solvers.linear import sinkhorn - -__all__ = [ - "sort", "ranks", "sort_with", "quantile", "quantile_normalization", - "quantize", "topk_mask", "multivariate_cdf_quantile_maps" -] - -Func_t = Callable[[jnp.ndarray], jnp.ndarray] - - -def transport_for_sort( - inputs: jnp.ndarray, - weights: Optional[jnp.ndarray] = None, - target_weights: Optional[jnp.ndarray] = None, - squashing_fun: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, - epsilon: float = 1e-2, - **kwargs: Any, -) -> sinkhorn.SinkhornOutput: - r"""Solve reg. OT, from inputs to a weighted family of increasing values. - - Args: - inputs: Array[num_points]. Must be one dimensional. - weights: Array[num_points]. Weight vector `a` for input values. - target_weights: Array[num_targets]: Weight vector of the target - measure. It may be of different size than `weights`. - squashing_fun: function taking an array to squash all its entries in [0,1]. - sigmoid of whitened values by default. Can be set to be the identity by - passing ``squashing_fun = lambda x : x`` instead. - epsilon: the regularization parameter. - kwargs: keyword arguments for - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`. - - Returns: - A :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput` object. - """ - shape = inputs.shape - if len(shape) > 2 or (len(shape) == 2 and shape[1] != 1): - raise ValueError( - f"Shape ({shape}) not supported. The input should be one-dimensional." - ) - - x = jnp.expand_dims(jnp.squeeze(inputs), axis=1) - if squashing_fun is None: - squashing_fun = lambda z: jax.nn.sigmoid((z - jnp.mean(z)) / - (jnp.std(z) + 1e-10)) - x = squashing_fun(x) - - a = jnp.squeeze(weights) - b = jnp.squeeze(target_weights) - num_targets = inputs.shape[0] if b is None else b.shape[0] - y = jnp.linspace(0.0, 1.0, num_targets)[:, jnp.newaxis] - - geom = pointcloud.PointCloud(x, y, epsilon=epsilon) - prob = linear_problem.LinearProblem(geom, a=a, b=b) - - solver = sinkhorn.Sinkhorn(**kwargs) - - return solver(prob) - - -def apply_on_axis(op, inputs, axis, *args, **kwargs: Any) -> jnp.ndarray: - """Apply a differentiable operator on a given axis of the input. - - Args: - op: a differentiable operator (can be ranks, quantile, etc.) - inputs: Array of any shape. - axis: the axis (int) or tuple of ints on which to apply the operator. If - several axes are passed to the operator, those are merged as a single - dimension. - args: other positional arguments to the operator. - kwargs: other positional arguments to the operator. - - Returns: - An Array holding the output of the differentiable operator on the given - axis. - """ - op_inner = functools.partial(op, **kwargs) - axis = (axis,) if isinstance(axis, int) else axis - num_points = np.prod(np.array(inputs.shape)[(axis,)]) - permutation = np.arange(len(inputs.shape)) - axis = tuple(permutation[a] for a in axis) - permutation = tuple(sorted(set(permutation) - set(axis)) + sorted(axis)) - inputs = jnp.transpose(inputs, permutation) - - batch_fn = jax.vmap(op_inner, in_axes=(0,) + (None,) * len(args)) - result = batch_fn(jnp.reshape(inputs, (-1, num_points)), *args) - shrink = len(axis) - result = jnp.reshape(result, inputs.shape[:-shrink] + result.shape[-1:]) - - permutation = tuple(range(len(result.shape))) - rank = len(result.shape) - 1 - axis = min(axis) - permutation = permutation[:axis] + (rank,) + permutation[axis:-1] - return jnp.transpose(result, permutation) - - -def _sort( - inputs: jnp.ndarray, topk: int, num_targets: Optional[int], **kwargs: Any -) -> jnp.ndarray: - """Apply the soft sort operator on a one dimensional array.""" - num_points = inputs.shape[0] - a = jnp.ones((num_points,)) / num_points - if 0 < topk < num_points: - start_index = 1 - b = jnp.concatenate([ - jnp.array([(num_points - topk) / num_points]), - jnp.ones(topk) / num_points - ]) - else: - # Use the "sorting" initializer if default uniform weights of same size. - if num_targets is None or num_targets == num_points: - num_targets = num_points - # use sorting initializer in this case. - kwargs.setdefault("initializer", "sorting") - start_index = 0 - b = jnp.ones((num_targets,)) / num_targets - ot = transport_for_sort(inputs, a, b, **kwargs) - out = 1.0 / b * ot.apply(inputs, axis=0) - return out[start_index:] - - -def sort( - inputs: jnp.ndarray, - axis: int = -1, - topk: int = -1, - num_targets: Optional[int] = None, - **kwargs: Any, -) -> jnp.ndarray: - r"""Apply the soft sort operator on a given axis of the input. - - For instance: - - .. code-block:: python - - x = jax.random.uniform(rng, (100,)) - x_sorted = sort(x) - - will output sorted convex-combinations of values contained in ``x``, that are - differentiable approximations to the sorted vector of entries in ``x``. - These can be compared with the values produced by :func:`jax.numpy.sort`, - - .. code-block:: python - - x_sorted = jax.numpy.sort(x) - - Args: - inputs: Array of any shape. - axis: the axis on which to apply the soft-sorting operator. - topk: if set to a positive value, the returned vector will only contain - the top-k values. This also reduces the complexity of soft-sorting, since - the number of target points to which the slice of the ``inputs`` tensor - will be mapped to will be equal to ``topk + 1``. - num_targets: if ``topk`` is not specified, a vector of size``num_targets`` - is returned. This defines the number of (composite) sorted values computed - from the inputs (each value is a convex combination of values recorded in - the inputs, provided in increasing order). If neither ``topk`` nor - ``num_targets`` are specified, ``num_targets`` defaults to the size of the - slices of the input that are sorted, i.e. ``inputs.shape[axis]``, and the - number of composite sorted values is equal to the slice of the inputs that - are sorted. As a result, the output is of the same size as ``inputs``. - kwargs: keyword arguments passed on to lower level functions. Of interest - to the user are ``squashing_fun``, which will redistribute the values in - ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) - to solve the optimal transport problem; - :class:`cost_fn ` object of - :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground - 1D cost function to transport from ``inputs`` to the ``num_targets`` - target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` - are passed on to parameterize the - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. - - Returns: - An Array of the same shape as the input, except on ``axis``, where that size - will be equal to ``topk`` or ``num_targets``, with soft-sorted values on the - given axis. Same size as ``inputs`` if both these parameters are ``None``. - """ - return apply_on_axis(_sort, inputs, axis, topk, num_targets, **kwargs) - - -def _ranks( - inputs: jnp.ndarray, num_targets, target_weights, **kwargs: Any -) -> jnp.ndarray: - """Apply the soft ranks operator on a one dimensional array.""" - num_points = inputs.shape[0] - if target_weights is None: - num_targets = num_points if num_targets is None else num_targets - target_weights = jnp.ones((num_targets,)) / num_targets - else: - num_targets = target_weights.shape[0] - a = jnp.ones((num_points,)) / num_points - ot = transport_for_sort(inputs, a, target_weights, **kwargs) - out = 1.0 / a * ot.apply(jnp.arange(num_targets), axis=1) - out *= (num_points - 1.0) / (num_targets - 1.0) - return jnp.reshape(out, inputs.shape) - - -def ranks( - inputs: jnp.ndarray, - axis: int = -1, - num_targets: Optional[int] = None, - target_weights: Optional[jnp.ndarray] = None, - **kwargs: Any, -) -> jnp.ndarray: - r"""Apply the soft rank operator on input tensor. - - For instance: - - .. code-block:: python - - x = jax.random.uniform(rng, (100,)) - x_ranks = ranks(x) - - will output values that are differentiable approximations to the ranks of - entries in ``x``. These should be compared to the non-differentiable rank - vectors, namely the normalized inverse permutation produced by - :func:`jax.numpy.argsort`, which can be obtained as: - - .. code-block:: python - - x_ranks = jax.numpy.argsort(jax.numpy.argsort(x)) - - Args: - inputs: Array of any shape. - axis: the axis on which to apply the soft-sorting operator. - target_weights: This vector contains weights (summing to 1) that describe - amount of mass shipped to targets. - num_targets: If ``target_weights` is ``None``, ``num_targets`` is - considered to define the number of targets used to rank inputs. Each - rank in the output will be a convex combination of - ``{1, .., num_targets}``. The weight of each of these points - is assumed to be uniform. If neither ``num_targets`` nor - ``target_weights`` are specified, ``num_targets`` defaults to the size - of the slices of the input that are sorted, i.e. ``inputs.shape[axis]``. - kwargs: keyword arguments passed on to lower level functions. Of interest - to the user are ``squashing_fun``, which will redistribute the values in - ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) - to solve the optimal transport problem; - :class:`cost_fn ` object of - :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground - 1D cost function to transport from ``inputs`` to the ``num_targets`` - target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` - are passed on to parameterize the - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. - - Returns: - An Array of the same shape as the input with soft-rank values - normalized to be in :math:`[0, n-1]` where :math:`n` is - ``inputs.shape[axis]``. - """ - return apply_on_axis( - _ranks, inputs, axis, num_targets, target_weights, **kwargs - ) - - -def topk_mask( - inputs: jnp.ndarray, - axis: int = -1, - k: int = 1, - **kwargs: Any, -) -> jnp.ndarray: - r"""Soft :math:`\text{top-}k` selection mask. - - For instance: - - .. code-block:: python - - k = 5 - x = jax.random.uniform(rng, (100,)) - mask = topk_mask(x, k=k) - - will output a vector of shape ``x.shape``, with values in :math:`[0,1]`, that - are differentiable approximations to the binary mask selecting the top $k$ - entries in ``x``. These should be compared to the non-differentiable mask - obtained with :func:`jax.numpy.sort`, which can be obtained as: - - .. code-block:: python - - mask = x >= jax.numpy.sort(x).flip()[k-1] - - Args: - inputs: Array of any shape. - axis: the axis on which to apply the soft-sorting operator. - k: topk parameter. Should be smaller than ``inputs.shape[axis]``. - kwargs: keyword arguments passed on to lower level functions. Of interest - to the user are ``squashing_fun``, which will redistribute the values in - ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) - to solve the optimal transport problem; - :class:`cost_fn ` object of - :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground - 1D cost function to transport from ``inputs`` to the ``num_targets`` - target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` - are passed on to parameterize the - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. - - Returns: - The soft :math:`\text{top-}k` selection mask. - """ - num_points = inputs.shape[axis] - assert k < num_points, ( - f"`k` must be smaller than `inputs.shape[axis]`, yet {k} >= {num_points}." - ) - target_weights = jnp.array([1.0 - k / num_points, k / num_points]) - out = apply_on_axis( - _ranks, - inputs, - axis, - num_targets=None, - target_weights=target_weights, - **kwargs - ) - return out / (num_points - 1) - - -def quantile( - inputs: jnp.ndarray, - q: Optional[Union[float, jnp.ndarray]], - axis: Union[int, Tuple[int, ...]] = -1, - weight: Optional[Union[float, jnp.ndarray]] = None, - **kwargs: Any, -) -> jnp.ndarray: - r"""Apply the soft quantiles operator on the input tensor. - - For instance: - - .. code-block:: python - - x = jax.random.uniform(rng, (100,)) - x_quantiles = quantile(x, q=jnp.array([0.2, 0.8])) - - ``x_quantiles`` will hold an approximation to the 20 and 80 percentiles in - ``x``, computed as a convex combination (a weighted mean, with weights summing - to 1) of all values in ``x`` (and not, as for standard quantiles, the - values ``x_sorted[20]`` and ``x_sorted[80]`` if ``x_sorted=jnp.sort(x)``). - These values offer a trade-off between accuracy (closeness to the true - percentiles) and gradient (the Jacobian of ``x_quantiles`` w.r.t ``x`` will - impact all values listed in ``x``, not just those indexed at 20 and 80). - - The non-differentiable version is given by :func:`jax.numpy.quantile`, e.g. - - .. code-block:: python - - x_quantiles = jax.numpy.quantile(x, q=jnp.array([0.2, 0.8])) - - Args: - inputs: an Array of any shape. - q: values of the quantile level to be computed, e.g. [0.5] for median. - These values should all lie within the segment :math:`]0,1[`, excluding - boundary values :math:`0` and :math:`1`. - axis: the axis on which to apply the operator. - weight: the weight assigned to each quantile target value in the OT problem. - This weight should be small, typically of the order of ``1/n``, where ``n`` - is the size of ``x``. Note: Since the size of ``q`` times ``weight`` - must be strictly smaller than ``1``, in order to leave enough mass to set - other target values in the transport problem, the algorithm might ensure - this by setting, when needed, a lower value. - kwargs: keyword arguments passed on to lower level functions. Of interest - to the user are ``squashing_fun``, which will redistribute the values in - ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) - to solve the optimal transport problem; - :class:`cost_fn ` object of - :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground - 1D cost function to transport from ``inputs`` to the ``num_targets`` - target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` - are passed on to parameterize the - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. - - Returns: - An Array, which has the same shape as ``inputs``, except on the ``axis`` - that is passed, which has size ``q.shape[0]``, to collect soft-quantile - values. - """ - - def _quantile( - inputs: jnp.ndarray, q: float, weight: float, **kwargs - ) -> jnp.ndarray: - num_points = inputs.shape[0] - q = jnp.array([0.2, 0.5, 0.8]) if q is None else jnp.atleast_1d(q) - num_quantiles = q.shape[0] - a = jnp.ones((num_points,)) / num_points - idx = jnp.argsort(q) - q = q[idx] - - extended_q = jnp.concatenate([jnp.array([0.0]), q, jnp.array([1.0])]) - filler_weights = extended_q[1:] - extended_q[:-1] - safe_weight = 0.5 * jnp.concatenate([ - jnp.array([1.0 / num_quantiles]), filler_weights - ]) - if weight is None: - # Populate with other options. - safe_weight = jnp.concatenate([ - safe_weight, - jnp.array( - [0.02] - ), # reasonable mass per quantile for a small number of points - jnp.array( - [1.5 / num_points] - ), # this means each quantile would be ~ assigned 1.5 points. - ]) - else: - safe_weight = jnp.concatenate([safe_weight, jnp.atleast_1d(weight)]) - weight = jnp.min(safe_weight) - weights = jnp.ones(filler_weights.shape) * weight - - # Takes into account quantile_width in the definition of weights - shift = -jnp.ones(filler_weights.shape) - shift = shift + 0.5 * ( - jax.nn.one_hot(0, num_quantiles + 1) + - jax.nn.one_hot(num_quantiles, num_quantiles + 1) - ) - filler_weights = filler_weights + weights * shift - - # Adds one more value to have tensors of the same shape to interleave them. - quantile_weights = jnp.ones(num_quantiles + 1) * weights - - # Interleaves the filler_weights with the quantile weights. - weights = jnp.reshape( - jnp.stack([filler_weights, quantile_weights], axis=1), (-1,) - )[:-1] - - ot = transport_for_sort(inputs, a, weights, **kwargs) - out = 1.0 / weights * ot.apply(jnp.squeeze(inputs), axis=0) - - # Recover odd indices corresponding to the desired quantiles. - odds = np.concatenate([ - np.zeros((num_quantiles + 1, 1), dtype=bool), - np.ones((num_quantiles + 1, 1), dtype=bool) - ], - axis=1).ravel()[:-1] - return out[odds][idx] - - return apply_on_axis(_quantile, inputs, axis, q, weight, **kwargs) - - -def multivariate_cdf_quantile_maps( - inputs: jnp.ndarray, - target_sampler: Optional[Callable[[jax.Array, Tuple[int, int]], - jnp.ndarray]] = None, - rng: Optional[jax.Array] = None, - num_target_samples: Optional[int] = None, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, - input_weights: Optional[jnp.ndarray] = None, - target_weights: Optional[jnp.ndarray] = None, - **kwargs: Any -) -> Tuple[Func_t, Func_t]: - r"""Returns multivariate CDF and quantile maps, given input samples. - - Implements the multivariate generalizations for CDF and quantiles proposed in - :cite:`chernozhukov:17`. The reference measure is assumed to be the uniform - measure by default, but can be modified. For consistency, the reference - measure should be symmetrically centered around - :math:`(\tfrac{1}{2},\cdots,\tfrac{1}{2})` and supported on :math:`[0, 1]^d`. - - The implementation return two entropic map estimators, one for the CDF map, - the other for the quantiles map. - - Args: - inputs: 2D array of ``[n, d]`` vectors. - target_sampler: Callable that takes a ``rng`` and ``[m, d]`` shape. - ``m`` is passed on as ``target_num_samples``, dimension ``d`` is inferred - directly from the shape passed in ``inputs``. This is assumed by default - to be :func:`~jax.random.uniform`, and could be any other random sampler - properly wrapped to have the signature above. - rng: rng key used by ``target_sampler``. - num_target_samples: number ``m`` of points generated in the target - distribution. - cost_fn: Cost function, used to compare ``inputs`` and ``targets``. - Passed on to instantiate a - :class:`~ott.geometry.pointcloud.PointCloud` object. If :obj:`None`, - :class:`~ott.geometry.costs.SqEuclidean` is used. - epsilon: entropic regularization parameter used to instantiate the - :class:`~ott.geometry.pointcloud.PointCloud` object. - input_weights: ``[n,]`` vector of weights for input measure. Assumed to - be uniform by default. - target_weights: ``[m,]`` vector of weights for target measure. Assumed - to be uniform by default. - kwargs: keyword arguments passed on to the :func:`~ott.solvers.linear.solve` - function, which solves the OT problem between ``inputs`` and ``targets`` - using the :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` algorithm. - - Returns: - - The multivariate CDF map, taking a ``[b, d]`` batch of vectors in the - range of the ``inputs``, and mapping each vector within the range - of the reference measure (assumed by default to be :math:`[0, 1]^d`). - - The quantile map, mapping a batch ``[b, d]`` of multivariate quantile - vectors onto ``[b, d]`` vectors in :math:`[0, 1]^d`, the range of - the reference measure. - """ - n, d = inputs.shape - rng = utils.default_prng_key(rng) - - if num_target_samples is None: - num_target_samples = n - if target_sampler is None: - target_sampler = jax.random.uniform - - targets = target_sampler(rng, (num_target_samples, d)) - geom = pointcloud.PointCloud( - inputs, targets, cost_fn=cost_fn, epsilon=epsilon - ) - - out = linear.solve(geom, a=input_weights, b=target_weights, **kwargs) - potentials = out.to_dual_potentials() - - cdf_map = jtu.Partial(lambda x, p: p.transport(x), p=potentials) - quantile_map = jtu.Partial( - lambda x, p: p.transport(x, forward=False), p=potentials - ) - return cdf_map, quantile_map - - -def _quantile_normalization( - inputs: jnp.ndarray, targets: jnp.ndarray, weights: float, **kwargs: Any -) -> jnp.ndarray: - """Apply soft quantile normalization on a one dimensional array.""" - num_points = inputs.shape[0] - a = jnp.ones((num_points,)) / num_points - ot = transport_for_sort(inputs, a, weights, **kwargs) - return 1.0 / a * ot.apply(targets, axis=1) - - -def quantile_normalization( - inputs: jnp.ndarray, - targets: jnp.ndarray, - weights: Optional[jnp.ndarray] = None, - axis: int = -1, - **kwargs: Any, -) -> jnp.ndarray: - r"""Re-normalize inputs so that its quantiles match those of targets/weights. - - Quantile normalization rearranges the values in inputs to values that match - the distribution of values described in the discrete distribution ``targets`` - weighted by ``weights``. This transformation preserves the order of values - in ``inputs`` along the specified ``axis``. - - Args: - inputs: array of any shape whose values will be changed to match those in - ``targets``. - targets: sorted array (in ascending order) of dimension 1 describing a - discrete distribution. Note: the ``targets`` values must be provided as - a sorted vector. - weights: vector of non-negative weights, summing to :math:`1`, of the same - size as ``targets``. When not set, this defaults to the uniform - distribution. - axis: the axis along which the quantile transformation is applied. - kwargs: keyword arguments passed on to lower level functions. Of interest - to the user are ``squashing_fun``, which will redistribute the values in - ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) - to solve the optimal transport problem; - :class:`cost_fn ` object of - :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground - 1D cost function to transport from ``inputs`` to the ``num_targets`` - target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` - are passed on to parameterize the - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. - - Returns: - An array, which has the same shape as the input, except on the give axis on - which the dimension is 1. - - Raises: - A ValueError in case the weights and the targets are both set and not of - compatible shapes. - """ - if weights is not None and weights.shape != targets.shape: - raise ValueError( - "The target weights and targets values should have the " - f"same shape: {targets.shape} != {weights.shape}" - ) - if weights is None: - num_targets = targets.shape[0] - weights = jnp.ones((num_targets,)) / num_targets - - op = _quantile_normalization - return apply_on_axis(op, inputs, axis, targets, weights, **kwargs) - - -def sort_with( - inputs: jnp.ndarray, - criterion: jnp.ndarray, - topk: int = -1, - **kwargs: Any, -) -> jnp.ndarray: - r"""Sort a multidimensional array according to a real valued criterion. - - Given ``batch`` vectors of dimension `dim`, to which, for each, a real value - ``criterion`` is associated, this function produces ``topk`` (or - ``batch`` if ``topk`` is set to -1, as by default) composite vectors of size - ``dim`` that will be convex combinations of all vectors, ranked from smallest - to largest criterion. Composite vectors with the largest indices will contain - convex combinations of those vectors with highest criterion, vectors with - smaller indices will contain combinations of vectors with smaller criterion. - - Args: - inputs: Array of size [batch, dim]. - criterion: the values according to which to sort the inputs. It has shape - [batch, 1]. - topk: The number of outputs to keep. - kwargs: keyword arguments passed on to lower level functions. Of interest - to the user are ``squashing_fun``, which will redistribute the values in - ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) - to solve the optimal transport problem; - :class:`cost_fn ` object of - :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground - 1D cost function to transport from ``inputs`` to the ``num_targets`` - target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` - are passed on to parameterize the - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. - - Returns: - An Array of size [batch | topk, dim]. - """ - num_points = criterion.shape[0] - weights = jnp.ones(num_points) / num_points - if 0 < topk < num_points: - start_index = 1 - target_weights = jnp.concatenate([ - jnp.array([(num_points - topk) / num_points]), - jnp.ones(topk) / num_points - ]) - else: - start_index = 0 - target_weights = jnp.ones((num_points,)) / num_points - ot = transport_for_sort(criterion, weights, target_weights, **kwargs) - # Applies the topk on each of the dimensions of the inputs. - sort_fn = jax.vmap( - lambda x: (1.0 / target_weights * ot.apply(x, axis=0))[start_index:], - in_axes=(1,), - out_axes=1 - ) - return sort_fn(inputs) - - -def _quantize(inputs: jnp.ndarray, num_q: int, **kwargs: Any) -> jnp.ndarray: - """Apply the soft quantization operator on a one dimensional array.""" - num_points = inputs.shape[0] - a = jnp.ones((num_points,)) / num_points - b = jnp.ones((num_q,)) / num_q - ot = transport_for_sort(inputs, a, b, **kwargs) - return 1.0 / a * ot.apply(1.0 / b * ot.apply(inputs), axis=1) - - -def quantize( - inputs: jnp.ndarray, - num_levels: int = 10, - axis: int = -1, - **kwargs: Any, -) -> jnp.ndarray: - r"""Soft quantizes an input according using ``num_levels`` values along axis. - - The quantization operator consists in concentrating several values around - a few predefined ``num_levels``. The soft quantization operator proposed here - does so by carrying out a `soft` concentration that is differentiable. The - ``inputs`` values are first soft-sorted, resulting in ``num_levels`` values. - In a second step, the ``inputs`` values are then provided again a composite - value that is equal (for each) to a convex combination of those soft-sorted - values using the transportation matrix. As the regularization parameter - ``epsilon`` of regularized optimal transport goes to 0, this operator recovers - the expected behavior of quantization, namely each value in ``inputs`` is - assigned a single level. When using ``epsilon>0`` the behavior is similar but - differentiable. - - Args: - inputs: an Array of size [batch, dim]. - num_levels: number of quantiles available to quantize the signal. - axis: axis along which quantization is carried out. - kwargs: keyword arguments passed on to lower level functions. Of interest - to the user are ``squashing_fun``, which will redistribute the values in - ``inputs`` to lie in :math:`[0,1]` (sigmoid of whitened values by default) - to solve the optimal transport problem; - :class:`cost_fn ` object of - :class:`~ott.geometry.pointcloud.PointCloud`, which defines the ground - 1D cost function to transport from ``inputs`` to the ``num_targets`` - target values; ``epsilon`` regularization parameter. Remaining ``kwargs`` - are passed on to parameterize the - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver. - - Returns: - An Array of the same size as ``inputs``. - """ - return apply_on_axis(_quantize, inputs, axis, num_levels, **kwargs) diff --git a/ott/build/lib/ott/types.py b/ott/build/lib/ott/types.py deleted file mode 100644 index f6eefb7..0000000 --- a/ott/build/lib/ott/types.py +++ /dev/null @@ -1,42 +0,0 @@ -# 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. -from typing import Protocol, Union - -import jax.experimental.sparse as jesp -import jax.numpy as jnp - -__all__ = ["Transport", "Array_g"] - -# TODO(michalk8): introduce additional types here - -# Either a dense or sparse array. -Array_g = Union[jnp.ndarray, jesp.BCOO] - - -class Transport(Protocol): - """Interface for the solution of a transport problem. - - Classes implementing those function do not have to inherit from it, the - class can however be used in type hints to support duck typing. - """ - - @property - def matrix(self) -> jnp.ndarray: - ... - - def apply(self, inputs: jnp.ndarray, axis: int) -> jnp.ndarray: - ... - - def marginal(self, axis: int = 0) -> jnp.ndarray: - ... diff --git a/ott/build/lib/ott/utils.py b/ott/build/lib/ott/utils.py deleted file mode 100644 index e9aecb7..0000000 --- a/ott/build/lib/ott/utils.py +++ /dev/null @@ -1,210 +0,0 @@ -# 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 dataclasses -import functools -import io -import warnings -from typing import Any, Callable, NamedTuple, Optional, Tuple - -import jax -import numpy as np - -try: - from tqdm import tqdm -except ImportError: - tqdm = Any - -__all__ = [ - "register_pytree_node", - "deprecate", - "default_prng_key", - "default_progress_fn", - "tqdm_progress_fn", -] - -Status_t = Tuple[np.ndarray, np.ndarray, np.ndarray, NamedTuple] -IOCallback_t = Callable[[Status_t], None] - - -def register_pytree_node(cls: type) -> type: - """Register dataclasses as pytree_nodes.""" - cls = dataclasses.dataclass()(cls) - flatten = lambda obj: jax.tree_util.tree_flatten(dataclasses.asdict(obj)) - unflatten = lambda d, children: cls(**d.unflatten(children)) - jax.tree_util.register_pytree_node(cls, flatten, unflatten) - return cls - - -def deprecate( # noqa: D103 - *, - version: Optional[str] = None, - alt: Optional[str] = None, - func: Optional[Callable[[Any], Any]] = None -) -> Callable[[Any], Any]: - - def wrapper(*args: Any, **kwargs: Any) -> Any: - warnings.warn(msg, category=DeprecationWarning, stacklevel=2) - return func(*args, **kwargs) - - if func is None: - return lambda fn: deprecate(version=version, alt=alt, func=fn) - - msg = f"`{func.__name__}` will be removed in the " - msg += ("next" if version is None else f"`ott-jax=={version}`") + " release." - if alt: - msg += " " + alt - - return functools.wraps(func)(wrapper) - - -def default_prng_key(rng: Optional[jax.Array] = None) -> jax.Array: - """Get the default PRNG key. - - Args: - rng: PRNG key. - - Returns: - If ``rng = None``, returns the default PRNG key. - Otherwise, it returns the unmodified ``rng`` key. - """ - return jax.random.PRNGKey(0) if rng is None else rng - - -def default_progress_fn( - fmt: str = "{iter} / {max_iter} -- {error}", - stream: Optional[io.TextIOBase] = None, -) -> IOCallback_t: - """Return a callback that prints the progress when solving - :mod:`linear problems `. - - It prints the progress only when the error is computed, that is every - :attr:`~ott.solvers.linear.sinkhorn.Sinkhorn.inner_iterations`. - - Args: - fmt: Format used to print. It can format ``iter``, ``max_iter`` and - ``error`` values. - stream: Output IO stream. - - Returns: - A callback function accepting the following arguments - - - the current iteration number, - - the number of inner iterations after which the error is computed, - - the total number of iterations, and - - the current :class:`~ott.solvers.linear.sinkhorn.SinkhornState` or - :class:`~ott.solvers.linear.sinkhorn_lr.LRSinkhornState`. - - Examples: - .. code-block:: python - - import jax - import jax.numpy as jnp - - from ott import utils - from ott.geometry import pointcloud - from ott.solvers.linear import sinkhorn - - x = jax.random.normal(jax.random.PRNGKey(0), (100, 5)) - geom = pointcloud.PointCloud(x) - - progress_fn = utils.default_progress_fn() - solve_fn = jax.jit(sinkhorn.solve, static_argnames=["progress_fn"]) - out = solve_fn(geom, progress_fn=progress_fn) - """ # noqa: D205 - - def progress_callback(status: Status_t) -> None: - iteration, inner_iterations, total_iter, errors = _prepare_info(status) - # Avoid reporting error on each iteration, - # because errors are only computed every `inner_iterations`. - if iteration % inner_iterations == 0: - error_idx = max(0, iteration // inner_iterations - 1) - error = errors[error_idx] - - print( - fmt.format(iter=iteration, max_iter=total_iter, error=error), - file=stream - ) - - return progress_callback - - -def tqdm_progress_fn( - pbar: tqdm, - fmt: str = "error: {error:0.6e}", -) -> IOCallback_t: - """Return a callback that updates a progress bar when solving - :mod:`linear problems `. - - It updates the progress bar only when the error is computed, that is every - :attr:`~ott.solvers.linear.sinkhorn.Sinkhorn.inner_iterations`. - - Args: - pbar: `tqdm `_ progress bar. - fmt: Format used for the postfix. It can format ``iter``, ``max_iter`` and - ``error`` values. - - Returns: - A callback function accepting the following arguments - - - the current iteration number, - - the number of inner iterations after which the error is computed, - - the total number of iterations, and - - the current :class:`~ott.solvers.linear.sinkhorn.SinkhornState` or - :class:`~ott.solvers.linear.sinkhorn_lr.LRSinkhornState`. - - Examples: - .. code-block:: python - - import tqdm - - import jax - import jax.numpy as jnp - - from ott import utils - from ott.geometry import pointcloud - from ott.solvers.linear import sinkhorn - - x = jax.random.normal(jax.random.PRNGKey(0), (100, 5)) - geom = pointcloud.PointCloud(x) - - with tqdm.tqdm() as pbar: - progress_fn = utils.tqdm_progress_fn(pbar) - solve_fn = jax.jit(sinkhorn.solve, static_argnames=["progress_fn"]) - out = solve_fn(geom, progress_fn=progress_fn) - """ # noqa: D205 - - def progress_callback(status: Status_t) -> None: - iteration, inner_iterations, total_iter, errors = _prepare_info(status) - # Avoid reporting error on each iteration, - # because errors are only computed every `inner_iterations`. - if iteration % inner_iterations == 0: - error_idx = max(0, iteration // inner_iterations - 1) - error = errors[error_idx] - - postfix = fmt.format(iter=iteration, max_iter=total_iter, error=error) - pbar.set_postfix_str(postfix) - pbar.total = total_iter // inner_iterations - pbar.update() - - return progress_callback - - -def _prepare_info(status: Status_t) -> Tuple[int, int, int, np.ndarray]: - iteration, inner_iterations, total_iter, state = status - iteration = int(iteration) + 1 - inner_iterations = int(inner_iterations) - total_iter = int(total_iter) - errors = np.array(state.errors).ravel() - - return iteration, inner_iterations, total_iter, errors diff --git a/perturbot/build/lib/perturbot/__init__.py b/perturbot/build/lib/perturbot/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/perturbot/build/lib/perturbot/eval/.run_loo.py b/perturbot/build/lib/perturbot/eval/.run_loo.py deleted file mode 100755 index ddb2f82..0000000 --- a/perturbot/build/lib/perturbot/eval/.run_loo.py +++ /dev/null @@ -1,126 +0,0 @@ -# Run leave-one-out prediction & evaluation -from typing import Optional -import numpy as np -import argparse -import mudata as md -from sklearn.model_selection import LeaveOneOut -from software.perturbot.perturbot.ot.cot_labels import ( - predict_with_cot, - evaluate_training, -) - -method_dict = {"cot_labels": predict_with_cot} -get_train_metrics_dict = {"cot_labels": evaluate_training} - - -def parse_args(): - parser = argparse.ArgumentParser( - description="Run leave-one-out prediction & evaluation." - ) - parser.add_argument("input_h5mu", type=str, help="Path to the input .h5mu file.") - parser.add_argument( - "--label-col", - "-l", - type=str, - default="treatment", - help="Column label in .obs of MuData that encodes the perturbation labels.", - ) - parser.add_argument( - "--within-label-group", - "-g", - type=str, - default=None, - help="If provided, evaluation includes the concordance of the grouping.", - ) - parser.add_argument( - "--within-label-group", - "-g", - type=str, - default=None, - help="If provided, evaluation includes the concordance of the grouping.", - ) - parser.add_argument( - "--train-with-group", - "-tg", - type=bool, - action="store_true", - help="train with `within_label_group`.", - ) - parser.add_argument( - "--source-modality", - "-s", - type=str, - default="protein", - help="Source modality in MuData", - ) - parser.add_argument( - "--target-modality", - "-t", - type=str, - default="RNA", - help="Target modality in MuData", - ) - parser.add_argument( - "--rng-seed", - "-r", - type=int, - default=2024, - help="random number generator seed for NumPy", - ) - return parser.parse_args() - - -def get_data_with_group( - x, y, include_y, z: Optional[np.ndarray] = None, nbperclass: Optional[int] = None -): - # Randomly select nbperclass train datasets - xr = np.zeros((0, x.shape[1])) - yr = np.zeros((0)) - if z is not None: - zr = np.zeros((0)) - else: - zr = None - for i in include_y: - xi = x[y.ravel() == i, :] - xr = np.concatenate((xr, xi), 0) - yr = np.concatenate((yr, np.repeat(i, xi.shape[0]))) - if z is not None: - zi = z[y.ravel() == i] - zr = np.concatenate((zr, zi), 0) - return xr, yr, zr - - -def main(): - args = parse_args() - np.random.seed(args.rng_seed) - mdata = md.read_h5ad(args.input_h5mu) - Xtot1 = mdata[args.source_modality].X - Xtot2 = mdata[args.target_modality].X - labels = mdata.obs[args.label_col].values - within_label_group = mdata.obs[args.within_label_group].values - loo = LeaveOneOut() - train_data = {} - test_data = {} - for i, (train_idx, test_idx) in enumerate(loo.split(labels)): - Xs, Ys, Zs = get_data_with_group( - Xtot1, labels, labels[train_idx], within_label_group - ) - Xt, Yt, Zt = get_data_with_group( - Xtot2, labels, labels[train_idx], within_label_group - ) - Xs_test, Ys_test, Zs_test = get_data_with_group( - Xtot1, labels, labels[test_idx], within_label_group - ) - Xt_test, Yt_test, Zt_test = get_data_with_group( - Xtot1, labels, labels[test_idx], within_label_group - ) - for method_label, run_fn in method_dict.items(): - Xt_pred, log = run_fn(Xs, Ys, Xt, Yt, Xs_test) - train_eval_metrics = get_train_metrics_dict[method_label]( - Xs, Ys, Xt, Yt, Zt, log - ) - test_eval_metrics = get_test_metrics(Xt_pred, Xt_test, Zt_test) - - -if __name__ == "__main__": - main() diff --git a/perturbot/build/lib/perturbot/eval/__init__.py b/perturbot/build/lib/perturbot/eval/__init__.py deleted file mode 100755 index e69de29..0000000 diff --git a/perturbot/build/lib/perturbot/eval/all.py b/perturbot/build/lib/perturbot/eval/all.py deleted file mode 100755 index 67dca36..0000000 --- a/perturbot/build/lib/perturbot/eval/all.py +++ /dev/null @@ -1,190 +0,0 @@ -"""For 5-fold CV, run the single inner loop of train-validation to select hyperparameter.""" - -import os -import argparse -import pickle as pkl -from functools import partial -import numpy as np -import torch -from perturbot.eval.match import get_FOSCTTM, get_diag_fracs -from perturbot.eval.utils import get_Ts_from_nn_multKs -from perturbot.match.cot_labels import get_coupling_cotl_sinkhorn - -from perturbot.match.ott_egwl import ( - get_coupling_egw_labels_ott, - get_coupling_egw_all_ott, - get_coupling_eot_ott, - get_coupling_leot_ott, - get_coupling_egw_ott, -) -from perturbot.match.cot import ( - get_coupling_cot_sinkhorn, - get_coupling_each_cot_sinkhorn, -) -from perturbot.match.gw_labels import get_coupling_egw_labels -from perturbot.predict.scvi_vae import train_vae_model, infer_from_Xs, infer_from_Ys - -ot_method_map = { - "ECOOTL": get_coupling_cotl_sinkhorn, - "ECOOT": get_coupling_cot_sinkhorn, - "ECOOT_each": get_coupling_each_cot_sinkhorn, - "EGWL": get_coupling_egw_labels, - "EOT_ott": get_coupling_eot_ott, - "LEOT_ott": get_coupling_leot_ott, - "EGW_ott": get_coupling_egw_ott, - "EGW_all_ott": get_coupling_egw_all_ott, - "EGWL_ott": get_coupling_egw_labels_ott, - "VAE_label": train_vae_model, - "VAE": partial(train_vae_model, use_label=False), -} - - -def parse_args(): - parser = argparse.ArgumentParser( - "Run inner validation loop of CV", - "Run sample-matching OT in parallel and fit prediction model in leave-one-out mannter", - ) - parser.add_argument("method", type=str) - parser.add_argument("filepath", type=str) - parser.add_argument("eps", type=float) - - parser.add_argument( - "-l", - "--log-filepath", - type=str, - default=None, - help="Path to log.pkl if transport was already run. If not None, reevaluates but does not calculates the transport again.", - ) - parser.add_argument( - "-b", - "--baseline", - type=str, - default=None, - help="One of perfect, random, by_conc", - ) - return parser.parse_args() - - -ot_method_all_to_all = ["GW_all", "EGW_all_ott", "EOT_all_ott"] - - -def main(args): - # if args.log_filepath is not None and args.log_filepath != "None": - # with open(args.log_filepath, "rb") as f: - # prev_logs = pkl.load(f) - logs = { - "matching_evals": {}, - "pred_evals": {}, - "T": {}, - "pred": {}, - "log": {}, - "eps": {}, - } - try: - match_eps = args.eps - logs["eps"] = match_eps - with open(args.filepath, "rb") as f: - data_dict = pkl.load(f) - - X_dict = data_dict["Xs_dict"] - Y_dict = data_dict["Xt_dict"] - Zs_dict = data_dict["Zs_dict"]["dosage"] - Zt_dict = data_dict["Zt_dict"]["dosage"] - - labels = list(X_dict.keys()) - - train_X = X_dict - train_Y = Y_dict - train_Z = Zs_dict - train_data = (train_X, train_Y) - print(f"Calculating matching with {match_eps}") - if args.log_filepath is not None: - with open(args.log_filepath, "rb") as f: - d = pkl.load(f) - Ts_matching = d["T"] - log_matching = "" - else: - Ts_matching, log_matching = ot_method_map[args.method]( - train_data, match_eps - ) - if "VAE" in args.method: - dim_X = data_dict["Xs_dict"][labels[0]].shape[1] - dim_Y = data_dict["Xt_dict"][labels[0]].shape[1] - latent_Y = infer_from_Ys(train_Y, Ts_matching, dim_X) - latent_X = infer_from_Xs(train_X, Ts_matching, dim_Y) - _, mean_foscttm = get_FOSCTTM( - Ts_matching, - latent_X, - latent_Y, - use_agg="mean", - use_barycenter=False, - ) - ks = [1, 5, 10, 50, 100] - k_to_Ts = get_Ts_from_nn_multKs(latent_X, latent_Y, ks) # k -> T - dfracs = {} - rel_dfracs = {} - for k, Ts in k_to_Ts.items(): - dfracs[k], rel_dfracs[k] = get_diag_fracs( - Ts, train_X, train_Y, train_Z, train_Z - ) - - else: - if isinstance(Ts_matching, dict): - total_sum = 0 - for k, v in Ts_matching.items(): - total_sum += v.sum() - Ts_matching = { - k: v.astype(np.double) / total_sum for k, v in Ts_matching.items() - } - else: - Ts_matching = Ts_matching.astype(np.double) / Ts_matching.sum() - _, mean_foscttm = get_FOSCTTM(Ts_matching, train_X, train_Y, use_agg="mean") - dfracs, rel_dfracs = get_diag_fracs( - Ts_matching, train_X, train_Y, train_Z, train_Z - ) - - # if not all_to_all: - mean_mean_foscttm = mean_foscttm.mean() - # else: - # mean_mean_foscttm = mean_foscttm - print("foscttm", mean_foscttm) - logs["matching_evals"] = ( - { - "foscttm": mean_foscttm, - "mean_foscttm": mean_mean_foscttm, - "dfracs": dfracs, - "rel_dfracs": rel_dfracs, - }, - ) - logs["T"] = Ts_matching - logs["log"] = (log_matching,) - except Exception as e: - with open(f"all_{args.method}.{args.eps}.tmp.pkl", "wb") as f: - pkl.dump({"log": logs, "T": Ts_matching}, f) - raise e - - with open(f"all_{args.method}.{args.eps}.pkl", "wb") as f: - pkl.dump(logs, f) - - -def submit_all_run(data_path, ot_method_label, load_existing=False): - epsilons = [1e-2, 1e-3, 1e-4, 1e-5] - for eps in epsilons: - run_label = f"all.{ot_method_label}.{eps}" - f = open(f"{run_label}.bsub", "w") - f.write("#!/bin/bash\n") - f.write("source ~/.bashrc\n") - f.write("pwd\n") - f.write("conda activate ot \n") - command = f"python /gpfs/scratchfs01/site/u/ryuj6/OT/software/perturbot/perturbot/eval/all.py {ot_method_label} {data_path} {eps}" - f.write(f"echo {command}\n") - f.write(f"{command}\n") - f.close() - command = f"bsub -M 10G -n 10 -q long -J {run_label} -o log/log.{run_label}.o%J -e log/log.{run_label}.e%J < {run_label}.bsub" - print(command) - os.system(command) - - -if __name__ == "__main__": - args = parse_args() - main(args) diff --git a/perturbot/build/lib/perturbot/eval/cv.py b/perturbot/build/lib/perturbot/eval/cv.py deleted file mode 100755 index b035245..0000000 --- a/perturbot/build/lib/perturbot/eval/cv.py +++ /dev/null @@ -1,251 +0,0 @@ -import os -from typing import List, Dict, Any -from functools import partial -from itertools import product -import argparse -import os - -os.environ["OPENBLAS_NUM_THREADS"] = "1" -import numpy as np -import pandas as pd -from sklearn.model_selection import KFold -from perturbot.eval.prediction import get_evals_preds, get_evals -from multiprocessing import Pool -import pickle as pkl -from perturbot.eval.match import get_FOSCTTM, get_diag_fracs -from perturbot.eval.utils import get_Ts_from_nn_multKs, _pop_key -from perturbot.predict.scvi_vae import predict_from_model, infer_from_Xs, infer_from_Ys -from perturbot.preprocess.vae import ( - train_vae_rna, - train_vae_acc, - train_vae_prot, - SCVI_LATENT_KEY, -) -from perturbot.match.cot_labels import get_coupling_cotl, get_coupling_cotl_sinkhorn -from perturbot.match.gw import ( - get_coupling_gw_cg, - get_coupling_egw, - get_coupling_egw_all, - get_coupling_gw_all, -) -from perturbot.match.ott_egwl import ( - get_coupling_egw_labels_ott, - get_coupling_egw_all_ott, - get_coupling_eot_ott, - get_coupling_leot_ott, - get_coupling_egw_ott, -) -from perturbot.match.cot import ( - get_coupling_cot, - get_coupling_cot_sinkhorn, - get_coupling_each_cot_sinkhorn, -) -from perturbot.match.gw_labels import get_coupling_gw_labels, get_coupling_egw_labels -from perturbot.predict.scvi_vae import train_vae_model - -np.seterr(all="raise") - -vae_trainer_dict = { - "rna": train_vae_rna, - "accessibility": train_vae_acc, - "protein": train_vae_prot, -} - -ot_method_map = { - "COOTL": get_coupling_cotl, - "ECOOTL": get_coupling_cotl_sinkhorn, - "COOT": get_coupling_cot, - "ECOOT": get_coupling_cot_sinkhorn, - "ECOOT_each": get_coupling_each_cot_sinkhorn, - "GW_all": get_coupling_gw_all, - # "EGW_all": get_coupling_egw_all, - "GW": get_coupling_gw_cg, - # "EGW": get_coupling_egw, - "GWL": get_coupling_gw_labels, - "EGWL": get_coupling_egw_labels, - "EOT_ott": get_coupling_eot_ott, - "LEOT_ott": get_coupling_leot_ott, - "EGW_ott": get_coupling_egw_ott, - "EGW_all_ott": get_coupling_egw_all_ott, - "EGWL_ott": get_coupling_egw_labels_ott, - "VAE_label": train_vae_model, - "VAE": partial(train_vae_model, use_label=False), -} -ot_method_hyperparams = {} -for method in ["EGWL", "EOT_ott", "LEOT_ott", "EGW_ott", "EGW_all_ott", "EGWL_OTT"]: - ot_method_hyperparams[method] = [ - 1e-2, - 1e-3, - 1e-4, - 1e-5, - 1e-6, - ] -for method in ["ECOOTL", "ECOOT", "ECOOT_each"]: - ot_method_hyperparams[method] = [0.1, 0.05, 0.01, 0.005, 0.001] -ot_method_all_to_all = ["GW_all", "EGW_all_ott", "EOT_all_ott"] - - -def run_cv_models( - filepath, - ot_method_label, - data_label="k", - log_filepath=None, - full_filepath=None, - run_mlp=False, - mlp_baseline=False, - rerun=False, -): - if ot_method_label == "VAE": - assert full_filepath is not None - with open(filepath, "rb") as f: - data_dict = pkl.load(f) - labels = list(data_dict["Xs_dict"].keys()) - cv = KFold(n_splits=5) - # Outer loop will evaluate final performance. - for i, (train_val_idx, test_idx) in enumerate(cv.split(labels)): - # Inner loop will select the best-performing hyperparameter. - # Run CV for multiple hyperparameters (eps) - run_label = f"{'mlp.' if run_mlp else ''}{ot_method_label}.{i}" - - if ( - run_mlp - and os.path.isfile(f"val_mlp.{ot_method_label}.{i}.best_eps.pkl") - and not rerun - ): - print(f"val_mlp.{ot_method_label}.{i}.best_eps.pkl exists. Skipping") - continue - elif ( - (not run_mlp) - and os.path.isfile(f"val_CV_{ot_method_label}.{i}.best_eps.pkl") - and log_filepath is None - ): - print(f"val_CV_{ot_method_label}.{i}.best_eps.pkl exists. Skipping") - continue - # Write job script - f = open(f"{run_label}.bsub", "w") - f.write("#!/bin/bash\n") - f.write("source ~/.bashrc\n") - f.write("pwd\n") - f.write("conda activate ot \n") - if "VAE" in ot_method_label: - command = f"python /gpfs/scratchfs01/site/u/ryuj6/OT/software/perturbot/perturbot/eval/cv_inner_loop.py {ot_method_label} {i} {full_filepath} {'' if log_filepath is None else f'-l val_CV_{ot_method_label}.{i}.pkl'} {'--mlp' if run_mlp else ''} " - else: - command = f"python /gpfs/scratchfs01/site/u/ryuj6/OT/software/perturbot/perturbot/eval/cv_inner_loop.py {ot_method_label} {i} {full_filepath if run_mlp else filepath} {'' if log_filepath is None else f'-l val_CV_{ot_method_label}.{i}.pkl'} {'' if full_filepath is None else '-p ' + full_filepath} {'--mlp' if run_mlp else ''} {'--baseline' if mlp_baseline else ''}" - - f.write(f"echo {command}\n") - f.write(f"{command}\n") - f.close() - label = "mlp" if run_mlp else "cv" - if log_filepath is not None: - label = "reval" - command = f"bsub {'-M 20G -n 10' if 'VAE' in ot_method_label or run_mlp else '-M 10G -n 25'} -q {'short' if log_filepath is not None else 'long'} -J {data_label}{'mlp' if run_mlp else 'CV'}.{run_label} -o log/log_{label}.{run_label}.o%J -e log/log_{label}.{run_label}.e%J < {run_label}.bsub" - print(command) - os.system(command) - - -def get_test_eval( - filepath, - ot_method_label, - data_label="k", - log_filepath=None, - full_filepath=None, - run_mlp=False, - baseline=None, -): - if ot_method_label == "VAE": - assert full_filepath is not None - with open(filepath, "rb") as f: - data_dict = pkl.load(f) - labels = list(data_dict["Xs_dict"].keys()) - cv = KFold(n_splits=5) - # Outer loop will evaluate final performance. - if baseline is not None: - for i, (train_val_idx, test_idx) in enumerate(cv.split(labels)): - if log_filepath is not None: - log_filepath = f"test_{baseline}.{i}.pkl" - run_label = f"TM.{baseline}.{i}" - f = open(f"{run_label}.bsub", "w") - f.write("#!/bin/bash\n") - f.write("source ~/.bashrc\n") - f.write("pwd\n") - f.write("conda activate ot \n") - command = f"python /gpfs/scratchfs01/site/u/ryuj6/OT/software/perturbot/perturbot/eval/cv_outer_loop.py {baseline} {i} {filepath} 0,0,0 -p {full_filepath} -l {log_filepath} --baseline {baseline}" - f.write(f"echo {command}\n") - f.write(f"{command}\n") - f.close() - command = f"bsub -M 10G -n 10 -q long -J {data_label}{run_label} -o test_log/log.{run_label}.o%J -e test_log/log.{run_label}.e%J < {run_label}.bsub" - print(command) - os.system(command) - else: - for i, (train_val_idx, test_idx) in enumerate(cv.split(labels)): - # Evaluate matching - try: - with open(f"val_CV_{ot_method_label}.{i}.best_eps.pkl", "rb") as f: - log = pkl.load(f) - - except FileNotFoundError: - print(f"val_CV_{ot_method_label}.{i}.best_eps.pkl not found. Skipping") - continue - if "VAE" not in ot_method_label: - try: - with open(f"val_mlp.{ot_method_label}.{i}.best_eps.pkl", "rb") as f: - log2 = pkl.load(f) - eps_pred = log2["eps"] - except FileNotFoundError: - try: - with open(f"val_CV_MLP_{ot_method_label}.{i}.pkl", "rb") as f: - log2 = pkl.load(f) - eps_pred = log2["pred_evals"].T["Pearson_samples"].idxmax() - except FileNotFoundError: - print( - f"val_CV_MLP_{ot_method_label}.{i}.pkl and val_mlp.{ot_method_label}.{i}.best_eps.pkl not found. Skipping" - ) - continue - # elif os.path.exists(f"test_{ot_method_label}.{i}.pkl"): - # print(f"test_{ot_method_label}.{i}.pkl already exists. Skipping") - # elif os.path.exists(f"test_log/test_{ot_method_label}.{i}.tmp.pkl"): - # print( - # f"test_log/test_{ot_method_label}.{i}.tmp.pkl already exists. Skipping" - # ) - else: - eps_pred = np.nan - - allowed_eps = [1e-2, 1e-3, 1e-4, 1e-5] - eps_lin = log["pred"] - - eps_matching = log["matching"] - if isinstance(eps_lin, tuple): - eps_lin = eps_lin[0] - if isinstance(eps_matching, tuple): - eps_matching = eps_matching[0] - - if "VAE" not in ot_method_label and ( - eps_lin not in allowed_eps or eps_matching not in allowed_eps - ): - try: - with open(f"val_CV_{ot_method_label}.{i}.pkl", "rb") as f: - eval_log = pkl.load(f) - except FileNotFoundError: - print(f"val_CV_{ot_method_label}.{i}.pkl not found. Skipping") - eps_lin = eval_log["pred_evals"].loc["MSE", allowed_eps].idxmin() - eps_matching = eval_log["matching_evals"][allowed_eps].idxmin() - elif "VAE" in ot_method_label: - if isinstance(eps_lin, tuple): - eps_lin = eps_lin[0] - if isinstance(eps_matching, tuple): - eps_matching = eps_matching[0] - if log_filepath is not None: - log_filepath = f"test_{ot_method_label}.{i}.pkl" - run_label = f"TM.{ot_method_label}.{i}" - f = open(f"{run_label}.bsub", "w") - f.write("#!/bin/bash\n") - f.write("source ~/.bashrc\n") - f.write("pwd\n") - f.write("conda activate ot \n") - command = f"python /gpfs/scratchfs01/site/u/ryuj6/OT/software/perturbot/perturbot/eval/cv_outer_loop.py {ot_method_label} {i} {filepath} {eps_matching},{eps_lin},{eps_pred} -p {full_filepath} -l {log_filepath} {'' if baseline is None else '--baseline '+baseline}" - f.write(f"echo {command}\n") - f.write(f"{command}\n") - f.close() - command = f"bsub -M 10G -n 10 -q long -J {data_label}{run_label} -o test_log/log.{run_label}.o%J -e test_log/log.{run_label}.e%J < {run_label}.bsub" - print(command) - os.system(command) diff --git a/perturbot/build/lib/perturbot/eval/cv_inner_loop.py b/perturbot/build/lib/perturbot/eval/cv_inner_loop.py deleted file mode 100755 index a38316e..0000000 --- a/perturbot/build/lib/perturbot/eval/cv_inner_loop.py +++ /dev/null @@ -1,673 +0,0 @@ -"""For 5-fold CV, run the single inner loop of train-validation to select hyperparameter.""" - -from typing import Dict, Any -import sys -import argparse -import pickle as pkl -import jax -from functools import partial -from itertools import product -from multiprocessing import Pool -import numpy as np -import pandas as pd -from perturbot.eval.utils import get_Ts_from_nn_multKs -from sklearn.model_selection import KFold -import torch -from tqdm import tqdm -import psutil -import gc -from perturbot.eval.utils import _pop_keys, _pop_key -from perturbot.eval.match import get_FOSCTTM, get_diag_fracs -from perturbot.eval.prediction import get_evals_preds, get_evals -from perturbot.preprocess.vae import ( - train_vae_rna, - train_vae_acc, - train_vae_prot, - SCVI_LATENT_KEY, -) -from perturbot.match.cot_labels import get_coupling_cotl, get_coupling_cotl_sinkhorn -from perturbot.match.gw import ( - get_coupling_gw_cg, - get_coupling_egw, - get_coupling_egw_all, - get_coupling_gw_all, -) -from perturbot.match.ott_egwl import ( - get_coupling_egw_labels_ott, - get_coupling_egw_all_ott, - get_coupling_eot_ott, - get_coupling_leot_ott, - get_coupling_egw_ott, -) -from perturbot.match.cot import ( - get_coupling_cot, - get_coupling_cot_sinkhorn, - get_coupling_each_cot_sinkhorn, -) -from perturbot.match.gw_labels import get_coupling_gw_labels, get_coupling_egw_labels -from perturbot.predict.scvi_vae import train_vae_model -from perturbot.predict.linear_regression import ( - ols_normed, - weight_conc_normed, - weighted_ols_normed, - weight_1_ols_normed, - predict, -) -from perturbot.predict.mlp import train_mlp -from perturbot.predict.scvi_vae import predict_from_model, infer_from_Xs, infer_from_Ys - -ot_method_map = { - "ECOOTL": get_coupling_cotl_sinkhorn, - "ECOOT_each": get_coupling_each_cot_sinkhorn, - "ECOOT": get_coupling_cot_sinkhorn, - "EGWL": get_coupling_egw_labels, - "EOT_ott": get_coupling_eot_ott, - "LEOT_ott": get_coupling_leot_ott, - "EGW_ott": get_coupling_egw_ott, - "EGW_all_ott": get_coupling_egw_all_ott, - "EGWL_ott": get_coupling_egw_labels_ott, - "VAE_label": train_vae_model, - "VAE": partial(train_vae_model, use_label=False), -} - - -def parse_args(): - parser = argparse.ArgumentParser( - "Run inner validation loop of CV", - "Run sample-matching OT in parallel and fit prediction model in leave-one-out mannter", - ) - parser.add_argument("method", type=str) - parser.add_argument("test_idx", type=int) - parser.add_argument("filepath", type=str) - parser.add_argument("-m", "--mlp", action="store_true") - parser.add_argument("--baseline", action="store_true") - - parser.add_argument( - "-p", - "--pred-filepath", - type=str, - default=None, - help="Path to data.pkl file with full features.", - ) - parser.add_argument( - "-l", - "--log-filepath", - type=str, - default=None, - help="Path to log.pkl if transport was already run. If not None, reevaluates but does not calculates the transport again.", - ) - return parser.parse_args() - - -ot_method_hyperparams = {} -for method in [ - "EGWL", - "EOT_ott", - "LEOT_ott", - "EGW_ott", - "EGW_all_ott", - "EGWL_ott", - "ECOOT", - "ECOOT_each", - "ECOOTL", -]: - ot_method_hyperparams[method] = [ - 0.1, - 1e-2, - 1e-3, - 1e-4, - 1e-5, - ] -vae_adv_loss = [1, 5, 10, 50, 100] -vae_latent_dims = [128] -vae_learning_rates = [1e-4] -# vae_latent_dims = [32] -# vae_learning_rates = [1e-4] -for method in ["VAE", "VAE_label"]: - ot_method_hyperparams[method] = list( - product(vae_adv_loss, vae_latent_dims, vae_learning_rates) - ) - -ot_method_all_to_all = ["GW_all", "EGW_all_ott", "EOT_all_ott"] -pred_method = weighted_ols_normed -baseline_pred_methods = [ols_normed, weight_1_ols_normed, weight_conc_normed] -baseline_pred_method_labels = ["perfect", "random", "by_conc"] -pred_from_param = predict - - -def main(args): - epsilons = ot_method_hyperparams[args.method] - all_to_all = args.method in ot_method_all_to_all - with open(args.filepath, "rb") as f: - data_dict = pkl.load(f) - if args.pred_filepath is not None: - with open(args.pred_filepath, "rb") as f: - pred_data_dict = pkl.load(f) - - X_dict = data_dict["Xs_dict"] - Y_dict = data_dict["Xt_dict"] - Zs_dict = data_dict["Zs_dict"]["dosage"] - Zt_dict = data_dict["Zt_dict"]["dosage"] - - labels = list(X_dict.keys()) - dim_X = X_dict[labels[0]].shape[1] - dim_Y = Y_dict[labels[0]].shape[1] - cv = KFold(n_splits=5) - train_val_idx, test_idx = list(cv.split(labels))[args.test_idx] - cv_inner = KFold(n_splits=5) - - test_labels = [labels[tidx] for tidx in test_idx] - train_val_X = _pop_keys(X_dict, test_labels) - train_val_Y = _pop_keys(Y_dict, test_labels) - if args.pred_filepath is not None: - train_val_X_full = _pop_keys(pred_data_dict["Xs_dict"], test_labels) - train_val_Y_full = _pop_keys(pred_data_dict["Xt_dict"], test_labels) - train_val_Z = _pop_keys(Zs_dict, test_labels) - - train_val_idx_split = cv_inner.split(train_val_idx) - val_labels_all = [] - train_Zs = [] - train_data = [] - train_val_labels = [labels[tvidx] for tvidx in train_val_idx] - for train_idx, val_idx in train_val_idx_split: - val_labels = tuple([train_val_labels[vidx] for vidx in val_idx]) - val_labels_all.append(val_labels) - train_X = _pop_keys(train_val_X, val_labels) - train_Y = _pop_keys(train_val_Y, val_labels) - train_Z = _pop_keys(train_val_Z, val_labels) - train_Zs.append(train_Z) - train_data.append((train_X, train_Y)) - - # Run 5-fold inner CV for all eps - eps_prod, train_data_prod = zip(*product(epsilons, train_data)) - val_labels_prod = val_labels_all * len(epsilons) - train_Z_prod = train_Zs * len(epsilons) - print("Len eps", eps_prod) - if args.log_filepath is None and "VAE" not in args.method: - Ts_list = [] - logs = [] - # for i, (_train_data, _eps) in tqdm(enumerate(zip(train_data_prod, eps_prod))): - # print(f"{i}th iteration with {_eps}") - # _T, _log = ot_method_map[args.method](_train_data, _eps) - # Ts_list.append(_T) - # logs.append(_log) - # jax.clear_caches() - try: - with Pool(5 * len(epsilons)) as p: - # with Pool(5) as p: - Ts_list, logs = zip( - *p.starmap( - ot_method_map[args.method], zip(train_data_prod, eps_prod) - ) - ) - - except KeyboardInterrupt: - p.terminate() - finally: - p.join() - elif args.log_filepath is not None: - with open(args.log_filepath, "rb") as f: - prev_log = pkl.load(f) - Ts_list = [None] * len(train_data_prod) - logs = [None] * len(train_data_prod) - if "VAE" not in args.method: - epsilons = [ - e for e in epsilons if e in np.unique(list(prev_log["T"].keys())) - ] - eps_prod, train_data_prod = zip(*product(epsilons, train_data)) - val_labels_prod = val_labels_all * len(epsilons) - else: - Ts_list = [] - logs = [] - for train_data, eps in zip(train_data_prod, eps_prod): - t, l = ot_method_map[args.method](train_data, eps) - Ts_list.append(t) - logs.append(l) - - # Evaluate & select best eps - eps_to_pred_evals = {eps: [] for eps in epsilons} - eps_to_pred_full_evals = {eps: [] for eps in epsilons} - eps_to_matching_evals = {eps: [] for eps in epsilons} - eps_to_dfracs_evals = {eps: [] for eps in epsilons} - eps_to_val_to_T = {eps: {} for eps in epsilons} - eps_to_val_to_log = {eps: {} for eps in epsilons} - iters = zip(eps_prod, val_labels_prod, train_data_prod, Ts_list, logs) - assert len(list(iters)) == len(eps_prod) - print(f"Ts_list: (len {len(Ts_list)})") - for eps, val_labels, train_pair, Ts, train_Z, log in zip( - eps_prod, val_labels_prod, train_data_prod, Ts_list, train_Z_prod, logs - ): - if args.log_filepath is not None: - Ts = prev_log["T"][eps][val_labels] - log = prev_log["log"][eps][val_labels] - # Ts: Dict[Union[float,int], np.ndarray] for each label, transport plan - val_X_dict = {tidx: train_val_X[tidx] for tidx in val_labels} - val_Y_dict = {tidx: train_val_Y[tidx] for tidx in val_labels} - if args.pred_filepath is not None: - val_X_dict_full = {tidx: train_val_X_full[tidx] for tidx in val_labels} - val_Y_dict_full = {tidx: train_val_Y_full[tidx] for tidx in val_labels} - train_X, train_Y = train_pair - print(f"in the loop; {eps}, {val_labels}") - - if isinstance(Ts, int): - # COOT underflow - print("Underflow", eps, val_labels) - eps_to_matching_evals[eps].append(100) - for tidx in val_labels: - eps_to_pred_evals[eps].append( - pd.Series( - [np.nan, np.nan, np.nan, np.nan, np.nan], - index=[ - "Pearson_corr", - "Spearman_corr", - "Pearson_samples", - "Spearman_samples", - "MSE", - ], - ).to_frame() - ) - if args.pred_filepath is not None: - for tidx in val_labels: - eps_to_pred_full_evals[eps].append( - pd.Series( - [np.nan, np.nan, np.nan, np.nan, np.nan], - index=[ - "Pearson_corr", - "Spearman_corr", - "Pearson_samples", - "Spearman_samples", - "MSE", - ], - ).to_frame() - ) - eps_to_val_to_T[eps][val_labels] = Ts - eps_to_val_to_log[eps][val_labels] = {} - continue - # Evaluate matching - if "VAE" in args.method: - ks = [5, 10, 25, 50] - latent_Y = infer_from_Ys(train_Y, Ts, dim_X) - latent_X = infer_from_Xs(train_X, Ts, dim_Y) - _, mean_foscttm = get_FOSCTTM( - Ts, latent_X, latent_Y, use_agg="mean", use_barycenter=False - ) - - Ts_multK = get_Ts_from_nn_multKs(latent_X, latent_Y, ks) # k -> T - diag_fracs_dict = {} - for k, T_k in Ts_multK.items(): - dfracs, rel_dfracs = get_diag_fracs( - T_k, train_X, train_Y, train_Z, train_Z - ) - diag_fracs_dict[k] = rel_dfracs - eps_to_dfracs_evals[eps].append(diag_fracs_dict) - else: - _, mean_foscttm = get_FOSCTTM(Ts, train_X, train_Y, use_agg="mean") - dfracs, rel_dfracs = get_diag_fracs(Ts, train_X, train_Y, train_Z, train_Z) - - eps_to_dfracs_evals[eps].append(rel_dfracs) - eps_to_matching_evals[eps].append(mean_foscttm.mean()) - print("map", eps_to_matching_evals[eps]) - print("foscttm", mean_foscttm) - - # Evaluate prediction - for tidx in val_labels: - val_X = val_X_dict[tidx] - val_Y = val_Y_dict[tidx] - if "VAE" in args.method: - pred = predict_from_model(val_X, Ts, dim_Y) - else: - param = pred_method(train_X, train_Y, Ts) - pred = pred_from_param(val_X, param) - try: - df = get_evals( - val_Y, - pred, - prediction_id=(eps, val_labels), - full=False, - agg_method="mean", - ) - except: - df = pd.Series( - [np.nan, np.nan, np.nan, np.nan, np.nan], - index=[ - "Pearson_corr", - "Spearman_corr", - "Pearson_samples", - "Spearman_samples", - "MSE", - ], - ).to_frame() - eps_to_pred_evals[eps].append(df) - if args.pred_filepath is not None: - for tidx in val_labels: - train_X_full = _pop_keys(train_val_X_full, val_labels) - train_Y_full = _pop_keys(train_val_Y_full, val_labels) - val_X_full = val_X_dict_full[tidx] - val_Y_full = val_Y_dict_full[tidx] - if "VAE" in args.method: - pred = predict_from_model(val_X, Ts, dim_Y) - else: - param = pred_method(train_X_full, train_Y_full, Ts) - pred = pred_from_param(val_X_full, param) - eps_to_pred_full_evals[eps].append( - get_evals( - val_Y_full, - pred, - prediction_id=(eps, val_labels), - full=False, - agg_method="mean", - ) - ) - eps_to_val_to_T[eps][val_labels] = Ts - eps_to_val_to_log[eps][val_labels] = log - print(eps_to_matching_evals) - eps_to_matching_evals = pd.Series( - {k: np.nanmean(v) for k, v in eps_to_matching_evals.items()}, - ) - max_matching_eps = eps_to_matching_evals.idxmin() - eps_to_pred_evals = pd.DataFrame( - {k: pd.concat(v, axis=1).mean(axis=1) for k, v in eps_to_pred_evals.items()} - ) - max_pred_eps = eps_to_pred_evals.loc["MSE"].idxmin() - best_eps = {"matching": max_matching_eps, "pred": max_pred_eps} - if args.pred_filepath is not None: - eps_to_pred_full_evals = pd.DataFrame( - { - k: pd.concat(v, axis=1).mean(axis=1) - for k, v in eps_to_pred_full_evals.items() - } - ) - max_pred_full_eps = eps_to_pred_full_evals.loc["Pearson_samples"].idxmax() - best_eps["pred_full"] = eps_to_pred_full_evals.loc["Pearson_samples"].idxmax() - val_logs = { - "matching_evals": eps_to_matching_evals, - "dfracs": eps_to_dfracs_evals, - "pred_evals": eps_to_pred_evals, - "T": eps_to_val_to_T, - "log": eps_to_val_to_log, - "best_eps": best_eps, - } - if args.pred_filepath is not None: - val_logs["pred_full_evals"] = eps_to_pred_full_evals - - edited_label = "e." if args.log_filepath is not None else "" - with open( - f"val_CV_{args.method}.{args.test_idx}.best_eps.{edited_label}pkl", - "wb", - ) as f: - pkl.dump(best_eps, f) - - with open( - f"val_CV_{args.method}.{args.test_idx}.{edited_label}pkl", - "wb", - ) as f: - pkl.dump(val_logs, f) - - -def run_mlp(args): - epsilons = [e for e in ot_method_hyperparams[args.method] if e != 0.1 and e != 1e-6] - with open(args.log_filepath, "rb") as f: - log = pkl.load(f) - with open(args.filepath, "rb") as f: - data_dict = pkl.load(f) - - X_dict = data_dict["Xs_dict"] - Y_dict = data_dict["Xt_dict"] - Z_dict = data_dict["Zs_dict"]["dosage"] - - labels = list(X_dict.keys()) - dim_X = X_dict[labels[0]].shape[1] - dim_Y = Y_dict[labels[0]].shape[1] - cv = KFold(n_splits=5) - train_val_idx, test_idx = list(cv.split(labels))[args.test_idx] - cv_inner = KFold(n_splits=5) - - test_labels = [labels[tidx] for tidx in test_idx] - train_val_X = _pop_keys(X_dict, test_labels) - train_val_Y = _pop_keys(Y_dict, test_labels) - train_val_Z = _pop_keys(Z_dict, test_labels) - - train_val_idx_split = cv_inner.split(train_val_idx) - val_labels_all = [] - train_data = [] - val_data = [] - train_val_labels = [labels[tvidx] for tvidx in train_val_idx] - for train_idx, val_idx in train_val_idx_split: - val_labels = tuple([train_val_labels[vidx] for vidx in val_idx]) - val_labels_all.append(val_labels) - train_X = _pop_keys(train_val_X, val_labels) - train_Y = _pop_keys(train_val_Y, val_labels) - val_X = np.concatenate([train_val_X[l] for l in val_labels], axis=0) - val_Y = np.concatenate([train_val_Y[l] for l in val_labels], axis=0) - val_Z = np.concatenate([train_val_Z[l] for l in val_labels], axis=0) - train_data.append((train_X, train_Y)) - val_data.append((val_X, val_Y, val_Z)) - - # Run 5-fold inner CV for all eps - eps_prod, train_data_prod = zip(*product(epsilons, train_data)) - val_labels_prod = val_labels_all * len(epsilons) - val_data_prod = val_data * len(epsilons) - Ts_list = [] - for eps, val_labels in zip(eps_prod, val_labels_prod): - Ts_list.append(log["T"][eps][val_labels]) - - trained_models = [] - logs = [] - for train_data, Ts in zip(train_data_prod, Ts_list): - t, l = train_mlp(train_data, Ts) - trained_models.append(t) - logs.append(l) - - # Evaluate & select best eps - eps_to_pred_evals = {eps: [] for eps in epsilons} - eps_to_val_to_preds = {eps: {} for eps in epsilons} - eps_to_matching_evals = {eps: [] for eps in epsilons} - eps_to_val_to_model = {eps: {} for eps in epsilons} - eps_to_val_to_log = {eps: {} for eps in epsilons} - iters = zip(val_data_prod, train_data_prod, trained_models) - assert len(list(iters)) == len(eps_prod) - print(f"Ts_list: (len {len(Ts_list)})") - for eps, val_labels, val_pair, model in zip( - eps_prod, val_labels_prod, val_data_prod, trained_models - ): - val_X, val_Y, val_Z = val_pair - print(f"in the loop; {val_labels}") - - # Evaluate prediction - Y_pred = model(torch.tensor(val_X)).detach().numpy() - eval_df = get_evals( - val_Y, - Y_pred, - prediction_id=(eps, val_labels), - full=False, - agg_method="mean", - ) - print(eval_df) - eps_to_pred_evals[eps] = eps_to_pred_evals[eps] + [eval_df] - eps_to_val_to_preds[eps][val_labels] = (Y_pred, val_Y, val_Z) - eps_to_val_to_model[eps][val_labels] = model - - def concat_if_possible(v): - if len(v) > 0: - return pd.concat(v, axis=1).mean(axis=1) - else: - return v - - eps_to_pred_evals = pd.DataFrame( - {k: concat_if_possible(v) for k, v in eps_to_pred_evals.items()} - ) - # max_pred_eps = eps_to_pred_evals.loc["Pearson_c"].idxmax() - val_logs = { - "pred": eps_to_val_to_preds, - "pred_evals": eps_to_pred_evals, - "model": eps_to_val_to_model, - } - # with open(f"val_CV_MLP_{args.method}.{args.test_idx}.best_eps.pkl", "wb") as f: - # pkl.dump({"eps": max_pred_eps}, f) - - with open(f"val_CV_MLP_{args.method}.{args.test_idx}.pkl", "wb") as f: - pkl.dump(val_logs, f) - - -def run_mlp_control(args): - with open(args.filepath, "rb") as f: - data_dict = pkl.load(f) - - X_dict = data_dict["Xs_dict"] - Y_dict = data_dict["Xt_dict"] - Z_dict = data_dict["Zs_dict"]["dosage"] - - labels = list(X_dict.keys()) - dim_X = X_dict[labels[0]].shape[1] - dim_Y = Y_dict[labels[0]].shape[1] - cv = KFold(n_splits=5) - train_val_idx, test_idx = list(cv.split(labels))[args.test_idx] - cv_inner = KFold(n_splits=5) - - test_labels = [labels[tidx] for tidx in test_idx] - train_val_X = _pop_keys(X_dict, test_labels) - train_val_Y = _pop_keys(Y_dict, test_labels) - train_val_Z = _pop_keys(Z_dict, test_labels) - - train_val_idx_split = cv_inner.split(train_val_idx) - val_labels_all = [] - train_data = [] - val_data = [] - train_Z = [] - train_val_labels = [labels[tvidx] for tvidx in train_val_idx] - for train_idx, val_idx in train_val_idx_split: - val_labels = tuple([train_val_labels[vidx] for vidx in val_idx]) - val_labels_all.append(val_labels) - train_X = _pop_keys(train_val_X, val_labels) - train_Y = _pop_keys(train_val_Y, val_labels) - val_X = np.concatenate([train_val_X[l] for l in val_labels], axis=0) - val_Y = np.concatenate([train_val_Y[l] for l in val_labels], axis=0) - val_Z = np.concatenate([train_val_Z[l] for l in val_labels], axis=0) - train_data.append((train_X, train_Y)) - train_Z.append(train_val_Z) - val_data.append((val_X, val_Y, val_Z)) - - # Make baseline Ts - perf_T = {k: np.diag(np.ones(v.shape[0])) for k, v in X_dict.items()} - - def make_G(size, label, k): - assert size == len(label), (k, size, len(label)) - G = np.zeros((size, size)) - for l in np.unique(label): - l_idx = np.where(label == l)[0] - for i in l_idx: - for j in l_idx: - G[i, j] = 1.0 - assert (G.sum(axis=0) > 0).all(), (k, (G.sum(axis=0) == 0).sum()) - return G - - cond_T = {k: make_G(X_dict[k].shape[0], Z_dict[k], k) for k in X_dict.keys()} - rand_T = {k: np.ones((v.shape[0], v.shape[0])) for k, v in X_dict.items()} - Ts_list = [perf_T, cond_T, rand_T] - Ts_prod, train_data_prod = zip(*product(Ts_list, train_data)) - Ts_label = ["perfect", "dosage", "random"] - Ts_label_prod = np.repeat(Ts_label, len(train_data)) - train_Z_prod = train_Z * len(Ts_label) - val_labels_prod = val_labels_all * len(Ts_label) - val_data_prod = val_data * len(Ts_label) - - trained_models = [] - logs = [] - for train_data, val_labels, Ts, tlab in zip( - train_data_prod, val_labels_prod, Ts_prod, Ts_label_prod - ): - T = _pop_keys(Ts, test_labels) - T = _pop_keys(T, val_labels) - try: - t, l = train_mlp(train_data, T) - except Exception as e: - print(f"error at {tlab}, {val_labels}") - raise e - trained_models.append(t) - logs.append(l) - - tlab_to_pred_evals = {tlab: [] for tlab in Ts_label} - tlab_to_matching_evals = {tlab: [] for tlab in Ts_label} - tlab_to_val_to_model = {tlab: {} for tlab in Ts_label} - tlab_to_val_to_log = {tlab: {} for tlab in Ts_label} - tlab_to_val_to_pred = {tlab: {} for tlab in Ts_label} - iters = zip(val_data_prod, train_data_prod, trained_models) - assert len(list(iters)) == len(Ts_label_prod) - print(f"Ts_list: (len {len(Ts_list)})") - for Ts, tlab, val_labels, val_pair, train_data, model, train_Z, log in zip( - Ts_prod, - Ts_label_prod, - val_labels_prod, - val_data_prod, - train_data_prod, - trained_models, - train_Z_prod, - logs, - ): - val_X, val_Y, val_Z = val_pair - train_X, train_Y = train_data - print(f"in the loop; {val_labels}") - # Evaluate matching - T = _pop_keys(Ts, test_labels) - T = _pop_keys(T, val_labels) - try: - _, mean_foscttm = get_FOSCTTM(T, train_X, train_Y, use_agg="mean") - except KeyError as e: - print(Ts.keys()) - print(T.keys()) - print(train_X.keys()) - raise e - dfracs, rel_dfracs = get_diag_fracs(T, train_X, train_Y, train_Z, train_Z) - tlab_to_matching_evals[tlab] = { - "foscttm": mean_foscttm, - "dfracs": dfracs, - "rel_dfracs": rel_dfracs, - } - # Evaluate prediction - Y_pred = model(torch.tensor(val_X)).detach().numpy() - eval_df = get_evals( - val_Y, - Y_pred, - prediction_id=(tlab, val_labels), - full=False, - agg_method="mean", - norm_Y=Y_dict[3].mean(axis=0), - ) - tlab_to_val_to_log[tlab] = log - tlab_to_val_to_pred[tlab][val_labels] = (val_Y, Y_pred, val_Z) - print(eval_df) - tlab_to_pred_evals[tlab] = tlab_to_pred_evals[tlab] + [eval_df] - tlab_to_val_to_model[tlab][val_labels] = model - - def concat_if_possible(v): - if len(v) > 0: - return pd.concat(v, axis=1).mean(axis=1) - else: - return v - - tlab_to_pred_evals = pd.DataFrame( - {k: concat_if_possible(v) for k, v in tlab_to_pred_evals.items()} - ) - val_logs = { - "pred_evals": tlab_to_pred_evals, - "match_evals": tlab_to_matching_evals, - "model": tlab_to_val_to_model, - "logs": tlab_to_val_to_log, - "preds": tlab_to_val_to_pred, - } - - with open(f"val_CV_MLP_baseline.{args.test_idx}.pkl", "wb") as f: - pkl.dump(val_logs, f) - - -if __name__ == "__main__": - args = parse_args() - if args.mlp: - if args.baseline: - run_mlp_control(args) - else: - if args.log_filepath is None: - args.log_filepath = f"val_CV_{args.method}.{args.test_idx}.pkl" - run_mlp(args) - else: - main(args) diff --git a/perturbot/build/lib/perturbot/eval/cv_outer_loop.py b/perturbot/build/lib/perturbot/eval/cv_outer_loop.py deleted file mode 100755 index 8acbe10..0000000 --- a/perturbot/build/lib/perturbot/eval/cv_outer_loop.py +++ /dev/null @@ -1,325 +0,0 @@ -"""For 5-fold CV, run the single inner loop of train-validation to select hyperparameter.""" - -from typing import Dict, Any -import sys -import argparse -import pickle as pkl -from functools import partial -from itertools import product -from multiprocessing import Pool -import numpy as np -import pandas as pd -from sklearn.model_selection import KFold -import torch - -from perturbot.eval.utils import _pop_keys, _pop_key, get_Ts_from_nn_multKs, make_G -from perturbot.eval.match import get_FOSCTTM, get_diag_fracs -from perturbot.eval.prediction import get_evals_preds, get_evals -from perturbot.preprocess.vae import ( - train_vae_rna, - train_vae_acc, - train_vae_prot, - SCVI_LATENT_KEY, -) -from perturbot.match.cot_labels import get_coupling_cotl, get_coupling_cotl_sinkhorn -from perturbot.match.gw import ( - get_coupling_gw_cg, - get_coupling_egw, - get_coupling_egw_all, - get_coupling_gw_all, -) -from perturbot.match.ott_egwl import ( - get_coupling_egw_labels_ott, - get_coupling_egw_all_ott, - get_coupling_eot_ott, - get_coupling_leot_ott, - get_coupling_egw_ott, -) -from perturbot.match.cot import get_coupling_cot, get_coupling_cot_sinkhorn -from perturbot.match.gw_labels import get_coupling_gw_labels, get_coupling_egw_labels -from perturbot.predict.scvi_vae import train_vae_model -from perturbot.predict.linear_regression import ( - ols_normed, - weight_conc_normed, - weighted_ols_normed, - weight_1_ols_normed, - predict, -) -from perturbot.predict.mlp import train_mlp -from perturbot.predict.scvi_vae import predict_from_model, infer_from_Xs, infer_from_Ys - -ot_method_map = { - "ECOOTL": get_coupling_cotl_sinkhorn, - "ECOOT": get_coupling_cot_sinkhorn, - "EGWL": get_coupling_egw_labels, - "EOT_ott": get_coupling_eot_ott, - "LEOT_ott": get_coupling_leot_ott, - "EGW_ott": get_coupling_egw_ott, - "EGW_all_ott": get_coupling_egw_all_ott, - "EGWL_ott": get_coupling_egw_labels_ott, - "VAE_label": train_vae_model, - "VAE": partial(train_vae_model, use_label=False), -} - - -def parse_args(): - parser = argparse.ArgumentParser( - "Run inner validation loop of CV", - "Run sample-matching OT in parallel and fit prediction model in leave-one-out mannter", - ) - parser.add_argument("method", type=str) - parser.add_argument("test_idx", type=int) - parser.add_argument("filepath", type=str) - parser.add_argument("eps", type=str) - parser.add_argument( - "-p", - "--pred-filepath", - type=str, - default=None, - help="Path to data.pkl file with full features.", - ) - parser.add_argument( - "-l", - "--log-filepath", - type=str, - default=None, - help="Path to log.pkl if transport was already run. If not None, reevaluates but does not calculates the transport again.", - ) - parser.add_argument( - "-b", - "--baseline", - type=str, - default=None, - help="One of perfect, random, by_conc", - ) - return parser.parse_args() - - -ot_method_all_to_all = ["GW_all", "EGW_all_ott", "EOT_all_ott"] -pred_from_param = predict - - -def main(args): - if args.log_filepath is not None and args.log_filepath != "None": - with open(args.log_filepath, "rb") as f: - prev_logs = pkl.load(f) - logs = { - "matching_evals": {}, - "pred_evals": {}, - "T": {}, - "pred": {}, - "log": {}, - "eps": {}, - } - try: - match_eps, lin_eps, pred_eps = tuple(map(float, args.eps.split(","))) - logs["eps"]["match"] = match_eps - logs["eps"]["lin"] = lin_eps - logs["eps"]["pred"] = pred_eps - all_to_all = args.method in ot_method_all_to_all - with open(args.filepath, "rb") as f: - data_dict = pkl.load(f) - if args.pred_filepath is not None: - with open(args.pred_filepath, "rb") as f: - pred_data_dict = pkl.load(f) - - X_dict = data_dict["Xs_dict"] - Y_dict = data_dict["Xt_dict"] - Zs_dict = data_dict["Zs_dict"]["dosage"] - Zt_dict = data_dict["Zt_dict"]["dosage"] - - labels = list(X_dict.keys()) - dim_X = pred_data_dict["Xs_dict"][labels[0]].shape[1] - dim_Y = pred_data_dict["Xt_dict"][labels[0]].shape[1] - cv = KFold(n_splits=5) - train_idx, test_idx = list(cv.split(labels))[args.test_idx] - - test_labels = [labels[tidx] for tidx in test_idx] - train_X = _pop_keys(X_dict, test_labels) - train_Y = _pop_keys(Y_dict, test_labels) - train_Z = _pop_keys(Zs_dict, test_labels) - if args.pred_filepath is not None: - train_X_full = _pop_keys(pred_data_dict["Xs_dict"], test_labels) - train_Y_full = _pop_keys(pred_data_dict["Xt_dict"], test_labels) - train_Z = _pop_keys(Zs_dict, test_labels) - train_data = (train_X, train_Y) - train_data_full = (train_X_full, train_Y_full) - print(f"Calculating matching with {match_eps}") - - # Get matching - if args.log_filepath is not None and args.log_filepath != "None": - Ts_matching = prev_logs["T"]["match"] - Ts_pred = prev_logs["T"]["pred"] - try: - log_matching = prev_logs["log"]["match"] - except: - log_matching = None - try: - log_matching_predeps = prev_logs["log"]["match_pred"] - except: - log_matching_predeps = None - elif args.baseline is not None: - if args.baseline == "perfect": - Ts_matching = Ts_pred = { - k: np.diag(np.ones(v.shape[0])) for k, v in train_X.items() - } - elif args.baseline == "random": - Ts_matching = Ts_pred = { - k: np.ones((v.shape[0], v.shape[0])) for k, v in train_X.items() - } - elif args.baseline == "by_conc": - Ts_matching = Ts_pred = { - k: make_G(train_X[k].shape[0], train_Z[k], k) - for k in train_X.keys() - } - log_matching = None - log_matching_predeps = None - else: - if "VAE" in args.method: - Ts_matching, log_matching = ot_method_map[args.method]( - train_data_full, match_eps - ) - else: - Ts_matching, log_matching = ot_method_map[args.method]( - train_data, match_eps - ) - - if "VAE" in args.method: - if match_eps == pred_eps: - Ts_pred = Ts_matching - log_matching_predeps = None - else: - Ts_pred, log_matching_predeps = ot_method_map[args.method]( - train_data_full, pred_eps - ) - else: - if match_eps != pred_eps: - print(f"Calculating pred matching with {pred_eps}") - Ts_pred, log_matching_predeps = ot_method_map[args.method]( - train_data, pred_eps - ) - else: - Ts_pred = Ts_matching - log_matching_predeps = None - logs["T"]["match"] = Ts_matching - logs["T"]["pred"] = Ts_pred - - if "VAE" in args.method: - latent_Y = infer_from_Ys(train_Y_full, Ts_matching, dim_X) - latent_X = infer_from_Xs(train_X_full, Ts_matching, dim_Y) - _, mean_foscttm = get_FOSCTTM( - Ts_matching, - latent_X, - latent_Y, - use_agg="mean", - use_barycenter=False, - ) - ks = [1, 5, 10, 50, 100] - k_to_Ts = get_Ts_from_nn_multKs(latent_X, latent_Y, ks) # k -> T - dfracs = {} - rel_dfracs = {} - for k, Ts in k_to_Ts.items(): - dfracs[k], rel_dfracs[k] = get_diag_fracs( - Ts, train_X, train_Y, train_Z, train_Z - ) - - else: - # normalize - if isinstance(Ts_matching, dict): - total_sum = 0 - for k, v in Ts_matching.items(): - total_sum += v.sum() - Ts_matching = { - k: v.astype(np.double) / total_sum for k, v in Ts_matching.items() - } - else: - Ts_matching = Ts_matching.astype(np.double) / Ts_matching.sum() - _, mean_foscttm = get_FOSCTTM(Ts_matching, train_X, train_Y, use_agg="mean") - dfracs, rel_dfracs = get_diag_fracs( - Ts_matching, train_X, train_Y, train_Z, train_Z - ) - - # if not all_to_all: - mean_mean_foscttm = mean_foscttm.mean() - # else: - # mean_mean_foscttm = mean_foscttm - print("foscttm", mean_foscttm) - logs["matching_evals"] = ( - { - "foscttm": mean_foscttm, - "mean_foscttm": mean_mean_foscttm, - "dfracs": dfracs, - "rel_dfracs": rel_dfracs, - }, - ) - # Eval prediction: full features - test_X_full = np.concatenate( - [pred_data_dict["Xs_dict"][l] for l in test_labels], axis=0 - ) - test_Y_full = np.concatenate( - [pred_data_dict["Xt_dict"][l] for l in test_labels], axis=0 - ) - if "VAE" not in args.method: - trained_model, log_pred = train_mlp((train_X_full, train_Y_full), Ts_pred) - Y_pred_full = trained_model(torch.tensor(test_X_full)).detach().numpy() - else: - Y_pred_full = predict_from_model(test_X_full, Ts_pred, dim_Y) - logs["pred"] = { - "model": trained_model if "VAE" not in args.method else None, - "Y_pred": Y_pred_full, - "Y_true": test_Y_full, - "train_Y": train_Y_full, - "test_Z": {k: Zs_dict[k] for k in test_labels}, - } - eval_df_full = get_evals( - test_Y_full, - Y_pred_full, - prediction_id="eval", - full=False, - agg_method="mean", - ) - print(eval_df_full) - logs["pred_evals"]["full"] = eval_df_full - - # # Eval prediction - PC - # test_X = np.concatenate([X_dict[l] for l in test_labels], axis=0) - # test_Y = np.concatenate([Y_dict[l] for l in test_labels], axis=0) - # if "VAE" in args.method: - # pred_Ys = predict_from_model(test_X, Ts_lin, dim_Y) - - # else: - # params = [ - # lin_method(train_X, train_Y, Ts_lin) - # for lin_method in [ols_normed, weight_1_ols_normed, weight_conc_normed] - # ] - # pred_Ys = [pred_from_param(test_X, param) for param in params] - # pred_labels = ["ot", "perfect", "random", "by_conc"] - # eval_df = get_evals_preds( - # test_Y, - # pred_Ys, - # pred_labels=pred_labels, - # full=False, - # ) - # logs["pred_evals"]["PC"] = eval_df - # print(eval_df) - - logs["log"] = ( - { - "match": log_matching, - "match_pred": log_matching_predeps, - # "match_lin": log_matching_lin, - "mlp": log_pred if "VAE" not in args.method else None, - }, - ) - except Exception as e: - with open(f"test_{args.method}.{args.test_idx}.tmp.pkl", "wb") as f: - pkl.dump(logs, f) - raise e - - with open(f"test_{args.method}.{args.test_idx}.e.pkl", "wb") as f: - pkl.dump(logs, f) - - -if __name__ == "__main__": - args = parse_args() - main(args) diff --git a/perturbot/build/lib/perturbot/eval/feature_matching.py b/perturbot/build/lib/perturbot/eval/feature_matching.py deleted file mode 100755 index 2bf61d7..0000000 --- a/perturbot/build/lib/perturbot/eval/feature_matching.py +++ /dev/null @@ -1,160 +0,0 @@ -"""For 5-fold CV, run the single inner loop of train-validation to select hyperparameter.""" - -import os -import argparse -import pickle as pkl -from functools import partial -import numpy as np -from perturbot.match.fot import get_coupling_fot -from perturbot.eval.utils import make_G -from perturbot.predict.scvi_vae import train_vae_model, infer_from_Xs, infer_from_Ys -from perturbot.eval.utils import get_Ts_from_nn_multKs - - -def parse_args(): - parser = argparse.ArgumentParser( - "Run inner validation loop of CV", - "Run sample-matching OT in parallel and fit prediction model in leave-one-out mannter", - ) - parser.add_argument("method", type=str) - parser.add_argument("filepath", type=str) - parser.add_argument("best_eps", type=float) - parser.add_argument("eps", type=float) - - parser.add_argument( - "-l", - "--log-filepath", - type=str, - default=None, - help="Path to log.pkl if transport was already run. If not None, reevaluates but does not calculates the transport again.", - ) - parser.add_argument( - "-b", - "--baseline", - type=str, - default=None, - help="One of perfect, random, by_conc", - ) - parser.add_argument( - "--best-k", - type=int, - default=None, - help="Best k to use for VAE", - ) - return parser.parse_args() - - -ot_method_all_to_all = ["GW_all", "EGW_all_ott", "EOT_all_ott"] - - -def main(args): - logs = { - "matching_evals": {}, - "pred_evals": {}, - "T": {}, - "pred": {}, - "log": {}, - "eps": {}, - } - try: - if args.best_eps != 0: - with open(f"all_{args.method}.{args.best_eps}.pkl", "rb") as f: - d = pkl.load(f) - Ts = d["T"] - logs["sample_eps"] = args.best_eps - # logs['min_foscttm'] = d[] - else: - Ts = None - with open(args.filepath, "rb") as f: - data_dict = pkl.load(f) # full features - - X_dict = data_dict["Xs_dict"] - Y_dict = data_dict["Xt_dict"] - Z_dict = data_dict["Zs_dict"]["dosage"] - - if "VAE" in args.method: - labels = list(X_dict.keys()) - dim_X = data_dict["Xs_dict"][labels[0]].shape[1] - dim_Y = data_dict["Xt_dict"][labels[0]].shape[1] - latent_Y = infer_from_Ys(Y_dict, Ts, dim_X) - latent_X = infer_from_Xs(X_dict, Ts, dim_Y) - Ts = get_Ts_from_nn_multKs(latent_X, latent_Y, [args.best_k])[args.best_k] - - if Ts is None: - if args.method == "random": - Ts = { - k: np.ones((X_dict[k].shape[0], Y_dict[k].shape[0])) - / (X_dict[k].shape[0] * Y_dict[k].shape[0]) - for k in X_dict.keys() - } - elif args.method == "perfect": - Ts = { - k: np.diag( - np.ones(X_dict[k].shape[0]), - ) - / X_dict[k].shape[0] - for k in X_dict.keys() - } - elif args.method == "by_conc": - Ts = { - k: make_G(X_dict[k].shape[0], Z_dict[k], k) for k in X_dict.keys() - } - Tv, log = get_coupling_fot((X_dict, Y_dict), Ts, args.eps) - logs["log"] = log - logs["Tv"] = Tv - except Exception as e: - with open(f"features_{args.method}.{args.eps}.tmp.pkl", "wb") as f: - pkl.dump(logs, f) - raise e - - with open(f"features_{args.method}.{args.eps}.pkl", "wb") as f: - pkl.dump(logs, f) - - -def submit_feature_run(data_path, ot_method_label): - epsilons = [1e-2, 1e-3, 1e-4, 1e-5] - rel_dfracs = [] - if ot_method_label in ["perfect", "random", "by_conc"]: - best_eps = 0 - else: - best_k_dict = {} - for eps in epsilons: - # try: - with open(f"all_{ot_method_label}.{eps}.pkl", "rb") as f: - d = pkl.load(f) - _rel_dfracs = d["matching_evals"][0]["rel_dfracs"] - if isinstance(_rel_dfracs, dict): - max_rel_dfracs = -10 - for k, v in _rel_dfracs.items(): - if max_rel_dfracs < v: - max_rel_dfracs = v - best_k_dict[eps] = k - _rel_dfracs = max_rel_dfracs - rel_dfracs.append(_rel_dfracs) - # except Exception as e: - # print(e) - # rel_dfracs.append(-1) - best_eps = epsilons[rel_dfracs.index(max(rel_dfracs))] - if "VAE" in ot_method_label: - best_k = best_k_dict[best_eps] - for eps in epsilons: - run_label = f"FM.{ot_method_label}.{eps}" - f = open(f"{run_label}.bsub", "w") - f.write("source ~/.bashrc\n") - f.write("pwd\n") - f.write("conda activate ot \n") - command = f"python /gpfs/scratchfs01/site/u/ryuj6/OT/software/perturbot/perturbot/eval/feature_matching.py {ot_method_label} {data_path} {best_eps} {eps}" - - if "VAE" in ot_method_label: - command += f" --best-k {best_k}" - f.write(f"echo {command}\n") - f.write(f"{command}\n") - f.close() - command = f"bsub -M 10G -n 1 -q short -J {run_label} -o log/log.{run_label}.o%J -e log/log.{run_label}.e%J < {run_label}.bsub" - print(command) - os.system(command) - - -if __name__ == "__main__": - args = parse_args() - main(args) diff --git a/perturbot/build/lib/perturbot/eval/loo.py b/perturbot/build/lib/perturbot/eval/loo.py deleted file mode 100755 index a3d0e28..0000000 --- a/perturbot/build/lib/perturbot/eval/loo.py +++ /dev/null @@ -1,432 +0,0 @@ -from typing import List, Dict -import numpy as np -import pandas as pd -from sklearn.model_selection import LeaveOneOut -from .prediction import get_evals_preds -from multiprocessing import Pool -import pickle as pkl -from .match import get_FOSCTTM, get_diag_fracs -from .utils import get_Ts_from_nn_multKs -from perturbot.predict.scvi_vae import predict_from_model, infer_from_Xs, infer_from_Ys -from perturbot.preprocess.vae import ( - train_vae_rna, - train_vae_acc, - train_vae_prot, - SCVI_LATENT_KEY, -) -from .utils import _pop_key - -np.seterr(all="raise") - -vae_trainer_dict = { - "rna": train_vae_rna, - "accessibility": train_vae_acc, - "protein": train_vae_prot, -} - - -def run_models( - X_dict, - Y_dict, - ot_method, - pred_method, - pred_method_label: str, - baseline_pred_methods: List, - baseline_pred_method_labels: List[str], - pred_from_param, - Ts_dict=None, - Z_dict=None, - *args, - **kwargs, -): - if pred_method_label == "VAE": - return run_models_vae( - X_dict, - Y_dict, - ot_method, - Ts_dict=None, - *args, - **kwargs, - ) - labels = list(X_dict.keys()) - loo = LeaveOneOut() - log = {"ot_couplings": {}, "params": {}, "preds": {}, "logs": {}} - eval_dfs = [] - train_data = [] - test_data = {} - - for i, (train_idx, test_idx) in enumerate(loo.split(labels)): - test_idx = test_idx.item() - test_X = X_dict[test_idx] - test_Y = Y_dict[test_idx] - train_X = X_dict.copy() - train_Y = Y_dict.copy() - del train_X[test_idx] - del train_Y[test_idx] - train_data.append((train_X, train_Y)) - test_data[test_idx] = (test_X, test_Y) - if Ts_dict is None: - try: - with Pool(10) as p: - Ts_list, logs = zip(*p.map(ot_method, train_data)) - except KeyboardInterrupt: - p.terminate() - finally: - p.join() - - for i, (test_idx, (test_X, test_Y)) in enumerate(test_data.items()): - if Ts_dict is None: - Ts = Ts_list[i] - log["logs"][i] = logs[i] - else: - Ts = Ts_dict[test_idx] - log["ot_couplings"][test_idx] = Ts - train_X = X_dict.copy() - train_Y = Y_dict.copy() - if Z_dict is not None: - train_Z = Z_dict.copy() - del train_X[test_idx] - del train_Y[test_idx] - params = [pred_method(train_X, train_Y, Ts)] - if Z_dict is not None: - params += [ - baseline_method(train_X, train_Y, train_Z) - for baseline_method in baseline_pred_methods - ] - else: - params += [ - baseline_method(train_X, train_Y) - for baseline_method in baseline_pred_methods - ] - log["params"][test_idx] = params - - preds = [pred_from_param(test_X, param) for param in params] - log["preds"][test_idx] = preds - method_labels = [pred_method_label] + baseline_pred_method_labels - assert len(preds) == len(method_labels) - eval_df = get_evals_preds(test_Y, preds, method_labels) - eval_df["loo_test_idx"] = test_idx - eval_dfs.append(eval_df) - - return pd.concat(eval_dfs, axis=0), log - - -def run_models_vae( - X_dict, - Y_dict, - ot_method, - Ts_dict=None, - *args, - **kwargs, -): - labels = list(X_dict.keys()) - loo = LeaveOneOut() - log = { - "models": {}, - "params": {}, - "preds": {}, - "latent_X": {}, - "latent_Y": {}, - "logs": {}, - } - eval_dfs = [] - train_data = [] - test_data = {} - - for i, (train_idx, test_idx) in enumerate(loo.split(labels)): - test_idx = test_idx.item() - test_X = X_dict[test_idx] - test_Y = Y_dict[test_idx] - train_X = X_dict.copy() - train_Y = Y_dict.copy() - del train_X[test_idx] - del train_Y[test_idx] - train_data.append((train_X, train_Y)) - test_data[test_idx] = (test_X, test_Y) - if Ts_dict is None: - try: - with Pool(10) as p: - model_list, logs = zip( - *p.map( - ot_method, - train_data, - ) - ) - except KeyboardInterrupt: - p.terminate() - finally: - p.join() - ks = [5, 10, 50, 100] - for k in ks: - log[f"pred_T_k{k}"] = {} # test_idx -> train_idx -> T - for i, (test_idx, (test_X, test_Y)) in enumerate(test_data.items()): - if Ts_dict is None: - model = model_list[i] - log["models"][test_idx] = model - log["logs"][test_idx] = logs[i] - else: - model = Ts_dict[test_idx] - log["models"][test_idx] = model - dX = test_X.shape[1] - dY = test_Y.shape[1] - latent_Y = infer_from_Ys(_pop_key(Y_dict, test_idx), model, dX) - latent_X = infer_from_Xs(_pop_key(X_dict, test_idx), model, dY) - pred_Y = predict_from_model(test_X, model, test_Y.shape[1]) - log["preds"][test_idx] = pred_Y - log["latent_X"][test_idx] = latent_X - log["latent_Y"][test_idx] = latent_Y - Ts = get_Ts_from_nn_multKs(latent_X, latent_Y, ks) # k -> T - for k, v in Ts.items(): - log[f"pred_T_k{k}"][test_idx] = v - eval_df = get_evals_preds(test_Y, [pred_Y], ["VAE"]) - eval_df["loo_test_idx"] = test_idx - eval_dfs.append(eval_df) - - return pd.concat(eval_dfs, axis=0), log - - -def run_models_vae_then_ot( - source_adata, - target_adata, - source_modality, - target_modality, - label_col, - ot_method, - pred_method, - pred_method_label: str, - baseline_pred_methods: List, - baseline_pred_method_labels: List[str], - pred_from_param, - T_dict=None, - model_dict=None, - Z_dict=None, - n_threads=1, - *args, - **kwargs, -): - labels = source_adata.obs[label_col].unique().values - loo = LeaveOneOut() - log = {"ot_couplings": {}, "params": {}, "preds": {}, "logs": {}} - eval_dfs = [] - test_data = {} - source_adatas_train = [] - target_adatas_train = [] - for i, (train_idx, test_idx) in enumerate(loo.split(labels)): - test_idx = test_idx.item() - source_adata_test = source_adata.obs[label_col == test_idx].copy() - target_adata_test = target_adata.obs[label_col == test_idx].copy() - source_adata_train = source_adata.obs[label_col != test_idx].copy() - target_adata_train = target_adata.obs[label_col != test_idx].copy() - source_adatas_train.append(source_adata_train) - target_adatas_train.append(target_adata_train) - test_data[test_idx] = (source_adata_test, target_adata_test) - - # Train VAE - if model_dict is None: - try: - with Pool(n_threads) as p: - source_adatas, source_models = zip( - *p.map(vae_trainer_dict[source_modality], source_adatas_train) - ) - except KeyboardInterrupt: - p.terminate() - finally: - p.join() - try: - with Pool(n_threads) as p: - target_adatas, target_models = zip( - *p.map(vae_trainer_dict[target_modality], target_adatas_train) - ) - except KeyboardInterrupt: - p.terminate() - finally: - p.join() - - # Train OT - train_data = [] - for i, (test_idx, (test_X, test_Y)) in enumerate(test_data.items()): - train_X = source_adatas[i].obsm[SCVI_LATENT_KEY] - train_Y = target_adatas[i].obsm[SCVI_LATENT_KEY] - train_data.append((train_X, train_Y)) - - if T_dict is None: - try: - with Pool(1) as p: - T_list, logs = zip(*p.map(ot_method, train_data)) - except KeyboardInterrupt: - p.terminate() - finally: - p.join() - - for i, (test_idx, (test_X, test_Y)) in enumerate(test_data.items()): - if T_dict is None: - Ts = T_list[i] - log["logs"][i] = logs[i] - else: - Ts = T_dict[test_idx] - log["ot_couplings"][test_idx] = Ts - - # @TODO see from here - pred_Y = predict_from_model_with_OT( - test_X, source_models[i], target_models[i], Ts, test_Y.shape[1] - ) - - preds = [pred_from_param(test_X, param) for param in params] - log["preds"][test_idx] = preds - method_labels = [pred_method_label] + baseline_pred_method_labels - assert len(preds) == len(method_labels) - eval_df = get_evals_preds(test_Y, preds, method_labels) - eval_df["loo_test_idx"] = test_idx - eval_dfs.append(eval_df) - - return pd.concat(eval_dfs, axis=0), log - - -def evaluate_loo_vae( - data_dict, logfile_name, test_idx_file, use_z_key="dosage", ks=[5, 10, 50, 100] -): - Xs_dict = data_dict["Xs_dict"] - Xt_dict = data_dict["Xt_dict"] - Zs_dict = data_dict["Zs_dict"] - Zt_dict = data_dict["Zt_dict"] - with open(logfile_name, "rb") as f: - log = pkl.load(f) - dfracs = {k: [] for k in ks} - rel_dfracs = {k: [] for k in ks} - med_foscttms = [] - idx_name = pd.read_csv(test_idx_file, header=None) - for test_idx in log["preds"].keys(): - model = log["models"] - Xs_train = _pop_key(Xs_dict, test_idx) - Xt_train = _pop_key(Xt_dict, test_idx) - for k in ks: - dfrac, rel_dfrac = get_diag_fracs( - log[f"pred_T_k{k}"][test_idx], - Xs_train, - Xt_train, - Zs_dict[use_z_key], - Zt_dict[use_z_key], - ) - dfracs[k].append(dfrac) - rel_dfracs[k].append(rel_dfrac) - _, median_foscttm = get_FOSCTTM( - model, - log["latent_X"][test_idx], - log["latent_Y"][test_idx], - use_barycenter=False, - ) - med_foscttms.append(median_foscttm) - med_foscttms = pd.concat(med_foscttms, axis=1).mean(axis=1) - rel_dfracs_ks = {} - dfracs_ks = {} - for k in ks: - rel_dfracs_ks[k] = pd.concat(rel_dfracs[k], axis=1).mean(axis=1) - dfracs_ks[k] = pd.concat(dfracs[k], axis=1).mean(axis=1) - main_k = 10 - other_ks = [k for k in ks if k != 3] - metrics_list = [ - med_foscttms, - dfracs_ks[main_k], - rel_dfracs_ks[main_k], - ] - labels = ["treat_idx", "FOSCTTM", "DFracs", "Relative DFracs"] - for k in other_ks: - metrics_list = metrics_list + [dfracs_ks[k], rel_dfracs_ks[k]] - labels = labels + [f"DFracs_{k}", f"Relative DFracs_{k}"] - metrics = pd.concat( - metrics_list, - axis=1, - ).reset_index() - - metrics.columns = labels - metrics["treatment"] = idx_name.loc[metrics["treat_idx"].astype(int), 0].tolist() - return metrics.set_index("treatment", drop=True)[labels[1:]].to_dict("series") - - -def evaluate_loo_run( - data_dict, - logfile_name, - test_idx_file="/home/ryuj6/scratch/OT/data/chemical_screen/chemical_screen_pca_idx.txt", - use_z_key="dosage", -): - Xs_dict = data_dict["Xs_dict"] - Xt_dict = data_dict["Xt_dict"] - Zs_dict = data_dict["Zs_dict"] - Zt_dict = data_dict["Zt_dict"] - with open(logfile_name, "rb") as f: - log = pkl.load(f) - dfracs = [] - rel_dfracs = [] - med_foscttms = [] - idx_name = pd.read_csv(test_idx_file, header=None) - for test_idx, Ts in log["ot_couplings"].items(): - Xs_train = _pop_key(Xs_dict, test_idx) - Xt_train = _pop_key(Xt_dict, test_idx) - - dfrac, rel_dfrac = get_diag_fracs( - Ts, Xs_train, Xt_train, Zs_dict[use_z_key], Zt_dict[use_z_key] - ) - _, median_foscttm = get_FOSCTTM(Ts, Xs_train, Xt_train) - med_foscttms.append(median_foscttm) - dfracs.append(dfrac) - rel_dfracs.append(rel_dfrac) - if isinstance(Ts, dict): - med_foscttms = pd.concat(med_foscttms, axis=1).mean(axis=1) - rel_dfracs = pd.concat(rel_dfracs, axis=1).mean(axis=1) - dfracs_df = pd.concat(dfracs, axis=1).mean(axis=1) - metrics = pd.concat( - [ - med_foscttms, - dfracs_df, - rel_dfracs, - ], - axis=1, - ).reset_index() - else: - med_foscttms = pd.Series(med_foscttms, index=log["ot_couplings"].keys()) - rel_dfracs = pd.Series(rel_dfracs, index=log["ot_couplings"].keys()) - dfracs_df = pd.Series(dfracs_df, index=log["ot_couplings"].keys()) - metrics = pd.concat( - [ - med_foscttms, - dfracs_df, - rel_dfracs, - ], - axis=1, - ).reset_index() - - metrics.columns = ["treat_idx", "FOSCTTM", "DFracs", "Relative DFracs"] - metrics["treatment"] = idx_name.loc[metrics["treat_idx"].astype(int), 0].tolist() - return metrics.set_index("treatment", drop=True)[ - ["FOSCTTM", "DFracs", "Relative DFracs"] - ].to_dict("series") - - -def evaluate_loo_runs_OT( - data_dict, - logfile_paths: List[str], - labels: List[str], - test_idx_file="/home/ryuj6/scratch/OT/data/chemical_screen/chemical_screen_pca_idx.txt", -): - metric_names = ["FOSCTTM", "DFracs", "Relative DFracs"] - metrics = {m: [] for m in metric_names} - non_foscttm_labels = [] - for logfile_path, label in zip(logfile_paths, labels): - if "vae" in label: - evals = evaluate_loo_vae(data_dict, logfile_path, test_idx_file) - else: - non_foscttm_labels.append(label) - evals = evaluate_loo_run(data_dict, logfile_path, test_idx_file) - for m in metric_names: - if m in evals: - metrics[m].append(evals[m]) - else: - metrics[m].append(pd.Series(np.zeros(evals["FOSCTTM"].shape))) - - def collect_series(list_series): - df = pd.concat(list_series, axis=1) - df.columns = labels - return df - - metrics = {m: collect_series(metrics[m]) for m in metrics.keys()} - return metrics diff --git a/perturbot/build/lib/perturbot/eval/match.py b/perturbot/build/lib/perturbot/eval/match.py deleted file mode 100755 index be475a2..0000000 --- a/perturbot/build/lib/perturbot/eval/match.py +++ /dev/null @@ -1,242 +0,0 @@ -"""Evaluate the quality of OT""" - -from typing import Dict, Tuple, Union, Sequence, List -import numpy as np -import pandas as pd -from perturbot.eval.utils import foscttm -from perturbot.utils import mdict_to_matrix - - -def get_rel_mse(T_dict): - rel_err = {} - for k, T in T_dict.items(): - T = T / T.sum() - perfect_match = np.identity( - T_dict[k].shape[0], - ) - perfect_match /= perfect_match.sum() - err = np.mean((np.diag(T_dict[k]) - np.diag(perfect_match)) ** 2) - - all2all = np.ones((T_dict[k].shape[0], T_dict[k].shape[1])) - all2all /= all2all.sum() - worst_err = np.mean((np.diag(all2all) - np.diag(perfect_match)) ** 2) - - rel_err[k] = err / worst_err - - return rel_err - - -def get_confusion_matrix( - T_dict, - Xs_dict, - Xt_dict, - Zs_dict, - Zt_dict, - norm=True, -) -> Tuple[Dict[Union[float, int], np.ndarray], pd.Series]: - """ - Ts: Dictionary of n_l x n'_l matrices for each label l - zs: m x m Confusion matrix is calculated where k is the number of unique values in z. - Must contain integer values (0, ..., m-1) - """ - labels = list(Xs_dict.keys()) - if not isinstance(T_dict, dict): - return get_confusion_matrix_single(T_dict, Xs_dict, Xt_dict, Zs_dict, Zt_dict) - classes = set(val for vals in Zs_dict.values() for val in vals) - Cmat_dict = np.zeros((len(classes), len(classes))) - diag_frac = {} - for k in labels: - # Cmat_dict[k] = np.zeros((len(classes), len(classes))) - try: - Zs = Zs_dict[k] - except KeyError as e: - print(e) - print(Zs_dict.keys()) - raise e - Zt = Zt_dict[k] - # if norm: - # T = T_dict[k] / T_dict[k].sum() - # else: - T = T_dict[k] - for i in range(T.shape[0]): - for j in range(T.shape[1]): - if T[i, j]: - Cmat_dict[int(Zs[i]), int(Zt[j])] += T[i, j] - diag_frac = np.diag(Cmat_dict).sum() - return Cmat_dict, diag_frac - - -def get_confusion_matrix_single( - T, - Xs_dict, - Xt_dict, - Zs_dict, - Zt_dict, -) -> Tuple[np.ndarray, float]: - classes = set(val for vals in Zs_dict.values() for val in vals) - Cmat = np.zeros((len(classes), len(classes))) - Zs = np.concatenate([Zs_dict[k] for k in Xs_dict.keys()]) - Zt = np.concatenate([Zt_dict[k] for k in Xs_dict.keys()]) - T = T / T.sum() - for i in range(T.shape[0]): - for j in range(T.shape[1]): - if T[i, j]: - Cmat[int(Zs[i]), int(Zt[j])] += T[i, j] - diag_frac = np.diag(Cmat).sum() - return Cmat, diag_frac - - -def get_diag_fracs( - T_dict: Union[np.ndarray, Dict[float, np.ndarray]], - Xs_dict, - Xt_dict, - Zs_dict, - Zt_dict, -) -> Union[Tuple[pd.Series, pd.Series], Tuple[float, float]]: - if not isinstance(T_dict, dict): - T = T_dict.copy() - sidx = 0 - tidx = 0 - T_dict = {} - for k, v in Xs_dict.items(): - T_dict[k] = T[ - sidx : (sidx + v.shape[0]), tidx : (tidx + Xt_dict[k].shape[0]) - ] - sidx += v.shape[0] - tidx += Xt_dict[k].shape[0] - # return get_diag_fracs_single(T_dict, Xs_dict, Xt_dict, Zs_dict, Zt_dict) - Cmat_dict, diag_fracs = get_confusion_matrix( - T_dict, Xs_dict, Xt_dict, Zs_dict, Zt_dict - ) - total_size_perfect = sum([T_dict[k].shape[0] for k in T_dict.keys()]) - T_perfect = {} - for k in T_dict.keys(): - T_perfect[k] = ( - np.identity( - T_dict[k].shape[0], - ) - / total_size_perfect - ) - Cmats_perfect, perfect_diag_fracs = get_confusion_matrix( - T_perfect, Xs_dict, Xt_dict, Zs_dict, Zt_dict - ) - total_size_random = sum([T_dict[k].size for k in T_dict.keys()]) - T_random = {k: np.ones(T_dict[k].shape) / total_size_random for k in T_dict.keys()} - Cmats_random, random_diag_fracs = get_confusion_matrix( - T_random, Xs_dict, Xt_dict, Zs_dict, Zt_dict - ) - - return diag_fracs, (diag_fracs - random_diag_fracs) / ( - perfect_diag_fracs - random_diag_fracs - ) - - -def get_diag_fracs_single( - T, - Xs_dict, - Xt_dict, - Zs_dict, - Zt_dict, -): - Cmat, diag_fracs = get_confusion_matrix_single( - T, Xs_dict, Xt_dict, Zs_dict, Zt_dict - ) - T_perfect = {} - for k in Xs_dict.keys(): - T_perfect[k] = ( - np.identity( - Xs_dict[k].shape[0], - ) - / Xs_dict[k].shape[0] - ) - T_perfect = mdict_to_matrix( - T_perfect, - np.concatenate([np.ones(Xs_dict[l].shape[0]) * l for l in Xs_dict.keys()]), - np.concatenate([np.ones(Xt_dict[l].shape[0]) * l for l in Xs_dict.keys()]), - ) - Cmats_perfect, perfect_diag_fracs = get_confusion_matrix_single( - T_perfect, Xs_dict, Xt_dict, Zs_dict, Zt_dict - ) - - T_random = { - k: np.ones((Xs_dict[k].shape[0], Xt_dict[k].shape[0])) for k in Xs_dict.keys() - } - T_random = mdict_to_matrix( - T_random, - np.concatenate([np.ones(Xs_dict[l].shape[0]) * l for l in Xs_dict.keys()]), - np.concatenate([np.ones(Xt_dict[l].shape[0]) * l for l in Xs_dict.keys()]), - ) - Cmats_random, random_diag_fracs = get_confusion_matrix_single( - T_random, Xs_dict, Xt_dict, Zs_dict, Zt_dict - ) - - return diag_fracs, (diag_fracs - random_diag_fracs) / ( - perfect_diag_fracs - random_diag_fracs - ) - - -def get_FOSCTTM( - T_dict, Xs_dict, Xt_dict, use_barycenter=True, use_agg="mean" -) -> Union[Tuple[Dict[float, Sequence], pd.Series], Tuple[List[float], float]]: - """Obtain the fraction of cells with higher assignment probability than the - true match (assume the diagonal is the true match.). Barycenter is used as the prediction. - - Returns - foscttm_dict: label to the list of FOSCTTM per cells - median_foscttm_dict: label to the aggregated FOSCTTM - """ - agg_fn = np.nanmedian if use_agg == "median" else np.nanmean - if isinstance(T_dict, dict): - T = mdict_to_matrix( - T_dict, - np.concatenate([np.ones(Xs_dict[l].shape[0]) * l for l in Xs_dict.keys()]), - np.concatenate([np.ones(Xt_dict[l].shape[0]) * l for l in Xs_dict.keys()]), - ) - else: - T = T_dict - - foscttm_dict = {} - median_foscttm_dict = {} - Xs_true = np.concatenate([Xs_dict[l] for l in Xs_dict.keys()]) - Xt_true = np.concatenate([Xt_dict[l] for l in Xt_dict.keys()]) - if use_barycenter: - marg = T.sum(axis=-1) - marg[marg == 0] = 1e-30 - Xt_pred = (T / marg[:, None]) @ Xt_true - foscttm_ = foscttm(Xt_pred, Xt_true) - else: - foscttm_ = foscttm(Xs_true, Xt_true) - return foscttm_, agg_fn(foscttm_) - for l, Xt_true in Xt_dict.items(): - Xs_true = Xs_dict[l] - if use_barycenter: - T = T_dict[l] - marg = T.sum(axis=-1) - marg[marg == 0] = 1e-30 - Xt_pred = (T / marg[:, None]) @ Xt_true - foscttm_ = foscttm(Xt_pred, Xt_true) - else: - foscttm_ = foscttm(Xs_true, Xt_true) - foscttm_dict[l] = foscttm_ - median_foscttm_dict[l] = agg_fn(foscttm_) - return foscttm_dict, pd.Series(median_foscttm_dict) - - -def get_FOSCTTM_single( - T, Xs_dict, Xt_dict, use_barycenter=True -) -> Tuple[List[float], float]: - """Obtain the fraction of cells with higher assignment probability than the - true match (assume the diagonal is the true match.). Barycenter is used as the prediction. - """ - Xs = np.concatenate([Xs_dict[l] for l in Xs_dict.keys()], axis=0) - Xt = np.concatenate([Xt_dict[l] for l in Xt_dict.keys()], axis=0) - - if use_barycenter: - marg = T.sum(axis=-1) - marg[marg == 0] = 1e-30 - Xt_pred = (T / marg[:, None]) @ Xt - foscttm_ = foscttm(Xt_pred, Xt) - else: - foscttm_ = foscttm(Xs, Xt) - - return foscttm_, np.nanmedian(foscttm_) diff --git a/perturbot/build/lib/perturbot/eval/prediction.py b/perturbot/build/lib/perturbot/eval/prediction.py deleted file mode 100755 index 87d4111..0000000 --- a/perturbot/build/lib/perturbot/eval/prediction.py +++ /dev/null @@ -1,210 +0,0 @@ -import numpy as np -import pandas as pd -from scipy.stats import spearmanr -from perturbot.eval.utils import foscttm - - -def _pearson_rowwise(A, B, eps=1e-8): - assert A.shape[0] == B.shape[0] - A_mA = A - A.mean(1)[:, None] - B_mB = B - B.mean(1)[:, None] - ssA = np.einsum("ij,ij->i", A_mA, A_mA) - ssB = np.einsum("ij,ij->i", B_mB, B_mB) - return np.einsum("ij,ij->i", A_mA, B_mB) / (np.sqrt(ssA * ssB) + eps) - - -def _spearman_rowwise(A, B): - assert A.shape[0] == B.shape[0] - scorr = [] - for i in range(A.shape[0]): - scorr.append( - spearmanr( - A[i, :], - B[i, :], - )[0] - ) - return scorr - - -def get_corrs(Y_pred, Y_true, idx=None): - if idx is not None: - Y_pred = Y_pred[:, idx] - Y_true = Y_true[:, idx] - pearson = _pearson_rowwise(Y_pred, Y_true) - spearman = _spearman_rowwise(Y_pred, Y_true) - return pearson, spearman - - -def mse(Y_pred, Y_true, idx=None): - if idx is None: - return (np.abs(Y_pred - Y_true) ** 2).mean(axis=1) - else: - return (np.abs(Y_pred[:, idx] - Y_true[:, idx]) ** 2).mean(axis=1) - - -def get_evals( - Y_pred, - Y_true, - idx=None, - idx_id="subset", - prediction_id="pred", - full=True, - agg_method="mean", - norm_Y: np.ndarray = None, -): - if agg_method == "median": - agg = np.median - elif agg_method == "mean": - agg = np.mean - else: - raise ValueError() - if norm_Y is not None: - pearson, spearman = get_corrs( - Y_pred / norm_Y[None, :], - Y_true / norm_Y[None, :], - ) - else: - pearson, spearman = get_corrs(Y_pred, Y_true) - pearson_c, spearman_c = get_corrs(Y_pred.T, Y_true.T) - _mse = mse(Y_pred, Y_true) - metrics = pd.Series( - [agg(pearson), agg(spearman), agg(pearson_c), agg(spearman_c), agg(_mse)], - index=[ - "Pearson_corr", - "Spearman_corr", - "Pearson_samples", - "Spearman_samples", - "MSE", - ], - ).to_frame() - if idx is not None: - if norm_Y is not None: - pearson_idx, spearman_idx = get_corrs( - Y_pred / norm_Y[None, :], Y_true / norm_Y[None, :], idx - ) - else: - pearson_idx, spearman_idx = get_corrs(Y_pred, Y_true, idx) - pearson_idx_c, spearman_idx_c = get_corrs(Y_pred, Y_true, idx) - _foscttm_idx = foscttm(Y_pred, Y_true, idx) - _mse_idx = mse(Y_pred, Y_true, idx) - metrics = pd.concat( - [ - metrics, - pd.Series( - [ - agg(pearson_idx), - agg(spearman_idx), - agg(pearson_idx_c), - agg(spearman_idx_c), - agg(_mse_idx), - ], - index=[ - f"Pearson_corr_{idx_id}", - f"Spearman_corr_{idx_id}", - f"Pearson_samples_{idx_id}", - f"Spearman_samples_{idx_id}", - f"MSE_{idx_id}", - ], - ).to_frame(), - ], - axis=0, - ) - metrics.columns = [prediction_id] - if full: - if idx is not None: - full_metrics = pd.DataFrame( - { - "metric": np.concatenate( - [ - pearson, - spearman, - pearson_c, - spearman_c, - _mse, - pearson_idx, - spearman_idx, - pearson_idx_c, - spearman_idx_c, - _mse_idx, - ] - ), - "group": np.repeat( - [ - "Pearson_corr", - "Spearman_corr", - "Pearson_samples", - "Spearman_samples", - "MSE", - f"Pearson_corr_{idx_id}", - f"Spearman_corr_{idx_id}", - f"Pearson_samples_{idx_id}", - f"Spearman_samples_{idx_id}", - f"MSE_{idx_id}", - ], - len(pearson), - ), - "pred_label": prediction_id, - } - ) - else: - full_metrics = pd.DataFrame( - { - "metric": np.concatenate( - [ - pearson, - spearman, - pearson_c, - spearman_c, - _mse, - ] - ), - "group": np.repeat( - [ - "Pearson_corr", - "Spearman_corr", - "Pearson_samples", - "Spearman_samples", - "MSE", - ], - len(pearson), - ), - "pred_label": prediction_id, - } - ) - - return metrics, full_metrics - return metrics - - -def get_evals_preds( - Y_true, Y_preds, pred_labels, idx=None, idx_id="subset", full=False -): - metric_dfs = [] - if full: - foscttm_dfs = [] - for Y_pred, pred_label in zip(Y_preds, pred_labels): - if full: - metrics, foscttm_df = get_evals( - Y_true, - Y_pred, - idx=idx, - idx_id=idx_id, - full=full, - prediction_id=pred_label, - ) - metric_dfs.append(metrics) - foscttm_dfs.append(foscttm_df) - else: - metric_dfs.append( - get_evals( - Y_true, - Y_pred, - idx=idx, - idx_id=idx_id, - full=full, - prediction_id=pred_label, - ) - ) - if full: - return pd.concat(metric_dfs, axis=1), pd.conat(foscttm_dfs, axis=1) - return pd.concat(metric_dfs, axis=1) diff --git a/perturbot/build/lib/perturbot/eval/utils.py b/perturbot/build/lib/perturbot/eval/utils.py deleted file mode 100755 index 9680f52..0000000 --- a/perturbot/build/lib/perturbot/eval/utils.py +++ /dev/null @@ -1,105 +0,0 @@ -from typing import Sequence, List -import numpy as np -from sklearn.neighbors import NearestNeighbors - - -def make_G(size, label, k): - assert size == len(label), (k, size, len(label)) - G = np.zeros((size, size)) - for l in np.unique(label): - l_idx = np.where(label == l)[0] - for i in l_idx: - for j in l_idx: - G[i, j] = 1.0 - assert (G.sum(axis=0) > 0).all(), (k, (G.sum(axis=0) == 0).sum()) - return G - - -def foscttm(Y_pred, Y_true, idx=None) -> List[float]: - """ - From https://github.com/rsinghlab/SCOT/blob/9787d0a6a1a494059cb3fb49e8ceda9318273c06/src/.ipynb_checkpoints/evals-checkpoint.py#L32 - Returns fraction closer than true match for each sample (as an array) - """ - if idx is not None: - Y_pred = Y_pred[:, idx] - Y_true = Y_true[:, idx] - fracs = [] - nsamp = Y_pred.shape[0] - rank = 0 - for row_idx in range(nsamp): - euc_dist = np.sqrt( - np.sum(np.square(np.subtract(Y_pred[row_idx, :], Y_true)), axis=1) - ) - true_nbr = euc_dist[row_idx] - sort_euc_dist = sorted(euc_dist) - try: - rank = np.where(sort_euc_dist == true_nbr)[0].mean() - except FloatingPointError: - print(true_nbr) - print(sort_euc_dist) - np.where(np.abs(sort_euc_dist - true_nbr < 1e-8))[0] - rank = np.where(np.abs(sort_euc_dist - true_nbr < 1e-8))[0].mean() - - frac = float(rank) / (nsamp - 1) - fracs.append(frac) - return fracs - - -def get_T_from_nn(X, Y, k): - """Obtain kNN label of observation in X from the k nearest neighbors in Y""" - nsamp = X.shape[0] - T = np.zeros((X.shape[0], Y.shape[0])) - for row_idx in range(nsamp): - euc_dist = np.sqrt(np.sum(np.square(np.subtract(X[row_idx, :], Y)), axis=1)) - smallest_k_idx = np.argpartition(euc_dist, k)[:k] - T[row_idx, smallest_k_idx] = 1.0 / (nsamp * k) - return T - - -def get_Ts_from_nn_multKs(X_dict, Y_dict, ks): - """Obtain kNN label of observation in X from the k nearest neighbors in Y""" - X = np.concatenate([X_dict[l] for l in X_dict.keys()]) - Y = np.concatenate([Y_dict[l] for l in X_dict.keys()]) - k_to_T = {k: np.zeros((X.shape[0], Y.shape[0])) for k in ks} - nsamp = X.shape[0] - for row_idx in range(nsamp): - euc_dist = np.sqrt(np.sum(np.square(np.subtract(X[row_idx, :], Y)), axis=1)) - for k in ks: - smallest_k_idx = np.argpartition(euc_dist, k)[:k] - k_to_T[k][row_idx, smallest_k_idx] = 1.0 / (nsamp * k) - k_to_Tdict = {} - for k, T in k_to_T.items(): - i = 0 - j = 0 - k_to_Tdict[k] = { - l: np.zeros((X_dict[l].shape[0], Y_dict[l].shape[0])) for l in X_dict.keys() - } - for l, X in X_dict.items(): - k_to_Tdict[k][l] = T[ - i : (i + X_dict[l].shape[0]), j : (j + Y_dict[l].shape[0]) - ] - i += X_dict[l].shape[0] - j += Y_dict[l].shape[0] - # k_to_Tdict[k][l] = k_to_Tdict[k][l] / k_to_Tdict[k][l].sum() - - return k_to_Tdict - - -def _pop_key(d, k, sub_key=None): - d = d.copy() - if sub_key is None: - del d[k] - else: - del d[sub_key][k] - return d - - -def _pop_keys(d, ks, sub_key=None): - d = d.copy() - if sub_key is None: - for k in ks: - del d[k] - else: - for k in ks: - del d[sub_key][k] - return d diff --git a/perturbot/build/lib/perturbot/match/__init__.py b/perturbot/build/lib/perturbot/match/__init__.py deleted file mode 100755 index 4284d03..0000000 --- a/perturbot/build/lib/perturbot/match/__init__.py +++ /dev/null @@ -1,27 +0,0 @@ -from .cot_labels import get_coupling_cotl, get_coupling_cotl_sinkhorn, cotl_numpy -from .ott_egwl import ( - get_coupling_egw_labels_ott, - get_coupling_egw_all_ott, - get_coupling_eot_ott, - get_coupling_leot_ott, - get_coupling_egw_ott, -) -from .cot import get_coupling_cot, get_coupling_cot_sinkhorn, cot_numpy -from .gw_labels import get_coupling_gw_labels -from .fot import get_coupling_fot - -__all__ = [ - "cotl_numpy", - "cot_numpy", - "get_coupling_eot_ott", - "get_coupling_leot_ott", - "get_coupling_cot", - "get_coupling_cot_sinkhorn", - "get_coupling_cotl", - "get_coupling_cotl_sinkhorn", - "get_coupling_egw_ott", - "get_coupling_gw_labels", - "get_coupling_egw_all_ott", - "get_coupling_egw_labels_ott", - "get_coupling_fot", -] diff --git a/perturbot/build/lib/perturbot/match/cot.py b/perturbot/build/lib/perturbot/match/cot.py deleted file mode 100755 index 900c525..0000000 --- a/perturbot/build/lib/perturbot/match/cot.py +++ /dev/null @@ -1,389 +0,0 @@ -from typing import Union, Dict, Tuple, Optional -from numbers import Number -import numpy as np -import time -import ot -from ott.solvers import linear -from ott.geometry import geometry -from .utils import random_gamma_init, init_matrix_np -import jax - - -def cot_numpy( - X1, - X2, - w1=None, - w2=None, - v1=None, - v2=None, - niter=10, - algo="emd", - reg=0, - algo2="emd", - reg2=0, - verbose=True, - log=False, - random_init=False, - C_lin=None, -): - """Returns COOT between two datasets X1,X2 (see [1]), Sinkhorn reimplemented with OTT - - The function solves the following optimization problem: - - .. math:: - - COOT = \min_{Ts,Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}*Tv_{k,l} - - Where : - - X1 : The source dataset - - X2 : The target dataset - - w1,w2 : weights (histograms) on the samples (rows) of resp. X1 and X2 - - v1,v2 : weights (histograms) on the features (columns) of resp. X1 and X2 - - Parameters - ---------- - X1 : numpy array, shape (n, d) - Source dataset - X2 : numpy array, shape (n', d') - Target dataset - w1 : numpy array, shape (n,) - Weight (histogram) on the samples of X1. If None uniform distribution is considered. - w2 : numpy array, shape (n',) - Weight (histogram) on the samples of X2. If None uniform distribution is considered. - v1 : numpy array, shape (d,) - Weight (histogram) on the features of X1. If None uniform distribution is considered. - v2 : numpy array, shape (d',) - Weight (histogram) on the features of X2. If None uniform distribution is considered. - niter : integer - Number max of iterations of the BCD for solving COOT. - algo : string - Choice of algorithm for solving OT problems on samples each iteration. Choice ['emd','sinkhorn']. - If 'emd' returns sparse solution - If 'sinkhorn' returns regularized solution - algo2 : string - Choice of algorithm for solving OT problems on features each iteration. Choice ['emd','sinkhorn']. - If 'emd' returns sparse solution - If 'sinkhorn' returns regularized solution - reg : float - Regularization parameter for samples coupling matrix. Ignored if algo='emd' - reg2 : float - Regularization parameter for features coupling matrix. Ignored if algo='emd' - eps : float - Threshold for the convergence - random_init : bool - Wether to use random initialization for the coupling matrices. If false identity couplings are considered. - log : bool, optional - record log if True - C_lin : numpy array, shape (n, n') - Prior on the sample correspondences. Added to the cost for the samples transport - - Returns - ------- - Ts : numpy array, shape (n,n') - Optimal Transport coupling between the samples - Tv : numpy array, shape (d,d') - Optimal Transport coupling between the features - cost : float - Optimization value after convergence - log : dict - convergence information and coupling marices - References - ---------- - .. [1] Redko Ievgen, Vayer Titouan, Flamary R{\'e}mi and Courty Nicolas - "CO-Optimal Transport" - - Examples - ---------- - .. code-block:: python - - import numpy as np - from cot import cot_numpy - - n_samples=300 - Xs=np.random.rand(n_samples,2) - Xt=np.random.rand(n_samples,1) - cot_numpy(Xs,Xt) - """ - if v1 is None: - v1 = np.ones(X1.shape[1]) / X1.shape[1] # is (d,) - if v2 is None: - v2 = np.ones(X2.shape[1]) / X2.shape[1] # is (d',) - if w1 is None: - w1 = np.ones(X1.shape[0]) / X1.shape[0] # is (n',) - if w2 is None: - w2 = np.ones(X2.shape[0]) / X2.shape[0] # is (n,) - - if not random_init: - Ts = np.ones((X1.shape[0], X2.shape[0])) / ( - X1.shape[0] * X2.shape[0] - ) # is (n,n') - Tv = np.ones((X1.shape[1], X2.shape[1])) / ( - X1.shape[1] * X2.shape[1] - ) # is (d,d') - else: - Ts = random_gamma_init(w1, w2) - Tv = random_gamma_init(v1, v2) - - constC_s, hC1_s, hC2_s = init_matrix_np(X1, X2, v1, v2) - - constC_v, hC1_v, hC2_v = init_matrix_np(X1.T, X2.T, w1, w2) - cost = np.inf - - log_out = {} - log_out["cost"] = [] - - for i in range(niter): - Tsold = Ts - Tvold = Tv - costold = cost - - M = constC_s - np.dot(hC1_s, Tv).dot(hC2_s.T) - if C_lin is not None: - M = M + C_lin - if algo == "emd": - Ts = ot.emd(w1, w2, M, numItermax=1e7) - elif algo == "sinkhorn": - Ts = np.array( - linear.solve( - geometry.Geometry( - cost_matrix=M, epsilon=reg, scale_cost="max_cost" - ), - max_iterations=2000, - ).matrix - ) - - M = constC_v - np.dot(hC1_v, Ts).dot(hC2_v.T) - - if algo2 == "emd": - Tv = ot.emd(v1, v2, M, numItermax=1e7) - elif algo2 == "sinkhorn": - Tv = np.array( - linear.solve( - geometry.Geometry( - cost_matrix=M, epsilon=reg2, scale_cost="max_cost" - ), - max_iterations=2000, - ).matrix - ) - - delta = np.linalg.norm(Ts - Tsold) + np.linalg.norm(Tv - Tvold) - cost = np.sum(M * Tv) - - if log: - log_out["cost"].append(cost) - - if verbose: - print("Delta: {0} Loss: {1}".format(delta, cost)) - - if delta < 1e-16 or np.abs(costold - cost) < 1e-7: - if verbose: - print("converged at iter ", i) - break - jax.clear_caches() - if log: - return Ts, Tv, cost, log_out - else: - return Ts, Tv, cost - - -def predict_with_cot(Xs, ys, Xt, yt, Xs_test, log=True): - """Learn mapping from Xs, Ys, Yt, Yt and predict Xt_test from Xs_test.""" - vs = Xs.sum(axis=0) # set the weights on the features - vs /= vs.sum() - vt = Xt.sum(axis=0) - vt /= vt.sum() - - Ts, Tv, _, log = cot_numpy(Xs, Xt, y1=ys, y2=yt, v1=vs, v2=vt, niter=100, log=True) - - log["Ts"] = Ts - log["Tv"] = Tv - Xt_pred = Xs_test @ (Tv / Tv.sum(axis=-1)[:, None]) - return Xt_pred, log - - -def get_coupling_cot( - data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], -) -> Tuple[Union[int, Dict[Number, np.array]], Union[int, Dict]]: - """Returns sample coupling between two datasets X, Y given the labels, disregarding label information. - - The function solves the following optimization problem: - - .. math:: - - COOT = \min_{Ts,Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - eps: - Regularization parameter, relative to the max cost. - - Returns - ------- - T_dict : - Optimal Transport coupling between the samples per label - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_cot - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,2) for k in labels} - get_coupling_cot((Xs_dict, Xt_dict)) - """ - X_dict = data[0] - Y_dict = data[1] - X = np.concatenate([X_dict[l] for l in X_dict.keys()]) - Y = np.concatenate([Y_dict[l] for l in X_dict.keys()]) - start = time.time() - try: - T, Tv, cost, log = cot_numpy(X, Y, log=True, niter=2000) - except FloatingPointError: - return -1, -1 - log["time"] = time.time() - start - return T, log - - -def get_coupling_cot_sinkhorn( - data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], - eps: float = 5e-3, - eps2: Optional[float] = None, -) -> Tuple[Union[int, Dict[Number, np.array]], Union[int, Dict]]: - """Returns sample coupling between two datasets X, Y given the labels, disregarding label information. - - The function solves the following optimization problem: - - .. math:: - - ECOOT = \min_{Ts,Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} - \epsilon H(T_s) - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - eps: - Regularization parameter, relative to the max cost. - - Returns - ------- - T_dict : - Optimal Transport coupling between the samples per label - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_cot_sinkhorn - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,2) for k in labels} - get_coupling_cot_sinkhorn((Xs_dict, Xt_dict), 0.05) - """ - print(f"calculating with eps {eps}") - X_dict = data[0] - Y_dict = data[1] - X = np.concatenate([X_dict[l] for l in X_dict.keys()]) - Y = np.concatenate([Y_dict[l] for l in X_dict.keys()]) - if eps2 is None: - eps2 = eps - start = time.time() - try: - T, Tv, cost, log = cot_numpy( - X, - Y, - algo="sinkhorn", - reg=eps, - algo2="sinkhorn", - reg2=eps2, - log=True, - niter=2000, - ) - except FloatingPointError: - return -1, -1 - log["time"] = time.time() - start - return T, log - - -def get_coupling_each_cot_sinkhorn( - data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], - eps: float = 5e-3, - eps2: Optional[float] = None, -) -> Tuple[Union[int, Dict[Number, np.array]], Union[int, Dict]]: - """Returns sample coupling between two datasets X, Y given the labels, disregarding label information. - - The function solves the following optimization problem: - - .. math:: - - ECOOT = \min_{Ts,Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} - \epsilon H(T_s) - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - eps: - Regularization parameter, relative to the max cost. - - Returns - ------- - T_dict : - Optimal Transport coupling between the samples per label - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_cot_sinkhorn - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,2) for k in labels} - get_coupling_cot_sinkhorn((Xs_dict, Xt_dict), 0.05) - """ - print(f"calculating with eps {eps}") - X_dict = data[0] - Y_dict = data[1] - X = np.concatenate([X_dict[l] for l in X_dict.keys()]) - Y = np.concatenate([Y_dict[l] for l in X_dict.keys()]) - if eps2 is None: - eps2 = eps - start = time.time() - try: - T_dict = {} - for l in X_dict.keys(): - T, Tv, cost, log = cot_numpy( - X_dict[l], - Y_dict[l], - algo="sinkhorn", - reg=eps, - algo2="sinkhorn", - reg2=eps2, - log=True, - niter=2000, - ) - T_dict[l] = T - print(f"Done calculating with eps {eps}") - except FloatingPointError: - return -1, -1 - log["time"] = time.time() - start - return T_dict, log diff --git a/perturbot/build/lib/perturbot/match/cot_labels.py b/perturbot/build/lib/perturbot/match/cot_labels.py deleted file mode 100755 index 43fbf15..0000000 --- a/perturbot/build/lib/perturbot/match/cot_labels.py +++ /dev/null @@ -1,340 +0,0 @@ -# Adapted from https://github.com/PythonOT/COOT -from typing import Dict, Optional, Tuple, Union -from numbers import Number -import pandas as pd -import numpy as np -import time -import ot as pot -import matplotlib.pyplot as plt -from ott.solvers import linear -from ott.geometry import geometry -from .utils import random_gamma_init, init_matrix_np - - -def cotl_numpy( - X_dict: Dict[Number, np.ndarray], - Y_dict: Dict[Number, np.ndarray], - w1: Dict[Number, np.ndarray] = None, - w2: Dict[Number, np.ndarray] = None, - v1: Optional[np.ndarray] = None, - v2: Optional[np.ndarray] = None, - niter: int = 100, - algo: str = "emd", - reg: float = 0.1, - algo2: str = "emd", - reg2: float = 10.0, - verbose: bool = True, - log: bool = False, - random_init: bool = False, - C_lin: bool = None, -): - r"""Returns COOT between two datasets X, Y given labels. - - The function solves the following optimization problem: - - .. math:: - - COOTL = \min_{Ts^1,..Ts^{L},Tv} \sum_{t=1}^L \sum_{i,k \in {i|l_{x_i}=t}, j,l \in {j|l_{y_j}=t}} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} - - - Parameters - ---------- - X1 : numpy array, shape (n, d) - Source dataset - X2 : numpy array, shape (n', d') - Target dataset - y1 : numpy array, shape (n,) - y2 : numpy array, shape (n',) - w1 : numpy array, shape (n,) - Weight (histogram) on the samples of X1. If None uniform distribution is considered. - w2 : Ditionary of numpy array, shape (n',) - Weight (histogram) on the samples of X2. If None uniform distribution is considered. - v1 : numpy array, shape (d,) - Weight (histogram) on the features of X1. If None uniform distribution is considered. - v2 : numpy array, shape (d',) - Weight (histogram) on the features of X2. If None uniform distribution is considered. - niter : integer - Number max of iterations of the BCD for solving COOT. - algo : string - Choice of algorithm for solving OT problems on samples each iteration. Choice ['emd','sinkhorn']. - If 'emd' returns sparse solution - If 'sinkhorn' returns regularized solution - algo2 : string - Choice of algorithm for solving OT problems on features each iteration. Choice ['emd','sinkhorn']. - If 'emd' returns sparse solution - If 'sinkhorn' returns regularized solution - reg : float - Regularization parameter for samples coupling matrix. Ignored if algo='emd' - reg2 : float - Regularization parameter for features coupling matrix. Ignored if algo='emd' - eps : float - Threshold for the convergence - random_init : bool - Wether to use random initialization for the coupling matrices. If false identity couplings are considered. - log : bool, optional - record log if True - C_lin : numpy array, shape (n, n') - Prior on the sample correspondences. Added to the cost for the samples transport - - Returns - ------- - Ts : numpy array, shape (n,n') - Optimal Transport coupling between the samples - Tv : numpy array, shape (d,d') - Optimal Transport coupling between the features - cost : float - Optimization value after convergence - log : dict - convergence information and coupling marices - References - ---------- - .. [1] Redko Ievgen, Vayer Titouan, Flamary R{\'e}mi and Courty Nicolas - "CO-Optimal Transport" - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import cotl_numpy - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} - cotl_numpy(Xs_dict, Xt_dict) - """ - assert sorted(X_dict.keys()) == sorted( - Y_dict.keys() - ), "Labels don't match in y1 & y2." - labels = list(X_dict.keys()) - if v1 is None: - X = np.concatenate([X_dict[k] for k in labels], axis=0) - if (X >= 0).all(): - v1 = X.sum(axis=0) / X.sum() - else: - v1 = np.ones(X.shape[1]) / X.shape[1] - - if v2 is None: - Y = np.concatenate([Y_dict[k] for k in labels], axis=0) - if (Y >= 0).all(): - v2 = Y.sum(axis=0) / Y.sum() - else: - v2 = np.ones(Y.shape[1]) / Y.shape[1] - - if w1 is None: - w1 = { - k: np.ones(X_dict[k].shape[0]) / X_dict[k].shape[0] for k in labels - } # is (n',) - if w2 is None: - w2 = { - k: np.ones(Y_dict[k].shape[0]) / Y_dict[k].shape[0] for k in labels - } # is (n,) - - if not random_init: - Ts = { - k: np.ones((X_dict[k].shape[0], Y_dict[k].shape[0])) - / (X_dict[k].shape[0] * Y_dict[k].shape[0]) - for k in labels - } # is (n,n') - Tv = np.ones((X_dict[labels[0]].shape[1], Y_dict[labels[0]].shape[1])) / ( - X_dict[labels[0]].shape[1] * Y_dict[labels[0]].shape[1] - ) # is (d,d') - else: - Ts = {k: random_gamma_init(w1[k], w2[k]) for k in labels} - Tv = random_gamma_init(v1, v2) - - constC_s_dict = {} - hC1_s_dict = {} - hC2_s_dict = {} - - constC_v_dict = {} - hC1_v_dict = {} - hC2_v_dict = {} - - for k in labels: - constC_s_dict[k], hC1_s_dict[k], hC2_s_dict[k] = init_matrix_np( - X_dict[k], Y_dict[k], v1, v2 - ) - constC_v_dict[k], hC1_v_dict[k], hC2_v_dict[k] = init_matrix_np( - X_dict[k].T, Y_dict[k].T, w1[k], w2[k] - ) - cost = np.inf - - log_out = {} - log_out["cost"] = [] - - for i in range(niter): - Tsold = Ts - Tvold = Tv - costold = cost - - # Sample OT for each label - for k in labels: - M_k = constC_s_dict[k] - np.dot(hC1_s_dict[k], Tv).dot(hC2_s_dict[k].T) - print(f"M_{k}:{M_k.min()} - {M_k.max()}") - if C_lin is not None: - M_k = M_k + C_lin - if algo == "emd": - Ts[k] = pot.emd(w1[k], w2[k], M_k, numItermax=1e7) - elif algo == "sinkhorn": - Ts[k] = np.array( - linear.solve( - geometry.Geometry( - cost_matrix=M_k, epsilon=reg, scale_cost="max_cost" - ), - max_iterations=2000, - ).matrix - ) - - # Global feature OT - M = 0 - for k in labels: - M += constC_v_dict[k] - np.dot(hC1_v_dict[k], Ts[k]).dot(hC2_v_dict[k].T) - print(f"M:{M.min()} - {M.max()}") - if algo2 == "emd": - Tv = pot.emd(v1, v2, M, numItermax=1e7) - elif algo2 == "sinkhorn": - Tv = np.array( - linear.solve( - geometry.Geometry( - cost_matrix=M, epsilon=reg, scale_cost="max_cost" - ), - max_iterations=2000, - ).matrix - ) - if not np.abs(Tv.sum() - 1.0) < 1e-8: - Tv = Tv / Tv.sum() - delta = sum( - [np.linalg.norm(Ts[k] - Tsold[k]) for k in labels] - ) + np.linalg.norm(Tv - Tvold) - cost = np.sum(M * Tv) - - if log: - log_out["cost"].append(cost) - - if verbose: - print(f"It {i} Delta: {delta} Loss: {cost}") - - if delta < 1e-16 or np.abs(costold - cost) < 1e-7: - if verbose: - print("converged at iter ", i) - break - if log: - return Ts, Tv, cost, log_out - else: - return Ts, Tv, cost - - -def get_coupling_cotl( - data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], -) -> Tuple[Union[int, Dict[Number, np.array]], Union[int, Dict]]: - """Returns sample coupling between two datasets X, Y given the labels, disregarding label information. - - The function solves the following optimization problem: - - .. math:: - - COOTL = \min_{Ts^1,..Ts^{L},Tv} \sum_{t=1}^L \sum_{i,k \in {i|l_{x_i}=t}, j,l \in {j|l_{y_j}=t}} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} - - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - eps: - Regularization parameter, relative to the max cost. - - Returns - ------- - T_dict : - Optimal Transport coupling between the samples per label - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_eot_ott - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} - get_coupling_eot_ott((Xs_dict, Xt_dict), 0.05) - """ - X_dict = data[0] - Y_dict = data[1] - start = time.time() - try: - Ts, Tv, cost, log = cotl_numpy(X_dict, Y_dict, log=True, niter=2000) - except FloatingPointError: - return -1, -1 - log["time"] = time.time() - start - return Ts, log - - -def get_coupling_cotl_sinkhorn( - data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], - eps: float = 5e-3, - eps2: float = None, -) -> Tuple[Dict[Number, np.array], Dict]: - """Returns sample coupling between two datasets X, Y given the labels. - - The function solves the following optimization problem: - - .. math:: - - COOTL = \min_{Ts^1,..Ts^{L},Tv} \sum_{t=1}^L \sum_{i,k \in {i|l_{x_i}=t}, j,l \in {j|l_{y_j}=t}} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} - \epsilon_1 H(Ts) -\epsilon_2 H(Tv) - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - eps: - Regularization parameter, relative to the max cost. - - Returns - ------- - T_dict : - Optimal Transport coupling between the samples per label - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_eot_ott - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} - get_coupling_eot_ott((Xs_dict, Xt_dict), 0.05) - """ - print(f"calculating with eps {eps}") - X_dict = data[0] - Y_dict = data[1] - start = time.time() - if eps2 is None: - eps2 = eps - try: - Ts, Tv, cost, log = cotl_numpy( - X_dict, - Y_dict, - algo="sinkhorn", - reg=eps, - algo2="sinkhorn", - reg2=eps2, - log=True, - niter=2000, - ) - except FloatingPointError: - return -1, -1 - log["time"] = time.time() - start - return Ts, log diff --git a/perturbot/build/lib/perturbot/match/fot.py b/perturbot/build/lib/perturbot/match/fot.py deleted file mode 100755 index 7f25929..0000000 --- a/perturbot/build/lib/perturbot/match/fot.py +++ /dev/null @@ -1,220 +0,0 @@ -from typing import Dict, Tuple, Union -from numbers import Number -import numpy as np -import time -import ot -from scipy import stats -from scipy.sparse import random -from ott.solvers import linear -from ott.geometry import geometry -from .utils import init_matrix_np, random_gamma_init -from perturbot.utils import mdict_to_matrix - - -def fot_numpy( - X1, - X2, - Ts, - v1=None, - v2=None, - niter=10, - algo="emd", - reg=0, - algo2="emd", - reg2=0, - verbose=True, - log=False, - random_init=False, - C_lin=None, -): - """Returns COOT between two datasets X1,X2 (see [1]) - - The function solves the following optimization problem: - .. math:: - - FOT = \min_{Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}*Tv_{k,l} - - Where : - - X1 : The source dataset - - X2 : The target dataset - - w1,w2 : weights (histograms) on the samples (rows) of resp. X1 and X2 - - v1,v2 : weights (histograms) on the features (columns) of resp. X1 and X2 - - Parameters - ---------- - X1 : numpy array, shape (n, d) - Source dataset - X2 : numpy array, shape (n', d') - Target dataset - w1 : numpy array, shape (n,) - Weight (histogram) on the samples of X1. If None uniform distribution is considered. - w2 : numpy array, shape (n',) - Weight (histogram) on the samples of X2. If None uniform distribution is considered. - v1 : numpy array, shape (d,) - Weight (histogram) on the features of X1. If None uniform distribution is considered. - v2 : numpy array, shape (d',) - Weight (histogram) on the features of X2. If None uniform distribution is considered. - niter : integer - Number max of iterations of the BCD for solving COOT. - algo : string - Choice of algorithm for solving OT problems on samples each iteration. Choice ['emd','sinkhorn']. - If 'emd' returns sparse solution - If 'sinkhorn' returns regularized solution - algo2 : string - Choice of algorithm for solving OT problems on features each iteration. Choice ['emd','sinkhorn']. - If 'emd' returns sparse solution - If 'sinkhorn' returns regularized solution - reg : float - Regularization parameter for samples coupling matrix. Ignored if algo='emd' - reg2 : float - Regularization parameter for features coupling matrix. Ignored if algo='emd' - eps : float - Threshold for the convergence - random_init : bool - Wether to use random initialization for the coupling matrices. If false identity couplings are considered. - log : bool, optional - record log if True - C_lin : numpy array, shape (n, n') - Prior on the sample correspondences. Added to the cost for the samples transport - - Returns - ------- - Ts : numpy array, shape (n,n') - Optimal Transport coupling between the samples - Tv : numpy array, shape (d,d') - Optimal Transport coupling between the features - cost : float - Optimization value after convergence - log : dict - convergence information and coupling marices - References - ---------- - .. [1] Redko Ievgen, Vayer Titouan, Flamary R{\'e}mi and Courty Nicolas - "CO-Optimal Transport" - Example - ---------- - import numpy as np - from cot import cot_numpy - - n_samples=300 - Xs=np.random.rand(n_samples,2) - Xt=np.random.rand(n_samples,1) - cot_numpy(Xs,Xt) - """ - if v1 is None: - v1 = np.ones(X1.shape[1]) / X1.shape[1] # is (d,) - if v2 is None: - v2 = np.ones(X2.shape[1]) / X2.shape[1] # is (d',) - Ts = Ts / Ts.sum() - w1 = Ts.sum(axis=0) - w2 = Ts.sum(axis=1) - if not random_init: - Tv = np.ones((X1.shape[1], X2.shape[1])) / ( - X1.shape[1] * X2.shape[1] - ) # is (d,d') - else: - Tv = random_gamma_init(v1, v2) - - constC_v, hC1_v, hC2_v = init_matrix_np(X1.T, X2.T, w1, w2) - cost = np.inf - - log_out = {} - log_out["cost"] = [] - - for i in range(niter): - Tvold = Tv - costold = cost - - M = constC_v - np.dot(hC1_v, Ts).dot(hC2_v.T) - Tv = np.array( - linear.solve( - geometry.Geometry(cost_matrix=M, epsilon=reg2, scale_cost="max_cost"), - max_iterations=2000, - ).matrix - ) - - delta = np.linalg.norm(Tv - Tvold) - cost = np.sum(M * Tv) - - if log: - log_out["cost"].append(cost) - - if verbose: - print("Delta: {0} Loss: {1}".format(delta, cost)) - - if delta < 1e-16 or np.abs(costold - cost) < 1e-7: - if verbose: - print("converged at iter ", i) - break - if log: - return Tv, cost, log_out - else: - return Tv, cost - - -def get_coupling_fot( - data: Tuple[Dict[Number, np.ndarray], Dict[Number, np.ndarray]], - Ts: Union[Dict[Number, np.ndarray], np.ndarray], - eps=5e-3, -): - r"""Returns GW coupling between features given two datasets X, Y and the sample coupling. - - The function solves the following optimization problem: - - .. math:: - - FOT = \min_{Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}^l*Tv_{k,l} - \epsilon H(T_v) - - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - Ts: - Sample-to-sample transport. - Per-label transport matched with source dataset, target dataset - or a global coupling matrix where the samples are concatenated by - the order of labels in data[0].keys(). - eps: - Regularization parameter, relative to the max cost. - - Returns - ------- - Tv : - Feature-to-feature coupling. - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_egw_labels_ott, get_coupling_fot - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} - Ts, log = get_coupling_egw_labels_ott((Xs_dict, Xt_dict), 0.05) - Tv, feature_matching_log = get_coupling_fot((Xs_dict, Xt_dict), Ts, 0.05) - """ - - X_dict = data[0] - Y_dict = data[1] - if isinstance(Ts, dict): - Ts = mdict_to_matrix( - Ts, - np.concatenate([np.ones(X_dict[l].shape[0]) * l for l in X_dict.keys()]), - np.concatenate([np.ones(Y_dict[l].shape[0]) * l for l in X_dict.keys()]), - ) - X = np.concatenate([X_dict[l] for l in X_dict.keys()]) - Y = np.concatenate([Y_dict[l] for l in X_dict.keys()]) - start = time.time() - try: - Tv, cost, log = fot_numpy(X, Y, Ts, log=True, reg=eps, reg2=eps, niter=2000) - except FloatingPointError: - return -1, -1 - log["time"] = time.time() - start - return Tv, log diff --git a/perturbot/build/lib/perturbot/match/gw.py b/perturbot/build/lib/perturbot/match/gw.py deleted file mode 100755 index cbcec81..0000000 --- a/perturbot/build/lib/perturbot/match/gw.py +++ /dev/null @@ -1,137 +0,0 @@ -import numpy as np -import scipy as sp -import ot -import time - - -def gw_cg(X_dict, Y_dict): - """GW with conditional gradient algorithm""" - labels = X_dict.keys() - Ts = {} - log = {"cost_time": {}, "time": {}} - for l in labels: - start = time.time() - C1 = sp.spatial.distance.cdist(X_dict[l], X_dict[l]) - C2 = sp.spatial.distance.cdist(Y_dict[l], Y_dict[l]) - - C1 /= C1.max() - C2 /= C2.max() - log["cost_time"][l] = time.time() - start - p = ot.unif(C1.shape[0]) - q = ot.unif(C2.shape[0]) - - start = time.time() - Ts[l], log[l] = ot.gromov.gromov_wasserstein( - C1, C2, p, q, "square_loss", verbose=True, log=True - ) - log["time"][l] = time.time() - start - return Ts, log - - -def egw_pgd(X_dict, Y_dict, epsilon=5e-3): - """EGW with projected gradient descent algorithm""" - labels = X_dict.keys() - Ts = {} - log = {"cost_time": {}, "time": {}} - for l in labels: - start = time.time() - C1 = sp.spatial.distance.cdist(X_dict[l], X_dict[l]) - C2 = sp.spatial.distance.cdist(Y_dict[l], Y_dict[l]) - - C1 /= C1.max() - C2 /= C2.max() - log["cost_time"][l] = time.time() - start - p = ot.unif(C1.shape[0]) - q = ot.unif(C2.shape[0]) - start = time.time() - Ts[l], log[l] = ot.gromov.entropic_gromov_wasserstein( - C1, - C2, - p, - q, - "square_loss", - epsilon=epsilon, - solver="PGD", - log=True, - verbose=True, - ) - log["time"][l] = time.time() - start - return Ts, log - - -def gw_all(Xtot, Ytot): - start = time.time() - C1 = sp.spatial.distance.cdist(Xtot, Xtot) - C2 = sp.spatial.distance.cdist(Ytot, Ytot) - C1 /= C1.max() - C2 /= C2.max() - cost_time = time.time() - start - p = ot.unif(C1.shape[0]) - q = ot.unif(C2.shape[0]) - start = time.time() - Ts, log = ot.gromov.gromov_wasserstein( - C1, C2, p, q, "square_loss", verbose=True, log=True, max_iter=1e8 - ) - log["time"] = time.time() - start - log["cost_time"] = cost_time - return Ts, log - - -def egw_all(Xtot, Ytot, epsilon=5e-3): - start = time.time() - C1 = sp.spatial.distance.cdist(Xtot, Xtot) - C2 = sp.spatial.distance.cdist(Ytot, Ytot) - C1 /= C1.max() - C2 /= C2.max() - cost_time = time.time() - start - p = ot.unif(C1.shape[0]) - q = ot.unif(C2.shape[0]) - start = time.time() - Ts, log = ot.gromov.entropic_gromov_wasserstein( - C1, - C2, - p, - q, - "square_loss", - verbose=True, - log=True, - epsilon=epsilon, - ) - log["time"] = time.time() - start - log["cost_time"] = cost_time - return Ts, log - - -def get_coupling_gw_cg(data): - X_dict = data[0] - Y_dict = data[1] - Ts, log = gw_cg(X_dict, Y_dict) - return Ts, log - - -def get_coupling_egw(data, eps=5e-3): - X_dict = data[0] - Y_dict = data[1] - Ts, log = egw_pgd(X_dict, Y_dict, epsilon=eps) - return Ts, log - - -def get_coupling_gw_all(data): - """Run GW all-to-all, ignore labels.""" - X_dict = data[0] - Y_dict = data[1] - Xtot = np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0) - Ytot = np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0) - T, log = gw_all(Xtot, Ytot) - return T, log - - -def get_coupling_egw_all(data, eps=5e-3): - """Run GW all-to-all, ignore labels.""" - X_dict = data[0] - Y_dict = data[1] - Xtot = np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0) - Ytot = np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0) - T, log = egw_all(Xtot, Ytot, eps) - - return T, log diff --git a/perturbot/build/lib/perturbot/match/gw_labels.py b/perturbot/build/lib/perturbot/match/gw_labels.py deleted file mode 100755 index 5a59b92..0000000 --- a/perturbot/build/lib/perturbot/match/gw_labels.py +++ /dev/null @@ -1,148 +0,0 @@ -from typing import Tuple, Dict -from numbers import Number -import numpy as np -import scipy as sp -import ot -import time - - -def get_coupling_gw_labels( - data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], -) -> Tuple[Dict[Number, np.array], Dict]: - r"""Returns GW coupling between two datasets X, Y given the labels. - - The function solves the following optimization problem: - - .. math:: - - GWL = \min_{T\in C_{p,q}^\ell} \sum_{i,k \in \{i|l_{x_i}=t\}, j,l \in \{j|l_{y_j}=t}\} |(x_i-x_k)^2 - (y_j-y_l)^2|^{2}*T_{i,j}T_{k,l} \\ - C_{p,q}^\ell = \{T | T \in C{p,q}, T_{ij} > 0 \implies l_{x_i} = l_{y_j}\} - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - - Returns - ------- - T_dict : - Optimal Transport coupling between the samples per label - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_gw_labels - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} - get_coupling_gw_labels((Xs_dict, Xt_dict)) - """ - X_dict = data[0] - Y_dict = data[1] - Xs_tot = np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0) - Xt_tot = np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0) - source_labels = np.concatenate( - [np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()] - ) - target_labels = np.concatenate( - [np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()] - ) - start = time.time() - C1_tot = sp.spatial.distance.cdist(Xs_tot, Xs_tot) - C2_tot = sp.spatial.distance.cdist(Xt_tot, Xt_tot) - C1_tot /= C1_tot.max() - C2_tot /= C2_tot.max() - cost_time = time.time() - start - start = time.time() - T, log = ot.gromov.gromov_wasserstein_labeled( - C1_tot, C2_tot, source_labels, target_labels, log=True - ) - end = time.time() - log["time"] = end - start - log["cosst_time"] = cost_time - T_dict = {} - for l in np.unique(source_labels): - T_dict[l] = T[source_labels == l, :][:, target_labels == l] - return T_dict, log - - -def get_coupling_egw_labels(data, eps=5e-3): - """Returns GW coupling between two datasets X, Y given the labels. - - The function solves the following optimization problem: - - .. math:: - - GWL = \min_{T\in C_{p,q}^\ell} \sum_{i,k \in \{i|l_{x_i}=t\} j,l \in \{j|l_{y_j}=t}\} |(x_i-x_k)^2 - (y_j-y_l)^2|^{2}*T_{i,j}T_{k,l} - \epsilon H(T)\\ - C_{p,q}^\ell = \{T | T \in C{p,q}, T_{ij} > 0 \implies l_{x_i} = l_{y_j}\} - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - eps: - Regularization parameter, relative to the max cost. - - Returns - ------- - T_dict : - Optimal Transport coupling between the samples per label - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_egw_labels - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} - get_coupling_egw_labels((Xs_dict, Xt_dict), 0.05) - """ - X_dict = data[0] - Y_dict = data[1] - Xs_tot = np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0) - Xt_tot = np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0) - source_labels = np.concatenate( - [np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()] - ) - target_labels = np.concatenate( - [np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()] - ) - start = time.time() - C1_tot = sp.spatial.distance.cdist(Xs_tot, Xs_tot) - C2_tot = sp.spatial.distance.cdist(Xt_tot, Xt_tot) - C1_tot /= C1_tot.max() - C2_tot /= C2_tot.max() - cost_time = time.time() - start - start = time.time() - print("running LEOT") - T, log = ot.gromov.entropic_gromov_wasserstein_labeled( - C1_tot, - C2_tot, - source_labels, - target_labels, - epsilon=eps, - log=True, - verbose=True, - ) - end = time.time() - print("Done running LEOT") - log["time"] = end - start - log["cost_time"] = cost_time - T_dict = {} - for l in np.unique(source_labels): - T_dict[l] = T[source_labels == l, :][:, target_labels == l] - return T_dict, log diff --git a/perturbot/build/lib/perturbot/match/ot_labels.py b/perturbot/build/lib/perturbot/match/ot_labels.py deleted file mode 100755 index 4b474a1..0000000 --- a/perturbot/build/lib/perturbot/match/ot_labels.py +++ /dev/null @@ -1,78 +0,0 @@ -import numpy as np -import scipy as sp -import ot -import time - - -def get_coupling_ot_labels(data, eps=5e-3): - X_dict = data[0] - Y_dict = data[1] - Xs_tot = np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0) - Xt_tot = np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0) - source_labels = np.concatenate( - [np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()] - ) - target_labels = np.concatenate( - [np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()] - ) - start = time.time() - C = sp.spatial.distance.cdist(Xs_tot, Xs_tot) - C = C / C.max() - cost_time = time.time() - start - start = time.time() - p = np.ones(C.shape[0]) / C.shape[0] - q = np.ones(C.shape[1]) / C.shape[1] - T, log = ot.bergman.sinkhorn_labeled( - p, - q, - T, - source_labels, - target_labels, - reg=eps, - log=True, - verbose=True, - ) - end = time.time() - log["time"] = end - start - log["cost_time"] = cost_time - T_dict = {} - for l in np.unique(source_labels): - T_dict[l] = T[source_labels == l, :][:, target_labels == l] - return T_dict, log - - -def get_coupling_eot_labels(data, eps=5e-3): - X_dict = data[0] - Y_dict = data[1] - Xs_tot = np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0) - Xt_tot = np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0) - source_labels = np.concatenate( - [np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()] - ) - target_labels = np.concatenate( - [np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()] - ) - start = time.time() - C = sp.spatial.distance.cdist(Xs_tot, Xs_tot) - C = C / C.max() - cost_time = time.time() - start - start = time.time() - p = np.ones(C.shape[0]) / C.shape[0] - q = np.ones(C.shape[1]) / C.shape[1] - T, log = ot.bergman.sinkhorn_labeled( - p, - q, - T, - source_labels, - target_labels, - reg=eps, - log=True, - verbose=True, - ) - end = time.time() - log["time"] = end - start - log["cost_time"] = cost_time - T_dict = {} - for l in np.unique(source_labels): - T_dict[l] = T[source_labels == l, :][:, target_labels == l] - return T_dict, log diff --git a/perturbot/build/lib/perturbot/match/ott_egwl.py b/perturbot/build/lib/perturbot/match/ott_egwl.py deleted file mode 100755 index 7f9e4df..0000000 --- a/perturbot/build/lib/perturbot/match/ott_egwl.py +++ /dev/null @@ -1,450 +0,0 @@ -# Modified OTT with labels -from typing import Dict, Tuple -from numbers import Number -import numpy as np -import jax.numpy as jnp -import time -import ott -from ott.geometry import pointcloud, geometry -from ott.problems.linear import linear_problem -from ott.problems.quadratic import quadratic_problem -from ott.solvers.quadratic import gromov_wasserstein -from ott.solvers import linear -from ott.solvers.linear import acceleration, sinkhorn - - -def create_block_diag_mat(labels_a, labels_b): - block_diag_mat = np.zeros((len(labels_a), len(labels_b))) - for l in np.unique(labels_a): - block_diag_mat[ - np.ix_(np.where(labels_a == l)[0], np.where(labels_b == l)[0]) - ] = 1.0 - return jnp.array(block_diag_mat) - - -def get_coupling_egw_labels_ott( - data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], eps: float = 5e-3 -) -> Tuple[Dict[Number, np.array], Dict]: - r"""Returns GW coupling between two datasets X, Y given the labels. - - The function solves the following optimization problem: - - .. math:: - - EGWL = \min_{T\in C_{p,q}^\ell} \sum_{i,k \in \{i|l_{x_i}=t\}, j,l \in \{j|l_{y_j}=t\}} |(x_i-x_k)^2 - (y_j-y_l)^2|^{2}*T_{i,j}T_{k,l} - \epsilon H(T)\\ - C_{p,q}^\ell = \{T | T \in C{p,q}, T_{ij} > 0 \implies l_{x_i} = l_{y_j}\} - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - eps: - Regularization parameter, relative to the max cost. - - Returns - ------- - T_dict : - Optimal Transport coupling between the samples per label - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_egw_labels_ott - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} - get_coupling_egw_labels_ott((Xs_dict, Xt_dict), 0.05) - """ - X_dict = data[0] - Y_dict = data[1] - Xs_tot = jnp.array(np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0)) - Xt_tot = jnp.array(np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0)) - source_labels = jnp.array( - np.concatenate([np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()]) - ) - target_labels = jnp.array( - np.concatenate([np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()]) - ) - start = time.time() - - geom_xx = pointcloud.PointCloud(x=Xs_tot, y=Xs_tot, scale_cost="max_cost") - geom_yy = pointcloud.PointCloud(x=Xt_tot, y=Xt_tot, scale_cost="max_cost") - bdm = create_block_diag_mat(source_labels, target_labels) - - cost_time = time.time() - start - start = time.time() - print("running EGWL with ott") - - prob = quadratic_problem.QuadraticProblem( - geom_xx, - geom_yy, - labels_a=source_labels, - labels_b=target_labels, - n_labels=len(np.unique(source_labels)), - block_diag_mat=bdm, - ) - - solver = gromov_wasserstein.GromovWasserstein( - epsilon=eps, - store_inner_errors=True, - max_iterations=2000, - kwargs_sinkhorn={"max_iterations": 2000}, - ) - - out = solver(prob) - - has_converged = bool(out.linear_convergence[out.n_iters - 1]) - log = {} - log["n_iters_outer"] = out.n_iters - log["converged_inner"] = has_converged - log["converged_outer"] = out.converged - log["GW cost"] = out.reg_gw_cost - T = np.array(out.matrix) - print(f"{out.n_iters} outer iterations were needed.") - print(f"The last Sinkhorn iteration has converged: {has_converged}") - print(f"The outer loop of Gromov Wasserstein has converged: {out.converged}") - print(f"The final regularized GW cost is: {out.reg_gw_cost:.3f}") - - end = time.time() - print("Done running LEGWOT with ott") - log["time"] = end - start - log["cost_time"] = cost_time - T_dict = {} - for l in np.unique(source_labels): - T_dict[l] = T[np.array(source_labels) == l, :][:, np.array(target_labels) == l] - return T_dict, log - - -def get_coupling_egw_ott( - data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], eps: float = 5e-3 -) -> Tuple[Dict[Number, np.array], Dict]: - r"""Returns GW coupling between two datasets X, Y per label. - - The function solves the following optimization problem: - - .. math:: - - GW^l = \min_{T^l} \sum_{i,j,k,l} |(x_i-x_k)^2 - (y_j-y_l)^2|^{2}*T^l_{i,j}T^l_{k,l} - \epsilon H(T^l)\\ - - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - eps: - Regularization parameter, relative to the max cost. - - Returns - ------- - T_dict : - Optimal Transport coupling between the samples per label - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_egw_ott - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} - get_coupling_egw_ott((Xs_dict, Xt_dict), 0.05) - """ - X_dict = data[0] - Y_dict = data[1] - labels = X_dict.keys() - Ts = {} - log = {} - for l in labels: - log[l] = {} - start = time.time() - - geom_xx = pointcloud.PointCloud(x=X_dict[l], y=X_dict[l], scale_cost="max_cost") - geom_yy = pointcloud.PointCloud(x=Y_dict[l], y=Y_dict[l], scale_cost="max_cost") - - cost_time = time.time() - start - start = time.time() - - prob = quadratic_problem.QuadraticProblem( - geom_xx, - geom_yy, - ) - - # Instantiate a jitt'ed Gromov-Wasserstein solver - solver = gromov_wasserstein.GromovWasserstein( - epsilon=eps, store_inner_errors=True, max_iterations=1000 - ) - - out = solver(prob) - - end = time.time() - Ts[l] = np.array(out.matrix) - - has_converged = bool(out.linear_convergence[out.n_iters - 1]) - log[l]["n_iters_outer"] = out.n_iters - log[l]["converged_inner"] = has_converged - log[l]["converged_outer"] = out.converged - log[l]["GW cost"] = out.reg_gw_cost - log[l]["time"] = end - start - log[l]["cost_time"] = cost_time - return Ts, log - - -def get_coupling_egw_all_ott( - data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], eps: float = 5e-3 -) -> Tuple[Dict[Number, np.array], Dict]: - r"""Returns GW coupling between two datasets X, Y, all-to-all manner disregarding labels. - - The function solves the following optimization problem: - - .. math:: - - GW = \min_{T\in C_{p,q}} \sum_{i,j,k,l} |(x_i-x_k)^2 - (y_j-y_l)^2|^{2}*T_{i,j}T_{k,l} - \epsilon H(T) - - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - eps: - Regularization parameter, relative to the max cost. - - Returns - ------- - T_dict : - Optimal Transport coupling between the samples per label - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_egw_all_ott - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,2) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} - get_coupling_egw_all_ott((Xs_dict, Xt_dict), 0.05) - """ - X_dict = data[0] - Y_dict = data[1] - Xs_tot = jnp.array(np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0)) - Xt_tot = jnp.array(np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0)) - source_labels = jnp.array( - np.concatenate([np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()]) - ) - target_labels = jnp.array( - np.concatenate([np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()]) - ) - start = time.time() - - geom_xx = pointcloud.PointCloud(x=Xs_tot, y=Xs_tot, scale_cost="max_cost") - geom_yy = pointcloud.PointCloud(x=Xt_tot, y=Xt_tot, scale_cost="max_cost") - - cost_time = time.time() - start - start = time.time() - print("running EGWOT with ott") - - prob = quadratic_problem.QuadraticProblem( - geom_xx, - geom_yy, - ) - - # Instantiate a jitt'ed Gromov-Wasserstein solver - solver = gromov_wasserstein.GromovWasserstein( - epsilon=eps, store_inner_errors=True, max_iterations=1000 - ) - - out = solver(prob) - - has_converged = bool(out.linear_convergence[out.n_iters - 1]) - log = {} - log["n_iters_outer"] = out.n_iters - log["converged_inner"] = has_converged - log["converged_outer"] = out.converged - log["GW cost"] = out.reg_gw_cost - T = np.array(out.matrix) - print(f"{out.n_iters} outer iterations were needed.") - print(f"The last Sinkhorn iteration has converged: {has_converged}") - print(f"The outer loop of Gromov Wasserstein has converged: {out.converged}") - print(f"The final regularized GW cost is: {out.reg_gw_cost:.3f}") - - end = time.time() - print("Done running EGWOT with ott") - log["time"] = end - start - log["cost_time"] = cost_time - return T, log - - -def get_coupling_eot_ott( - data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], eps: float = 5e-3 -) -> Tuple[Dict[Number, np.array], Dict]: - r"""Returns OT coupling between two datasets X, Y given the labels, disregarding label information. - - The function solves the following optimization problem: - - .. math:: - - EOT = \min_{T\in C_{p,q}} \sum_{i,j} (x_i-y_j)^2 T_{i,j} - \epsilon H(T)\\ - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - eps: - Regularization parameter, relative to the max cost. - - Returns - ------- - T_dict : - Optimal Transport coupling between the samples per label - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_eot_ott - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} - get_coupling_eot_ott((Xs_dict, Xt_dict), 0.05) - """ - X_dict = data[0] - Y_dict = data[1] - Xs_tot = jnp.array(np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0)) - Xt_tot = jnp.array(np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0)) - source_labels = jnp.array( - np.concatenate([np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()]) - ) - target_labels = jnp.array( - np.concatenate([np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()]) - ) - start = time.time() - - geom = pointcloud.PointCloud(x=Xs_tot, y=Xt_tot, scale_cost="max_cost") - - cost_time = time.time() - start - start = time.time() - print("running LEOT with ott") - - out = linear.solve(geometry.Geometry(cost_matrix=geom.cost_matrix, epsilon=eps)) - - log = {} - log["n_iters_outer"] = out.n_iters - log["converged"] = out.converged - log["OT cost"] = out.reg_ot_cost - T = np.array(out.matrix) - print(f"{out.n_iters} outer iterations were needed.") - print(f"The last Sinkhorn iteration has converged: {out.converged}") - print(f"The final regularized OT cost is: {out.reg_ot_cost:.3f}") - - end = time.time() - print("Done running EOT with ott") - log["time"] = end - start - log["cost_time"] = cost_time - - return T, log - - -def get_coupling_leot_ott( - data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], eps: float = 5e-3 -) -> Tuple[Dict[Number, np.array], Dict]: - r"""Returns OT coupling between two datasets X, Y per label. - - The function solves the following optimization problem: - - .. math:: - - EOT^l = \min_{T^l} \sum_{i,j} (x_i-y_j)^2 T^l_{i,j} - \epsilon H(T^l)\\ - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - eps: - Regularization parameter, relative to the max cost. - - Returns - ------- - T_dict : - Optimal Transport coupling between the samples per label - log : - Running log - - Example - ---------- - .. code-block:: python - - import numpy as np - from perturbot.match import get_coupling_leot_ott - - n_samples = 300 - labels = [0,1,2,3] - Xs_dict = {k: np.random.rand(n_samples,1) for k in labels} - Xt_dict = {k: np.random.rand(n_samples,1) for k in labels} - get_coupling_leot_ott((Xs_dict, Xt_dict), 0.05) - """ - X_dict = data[0] - Y_dict = data[1] - Xs_tot = jnp.array(np.concatenate([X_dict[l] for l in X_dict.keys()], axis=0)) - Xt_tot = jnp.array(np.concatenate([Y_dict[l] for l in X_dict.keys()], axis=0)) - source_labels = jnp.array( - np.concatenate([np.repeat(l, X_dict[l].shape[0]) for l in X_dict.keys()]) - ) - target_labels = jnp.array( - np.concatenate([np.repeat(l, Y_dict[l].shape[0]) for l in X_dict.keys()]) - ) - start = time.time() - - geom = pointcloud.PointCloud(x=Xs_tot, y=Xt_tot, scale_cost="max_cost") - geom = geometry.Geometry(cost_matrix=geom.cost_matrix, epsilon=eps) - - cost_time = time.time() - start - start = time.time() - print("running LEOT with ott") - problem = linear_problem.LinearProblem( - geom, labels_a=source_labels, labels_b=target_labels - ) - solver = sinkhorn.Sinkhorn() - out = solver(problem) - - log = {} - log["n_iters_outer"] = out.n_iters - log["converged"] = out.converged - log["OT cost"] = out.reg_ot_cost - T = np.array(out.matrix) - print(f"{out.n_iters} outer iterations were needed.") - print(f"The last Sinkhorn iteration has converged: {out.converged}") - print(f"The final regularized OT cost is: {out.reg_ot_cost:.3f}") - - end = time.time() - print("Done running LEOT with ott") - log["time"] = end - start - log["cost_time"] = cost_time - T_dict = {} - for l in np.unique(source_labels): - T_dict[l] = T[np.array(source_labels) == l, :][:, np.array(target_labels) == l] - return T_dict, log diff --git a/perturbot/build/lib/perturbot/match/utils.py b/perturbot/build/lib/perturbot/match/utils.py deleted file mode 100755 index 3ea2951..0000000 --- a/perturbot/build/lib/perturbot/match/utils.py +++ /dev/null @@ -1,184 +0,0 @@ -import numpy as np -from scipy import stats -from scipy.sparse import random - - -def sinkhorn_scaling( - a, - b, - K, - numItermax=1000, - stopThr=1e-9, - verbose=False, - log=False, - always_raise=False, - **kwargs, -): - a = np.asarray(a, dtype=np.float64) - b = np.asarray(b, dtype=np.float64) - K = np.asarray(K, dtype=np.float64) - - # init data - Nini = len(a) - Nfin = len(b) - - if len(b.shape) > 1: - nbb = b.shape[1] - else: - nbb = 0 - - if log: - log = {"err": []} - - # we assume that no distances are null except those of the diagonal of - # distances - if nbb: - u = np.ones((Nini, nbb)) / Nini - v = np.ones((Nfin, nbb)) / Nfin - else: - u = np.ones(Nini) / Nini - v = np.ones(Nfin) / Nfin - - # print(reg) - # print(np.min(K)) - - Kp = (1 / a).reshape(-1, 1) * K - cpt = 0 - err = 1 - while err > stopThr and cpt < numItermax: - uprev = u - vprev = v - KtransposeU = np.dot(K.T, u) - v = np.divide(b, KtransposeU) - u = 1.0 / np.dot(Kp, v) - - zero_in_transp = np.any(KtransposeU == 0) - nan_in_dual = np.any(np.isnan(u)) or np.any(np.isnan(v)) - inf_in_dual = np.any(np.isinf(u)) or np.any(np.isinf(v)) - if zero_in_transp or nan_in_dual or inf_in_dual: - # we have reached the machine precision - # come back to previous solution and quit loop - print("Warning: numerical errors at iteration in sinkhorn_scaling", cpt) - # if zero_in_transp: - # print('Zero in transp : ',KtransposeU) - # if nan_in_dual: - # print('Nan in dual') - # print('u : ',u) - # print('v : ',v) - # print('KtransposeU ',KtransposeU) - # print('K ',K) - # print('M ',M) - - # if always_raise: - # raise NanInDualError - # if inf_in_dual: - # print('Inf in dual') - u = uprev - v = vprev - - break - if cpt % 10 == 0: - # we can speed up the process by checking for the error only all - # the 10th iterations - if nbb: - err = np.sum((u - uprev) ** 2) / np.sum((u) ** 2) + np.sum( - (v - vprev) ** 2 - ) / np.sum((v) ** 2) - else: - transp = u.reshape(-1, 1) * (K * v) - err = np.linalg.norm((np.sum(transp, axis=0) - b)) ** 2 - if log: - log["err"].append(err) - - if verbose: - if cpt % 200 == 0: - print("{:5s}|{:12s}".format("It.", "Err") + "\n" + "-" * 19) - print("{:5d}|{:8e}|".format(cpt, err)) - cpt = cpt + 1 - if log: - log["u"] = u - log["v"] = v - - if nbb: # return only loss - res = np.zeros((nbb)) - for i in range(nbb): - res[i] = np.sum(u[:, i].reshape((-1, 1)) * K * v[:, i].reshape((1, -1)) * M) - if log: - return res, log - else: - return res - - else: # return OT matrix - if log: - return u.reshape((-1, 1)) * K * v.reshape((1, -1)), log - else: - return u.reshape((-1, 1)) * K * v.reshape((1, -1)) - - -def random_gamma_init(p, q, **kwargs): - """Returns random coupling matrix with marginal p,q""" - rvs = stats.beta(1e-1, 1e-1).rvs - S = random(len(p), len(q), density=1, data_rvs=rvs) - return sinkhorn_scaling(p, q, S.A, **kwargs) - - -def init_matrix_np(X1, X2, v1, v2): - """Return loss matrices and tensors for COOT fast computation - Returns the value of |X1-X2|^{2} \otimes T as done in [1] based on [2] for the Gromov-Wasserstein distance. - Where : - - X1 : The source dataset of shape (n,d) - - X2 : The target dataset of shape (n',d') - - v1 ,v2 : weights (histograms) on the columns of resp. X1 and X2 - - T : Coupling matrix of shape (n,n') - Parameters - ---------- - X1 : numpy array, shape (n, d) - Source dataset - X2 : numpy array, shape (n', d') - Target dataset - v1 : numpy array, shape (d,) - Weight (histogram) on the features of X1. - v2 : numpy array, shape (d',) - Weight (histogram) on the features of X2. - - Returns - ------- - constC : ndarray, shape (n, n') - Constant C matrix (see paragraph 1.2 of supplementary material in [1]) - hC1 : ndarray, shape (n, d) - h1(X1) matrix (see paragraph 1.2 of supplementary material in [1]) - hC2 : ndarray, shape (n', d') - h2(X2) matrix (see paragraph 1.2 of supplementary material in [1]) - References - ---------- - .. [1] Redko Ievgen, Vayer Titouan, Flamary R{\'e}mi and Courty Nicolas - "CO-Optimal Transport" - .. [2] Peyré, Gabriel, Marco Cuturi, and Justin Solomon, - "Gromov-Wasserstein averaging of kernel and distance matrices." - International Conference on Machine Learning (ICML). 2016. - """ - - def f1(a): - return a**2 - - def f2(b): - return b**2 - - def h1(a): - return a - - def h2(b): - return 2 * b - - constC1 = np.dot( - np.dot(f1(X1), v1.reshape(-1, 1)), np.ones(f1(X2).shape[0]).reshape(1, -1) - ) - constC2 = np.dot( - np.ones(f1(X1).shape[0]).reshape(-1, 1), np.dot(v2.reshape(1, -1), f2(X2).T) - ) - - constC = constC1 + constC2 - hX1 = h1(X1) - hX2 = h2(X2) - - return constC, hX1, hX2 diff --git a/perturbot/build/lib/perturbot/predict/__init__.py b/perturbot/build/lib/perturbot/predict/__init__.py deleted file mode 100755 index 07310cd..0000000 --- a/perturbot/build/lib/perturbot/predict/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .mlp import train_mlp - -__all__ = ["train_mlp"] diff --git a/perturbot/build/lib/perturbot/predict/linear_regression.py b/perturbot/build/lib/perturbot/predict/linear_regression.py deleted file mode 100755 index 56c7fea..0000000 --- a/perturbot/build/lib/perturbot/predict/linear_regression.py +++ /dev/null @@ -1,155 +0,0 @@ -import numpy as np -from typing import Dict, Union -from numpy.linalg import inv -from tqdm.auto import tqdm - - -def ols(X, Y): - """Regular OLS where we know item-to. Note that given different number of perturbation, - each data points are weighted as the same.""" - X = np.concatenate([np.ones((X.shape[0], 1)), X], axis=1) - return inv(X.T @ X) @ (X.T @ Y) - - -def ols_label(X_dict, Y_dict): - assert X_dict.keys() == Y_dict.keys() - xtx = 0 - xty = 0 - for l, X in X_dict.items(): - X = np.concatenate([np.ones((X.shape[0], 1)), X], axis=1) - xtx += X.T @ X - xty += X.T @ Y_dict[l] - """Equivalent to ols""" - return inv(xtx) @ (xty) - - -def weighted_ols( - X_dict: Dict[int, np.ndarray], - Y_dict: Dict[int, np.ndarray], - G_dict: Dict[int, np.ndarray], -): - """OLS for weighted matching between X and Y provided as G. - Return B = (\sum_l{(X^l)'Diag(G_x)(X^l)})^{-1}(\sum_l{\sum_i{(X_i^l)'(G_i)(Y^l)}}). - Each sample from the source X has the same weight. - where - l: label (keys of the dictionaries), - i: sample index in X^l, - G_i: i th row of G, - G_x: Marginal distribution of x from G - """ - assert X_dict.keys() == Y_dict.keys() - xtx = 0 - xty = 0 - for l, X in tqdm(X_dict.items()): - X = np.concatenate([np.ones((X.shape[0], 1)), X], axis=1) - Y = Y_dict[l] - G = G_dict[l] - G = G / G.sum(axis=-1)[:, None] - assert X.shape[0] == Y.shape[0] - assert X.shape[0] == G.shape[0] - assert Y.shape[0] == G.shape[1] - xtx += X.T @ np.diag(G.sum(axis=-1)) @ X - for i in range(X.shape[0]): - xty += X[[i], :].T @ (G[[i], :] @ Y) - return inv(xtx) @ xty - - -def weight_1_ols(X_dict: Dict[int, np.ndarray], Y_dict: Dict[int, np.ndarray]): - """OLS for all-to-all matching given label. - Return B = (\sum_l{(X^l)'X^l)})^{-1}(\sum_l (X^l)'Y^l). - Each sample from the source X has the same weight. - """ - assert X_dict.keys() == Y_dict.keys() - xtx = 0 - xty = 0 - for l, X in X_dict.items(): - X = np.concatenate([np.ones((X.shape[0], 1)), X], axis=1) - Y = Y_dict[l] - xtx += X.T @ X - xty += ( - X.sum(axis=0, keepdims=True).T @ Y.sum(axis=0, keepdims=True) / Y.shape[0] - ) - return inv(xtx) @ xty - - -def predict(X, params): - X = np.concatenate([np.ones((X.shape[0], 1)), X], axis=1) - return X @ params - - -def weighted_ols_normed( - X_dict: Dict[int, np.ndarray], - Y_dict: Dict[int, np.ndarray], - G_dict: Union[np.ndarray, Dict[int, np.ndarray]], -) -> np.ndarray: - """OLS for weighted matching between X and Y provided as G. - Return B = (\sum_l{(X^l)'Diag(G_x)(X^l)})^{-1}(\sum_l{\sum_i{(X_i^l)'(G_i)(Y^l)}}). - Each sample from the source X has the same weight. - where - l: label (keys of the dictionaries), - i: sample index in X^l, - G_i: i th row of G, - G_x: Marginal distribution of x from G - """ - assert X_dict.keys() == Y_dict.keys() - xtx = 0 - xty = 0 - if isinstance(G_dict, dict): - for l, X in tqdm(X_dict.items()): - X = np.concatenate([np.ones((X.shape[0], 1)), X], axis=1) - Y = Y_dict[l] - G = G_dict[l] - G = G / G.sum() - assert X.shape[0] == Y.shape[0] - assert X.shape[0] == G.shape[0] - assert Y.shape[0] == G.shape[1] - xtx += X.T @ np.diag(G.sum(axis=-1)) @ X - for i in range(X.shape[0]): - xty += X[[i], :].T @ (G[[i], :] @ Y) - else: - try: - G = G_dict / G_dict.sum() - except FloatingPointError: - G = (G_dict + 1e-6) / (G_dict.sum() + 1e-6) - labels = X_dict.keys() - X = np.concatenate( - [ - np.ones((sum([X_dict[l].shape[0] for l in labels]), 1)), - np.concatenate([X_dict[l] for l in labels], axis=0), - ], - axis=1, - ) - Y = np.concatenate([Y_dict[l] for l in labels], axis=0) - xtx += X.T @ np.diag(G.sum(axis=-1)) @ X - for i in range(X.shape[0]): - xty += X[[i], :].T @ (G[[i], :] @ Y) - xty += X[[i], :].T @ (G[[i], :] @ Y) - return inv(xtx) @ xty - - -def weight_1_ols_normed(X_dict, Y_dict, *args): - G_dict = { - k: np.ones((X_dict[k].shape[0], Y_dict[k].shape[0])) for k in X_dict.keys() - } - return weighted_ols_normed(X_dict, Y_dict, G_dict) - - -def ols_normed(X_dict, Y_dict, *args): - G_dict = {k: np.identity(X_dict[k].shape[0]) for k in X_dict.keys()} - return weighted_ols_normed(X_dict, Y_dict, G_dict) - - -def weight_conc_normed(X_dict, Y_dict, Z_dict, z_key="dosage"): - def make_G(size, label): - G = np.zeros((size, size)) - for l in np.unique(label): - l_idx = np.where(label == l)[0] - for i in l_idx: - for j in l_idx: - G[i, j] = 1 - return G - - if z_key in Z_dict: - Z_dict = Z_dict[z_key] - G_dict = {k: make_G(X_dict[k].shape[0], Z_dict[k]) for k in X_dict.keys()} - return weighted_ols_normed(X_dict, Y_dict, G_dict) diff --git a/perturbot/build/lib/perturbot/predict/mlp.py b/perturbot/build/lib/perturbot/predict/mlp.py deleted file mode 100755 index e45cf42..0000000 --- a/perturbot/build/lib/perturbot/predict/mlp.py +++ /dev/null @@ -1,250 +0,0 @@ -from typing import Sequence, Dict, Union, Tuple -from numbers import Number -import sys -import time -import numpy as np -import torch - -from torch.utils.data import Dataset, DataLoader, random_split -from torch import optim, nn -import lightning as L -from lightning.pytorch.callbacks import TQDMProgressBar -from scvi.train._callbacks import LoudEarlyStopping -from scvi.train._progress import ProgressBar -from perturbot.utils import mdict_to_matrix - - -class MLP(L.LightningModule): - def __init__(self, n_input, hidden_layers: Sequence[int], n_output: int): - super().__init__() - model = nn.Sequential() - self.log_scale = nn.Parameter(torch.randn(n_output)) - prev_dim = n_input - for i, n_dim in enumerate(hidden_layers): - model.add_module(f"dense{i}", nn.Linear(prev_dim, n_dim)) - model.add_module(f"batchnorm{i}", nn.BatchNorm1d(n_dim)) - model.add_module(f"act{i}", nn.ReLU()) - prev_dim = n_dim - model.add_module(f"dense_out", nn.Linear(prev_dim, n_output)) - self.model = model - self.batch_losses = [] - - def nll(self, x, x_hat): - scale = torch.exp(self.log_scale) - mean = x_hat - dist = torch.distributions.Normal(mean, scale) - - # measure prob of seeing image under p(x|z) - log_pxz = dist.log_prob(x) - return -log_pxz.sum() - - def training_step(self, batch, batch_idx): - x = batch["source"] - y = batch["target"][:, 0, :] - y_hat = self.model(x).squeeze() - loss = nn.functional.mse_loss(y_hat, y) - # loss = self.nll(y, y_hat) - self.batch_losses.append(loss) - self.log("train_loss", loss, on_epoch=True) - return loss - - def validation_step(self, val_batch, val_batch_idx): - x = val_batch["source"] - y = val_batch["target"][:, 0, :].squeeze() - y_hat = self.model(x).squeeze() - # loss = self.nll(y, y_hat) - loss = nn.functional.mse_loss(y_hat, y) - self.log("val_loss", loss, on_epoch=True) - - def configure_optimizers(self): - optimizer = optim.Adam(self.parameters(), lr=1e-3) - return optimizer - - def on_train_epoch_end(self): - loss = torch.stack(self.batch_losses).mean() - self.log("train_loss_epoch", loss, on_step=False, on_epoch=True, prog_bar=True) - self.batch_losses.clear() - - -class LabeledDataset(Dataset): - def __init__( - self, - X_dict: Dict[Union[int, float], np.ndarray], - Y_dict: Dict[Union[int, float], np.ndarray], - ): - """ - For data (X, Y) and coupling (T) dictionaries (label -> np.ndarray), - for each X, sample Y according to the coupling probability. - Assume T has uniform marginal on X. - If T is not provided, uniform matrix is used. - """ - assert sorted(X_dict.keys()) == sorted(Y_dict.keys()) - - self.X_dict = X_dict - self.Y_dict = Y_dict - self.labels = sorted(X_dict.keys()) - self.X = np.concatenate([X_dict[l] for l in self.labels]) - self.Y = np.concatenate([Y_dict[l] for l in self.labels]) - self.label = np.concatenate( - [np.ones(X_dict[l].shape[0]) * l for l in self.labels] - ) - - def __len__(self): - return self.X.shape[0] - - def __getitem__(self, idx): - sample = {"source": self.X[idx, :], "target": self.Y[idx, :]} - return sample - - -class LabeledCouplingDataset(Dataset): - def __init__( - self, - X_dict: Dict[Union[int, float], np.ndarray], - Y_dict: Dict[Union[int, float], np.ndarray], - T_dict: Dict[Union[int, float], np.ndarray] = None, - ): - """ - For data (X, Y) and coupling (T) dictionaries (label -> np.ndarray), - for each X, sample Y according to the coupling probability. - Assume T has uniform marginal on X. - If T is not provided, uniform matrix is used. - """ - assert sorted(X_dict.keys()) == sorted(Y_dict.keys()) - - self.labels = sorted(X_dict.keys()) - self.X = torch.from_numpy(np.concatenate([X_dict[l] for l in self.labels])) - self.Y = torch.from_numpy(np.concatenate([Y_dict[l] for l in self.labels])) - - if T_dict is None: - T_dict = {} - for l in self.labels: - T = np.ones((X_dict[l].shape[0], Y_dict[l].shape[0])) - T_dict[l] = T / T.sum() - - else: - assert self.labels == sorted(T_dict.keys()) - self.T = torch.from_numpy( - mdict_to_matrix( - T_dict, - np.concatenate([np.ones(X_dict[l].shape[0]) * l for l in self.labels]), - np.concatenate([np.ones(Y_dict[l].shape[0]) * l for l in self.labels]), - ) - ) - self.T[self.T.sum(axis=-1) == 0, :] = 1e-8 - - def _sample_Y_flattened(self, x_idx): - try: - Y_idx = torch.multinomial(self.T[x_idx, :], 1) - except RuntimeError as e: - print(x_idx) - torch.set_printoptions(threshold=10_000) - print(self.T[x_idx, :]) - raise e - return self.Y[Y_idx, :] - - def __len__(self): - return self.X.shape[0] - - def __getitem__(self, idx): - sample = {"source": self.X[idx, :], "target": self._sample_Y_flattened(idx)} - return sample - - -def train( - model: L.LightningModule, - dataset: Dataset, - val_prop: float = 0.05, - batch_size: int = 256, - seed: int = 2024, - max_epochs: int = 2000, - checkpoint_dir: str = ".", - loader_kwargs={}, - log_dir="./", -): - if not torch.cuda.is_available(): - torch.set_num_threads(10) - # csv_logger = CSVLogger(save_dir=log_dir, name="csv_file") - torch_seed = torch.Generator().manual_seed(seed) - # seed_everything(seed, workers=True) - train_set_size = int(len(dataset) * (1 - val_prop)) - valid_set_size = len(dataset) - train_set_size - train_set, valid_set = random_split( - dataset, [train_set_size, valid_set_size], generator=torch_seed - ) - train_loader = DataLoader(train_set, batch_size=batch_size, **loader_kwargs) - valid_loader = DataLoader(valid_set, batch_size=batch_size, **loader_kwargs) - trainer = L.Trainer( - max_epochs=max_epochs, - accelerator="cpu", - callbacks=[ - ProgressBar(), - LoudEarlyStopping(monitor="val_loss", mode="min", patience=45), - ], # , EarlyStopping(monitor="val_loss", mode="min")], - default_root_dir=checkpoint_dir, - # logger=[csv_logger], - # deterministic=True, - ) - trainer.fit(model.double(), train_loader, valid_loader) - return model - - -def train_mlp( - train_data: Tuple[Dict[Number, np.array], Dict[Number, np.array]], - T_dict: Dict[Number, np.array], -) -> Tuple[nn.Module, Dict]: - """Trains MLP predicting Y from X given the labeled sample-matching T_dict. - - Parameters - ---------- - data : - (source dataset, target dataset) where source and target datasets - are the dictionaries mapping label to np.ndarray with matched labels. - T_dict : - Optimal Transport coupling between the samples per label - - Returns - ------- - model : - Trained predictor - log : - Training log - """ - Xs_dict, Xt_dict = train_data - dim_X = Xs_dict[list(Xs_dict.keys())[0]].shape[1] - dim_Y = Xt_dict[list(Xt_dict.keys())[0]].shape[1] - if isinstance(T_dict, np.ndarray): - # T is all-to-all mapping, make it into a dictionary - dataset = LabeledCouplingDataset( - {0: np.concatenate([Xs_dict[l] for l in Xs_dict.keys()], axis=0)}, - {0: np.concatenate([Xt_dict[l] for l in Xs_dict.keys()], axis=0)}, - {0: T_dict}, - ) - else: - dataset = LabeledCouplingDataset(Xs_dict, Xt_dict, T_dict) - - mlp = MLP(dim_X, [min(dim_X, 256), min(dim_X, 256)], dim_Y) - start_time = time.time() - model = train(mlp, dataset).model - log = {"time": start_time - time.time()} - return model, log - - -class MyProgressBar(TQDMProgressBar): - def init_validation_tqdm(self): - bar = super().init_validation_tqdm() - if not sys.stdout.isatty(): - bar.disable = True - return bar - - def init_predict_tqdm(self): - bar = super().init_predict_tqdm() - if not sys.stdout.isatty(): - bar.disable = True - return bar - - def init_test_tqdm(self): - bar = super().init_test_tqdm() - if not sys.stdout.isatty(): - bar.disable = True - return bar diff --git a/perturbot/build/lib/perturbot/predict/scvi_vae.py b/perturbot/build/lib/perturbot/predict/scvi_vae.py deleted file mode 100755 index a1bad76..0000000 --- a/perturbot/build/lib/perturbot/predict/scvi_vae.py +++ /dev/null @@ -1,177 +0,0 @@ -from typing import Dict, Union, Tuple -import numpy as np -import pandas as pd -import anndata as ad -import torch -import scvi -from tqdm import tqdm -import time - - -def make_adata(X_dict, Y_dict): - adata_X = ad.AnnData( - X=np.concatenate([v for k, v in X_dict.items()]), - obs=pd.DataFrame( - { - "modality": "source", - "labels": np.concatenate( - [np.repeat(k, v.shape[0]) for k, v in X_dict.items()] - ), - } - ), - var=pd.DataFrame( - index=[f"PC{n+1}_X" for n in range(X_dict[list(X_dict.keys())[0]].shape[1])] - ), - ) - adata_Y = ad.AnnData( - X=np.concatenate([v for k, v in Y_dict.items()]), - obs=pd.DataFrame( - { - "modality": "target", - "labels": np.concatenate( - [np.repeat(k, v.shape[0]) for k, v in Y_dict.items()] - ), - } - ), - var=pd.DataFrame( - index=[f"PC{n+1}_Y" for n in range(Y_dict[list(Y_dict.keys())[0]].shape[1])] - ), - ) - sim_multi = ad.concat([adata_X, adata_Y], axis=1)[:0, :].copy() - return sim_multi - - -def train_vae_model( - data_dict, - eps: Union[float, Tuple[float]] = None, - use_label=True, -): - """VAE with adversarial loss. Assumes normal likelihood for the observations.""" - if isinstance(eps, tuple): - eps, latent_dim, learning_rate = eps - else: - latent_dim = 50 - learning_rate = 1e-4 - if not torch.cuda.is_available(): - scvi.settings.num_threads = 10 - else: - torch.set_float32_matmul_precision("medium") - X_dict = data_dict[0] - Y_dict = data_dict[1] - dim_X = X_dict[list(X_dict.keys())[0]].shape[1] - dim_Y = Y_dict[list(Y_dict.keys())[0]].shape[1] - assert X_dict.keys() == Y_dict.keys() - adata_X = ad.AnnData( - X=np.concatenate([v for k, v in X_dict.items()]), - obs=pd.DataFrame( - { - "modality": "source", - "labels": np.concatenate( - [np.repeat(k, v.shape[0]) for k, v in X_dict.items()] - ), - } - ), - var=pd.DataFrame(index=[f"PC{n+1}_X" for n in range(dim_X)]), - ) - adata_Y = ad.AnnData( - X=np.concatenate([v for k, v in Y_dict.items()]), - obs=pd.DataFrame( - { - "modality": "target", - "labels": np.concatenate( - [np.repeat(k, v.shape[0]) for k, v in Y_dict.items()] - ), - } - ), - var=pd.DataFrame(index=[f"PC{n+1}_Y" for n in range(dim_Y)]), - ) - sim_multi = make_adata(X_dict, Y_dict) - if not use_label: - sim_multi.obs.labels = 0 - adata_mvi = scvi.data.organize_multiome_anndatas(sim_multi, adata_X, adata_Y) - scvi.model.MATCHVI.setup_anndata(adata_mvi, batch_key="modality") - model = scvi.model.MATCHVI( - adata_mvi, - n_genes=adata_X.n_vars, - n_regions=adata_Y.n_vars, - gene_likelihood="normal", - accessibility_likelihood="normal", - n_hidden=min(256, adata_X.n_vars), - n_latent=latent_dim, - ) - start_time = time.time() - model.train( - plan_kwargs={"match": True, "adversarial_classifier": True}, - scale_match_loss=0, - scale_adversarial_loss=eps, - max_epochs=2000, - lr=learning_rate, - ) - return model.module, {"time": time.time() - start_time, "history": model.history} - - -def infer_from_X(X, module, dY): - fake_acc = np.empty((X.shape[0], dY)) - fake_acc[:] = np.nan - X = torch.from_numpy(np.concatenate([X, fake_acc], axis=1)).float() - print("infer from X:") - print(X.shape) - print(module.n_input_genes, module.n_input_regions) - - inference_output = module.inference( - X, - y=torch.zeros((X.shape[0], 1), requires_grad=False), - batch_index=torch.zeros((X.shape[0], 1)), - cont_covs=None, - cat_covs=None, - label=None, - cell_idx=torch.range(0, X.shape[0] - 1).unsqueeze(-1), - ) - return inference_output - - -def infer_from_Y(Y, module, dX): - fake_gexp = np.empty((Y.shape[0], dX)) - fake_gexp[:] = np.nan - X = torch.from_numpy(np.concatenate([fake_gexp, Y], axis=1)).float() - print("infer from Y:") - print(f"input:{X.shape}={fake_gexp.shape} + {Y.shape}") - print(module.n_input_genes, module.n_input_regions) - inference_output = module.inference( - X, - y=torch.zeros((X.shape[0], 1), requires_grad=False), - batch_index=torch.zeros((X.shape[0], 1)), - cont_covs=None, - cat_covs=None, - label=None, - cell_idx=torch.range(0, X.shape[0] - 1).unsqueeze(-1), - ) - return inference_output - - -def infer_from_Xs(X_dict, model, dY): - res = {} - for k, X in X_dict.items(): - res[k] = infer_from_X(X, model, dY)["qz_m"].detach().cpu().numpy() - return res - - -def infer_from_Ys(Y_dict, model, dX): - res = {} - for k, Y in Y_dict.items(): - res[k] = infer_from_Y(Y, model, dX)["qz_m"].detach().cpu().numpy() - return res - - -def predict_from_model(X, module, dY): - # model = model.to("cpu") - inference_output = infer_from_X(X, module, dY) - generative_output = module.generative( - inference_output["z"], - inference_output["qz_m"], - batch_index=torch.ones((X.shape[0], 1)), - libsize_expr=inference_output["libsize_expr"], - libsize_acc=inference_output["libsize_acc"], - use_z_mean=True, - ) - return generative_output["p"].detach().numpy() diff --git a/perturbot/build/lib/perturbot/predict/vae.py b/perturbot/build/lib/perturbot/predict/vae.py deleted file mode 100755 index aef119e..0000000 --- a/perturbot/build/lib/perturbot/predict/vae.py +++ /dev/null @@ -1,493 +0,0 @@ -from typing import Sequence, Dict, Union, Tuple -import collections.abc -import numpy as np - -import torch -from torch import nn -from torch.nn.functional import one_hot -from torch.nn import CrossEntropyLoss -import lightning as L -from scvi.nn import FCLayers # , one_hot -from scvi.module import Classifier -from torch.utils.data import Dataset -import warnings - -warnings.filterwarnings( - "ignore", ".*Trying to infer the `batch_size` from an ambiguous collection.*" -) - - -class AbstractVAE(L.LightningModule): - # https://github.com/williamFalcon/pytorch-lightning-vae/blob/main/vae.py - def __init__( - self, - ): - super().__init__() - - self.save_hyperparameters() - - # for the gaussian likelihood - self.log_scale = nn.Parameter(torch.Tensor([0.0])) - - def configure_optimizers(self): - return torch.optim.Adam(self.parameters(), lr=1e-4) - - def gaussian_likelihood(self, x_hat, logscale, x): - scale = torch.exp(logscale) - mean = x_hat - dist = torch.distributions.Normal(mean, scale) - - # measure prob of seeing image under p(x|z) - log_pxz = dist.log_prob(x) - return log_pxz.sum(dim=1) - - def kl_divergence(self, z, mu, std): - # -------------------------- - # Monte carlo KL divergence - # -------------------------- - # 1. define the first two probabilities (in this case Normal for both) - p = torch.distributions.Normal(torch.zeros_like(mu), torch.ones_like(std)) - q = torch.distributions.Normal(mu, std) - - # 2. get the probabilities from the equation - log_qzx = q.log_prob(z) - log_pz = p.log_prob(z) - - # kl - kl = log_qzx - log_pz - kl = kl.sum(-1) - return kl - - def training_step(self, batch, batch_idx): - x, _ = batch - - # encode x to get the mu and variance parameters - x_encoded = self.encoder(x) - mu, log_var = self.fc_mu(x_encoded), self.fc_var(x_encoded) - - # sample z from q - std = torch.exp(log_var / 2) - q = torch.distributions.Normal(mu, std) - z = q.rsample() - - # decoded - x_hat = self.decoder(z) - - # reconstruction loss - recon_loss = self.gaussian_likelihood(x_hat, self.log_scale, x) - - # kl - kl = self.kl_divergence(z, mu, std) - - # elbo - elbo = kl - recon_loss - elbo = elbo.mean() - - self.log_dict( - { - "elbo": elbo, - "kl": kl.mean(), - "recon_loss": recon_loss.mean(), - "reconstruction": recon_loss.mean(), - "kl": kl.mean(), - } - ) - - return elbo - - -class MatchVAE(AbstractVAE): - def __init__( - self, - n_input1, - n_input2, - n_hidden, - n_latent, - n_layers, - n_labels, - mod2_loss_scale=None, - adv_loss_scale=1.0, - weight_decay=1e-6, - ): - super().__init__() - self.automatic_optimization = False - self.save_hyperparameters() - self.log_scale1 = nn.Parameter(torch.zeros(n_input1)) - self.log_scale2 = nn.Parameter(torch.zeros(n_input2)) - # encoder, decoder - self.encoder1 = FCLayers( - n_in=n_input1, - n_out=n_latent, - n_layers=n_layers, - n_hidden=n_hidden, - ) - self.decoder1 = FCLayers( - n_in=n_latent, - n_out=n_input1, - n_layers=n_layers, - n_hidden=n_hidden, - ) - self.encoder2 = FCLayers( - n_in=n_input2, - n_out=n_latent, - n_layers=n_layers, - n_hidden=n_hidden, - ) - self.decoder2 = FCLayers( - n_in=n_latent, - n_out=n_input2, - n_layers=n_layers, - n_hidden=n_hidden, - ) - - # distribution parameters - self.fc_mu1 = nn.Linear(n_latent, n_latent) - self.fc_var1 = nn.Linear(n_latent, n_latent) - self.fc_mu2 = nn.Linear(n_latent, n_latent) - self.fc_var2 = nn.Linear(n_latent, n_latent) - - self.adversarial_classifier = Classifier( - n_input=n_latent + n_labels, - n_hidden=32, - n_labels=n_labels, - n_layers=2, - logits=True, - ) - if mod2_loss_scale is None: - mod2_loss_scale = n_input1 / n_input2 - self.mod2_loss_scale = mod2_loss_scale - self.adv_loss_scale = adv_loss_scale - self.weight_decay = weight_decay - self.n_labels = n_labels - - def encode(self, x1, x2): - z1 = self.fc_mu1(self.encoder1(x1)) - z2 = self.fc_mu2(self.encoder2(x2)) - return z1, z2 - - def configure_optimizers(self): - """Configure optimizers for adversarial training.""" - # params1 = filter(lambda p: p.requires_grad, self.parameters()) - generators = [ - self.encoder1.parameters(), - self.decoder1.parameters(), - self.encoder2.parameters(), - self.decoder2.parameters(), - self.fc_mu1.parameters(), - self.fc_var1.parameters(), - self.fc_mu2.parameters(), - self.fc_var2.parameters(), - ] - param1 = [] - for g in generators: - param1 += list(g) - param1 + [self.log_scale1, self.log_scale2] - optimizer1 = torch.optim.Adam( - param1, - lr=1e-3, - eps=0.01, - weight_decay=self.weight_decay, - ) - config1 = {"optimizer": optimizer1} - - params2 = filter( - lambda p: p.requires_grad, self.adversarial_classifier.parameters() - ) - optimizer2 = torch.optim.Adam( - params2, lr=1e-3, eps=0.01, weight_decay=self.weight_decay - ) - config2 = {"optimizer": optimizer2} - - # pytorch lightning requires this way to return - opts = [config1.pop("optimizer"), config2["optimizer"]] - if "lr_scheduler" in config1: - scheds = [config1["lr_scheduler"]] - return opts, scheds - else: - return opts - - def loss_adversarial_classifier( - self, z, batch_index, label_index, predict_true_class=True - ): - """Loss for adversarial classifier.""" - n_classes = self.n_labels - onehot_labels = one_hot(label_index, n_classes) - logits = self.adversarial_classifier(torch.concat([z, onehot_labels], axis=-1)) - - if predict_true_class: - cls_target = one_hot(batch_index, n_classes).float() - else: - one_hot_batch = one_hot(batch_index, n_classes) - # place zeroes where true label is - cls_target = (~one_hot_batch.bool()).float() - cls_target = cls_target / (n_classes - 1) - - ce_loss = CrossEntropyLoss() - loss = ce_loss(logits, cls_target) - return loss - - def training_step(self, batch, batch_idx): - opt1, opt2 = self.optimizers() - x1, x2, x1_label, x2_label = batch - - # encode x to get the mu and variance parameters - z1 = self.encoder1(x1) - mu1, log_var1 = self.fc_mu1(z1), self.fc_var1(z1) - - # sample z from q - std1 = torch.exp(log_var1 / 2) - q1 = torch.distributions.Normal(mu1, std1) - z1 = q1.rsample() - - # decoded - x_hat1 = self.decoder1(z1) - - # reconstruction loss - recon_loss1 = self.gaussian_likelihood(x_hat1, self.log_scale1, x1) - - # kl - kl1 = self.kl_divergence(z1, mu1, std1) - - # elbo - elbo1 = kl1 - recon_loss1 - elbo1 = elbo1.sum() / (x1.shape[0] + x2.shape[0]) - - # encode x to get the mu and variance parameters - z2 = self.encoder1(x2) - mu2, log_var2 = self.fc_mu1(z2), self.fc_var1(z2) - - # sample z from q - std2 = torch.exp(log_var2 / 2) - q2 = torch.distributions.Normal(mu2, std2) - z2 = q2.rsample() - - # decoded - x_hat2 = self.decoder1(z2) - - # reconstruction loss - recon_loss2 = self.gaussian_likelihood(x_hat2, self.log_scale2, x2) - - # kl - kl2 = self.kl_divergence(z2, mu2, std2) - - # elbo - elbo2 = kl2 - recon_loss2 - elbo2 = elbo2.sum() / (x1.shape[0] + x2.shape[0]) - - z = torch.concat([z1, z2], axis=0) - modality = ( - torch.concat([torch.zeros(z1.shape[0]), torch.zeros(z2.shape[0])], axis=0) - .long() - .to(z.device) - ) - label = torch.concat([x1_label, x2_label], axis=0) - fool_loss = self.loss_adversarial_classifier(z, modality, label, False) - - loss = elbo1 + elbo2 * self.mod2_loss_scale + fool_loss * self.adv_loss_scale - opt1.zero_grad() - self.manual_backward(loss) - opt1.step() - - adv_classifier_loss = ( - self.loss_adversarial_classifier(z.detach(), modality, label, True) - * self.adv_loss_scale - ) - opt2.zero_grad() - self.manual_backward(adv_classifier_loss) - opt2.step() - - self.log_dict( - { - "elbo1": elbo1, - "kl1": kl1.mean(), - "recon_loss1": recon_loss1.mean(), - "kl1": kl1.mean(), - "elbo2": elbo2, - "kl2": kl2.mean(), - "recon_loss2": recon_loss2.mean(), - "kl2": kl2.mean(), - "adv_loss": fool_loss, - "train_loss": loss, - "train_adv_classifier_loss": adv_classifier_loss, - "batch_size": x1.shape[0] + x2.shape[0], - } - ) - # self.log(batch_size=x1.shape[0] + x2.shape[0]) - - def validation_step(self, batch, batch_idx): - opt1, opt2 = self.optimizers() - x1, x2, x1_label, x2_label = batch - - # encode x to get the mu and variance parameters - z1 = self.encoder1(x1) - mu1, log_var1 = self.fc_mu1(z1), self.fc_var1(z1) - - # sample z from q - std1 = torch.exp(log_var1 / 2) - q1 = torch.distributions.Normal(mu1, std1) - z1 = q1.rsample() - - # decoded - x_hat1 = self.decoder1(z1) - - # reconstruction loss - recon_loss1 = self.gaussian_likelihood(x_hat1, self.log_scale1, x1) - - # kl - kl1 = self.kl_divergence(z1, mu1, std1) - - # elbo - elbo1 = kl1 - recon_loss1 - elbo1 = elbo1.sum() / (x1.shape[0] + x2.shape[0]) - - # encode x to get the mu and variance parameters - z2 = self.encoder1(x2) - mu2, log_var2 = self.fc_mu1(z2), self.fc_var1(z2) - - # sample z from q - std2 = torch.exp(log_var2 / 2) - q2 = torch.distributions.Normal(mu2, std2) - z2 = q2.rsample() - - # decoded - x_hat2 = self.decoder1(z2) - - # reconstruction loss - recon_loss2 = self.gaussian_likelihood(x_hat2, self.log_scale2, x2) - - # kl - kl2 = self.kl_divergence(z2, mu2, std2) - - # elbo - elbo2 = kl2 - recon_loss2 - elbo2 = elbo2.sum() / (x1.shape[0] + x2.shape[0]) - - z = torch.concat([z1, z2], axis=0) - modality = ( - torch.concat([torch.zeros(z1.shape[0]), torch.zeros(z2.shape[0])], axis=0) - .long() - .to(z.device) - ) - label = torch.concat([x1_label, x2_label], axis=0) - fool_loss = self.loss_adversarial_classifier(z, modality, label, False) - - loss = elbo1 + elbo2 * self.mod2_loss_scale + fool_loss * self.adv_loss_scale - - adv_classifier_loss = ( - self.loss_adversarial_classifier(z.detach(), modality, label, True) - * self.adv_loss_scale - ) - self.log_dict( - { - "val_elbo1": elbo1, - "val_kl1": kl1.mean(), - "val_recon_loss1": recon_loss1.mean(), - "val_kl1": kl1.mean(), - "val_elbo2": elbo2, - "val_kl2": kl2.mean(), - "val_recon_loss2": recon_loss2.mean(), - "val_kl2": kl2.mean(), - "val_adv_loss": fool_loss, - "val_loss": loss, - "val_train_adv_classifier_loss": adv_classifier_loss, - } - ) - # self.log(batch_size=x1.shape[0] + x2.shape[0]) - - def transfer_batch_to_device(self, batch, device, dataloader_idx): - if isinstance(batch, Tuple): - # move all tensors in your custom data structure to the device - return tuple(item.to(device) for item in batch) - batch.samples = batch.samples.to(device) - batch.targets = batch.targets.to(device) - elif dataloader_idx == 0: - # skip device transfer for the first dataloader or anything you wish - pass - else: - print(type(batch)) - batch = super().transfer_batch_to_device(batch, device, dataloader_idx) - return batch - - -class LabeledBimodalDataset(Dataset): - def __init__( - self, - X_dict: Dict[Union[int, float], np.ndarray], - Y_dict: Dict[Union[int, float], np.ndarray], - **kwargs, - ): - """ - For data (X, Y) and coupling (T) dictionaries (label -> np.ndarray), - for each X, sample Y according to the coupling probability. - Assume T has uniform marginal on X. - If T is not provided, uniform matrix is used. - """ - super().__init__() - assert sorted(X_dict.keys()) == sorted(Y_dict.keys()) - self.labels = sorted(X_dict.keys()) - self.labels_to_id = {i: l for i, l in enumerate(self.labels)} - X = np.zeros((0, X_dict[self.labels[0]].shape[1])) - Y = np.zeros((0, Y_dict[self.labels[0]].shape[1])) - modality = np.zeros((0)) - labels_X = np.zeros((0)) - labels_Y = np.zeros((0)) - for l in self.labels: - lid = self.labels_to_id[l] - X = np.concatenate([X, X_dict[l]]) - Y = np.concatenate([Y, Y_dict[l]]) - modality = np.concatenate( - [ - modality, - np.zeros(X_dict[l].shape[0], dtype=int), - np.zeros(Y_dict[l].shape[0], dtype=int), - ] - ) - labels_X = np.concatenate( - [labels_X, np.ones(X_dict[l].shape[0], dtype=int) * lid] - ) - labels_Y = np.concatenate( - [labels_Y, np.ones(X_dict[l].shape[0], dtype=int) * lid] - ) - self.X = X - self.Y = Y - self.modality = modality - self.labels_X = labels_X - self.labels_Y = labels_Y - self.data_idx = np.arange(0, self.X.shape[0] + self.Y.shape[0]) - np.random.shuffle(self.data_idx) - self.idx_order = np.argsort(self.data_idx) - # data_idx = 3 4 0 | 1 2 - # idx_order= 2 3 4 | 0 1 - # idx = 0, 1, 2 - # idx_order[idx] = [2, 3, 4] - - def __len__(self): - return self.X.shape[0] - - def __getitem__(self, idx): - X_idx = self.idx_order[idx][self.idx_order[idx] < self.X.shape[0]] - Y_idx = ( - self.idx_order[idx][self.idx_order[idx] >= self.X.shape[0]] - - self.X.shape[0] - ) - X = self.X[X_idx, :] - Y = self.Y[Y_idx, :] - labels_X = self.labels_X[X_idx] - labels_Y = self.labels_Y[Y_idx] - return X, Y, labels_X, labels_Y - - -def collate_fn(data): - """ - data: is a list of tuples with (example, label, length) - where 'example' is a tensor of arbitrary shape - and label/length are scalars - """ - x, y, labels_x, labels_y = zip(*data) - - return ( - torch.from_numpy(np.vstack(x)), - torch.from_numpy(np.vstack(y)), - torch.from_numpy(np.concatenate(labels_x)).long(), - torch.from_numpy(np.concatenate(labels_y)).long(), - ) diff --git a/perturbot/build/lib/perturbot/preprocess/__init__.py b/perturbot/build/lib/perturbot/preprocess/__init__.py deleted file mode 100755 index e69de29..0000000 diff --git a/perturbot/build/lib/perturbot/preprocess/vae.py b/perturbot/build/lib/perturbot/preprocess/vae.py deleted file mode 100755 index f5da651..0000000 --- a/perturbot/build/lib/perturbot/preprocess/vae.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Preprocess raw data to learn the latent embedding and projections.""" - -import os -import anndata as ad -import scanpy as sc -import scvi -import mudata as md - -SCVI_LATENT_KEY = "X_scVI" - - -def train_vae_rna(adata: ad.AnnData, save_dir="./"): - sc.pp.filter_genes(adata, min_counts=3) - adata.layers["counts"] = adata.X.copy() # preserve counts - sc.pp.normalize_total(adata, target_sum=1e4) - sc.pp.log1p(adata) - adata.raw = adata # freeze the state in `.raw` - sc.pp.highly_variable_genes( - adata, - n_top_genes=1200, - subset=True, - layer="counts", - flavor="seurat_v3", - batch_key="cell_source", - ) - model = scvi.model.SCVI(adata, n_latent=50) - model.train() - model_dir = os.path.join(save_dir, "scvi_model") - model.save(model_dir, overwrite=True) - - latent = model.get_latent_representation() - adata.obsm[SCVI_LATENT_KEY] = latent - return adata, model - - -def train_vae_acc(adata: ad.AnnData, save_dir="./"): - # compute the threshold: 5% of the cells - min_cells = int(adata.shape[0] * 0.05) - # in-place filtering of regions - sc.pp.filter_genes(adata, min_cells=min_cells) - scvi.model.PEAKVI.setup_anndata(adata) - model = scvi.model.PEAKVI(adata, n_hidden=50) - model.train() - model_dir = os.path.join(save_dir.name, "peakvi_pbmc") - model.save(model_dir, overwrite=True) - latent = model.get_latent_representation() - adata.obsm[SCVI_LATENT_KEY] = latent - return adata, model - - -def train_vae_prot(adata: ad.AnnData, save_dir="./"): - mdata = md.MuData({"rna": adata[:, :0].copy(), "protein": adata}) - scvi.model.TOTALVI.setup_mudata( - mdata, - rna_layer=None, - protein_layer=None, - modalities={ - "rna_layer": "rna", - "protein_layer": "protein", - }, - ) - model = scvi.model.TOTALVI(mdata, n_latent=50) - model.train() - adata.obsm[SCVI_LATENT_KEY] = model.get_latent_representation() - return adata, model diff --git a/perturbot/build/lib/perturbot/utils.py b/perturbot/build/lib/perturbot/utils.py deleted file mode 100644 index e8abb3e..0000000 --- a/perturbot/build/lib/perturbot/utils.py +++ /dev/null @@ -1,10 +0,0 @@ -import numpy as np - - -def mdict_to_matrix(M_dict, source_labels, target_labels): - Mtot = np.zeros((len(source_labels), len(target_labels))) - for l, M in M_dict.items(): - Mtot[ - np.ix_(np.where(source_labels == l)[0], np.where(target_labels == l)[0]) - ] = M - return Mtot From 2b533e5c88a31a40f332bf76ea4c4e06cfe6090f Mon Sep 17 00:00:00 2001 From: Aelrach Date: Sat, 29 Nov 2025 17:31:40 +0100 Subject: [PATCH 7/7] Removed other files not in the original repo --- .python-version | 1 - ott/src/ott_jax.egg-info/PKG-INFO | 356 ------------------ ott/src/ott_jax.egg-info/SOURCES.txt | 92 ----- ott/src/ott_jax.egg-info/dependency_links.txt | 1 - ott/src/ott_jax.egg-info/requires.txt | 41 -- ott/src/ott_jax.egg-info/top_level.txt | 1 - perturbot/perturbot.egg-info/PKG-INFO | 17 - perturbot/perturbot.egg-info/SOURCES.txt | 8 - .../perturbot.egg-info/dependency_links.txt | 1 - perturbot/perturbot.egg-info/requires.txt | 5 - perturbot/perturbot.egg-info/top_level.txt | 1 - pyproject.toml | 7 - 12 files changed, 531 deletions(-) delete mode 100644 .python-version delete mode 100644 ott/src/ott_jax.egg-info/PKG-INFO delete mode 100644 ott/src/ott_jax.egg-info/SOURCES.txt delete mode 100644 ott/src/ott_jax.egg-info/dependency_links.txt delete mode 100644 ott/src/ott_jax.egg-info/requires.txt delete mode 100644 ott/src/ott_jax.egg-info/top_level.txt delete mode 100644 perturbot/perturbot.egg-info/PKG-INFO delete mode 100644 perturbot/perturbot.egg-info/SOURCES.txt delete mode 100644 perturbot/perturbot.egg-info/dependency_links.txt delete mode 100644 perturbot/perturbot.egg-info/requires.txt delete mode 100644 perturbot/perturbot.egg-info/top_level.txt delete mode 100644 pyproject.toml diff --git a/.python-version b/.python-version deleted file mode 100644 index e4fba21..0000000 --- a/.python-version +++ /dev/null @@ -1 +0,0 @@ -3.12 diff --git a/ott/src/ott_jax.egg-info/PKG-INFO b/ott/src/ott_jax.egg-info/PKG-INFO deleted file mode 100644 index fa26352..0000000 --- a/ott/src/ott_jax.egg-info/PKG-INFO +++ /dev/null @@ -1,356 +0,0 @@ -Metadata-Version: 2.4 -Name: ott-jax -Version: 0.4.6a0 -Summary: Optimal Transport Tools in JAX. -Author-email: OTT team -License: Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - 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. - -Project-URL: Source Code, https://github.com/ott-jax/ott -Project-URL: Documentation, https://ott-jax.readthedocs.io -Project-URL: Issue Tracker, https://github.com/ott-jax/ott/issues -Project-URL: Changelog, https://github.com/ott-jax/ott/releases -Keywords: optimal transport,gromov wasserstein,sinkhorn,low-rank sinkhorn,sinkhorn divergences,wasserstein,wasserstein barycenter,jax,autodiff,implicit differentiation -Classifier: Typing :: Typed -Classifier: Development Status :: 4 - Beta -Classifier: License :: OSI Approved :: Apache Software License -Classifier: Topic :: Scientific/Engineering :: Mathematics -Classifier: Natural Language :: English -Classifier: Intended Audience :: Developers -Classifier: Intended Audience :: Science/Research -Classifier: Operating System :: POSIX :: Linux -Classifier: Operating System :: MacOS :: MacOS X -Classifier: Operating System :: Microsoft :: Windows -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Classifier: Programming Language :: Python :: 3.11 -Classifier: Programming Language :: Python :: 3.12 -Requires-Python: >=3.8 -Description-Content-Type: text/markdown -License-File: LICENSE -Requires-Dist: jax>=0.4.0 -Requires-Dist: jaxopt>=0.8 -Requires-Dist: lineax>=0.0.1; python_version >= "3.9" -Requires-Dist: numpy>=1.20.0 -Provides-Extra: neural -Requires-Dist: flax>=0.6.6; extra == "neural" -Requires-Dist: optax>=0.1.1; extra == "neural" -Provides-Extra: dev -Requires-Dist: pre-commit>=2.16.0; extra == "dev" -Requires-Dist: tox>=4; extra == "dev" -Provides-Extra: test -Requires-Dist: pytest; extra == "test" -Requires-Dist: pytest-xdist; extra == "test" -Requires-Dist: pytest-cov; extra == "test" -Requires-Dist: pytest-memray; extra == "test" -Requires-Dist: coverage[toml]; extra == "test" -Requires-Dist: chex; extra == "test" -Requires-Dist: networkx>=2.5; extra == "test" -Requires-Dist: scikit-learn>=1.0; extra == "test" -Requires-Dist: tqdm; extra == "test" -Requires-Dist: tslearn>=0.5; python_version < "3.12" and extra == "test" -Requires-Dist: lineax; python_version >= "3.9" and extra == "test" -Requires-Dist: matplotlib; extra == "test" -Provides-Extra: docs -Requires-Dist: sphinx>=4.0; extra == "docs" -Requires-Dist: sphinx-book-theme>=1.0.1; extra == "docs" -Requires-Dist: sphinx_autodoc_typehints>=1.12.0; extra == "docs" -Requires-Dist: sphinx-copybutton>=0.5.1; extra == "docs" -Requires-Dist: sphinxcontrib-bibtex>=2.5.0; extra == "docs" -Requires-Dist: sphinxcontrib-spelling>=7.7.0; extra == "docs" -Requires-Dist: myst-nb>=0.17.1; extra == "docs" -Dynamic: license-file - -logo - -# Optimal Transport Tools (OTT) -[![Downloads](https://static.pepy.tech/badge/ott-jax)](https://pypi.org/project/ott-jax/) -[![Tests](https://img.shields.io/github/actions/workflow/status/ott-jax/ott/tests.yml?branch=main)](https://github.com/ott-jax/ott/actions/workflows/tests.yml) -[![Docs](https://img.shields.io/readthedocs/ott-jax/latest)](https://ott-jax.readthedocs.io/en/latest/) -[![Coverage](https://img.shields.io/codecov/c/github/ott-jax/ott/main)](https://app.codecov.io/gh/ott-jax/ott) - -**See the [full documentation](https://ott-jax.readthedocs.io/en/latest/).** - -## What is OTT-JAX? -A ``JAX`` powered library to compute optimal transport at scale and on accelerators, ``OTT-JAX`` includes the fastest -implementation of the Sinkhorn algorithm you will find around. We have implemented all tweaks (scheduling, momentum, acceleration, initializations) and extensions (low-rank, entropic maps). They can be used directly between two datasets, or within more advanced problems -(Gromov-Wasserstein, barycenters). Some of ``JAX`` features, including -[JIT](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html#Using-jit-to-speed-up-functions), -[auto-vectorization](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html#Auto-vectorization-with-vmap) and -[implicit differentiation](https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html) -work towards the goal of having end-to-end differentiable outputs. ``OTT-JAX`` is led by a team of researchers at Apple, with contributions from Google and Meta researchers, as well as many academic partners, including TU München, Oxford, ENSAE/IP Paris, ENS Paris and the Hebrew University. - -## Installation -Install ``OTT-JAX`` from [PyPI](https://pypi.org/project/ott-jax/) as: -```bash -pip install ott-jax -``` -or with ``conda`` via [conda-forge](https://anaconda.org/conda-forge/ott-jax) as: -```bash -conda install -c conda-forge ott-jax -``` - -## What is optimal transport? -Optimal transport can be loosely described as the branch of mathematics and optimization that studies -*matching problems*: given two families of points, and a cost function on pairs of points, find a "good" (low cost) way -to associate bijectively to every point in the first family another in the second. - -Such problems appear in all areas of science, are easy to describe, yet hard to solve. Indeed, while matching optimally -two sets of $n$ points using a pairwise cost can be solved with the -[Hungarian algorithm](https://en.wikipedia.org/wiki/Hungarian_algorithm), solving it costs an order of $O(n^3)$ -operations, and lacks flexibility, since one may want to couple families of different sizes. - -Optimal transport extends all of this, through faster algorithms (in $n^2$ or even linear in $n$) along with numerous -generalizations that can help it handle weighted sets of different size, partial matchings, and even more evolved -so-called quadratic matching problems. - -In the simple toy example below, we compute the optimal coupling matrix between two point clouds sampled randomly -(2D vectors, compared with the squared Euclidean distance): - -## Example -```python -import jax -import jax.numpy as jnp - -from ott.geometry import pointcloud -from ott.problems.linear import linear_problem -from ott.solvers.linear import sinkhorn - -# sample two point clouds and their weights. -rngs = jax.random.split(jax.random.PRNGKey(0), 4) -n, m, d = 12, 14, 2 -x = jax.random.normal(rngs[0], (n,d)) + 1 -y = jax.random.uniform(rngs[1], (m,d)) -a = jax.random.uniform(rngs[2], (n,)) -b = jax.random.uniform(rngs[3], (m,)) -a, b = a / jnp.sum(a), b / jnp.sum(b) -# Computes the couplings using the Sinkhorn algorithm. -geom = pointcloud.PointCloud(x, y) -prob = linear_problem.LinearProblem(geom, a, b) - -solver = sinkhorn.Sinkhorn() -out = solver(prob) -``` - -The call to `solver(prob)` above works out the optimal transport solution. The `out` object contains a transport matrix -(here of size $12\times 14$) that quantifies the association strength between each point of the first point cloud, to one or -more points from the second, as illustrated in the plot below. We provide more flexibility to define custom cost -functions, objectives, and solvers, as detailed in the [full documentation](https://ott-jax.readthedocs.io/en/latest/). - -![obtained coupling](https://raw.githubusercontent.com/ott-jax/ott/main/docs/_static/images/couplings.png) - -## Citation -If you have found this work useful, please consider citing this reference: - -``` -@article{cuturi2022optimal, - title={Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein}, - author={Cuturi, Marco and Meng-Papaxanthos, Laetitia and Tian, Yingtao and Bunne, Charlotte and - Davis, Geoff and Teboul, Olivier}, - journal={arXiv preprint arXiv:2201.12324}, - year={2022} -} -``` -## See also -The [moscot](https://moscot.readthedocs.io/en/latest/index.html) package for OT analysis of multi-omics data also uses OTT as a backbone. diff --git a/ott/src/ott_jax.egg-info/SOURCES.txt b/ott/src/ott_jax.egg-info/SOURCES.txt deleted file mode 100644 index cdce04a..0000000 --- a/ott/src/ott_jax.egg-info/SOURCES.txt +++ /dev/null @@ -1,92 +0,0 @@ -CITATION.bib -CODE_OF_CONDUCT.md -CONTRIBUTING.md -LICENSE -MANIFEST.in -README.md -environment.yml -pyproject.toml -src/ott/__init__.py -src/ott/_version.py -src/ott/datasets.py -src/ott/py.typed -src/ott/types.py -src/ott/utils.py -src/ott/geometry/__init__.py -src/ott/geometry/costs.py -src/ott/geometry/distrib_costs.py -src/ott/geometry/epsilon_scheduler.py -src/ott/geometry/geodesic.py -src/ott/geometry/geometry.py -src/ott/geometry/graph.py -src/ott/geometry/grid.py -src/ott/geometry/low_rank.py -src/ott/geometry/pointcloud.py -src/ott/geometry/segment.py -src/ott/initializers/__init__.py -src/ott/initializers/linear/__init__.py -src/ott/initializers/linear/initializers.py -src/ott/initializers/linear/initializers_lr.py -src/ott/initializers/quadratic/__init__.py -src/ott/initializers/quadratic/initializers.py -src/ott/math/__init__.py -src/ott/math/fixed_point_loop.py -src/ott/math/matrix_square_root.py -src/ott/math/unbalanced_functions.py -src/ott/math/utils.py -src/ott/neural/__init__.py -src/ott/neural/layers.py -src/ott/neural/losses.py -src/ott/neural/models.py -src/ott/neural/solvers/__init__.py -src/ott/neural/solvers/conjugate.py -src/ott/neural/solvers/map_estimator.py -src/ott/neural/solvers/neuraldual.py -src/ott/problems/__init__.py -src/ott/problems/linear/__init__.py -src/ott/problems/linear/barycenter_problem.py -src/ott/problems/linear/linear_problem.py -src/ott/problems/linear/potentials.py -src/ott/problems/quadratic/__init__.py -src/ott/problems/quadratic/gw_barycenter.py -src/ott/problems/quadratic/quadratic_costs.py -src/ott/problems/quadratic/quadratic_problem.py -src/ott/solvers/__init__.py -src/ott/solvers/was_solver.py -src/ott/solvers/linear/__init__.py -src/ott/solvers/linear/_solve.py -src/ott/solvers/linear/acceleration.py -src/ott/solvers/linear/continuous_barycenter.py -src/ott/solvers/linear/discrete_barycenter.py -src/ott/solvers/linear/implicit_differentiation.py -src/ott/solvers/linear/lineax_implicit.py -src/ott/solvers/linear/lr_utils.py -src/ott/solvers/linear/sinkhorn.py -src/ott/solvers/linear/sinkhorn_lr.py -src/ott/solvers/linear/univariate.py -src/ott/solvers/quadratic/__init__.py -src/ott/solvers/quadratic/_solve.py -src/ott/solvers/quadratic/gromov_wasserstein.py -src/ott/solvers/quadratic/gromov_wasserstein_lr.py -src/ott/solvers/quadratic/gw_barycenter.py -src/ott/solvers/quadratic/lower_bound.py -src/ott/tools/__init__.py -src/ott/tools/k_means.py -src/ott/tools/plot.py -src/ott/tools/segment_sinkhorn.py -src/ott/tools/sinkhorn_divergence.py -src/ott/tools/soft_sort.py -src/ott/tools/gaussian_mixture/__init__.py -src/ott/tools/gaussian_mixture/fit_gmm.py -src/ott/tools/gaussian_mixture/fit_gmm_pair.py -src/ott/tools/gaussian_mixture/gaussian.py -src/ott/tools/gaussian_mixture/gaussian_mixture.py -src/ott/tools/gaussian_mixture/gaussian_mixture_pair.py -src/ott/tools/gaussian_mixture/linalg.py -src/ott/tools/gaussian_mixture/probabilities.py -src/ott/tools/gaussian_mixture/scale_tril.py -src/ott_jax.egg-info/PKG-INFO -src/ott_jax.egg-info/SOURCES.txt -src/ott_jax.egg-info/dependency_links.txt -src/ott_jax.egg-info/requires.txt -src/ott_jax.egg-info/top_level.txt \ No newline at end of file diff --git a/ott/src/ott_jax.egg-info/dependency_links.txt b/ott/src/ott_jax.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/ott/src/ott_jax.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/ott/src/ott_jax.egg-info/requires.txt b/ott/src/ott_jax.egg-info/requires.txt deleted file mode 100644 index c93e9c5..0000000 --- a/ott/src/ott_jax.egg-info/requires.txt +++ /dev/null @@ -1,41 +0,0 @@ -jax>=0.4.0 -jaxopt>=0.8 -numpy>=1.20.0 - -[:python_version >= "3.9"] -lineax>=0.0.1 - -[dev] -pre-commit>=2.16.0 -tox>=4 - -[docs] -sphinx>=4.0 -sphinx-book-theme>=1.0.1 -sphinx_autodoc_typehints>=1.12.0 -sphinx-copybutton>=0.5.1 -sphinxcontrib-bibtex>=2.5.0 -sphinxcontrib-spelling>=7.7.0 -myst-nb>=0.17.1 - -[neural] -flax>=0.6.6 -optax>=0.1.1 - -[test] -pytest -pytest-xdist -pytest-cov -pytest-memray -coverage[toml] -chex -networkx>=2.5 -scikit-learn>=1.0 -tqdm -matplotlib - -[test:python_version < "3.12"] -tslearn>=0.5 - -[test:python_version >= "3.9"] -lineax diff --git a/ott/src/ott_jax.egg-info/top_level.txt b/ott/src/ott_jax.egg-info/top_level.txt deleted file mode 100644 index 6ec4bc0..0000000 --- a/ott/src/ott_jax.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -ott diff --git a/perturbot/perturbot.egg-info/PKG-INFO b/perturbot/perturbot.egg-info/PKG-INFO deleted file mode 100644 index 629638f..0000000 --- a/perturbot/perturbot.egg-info/PKG-INFO +++ /dev/null @@ -1,17 +0,0 @@ -Metadata-Version: 2.4 -Name: perturbot -Version: 0.0.1 -Summary: OT for aligning multimodal perturbation -Author-email: Jayoung Ryu -Project-URL: Homepage, https://github.com/Genentech/Perturb-OT -Project-URL: Issues, https://github.com/Genentech/Perturb-OT/issues -Classifier: Programming Language :: Python :: 3 -Classifier: License :: OSI Approved :: MIT License -Classifier: Operating System :: OS Independent -Requires-Python: >=3.8 -Description-Content-Type: text/markdown -Requires-Dist: scanpy>=1.6 -Requires-Dist: POT>=0.9 -Requires-Dist: lightning -Requires-Dist: ott-jax==0.4.6a -Requires-Dist: scvi-tools==1.1.0a diff --git a/perturbot/perturbot.egg-info/SOURCES.txt b/perturbot/perturbot.egg-info/SOURCES.txt deleted file mode 100644 index 1c08e58..0000000 --- a/perturbot/perturbot.egg-info/SOURCES.txt +++ /dev/null @@ -1,8 +0,0 @@ -pyproject.toml -perturbot/__init__.py -perturbot/utils.py -perturbot.egg-info/PKG-INFO -perturbot.egg-info/SOURCES.txt -perturbot.egg-info/dependency_links.txt -perturbot.egg-info/requires.txt -perturbot.egg-info/top_level.txt \ No newline at end of file diff --git a/perturbot/perturbot.egg-info/dependency_links.txt b/perturbot/perturbot.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/perturbot/perturbot.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/perturbot/perturbot.egg-info/requires.txt b/perturbot/perturbot.egg-info/requires.txt deleted file mode 100644 index 6d3a9eb..0000000 --- a/perturbot/perturbot.egg-info/requires.txt +++ /dev/null @@ -1,5 +0,0 @@ -scanpy>=1.6 -POT>=0.9 -lightning -ott-jax==0.4.6a -scvi-tools==1.1.0a diff --git a/perturbot/perturbot.egg-info/top_level.txt b/perturbot/perturbot.egg-info/top_level.txt deleted file mode 100644 index 97cc772..0000000 --- a/perturbot/perturbot.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -perturbot diff --git a/pyproject.toml b/pyproject.toml deleted file mode 100644 index a499371..0000000 --- a/pyproject.toml +++ /dev/null @@ -1,7 +0,0 @@ -[project] -name = "perturb-ot" -version = "0.1.0" -description = "Add your description here" -readme = "README.md" -requires-python = ">=3.12" -dependencies = []