Deps/update ott

pyproject.toml

@@ -54,7 +54,8 @@ dependencies = [
     "anndata>=0.9.1",
     "scanpy>=1.9.3",
     "wrapt>=1.13.2",
-    "ott-jax==0.4.0",
+    "docrep>=0.3.2",
+    "ott-jax>=0.4.1",
     "cloudpickle>=2.2.0",
 ]
 

src/moscot/backends/ott/_utils.py

@@ -1,9 +1,9 @@
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import jax
 import jax.numpy as jnp
 import scipy.sparse as sp
-from ott.geometry import geometry, pointcloud
+from ott.geometry import epsilon_scheduler, geometry, pointcloud
 from ott.tools import sinkhorn_divergence as sdiv
 
 from moscot._logging import logger
@@ -17,7 +17,7 @@ def sinkhorn_divergence(
     point_cloud_2: ArrayLike,
     a: Optional[ArrayLike] = None,
     b: Optional[ArrayLike] = None,
-    epsilon: Optional[float] = 1e-1,
+    epsilon: Optional[Union[float, epsilon_scheduler.Epsilon]] = 1e-1,
     scale_cost: ScaleCost_t = 1.0,
     **kwargs: Any,
 ) -> float:

src/moscot/backends/ott/solver.py

@@ -141,7 +141,7 @@ def __init__(
         self,
         jit: bool = True,
         rank: int = -1,
-        epsilon: float = 0.0,
+        epsilon: Optional[float] = 0.0,
         initializer: SinkhornInitializer_t = None,
         initializer_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
         **kwargs: Any,

src/moscot/problems/spatiotemporal/_spatio_temporal.py

@@ -1,6 +1,8 @@
 import types
 from typing import Any, Literal, Mapping, Optional, Tuple, Type, Union
 
+from ott.geometry import epsilon_scheduler
+
 from anndata import AnnData
 
 from moscot import _constants
@@ -157,7 +159,7 @@ def prepare(
     def solve(
         self,
         alpha: float = 0.5,
-        epsilon: Optional[float] = 1e-3,
+        epsilon: Optional[Union[float, epsilon_scheduler.Epsilon]] = 1e-3,
         tau_a: float = 1.0,
         tau_b: float = 1.0,
         rank: int = -1,

tests/backends/ott/test_backend.py

@@ -67,8 +67,6 @@ def test_solver_rank(self, y: Geom_t, rank: Optional[int], initializer: str):
         np.testing.assert_allclose(solver._problem.geom.cost_matrix, problem.geom.cost_matrix, rtol=RTOL, atol=ATOL)
         np.testing.assert_allclose(gt.matrix, pred.transport_matrix, rtol=RTOL, atol=ATOL)
 
-    # TODO(michalk8): remove when new ott-jax version comes out
-    @pytest.mark.xfail(reason="broken on ott-jax==0.4.0")
     @pytest.mark.parametrize(
         ("rank", "cost_fn"), [(2, costs.Euclidean()), (3, costs.SqPNorm(p=1.5)), (5, costs.ElasticL1(0.1))]
     )

tests/problems/base/test_general_problem.py

@@ -61,8 +61,7 @@ def test_set_xy(self, adata_x: AnnData, adata_y: AnnData, tag: Literal["cost_mat
         prob.set_xy(cost_matrix, tag=tag)
         assert isinstance(prob.xy.data_src, np.ndarray)
         assert prob.xy.data_tgt is None
-
-        prob = prob.solve(max_iterations=5)  # TODO(@MUCDK) once fixed in OTT-JAX test for scale_cost
+        prob = prob.solve(epsilon=1.0, max_iterations=5)  # TODO(@MUCDK) once fixed in OTT-JAX test for scale_cost
         np.testing.assert_equal(prob.xy.data_src, cost_matrix.to_numpy())
 
     @pytest.mark.parametrize("tag", ["cost_matrix", "kernel"])
@@ -81,7 +80,7 @@ def test_set_x(self, adata_x: AnnData, adata_y: AnnData, tag: Literal["cost_matr
         assert isinstance(prob.x.data_src, np.ndarray)
         assert prob.x.data_tgt is None
 
-        prob = prob.solve(max_iterations=5)  # TODO(@MUCDK) once fixed in OTT-JAX test for scale_cost
+        prob = prob.solve(epsilon=1.0, max_iterations=5)  # TODO(@MUCDK) once fixed in OTT-JAX test for scale_cost
         np.testing.assert_equal(prob.x.data_src, cost_matrix.to_numpy())
 
     @pytest.mark.parametrize("tag", ["cost_matrix", "kernel"])
@@ -100,7 +99,7 @@ def test_set_y(self, adata_x: AnnData, adata_y: AnnData, tag: Literal["cost_matr
         assert isinstance(prob.y.data_src, np.ndarray)
         assert prob.y.data_tgt is None
 
-        prob = prob.solve(max_iterations=5)
+        prob = prob.solve(epsilon=1.0, max_iterations=5)
         np.testing.assert_equal(prob.y.data_src, cost_matrix.to_numpy())
 
     def test_set_xy_change_problem_kind(self, adata_x: AnnData, adata_y: AnnData):

tests/problems/cross_modality/test_translation_problem.py

@@ -3,6 +3,7 @@
 import pytest
 
 import numpy as np
+from ott.geometry import epsilon_scheduler
 
 from anndata import AnnData
 
@@ -88,7 +89,7 @@ def test_prepare_external_star_policy(
     def test_solve_balanced(
         self,
         adata_translation_split: Tuple[AnnData, AnnData],
-        epsilon: float,
+        epsilon: Optional[float],
         alpha: float,
         rank: int,
         src_attr: Mapping[str, str],
@@ -148,7 +149,15 @@ def test_pass_arguments(self, adata_translation_split: Tuple[AnnData, AnnData],
 
         geom = quad_prob.geom_xx
         for arg, val in geometry_args.items():
-            assert getattr(geom, val) == args_to_check[arg], arg
+            assert hasattr(geom, val)
+            el = getattr(geom, val)[0] if isinstance(getattr(geom, val), tuple) else getattr(geom, val)
+            if arg == "epsilon":
+                eps_processed = getattr(geom, val)
+                assert isinstance(eps_processed, epsilon_scheduler.Epsilon)
+                assert eps_processed.target == args_to_check[arg], arg
+            else:
+                assert getattr(geom, val) == args_to_check[arg], arg
+                assert el == args_to_check[arg]
 
         geom = quad_prob.geom_xy
         for arg, val in pointcloud_args.items():

tests/problems/generic/test_fgw_problem.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import pandas as pd
+from ott.geometry import epsilon_scheduler
 from ott.geometry.costs import (
     Cosine,
     ElasticL1,
@@ -117,7 +118,15 @@ def test_pass_arguments(self, adata_space_rotate: AnnData, args_to_check: Mappin
 
         geom = quad_prob.geom_xx
         for arg, val in geometry_args.items():
-            assert getattr(geom, val, object()) == args_to_check[arg], arg
+            assert hasattr(geom, val)
+            el = getattr(geom, val)[0] if isinstance(getattr(geom, val), tuple) else getattr(geom, val)
+            if arg == "epsilon":
+                eps_processed = getattr(geom, val)
+                assert isinstance(eps_processed, epsilon_scheduler.Epsilon)
+                assert eps_processed.target == args_to_check[arg], arg
+            else:
+                assert getattr(geom, val) == args_to_check[arg], arg
+                assert el == args_to_check[arg]
 
         geom = quad_prob.geom_xy
         for arg, val in pointcloud_args.items():
@@ -159,8 +168,8 @@ def test_set_xy(self, adata_time: AnnData, tag: Literal["cost_matrix", "kernel"]
             ("cosine", Cosine, {}),
             ("pnorm_p", PNormP, {"p": 3}),
             ("sq_pnorm", SqPNorm, {"xy": {"p": 5}, "x": {"p": 3}, "y": {"p": 4}}),
-            ("elastic_l1", ElasticL1, {"gamma": 1.1}),
-            ("elastic_stvs", ElasticSTVS, {"gamma": 1.2}),
+            ("elastic_l1", ElasticL1, {"scaling_reg": 1.1}),
+            ("elastic_stvs", ElasticSTVS, {"scaling_reg": 1.2}),
         ],
     )
     def test_prepare_costs(self, adata_time: AnnData, cost_str: str, cost_inst: Any, cost_kwargs: CostKwargs_t):

tests/problems/generic/test_gw_problem.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import pandas as pd
+from ott.geometry import epsilon_scheduler
 from ott.geometry.costs import (
     Cosine,
     ElasticL1,
@@ -135,7 +136,14 @@ def test_pass_arguments(self, adata_space_rotate: AnnData, args_to_check: Mappin
         geom = quad_prob.geom_xx
         for arg, val in geometry_args.items():
             assert hasattr(geom, val)
-            assert getattr(geom, val) == args_to_check[arg]
+            el = getattr(geom, val)[0] if isinstance(getattr(geom, val), tuple) else getattr(geom, val)
+            if arg == "epsilon":
+                eps_processed = getattr(geom, val)
+                assert isinstance(eps_processed, epsilon_scheduler.Epsilon)
+                assert eps_processed.target == args_to_check[arg], arg
+            else:
+                assert getattr(geom, val) == args_to_check[arg], arg
+                assert el == args_to_check[arg]
 
     @pytest.mark.fast()
     @pytest.mark.parametrize(
@@ -146,8 +154,8 @@ def test_pass_arguments(self, adata_space_rotate: AnnData, args_to_check: Mappin
             ("cosine", Cosine, {}),
             ("pnorm_p", PNormP, {"p": 3}),
             ("sq_pnorm", SqPNorm, {"x": {"p": 3}, "y": {"p": 4}}),
-            ("elastic_l1", ElasticL1, {"x": {"gamma": 3}, "y": {"gamma": 4}}),
-            ("elastic_stvs", ElasticSTVS, {"x": {"gamma": 3}, "y": {"gamma": 4}}),
+            ("elastic_l1", ElasticL1, {"x": {"scaling_reg": 3}, "y": {"scaling_reg": 4}}),
+            ("elastic_stvs", ElasticSTVS, {"x": {"scaling_reg": 3}, "y": {"scaling_reg": 4}}),
         ],
     )
     def test_prepare_costs(self, adata_time: AnnData, cost_str: str, cost_inst: Any, cost_kwargs: CostKwargs_t):

tests/problems/generic/test_sinkhorn_problem.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import pandas as pd
+from ott.geometry import epsilon_scheduler
 from ott.geometry.costs import (
     Cosine,
     ElasticL1,
@@ -72,8 +73,8 @@ def test_solve_balanced(self, adata_time: AnnData):
             ("cosine", Cosine, {}),
             ("pnorm_p", PNormP, {"p": 3}),
             ("sq_pnorm", SqPNorm, {"p": 3}),
-            ("elastic_l1", ElasticL1, {"gamma": 1.1}),
-            ("elastic_stvs", ElasticSTVS, {"gamma": 1.2}),
+            ("elastic_l1", ElasticL1, {"scaling_reg": 1.1}),
+            ("elastic_stvs", ElasticSTVS, {"scaling_reg": 1.2}),
         ],
     )
     def test_prepare_costs(self, adata_time: AnnData, cost_str: str, cost_inst: Any, cost_kwargs: Mapping[str, int]):
@@ -153,9 +154,15 @@ def test_pass_arguments(self, adata_time: AnnData, args_to_check: Mapping[str, A
 
         geom = lin_prob.geom
         for arg, val in geometry_args.items():
-            assert hasattr(geom, val), val
+            assert hasattr(geom, val)
             el = getattr(geom, val)[0] if isinstance(getattr(geom, val), tuple) else getattr(geom, val)
-            assert el == args_to_check[arg], arg
+            if arg == "epsilon":
+                eps_processed = getattr(geom, val)
+                assert isinstance(eps_processed, epsilon_scheduler.Epsilon)
+                assert eps_processed.target == args_to_check[arg], arg
+            else:
+                assert getattr(geom, val) == args_to_check[arg], arg
+                assert el == args_to_check[arg]
 
         args = pointcloud_args if args_to_check["rank"] == -1 else lr_pointcloud_args
         for arg, val in args.items():

tests/problems/space/test_alignment_problem.py

@@ -4,6 +4,7 @@
 import pytest
 
 import numpy as np
+from ott.geometry import epsilon_scheduler
 
 from anndata import AnnData
 
@@ -76,7 +77,7 @@ def test_prepare_star(self, adata_space_rotate: AnnData, reference: str):
     def test_solve_balanced(
         self,
         adata_space_rotate: AnnData,
-        epsilon: float,
+        epsilon: Optional[float],
         alpha: float,
         rank: int,
         initializer: Optional[Literal["random", "rank2"]],
@@ -154,7 +155,14 @@ def test_pass_arguments(self, adata_space_rotate: AnnData, args_to_check: Mappin
         geom = quad_prob.geom_xx
         for arg, val in geometry_args.items():
             assert hasattr(geom, val)
-            assert getattr(geom, val) == args_to_check[arg]
+            el = getattr(geom, val)[0] if isinstance(getattr(geom, val), tuple) else getattr(geom, val)
+            if arg == "epsilon":
+                eps_processed = getattr(geom, val)
+                assert isinstance(eps_processed, epsilon_scheduler.Epsilon)
+                assert eps_processed.target == args_to_check[arg], arg
+            else:
+                assert getattr(geom, val) == args_to_check[arg], arg
+                assert el == args_to_check[arg]
 
         geom = quad_prob.geom_xy
         for arg, val in pointcloud_args.items():

tests/problems/space/test_mapping_problem.py

@@ -4,6 +4,7 @@
 import pytest
 
 import numpy as np
+from ott.geometry import epsilon_scheduler
 
 from anndata import AnnData
 
@@ -96,7 +97,7 @@ def test_prepare_varnames(self, adata_mapping: AnnData, var_names: Optional[List
     def test_solve_balanced(
         self,
         adata_mapping: AnnData,
-        epsilon: float,
+        epsilon: Optional[float],
         alpha: float,
         rank: int,
         sc_attr: Mapping[str, str],
@@ -159,7 +160,14 @@ def test_pass_arguments(self, adata_mapping: AnnData, args_to_check: Mapping[str
         geom = quad_prob.geom_xx
         for arg, val in geometry_args.items():
             assert hasattr(geom, val)
-            assert getattr(geom, val) == args_to_check[arg]
+            el = getattr(geom, val)[0] if isinstance(getattr(geom, val), tuple) else getattr(geom, val)
+            if arg == "epsilon":
+                eps_processed = getattr(geom, val)
+                assert isinstance(eps_processed, epsilon_scheduler.Epsilon)
+                assert eps_processed.target == args_to_check[arg], arg
+            else:
+                assert getattr(geom, val) == args_to_check[arg], arg
+                assert el == args_to_check[arg]
 
         geom = quad_prob.geom_xy
         for arg, val in pointcloud_args.items():

tests/problems/space/test_mixins.py

@@ -51,6 +51,7 @@ def test_analysis(self, adata_space_rotate: AnnData):
                     == adata_ref[adata_ref.obs.batch == c].obsm["spatial"].shape
                 )
 
+    @pytest.mark.skip("TODO(giovp): regenerate solutions")
     def test_regression_testing(self, adata_space_rotate: AnnData):
         ap = AlignmentProblem(adata=adata_space_rotate).prepare(batch_key="batch").solve(alpha=0.5, epsilon=1)
         # TODO(giovp): unnecessary assert

tests/problems/spatio_temporal/test_spatio_temporal_problem.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import pandas as pd
+from ott.geometry import epsilon_scheduler
 
 from anndata import AnnData
 
@@ -213,7 +214,14 @@ def test_pass_arguments(self, adata_spatio_temporal: AnnData, args_to_check: Map
         geom = quad_prob.geom_xx
         for arg, val in geometry_args.items():
             assert hasattr(geom, val)
-            assert getattr(geom, val) == args_to_check[arg]
+            el = getattr(geom, val)[0] if isinstance(getattr(geom, val), tuple) else getattr(geom, val)
+            if arg == "epsilon":
+                eps_processed = getattr(geom, val)
+                assert isinstance(eps_processed, epsilon_scheduler.Epsilon)
+                assert eps_processed.target == args_to_check[arg], arg
+            else:
+                assert getattr(geom, val) == args_to_check[arg], arg
+                assert el == args_to_check[arg]
 
         geom = quad_prob.geom_xy
         for arg, val in pointcloud_args.items():

tests/problems/time/test_lineage_problem.py

@@ -3,6 +3,7 @@
 import pytest
 
 import numpy as np
+from ott.geometry import epsilon_scheduler
 
 from anndata import AnnData
 
@@ -254,7 +255,14 @@ def test_pass_arguments(self, adata_time_barcodes: AnnData, args_to_check: Mappi
         geom = quad_prob.geom_xx
         for arg, val in geometry_args.items():
             assert hasattr(geom, val)
-            assert getattr(geom, val) == args_to_check[arg]
+            el = getattr(geom, val)[0] if isinstance(getattr(geom, val), tuple) else getattr(geom, val)
+            if arg == "epsilon":
+                eps_processed = getattr(geom, val)
+                assert isinstance(eps_processed, epsilon_scheduler.Epsilon)
+                assert eps_processed.target == args_to_check[arg], arg
+            else:
+                assert getattr(geom, val) == args_to_check[arg], arg
+                assert el == args_to_check[arg]
 
         geom = quad_prob.geom_xy
         for arg, val in pointcloud_args.items():

tests/problems/time/test_temporal_problem.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import pandas as pd
+from ott.geometry import epsilon_scheduler
 
 from anndata import AnnData
 
@@ -269,7 +270,13 @@ def test_pass_arguments(self, adata_time: AnnData, args_to_check: Mapping[str, A
         for arg, val in geometry_args.items():
             assert hasattr(geom, val)
             el = getattr(geom, val)[0] if isinstance(getattr(geom, val), tuple) else getattr(geom, val)
-            assert el == args_to_check[arg]
+            if arg == "epsilon":
+                eps_processed = getattr(geom, val)
+                assert isinstance(eps_processed, epsilon_scheduler.Epsilon)
+                assert eps_processed.target == args_to_check[arg], arg
+            else:
+                assert getattr(geom, val) == args_to_check[arg], arg
+                assert el == args_to_check[arg]
 
         args = pointcloud_args if args_to_check["rank"] == -1 else lr_pointcloud_args
         for arg, val in args.items():