set defaults for lr sinkhorn (#739)

* set defaults for lr sinkhorn

* raise warning also for linear problems

* add additional neural optional dependencies

* adapt LRSinkhorn test to new default gamma

pyproject.toml

@@ -64,6 +64,13 @@ dependencies = [
 spatial = [
     "squidpy>=1.2.3"
 ]
+
+neural = [
+    "optax",
+    "flax",
+    "diffrax",
+
+]
 dev = [
     "pre-commit>=3.0.0",
     "tox>=4",

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,

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)