changes to gw_solver_args

src/moscot/backends/ott/solver.py

@@ -248,6 +248,7 @@ def __init__(
     ):
         super().__init__(jit=jit)
         if rank > -1:
+            kwargs = {**linear_solver_kwargs, **kwargs}
             kwargs.setdefault("gamma", 10)
             kwargs.setdefault("gamma_rescale", True)
             initializer = "random" if initializer is None else initializer
@@ -260,6 +261,7 @@ def __init__(
 
         else:
             initializer = None
+            kwargs = {**linear_solver_kwargs, **kwargs}
             self._solver = gromov_wasserstein.GromovWasserstein(
                 rank=rank,
                 quad_initializer=initializer,

tests/problems/conftest.py

@@ -154,7 +154,7 @@ def adata_time_with_tmap(adata_time: AnnData) -> AnnData:
     "gw_unbalanced_correction": False,
     "ranks": 3,
     "tolerances": 3e-2,
-    "warm_start": True,
+    #"warm_start": True,
     "linear_solver_kwargs": linear_solver_kwargs2,
 }
 
@@ -175,6 +175,14 @@ def adata_time_with_tmap(adata_time: AnnData) -> AnnData:
     "initializer": "quad_initializer",
 }
 
+gw_lr_solver_args = {
+    "epsilon": "epsilon",
+    "rank": "rank",
+    "threshold": "threshold",
+    "initializer_kwargs": "kwargs_init",
+    "initializer": "initializer",
+}
+
 gw_linear_solver_args = {
     "lse_mode": "lse_mode",
     "inner_iterations": "inner_iterations",

tests/problems/generic/test_gw_problem.py

@@ -29,6 +29,7 @@
     gw_linear_solver_args,
     gw_lr_linear_solver_args,
     gw_solver_args,
+    gw_lr_solver_args,
     quad_prob_args,
 )
 
@@ -113,11 +114,11 @@ def test_pass_arguments(self, adata_space_rotate: AnnData, args_to_check: Mappin
         problem = problem.solve(**args_to_check)
         key = ("0", "1")
         solver = problem[key].solver.solver
-        for arg, val in gw_solver_args.items():
+        for arg, val in gw_solver_args.items() if args_to_check["rank"] == -1 else gw_lr_solver_args.items():
             assert hasattr(solver, val)
             assert getattr(solver, val) == args_to_check[arg]
 
-        sinkhorn_solver = solver.linear_ot_solver
+        sinkhorn_solver = solver.linear_ot_solver if args_to_check["rank"] == -1 else solver
         lin_solver_args = gw_linear_solver_args if args_to_check["rank"] == -1 else gw_lr_linear_solver_args
         for arg, val in lin_solver_args.items():
             assert hasattr(sinkhorn_solver, val)