diff --git a/pyproject.toml b/pyproject.toml index 014775ffe..ea8f1a6ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,13 @@ dependencies = [ spatial = [ "squidpy>=1.2.3" ] + +neural = [ + "optax", + "flax", + "diffrax", + +] dev = [ "pre-commit>=3.0.0", "tox>=4", diff --git a/src/moscot/backends/ott/solver.py b/src/moscot/backends/ott/solver.py index 732404dbc..b4408cf8b 100644 --- a/src/moscot/backends/ott/solver.py +++ b/src/moscot/backends/ott/solver.py @@ -34,6 +34,7 @@ from ott.solvers.quadratic import gromov_wasserstein, gromov_wasserstein_lr from ott.solvers.utils import uniform_sampler +from moscot._logging import logger from moscot._types import ( ArrayLike, ProblemKind_t, @@ -277,8 +278,11 @@ def __init__( ): super().__init__(jit=jit) if rank > -1: - kwargs.setdefault("gamma", 10) + kwargs.setdefault("gamma", 500) kwargs.setdefault("gamma_rescale", True) + eps = kwargs.get("epsilon") + if eps is not None and eps > 0.0: + logger.info(f"Found `epsilon`={eps}>0. We recommend setting `epsilon`=0 for the low-rank solver.") initializer = "rank2" if initializer is None else initializer self._solver = sinkhorn_lr.LRSinkhorn( rank=rank, epsilon=epsilon, initializer=initializer, kwargs_init=initializer_kwargs, **kwargs @@ -395,6 +399,9 @@ def __init__( if rank > -1: kwargs.setdefault("gamma", 10) kwargs.setdefault("gamma_rescale", True) + eps = kwargs.get("epsilon") + if eps is not None and eps > 0.0: + logger.info(f"Found `epsilon`={eps}>0. We recommend setting `epsilon`=0 for the low-rank solver.") initializer = "rank2" if initializer is None else initializer self._solver = gromov_wasserstein_lr.LRGromovWasserstein( rank=rank, diff --git a/tests/backends/ott/test_backend.py b/tests/backends/ott/test_backend.py index 161b67f3b..d0b8cbb60 100644 --- a/tests/backends/ott/test_backend.py +++ b/tests/backends/ott/test_backend.py @@ -51,7 +51,8 @@ def test_matches_ott(self, x: Geom_t, eps: Optional[float], jit: bool): @pytest.mark.parametrize("initializer", ["random", "rank2", "k-means"]) def test_solver_rank(self, y: Geom_t, rank: Optional[int], initializer: str): eps = 1e-2 - lr_sinkhorn = LRSinkhorn(rank=rank, initializer=initializer) + default_gamma_lr_sinhorn = 500 + lr_sinkhorn = LRSinkhorn(rank=rank, initializer=initializer, gamma=default_gamma_lr_sinhorn) problem = LinearProblem(PointCloud(y, epsilon=eps)) gt = lr_sinkhorn(problem)