moscot 0.3.2__tar.gz → 0.3.4__tar.gz

{moscot-0.3.2 → moscot-0.3.4}/.pre-commit-config.yaml

@@ -7,29 +7,29 @@ default_stages:
 minimum_pre_commit_version: 3.0.0
 repos:
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.4.1
+    rev: v1.8.0
     hooks:
       - id: mypy
         additional_dependencies: [numpy>=1.25.0]
         files: ^src
   - repo: https://github.com/psf/black
-    rev: 23.7.0
+    rev: 24.2.0
     hooks:
       - id: black
         additional_dependencies: [toml]
   - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: v3.0.0
+    rev: v4.0.0-alpha.8
     hooks:
       - id: prettier
         language_version: system
   - repo: https://github.com/PyCQA/isort
-    rev: 5.12.0
+    rev: 5.13.2
     hooks:
       - id: isort
         additional_dependencies: [toml]
         args: [--order-by-type]
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
+    rev: v4.5.0
     hooks:
       - id: check-merge-conflict
       - id: check-ast
@@ -42,17 +42,17 @@ repos:
       - id: check-yaml
       - id: check-toml
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.9.0
+    rev: v3.15.1
     hooks:
       - id: pyupgrade
         args: [--py3-plus, --py38-plus, --keep-runtime-typing]
   - repo: https://github.com/asottile/blacken-docs
-    rev: 1.15.0
+    rev: 1.16.0
     hooks:
       - id: blacken-docs
         additional_dependencies: [black==23.1.0]
   - repo: https://github.com/rstcheck/rstcheck
-    rev: v6.1.2
+    rev: v6.2.0
     hooks:
       - id: rstcheck
         additional_dependencies: [tomli]
@@ -63,7 +63,7 @@ repos:
       - id: doc8
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.0.280
+    rev: v0.2.2
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{moscot-0.3.2 → moscot-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: moscot
-Version: 0.3.2
+Version: 0.3.4
 Summary: Multi-omic single-cell optimal transport tools
 Author: Dominik Klein, Giovanni Palla, Michal Klein, Zoe Piran, Marius Lange
 Maintainer-email: Dominik Klein <dominik.klein@helmholtz-muenchen.de>, Giovanni Palla <giovanni.palla@helmholtz-muenchen.de>, Michal Klein <michal.klein@helmholtz-muenchen.de>
@@ -56,11 +56,41 @@ Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
 Classifier: Topic :: Scientific/Engineering :: Mathematics
 Requires-Python: >=3.8
 Description-Content-Type: text/x-rst
+License-File: LICENSE
+Requires-Dist: numpy>=1.20.0
+Requires-Dist: scipy>=1.7.0
+Requires-Dist: pandas>=2.0.1
+Requires-Dist: networkx>=2.6.3
+Requires-Dist: matplotlib>=3.5.0
+Requires-Dist: anndata>=0.9.1
+Requires-Dist: scanpy>=1.9.3
+Requires-Dist: wrapt>=1.13.2
+Requires-Dist: docrep>=0.3.2
+Requires-Dist: ott-jax>=0.4.5
+Requires-Dist: cloudpickle>=2.2.0
+Requires-Dist: rich>=13.5
 Provides-Extra: spatial
+Requires-Dist: squidpy>=1.2.3; extra == "spatial"
 Provides-Extra: dev
+Requires-Dist: pre-commit>=3.0.0; extra == "dev"
+Requires-Dist: tox>=4; extra == "dev"
 Provides-Extra: test
+Requires-Dist: pytest>=7; extra == "test"
+Requires-Dist: pytest-xdist>=3; extra == "test"
+Requires-Dist: pytest-mock>=3.5.0; extra == "test"
+Requires-Dist: pytest-cov>=4; extra == "test"
+Requires-Dist: coverage[toml]>=7; extra == "test"
 Provides-Extra: docs
-License-File: LICENSE
+Requires-Dist: sphinx>=5.1.1; extra == "docs"
+Requires-Dist: sphinx_copybutton>=0.5.0; extra == "docs"
+Requires-Dist: sphinxcontrib-bibtex>=2.3.0; extra == "docs"
+Requires-Dist: sphinxcontrib-spelling>=7.6.2; extra == "docs"
+Requires-Dist: sphinx-autodoc-typehints; extra == "docs"
+Requires-Dist: furo>=2022.09.29; extra == "docs"
+Requires-Dist: sphinx-tippy>=0.4.1; extra == "docs"
+Requires-Dist: myst-nb>=0.17.1; extra == "docs"
+Requires-Dist: ipython>=7.20.0; extra == "docs"
+Requires-Dist: sphinx_design>=0.3.0; extra == "docs"
 
 |PyPI| |Downloads| |CI| |Pre-commit| |Codecov| |Docs|
 
@@ -125,6 +155,6 @@ Our preprint "Mapping cells through time and space with moscot" can be found `he
     :target: https://moscot.readthedocs.io/en/stable/
     :alt: Documentation
 
-.. |Downloads| image:: https://pepy.tech/badge/moscot
+.. |Downloads| image:: https://static.pepy.tech/badge/moscot
     :target: https://pepy.tech/project/moscot
     :alt: Downloads

{moscot-0.3.2 → moscot-0.3.4}/README.rst

@@ -61,6 +61,6 @@ Our preprint "Mapping cells through time and space with moscot" can be found `he
     :target: https://moscot.readthedocs.io/en/stable/
     :alt: Documentation
 
-.. |Downloads| image:: https://pepy.tech/badge/moscot
+.. |Downloads| image:: https://static.pepy.tech/badge/moscot
     :target: https://pepy.tech/project/moscot
     :alt: Downloads

{moscot-0.3.2 → moscot-0.3.4}/pyproject.toml

@@ -55,8 +55,9 @@ dependencies = [
     "scanpy>=1.9.3",
     "wrapt>=1.13.2",
     "docrep>=0.3.2",
-    "ott-jax>=0.4.3",
+    "ott-jax>=0.4.5",
     "cloudpickle>=2.2.0",
+    "rich>=13.5",
 ]
 
 [project.optional-dependencies]
@@ -121,6 +122,8 @@ ignore = [
     "D107",
     # Missing docstring in magic method
     "D105",
+    # Use `X | Y` for type annotations
+    "UP007",
 ]
 line-length = 120
 select = [

{moscot-0.3.2 → moscot-0.3.4}/src/moscot/_types.py

@@ -8,7 +8,7 @@ import numpy as np
 try:
     from numpy.typing import DTypeLike, NDArray
 
-    ArrayLike = NDArray[np.float_]
+    ArrayLike = NDArray[np.float64]
 except (ImportError, TypeError):
     ArrayLike = np.ndarray  # type: ignore[misc]
     DTypeLike = np.dtype  # type: ignore[misc]
@@ -33,20 +33,20 @@ OttCostFn_t = Literal[
     "euclidean",
     "sq_euclidean",
     "cosine",
-    "PNormP",
-    "SqPNorm",
-    "Euclidean",
-    "SqEuclidean",
-    "Cosine",
-    "ElasticL1",
-    "ElasticSTVS",
-    "ElasticSqKOverlap",
+    "pnorm_p",
+    "sq_pnorm",
+    "cosine",
+    "elastic_l1",
+    "elastic_l2",
+    "elastic_stvs",
+    "elastic_sqk_overlap",
+    "geodesic",
 ]
 OttCostFnMap_t = Union[OttCostFn_t, Mapping[Literal["xy", "x", "y"], OttCostFn_t]]
 GenericCostFn_t = Literal["barcode_distance", "leaf_distance", "custom"]
 CostFn_t = Union[str, GenericCostFn_t, OttCostFn_t]
 CostFnMap_t = Union[Union[OttCostFn_t, GenericCostFn_t], Mapping[str, Union[OttCostFn_t, GenericCostFn_t]]]
-PathLike = Union[os.PathLike, str]
+PathLike = Union[os.PathLike, str]  # type: ignore[type-arg]
 Policy_t = Literal[
     "sequential",
     "star",

{moscot-0.3.2 → moscot-0.3.4}/src/moscot/backends/ott/__init__.py

@@ -1,11 +1,11 @@
 from ott.geometry import costs
 
 from moscot.backends.ott._utils import sinkhorn_divergence
-from moscot.backends.ott.output import OTTOutput
+from moscot.backends.ott.output import GraphOTTOutput, OTTOutput
 from moscot.backends.ott.solver import GWSolver, SinkhornSolver
 from moscot.costs import register_cost
 
-__all__ = ["OTTOutput", "GWSolver", "SinkhornSolver", "sinkhorn_divergence"]
+__all__ = ["OTTOutput", "GraphOTTOutput", "GWSolver", "SinkhornSolver", "sinkhorn_divergence"]
 
 register_cost("euclidean", backend="ott")(costs.Euclidean)
 register_cost("sq_euclidean", backend="ott")(costs.SqEuclidean)
@@ -13,4 +13,6 @@ register_cost("cosine", backend="ott")(costs.Cosine)
 register_cost("pnorm_p", backend="ott")(costs.PNormP)
 register_cost("sq_pnorm", backend="ott")(costs.SqPNorm)
 register_cost("elastic_l1", backend="ott")(costs.ElasticL1)
+register_cost("elastic_l2", backend="ott")(costs.ElasticL2)
 register_cost("elastic_stvs", backend="ott")(costs.ElasticSTVS)
+register_cost("elastic_sqk_overlap", backend="ott")(costs.ElasticSqKOverlap)

{moscot-0.3.2 → moscot-0.3.4}/src/moscot/backends/ott/_utils.py

@@ -1,14 +1,17 @@
-from typing import Any, Optional, Union
+from typing import Any, Literal, Optional, Tuple, Union
 
 import jax
 import jax.numpy as jnp
 import scipy.sparse as sp
-from ott.geometry import epsilon_scheduler, geometry, pointcloud
+from ott.geometry import epsilon_scheduler, geodesic, geometry, pointcloud
 from ott.tools import sinkhorn_divergence as sdiv
 
 from moscot._logging import logger
 from moscot._types import ArrayLike, ScaleCost_t
 
+Scale_t = Union[float, Literal["mean", "median", "max_cost", "max_norm", "max_bound"]]
+
+
 __all__ = ["sinkhorn_divergence"]
 
 
@@ -86,3 +89,23 @@ def ensure_2d(arr: ArrayLike, *, reshape: bool = False) -> jax.Array:
     if arr.ndim != 2:
         raise ValueError(f"Expected array to have 2 dimensions, found `{arr.ndim}`.")
     return arr
+
+
+def _instantiate_geodesic_cost(
+    arr: jax.Array,
+    problem_shape: Tuple[int, int],
+    t: Optional[float],
+    is_linear_term: bool,
+    epsilon: Union[float, epsilon_scheduler.Epsilon] = None,
+    relative_epsilon: Optional[bool] = None,
+    scale_cost: Scale_t = 1.0,
+    directed: bool = True,
+    **kwargs: Any,
+) -> geometry.Geometry:
+    n_src, n_tgt = problem_shape
+    if is_linear_term and n_src + n_tgt != arr.shape[0]:
+        raise ValueError(f"Expected `x` to have `{n_src + n_tgt}` points, found `{arr.shape[0]}`.")
+    t = epsilon / 4.0 if t is None else t
+    cm_full = geodesic.Geodesic.from_graph(arr, t=t, directed=directed, **kwargs).cost_matrix
+    cm = cm_full[:n_src, n_src:] if is_linear_term else cm_full
+    return geometry.Geometry(cm, epsilon=epsilon, relative_epsilon=relative_epsilon, scale_cost=scale_cost)

{moscot-0.3.2 → moscot-0.3.4}/src/moscot/backends/ott/output.py

@@ -6,7 +6,7 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 from ott.solvers.linear import sinkhorn, sinkhorn_lr
-from ott.solvers.quadratic import gromov_wasserstein
+from ott.solvers.quadratic import gromov_wasserstein, gromov_wasserstein_lr
 
 import matplotlib as mpl
 import matplotlib.pyplot as plt
@@ -14,7 +14,7 @@ import matplotlib.pyplot as plt
 from moscot._types import ArrayLike, Device_t
 from moscot.base.output import BaseSolverOutput
 
-__all__ = ["OTTOutput"]
+__all__ = ["OTTOutput", "GraphOTTOutput"]
 
 
 class OTTOutput(BaseSolverOutput):
@@ -29,7 +29,13 @@ class OTTOutput(BaseSolverOutput):
     _NOT_COMPUTED = -1.0  # sentinel value used in `ott`
 
     def __init__(
-        self, output: Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput, gromov_wasserstein.GWOutput]
+        self,
+        output: Union[
+            sinkhorn.SinkhornOutput,
+            sinkhorn_lr.LRSinkhornOutput,
+            gromov_wasserstein.GWOutput,
+            gromov_wasserstein_lr.LRGWOutput,
+        ],
     ):
         super().__init__()
         self._output = output
@@ -168,7 +174,10 @@ class OTTOutput(BaseSolverOutput):
     def _apply(self, x: ArrayLike, *, forward: bool) -> ArrayLike:
         if x.ndim == 1:
             return self._output.apply(x, axis=1 - forward)
-        return self._output.apply(x.T, axis=1 - forward).T  # convert to batch first
+        return self._output.apply(
+            x.T,
+            axis=1 - forward,
+        ).T  # convert to batch first
 
     @property
     def shape(self) -> Tuple[int, int]:  # noqa: D102
@@ -218,8 +227,72 @@ class OTTOutput(BaseSolverOutput):
 
     @property
     def rank(self) -> int:  # noqa: D102
-        lin_output = self._output if self.is_linear else self._output.linear_state
-        return len(lin_output.g) if isinstance(lin_output, sinkhorn_lr.LRSinkhornOutput) else -1
+        output = self._output.linear_state if isinstance(self._output, gromov_wasserstein.GWOutput) else self._output
+        return (
+            len(output.g)
+            if isinstance(output, (sinkhorn_lr.LRSinkhornOutput, gromov_wasserstein_lr.LRGWOutput))
+            else -1
+        )
 
     def _ones(self, n: int) -> ArrayLike:  # noqa: D102
         return jnp.ones((n,))
+
+
+class GraphOTTOutput(OTTOutput):
+    """Output of :term:`OT` problems with a graph geometry in the linear term.
+
+    Parameters
+    ----------
+    output
+        Output of the :mod:`ott` backend.
+    shape
+        Shape of the problem.
+    """
+
+    def __init__(
+        self,
+        output: Union[
+            sinkhorn.SinkhornOutput,
+            sinkhorn_lr.LRSinkhornOutput,
+            gromov_wasserstein.GWOutput,
+            gromov_wasserstein_lr.LRGWOutput,
+        ],
+        shape: Tuple[int, int],
+    ):
+        super().__init__(output)
+        self._shape = shape
+
+    @property
+    def shape(self) -> Tuple[int, int]:  # noqa: D102
+        return self._shape
+
+    def _expand_data(self, x: jnp.ndarray, forward: bool) -> jnp.ndarray:
+        if forward:
+            shape = (self.shape[1],) if x.ndim == 1 else (self.shape[1], x.shape[1])
+            return jnp.concatenate((x, jnp.zeros(shape)))
+        shape = (self.shape[0],) if x.ndim == 1 else (self.shape[0], x.shape[1])
+        return jnp.concatenate((jnp.zeros(shape), x))
+
+    def _apply(self, x: ArrayLike, *, forward: bool) -> ArrayLike:
+        x_expanded = self._expand_data(x, forward=forward)
+        # ott-jax only supports lse_mode=False with graph geometry
+        res = self._output.apply(x_expanded.T, axis=1 - forward, lse_mode=False).T
+        return res[len(x) :] if forward else res[: -len(x)]
+
+    def to(self, device: Optional[Device_t] = None) -> "GraphOTTOutput":  # noqa: D102
+        if device is None:
+            return GraphOTTOutput(jax.device_put(self._output, device=device), shape=self.shape)
+
+        if isinstance(device, str) and ":" in device:
+            device, ix = device.split(":")
+            idx = int(ix)
+        else:
+            idx = 0
+
+        if not isinstance(device, xla_ext.Device):
+            try:
+                device = jax.devices(device)[idx]
+            except IndexError:
+                raise IndexError(f"Unable to fetch the device with `id={idx}`.") from None
+
+        return GraphOTTOutput(jax.device_put(self._output, device), shape=self.shape)

{moscot-0.3.2 → moscot-0.3.4}/src/moscot/backends/ott/solver.py

@@ -4,22 +4,34 @@ import types
 from typing import Any, Literal, Mapping, Optional, Set, Tuple, Union
 
 import jax
-from ott.geometry import costs, epsilon_scheduler, geometry, pointcloud
+import jax.numpy as jnp
+from ott.geometry import costs, epsilon_scheduler, geodesic, geometry, pointcloud
 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
+from ott.solvers.quadratic import gromov_wasserstein, gromov_wasserstein_lr
 
 from moscot._types import ProblemKind_t, QuadInitializer_t, SinkhornInitializer_t
-from moscot.backends.ott._utils import alpha_to_fused_penalty, check_shapes, ensure_2d
-from moscot.backends.ott.output import OTTOutput
+from moscot.backends.ott._utils import (
+    _instantiate_geodesic_cost,
+    alpha_to_fused_penalty,
+    check_shapes,
+    ensure_2d,
+)
+from moscot.backends.ott.output import GraphOTTOutput, OTTOutput
+from moscot.base.problems._utils import TimeScalesHeatKernel
 from moscot.base.solver import OTSolver
 from moscot.costs import get_cost
 from moscot.utils.tagged_array import TaggedArray
 
 __all__ = ["SinkhornSolver", "GWSolver"]
 
-OTTSolver_t = Union[sinkhorn.Sinkhorn, sinkhorn_lr.LRSinkhorn, gromov_wasserstein.GromovWasserstein]
+OTTSolver_t = Union[
+    sinkhorn.Sinkhorn,
+    sinkhorn_lr.LRSinkhorn,
+    gromov_wasserstein.GromovWasserstein,
+    gromov_wasserstein_lr.LRGromovWasserstein,
+]
 OTTProblem_t = Union[linear_problem.LinearProblem, quadratic_problem.QuadraticProblem]
 Scale_t = Union[float, Literal["mean", "median", "max_cost", "max_norm", "max_bound"]]
 
@@ -38,14 +50,21 @@ class OTTJaxSolver(OTSolver[OTTOutput], abc.ABC):
         self._solver: Optional[OTTSolver_t] = None
         self._problem: Optional[OTTProblem_t] = None
         self._jit = jit
+        self._a: Optional[jnp.ndarray] = None
+        self._b: Optional[jnp.ndarray] = None
 
     def _create_geometry(
         self,
         x: TaggedArray,
+        *,
+        is_linear_term: bool,
         epsilon: Union[float, epsilon_scheduler.Epsilon] = None,
         relative_epsilon: Optional[bool] = None,
         scale_cost: Scale_t = 1.0,
         batch_size: Optional[int] = None,
+        problem_shape: Optional[Tuple[int, int]] = None,
+        t: Optional[float] = None,
+        directed: bool = True,
         **kwargs: Any,
     ) -> geometry.Geometry:
         if x.is_point_cloud:
@@ -83,17 +102,80 @@ class OTTJaxSolver(OTSolver[OTTOutput], abc.ABC):
             return geometry.Geometry(
                 kernel_matrix=arr, epsilon=epsilon, relative_epsilon=relative_epsilon, scale_cost=scale_cost
             )
+        if x.is_graph:  # we currently only support this for the linear term.
+            return self._create_graph_geometry(
+                is_linear_term=is_linear_term,
+                x=x,
+                arr=arr,
+                problem_shape=problem_shape,
+                t=t,
+                epsilon=epsilon,
+                relative_epsilon=relative_epsilon,
+                scale_cost=scale_cost,
+                directed=directed,
+                **kwargs,
+            )
         raise NotImplementedError(f"Creating geometry from `tag={x.tag!r}` is not yet implemented.")
 
     def _solve(  # type: ignore[override]
         self,
         prob: OTTProblem_t,
         **kwargs: Any,
-    ) -> OTTOutput:
+    ) -> Union[OTTOutput, GraphOTTOutput]:
         solver = jax.jit(self.solver) if self._jit else self.solver
         out = solver(prob, **kwargs)
+        if isinstance(prob, linear_problem.LinearProblem) and isinstance(prob.geom, geodesic.Geodesic):
+            return GraphOTTOutput(out, shape=(len(self._a), len(self._b)))  # type: ignore[arg-type]
         return OTTOutput(out)
 
+    def _create_graph_geometry(
+        self,
+        is_linear_term: bool,
+        x: TaggedArray,
+        arr: jax.Array,
+        problem_shape: Optional[Tuple[int, int]],
+        t: Optional[float],
+        epsilon: Union[float, epsilon_scheduler.Epsilon] = None,
+        relative_epsilon: Optional[bool] = None,
+        scale_cost: Scale_t = 1.0,
+        directed: bool = True,
+        **kwargs: Any,
+    ) -> geometry.Geometry:
+        if x.cost == "geodesic":
+            if self.problem_kind == "linear":
+                if t is None:
+                    if epsilon is None:
+                        raise ValueError("`epsilon` cannot be `None`.")
+                    return geodesic.Geodesic.from_graph(arr, t=epsilon / 4.0, directed=directed, **kwargs)
+
+                return _instantiate_geodesic_cost(
+                    arr=arr,
+                    problem_shape=problem_shape,  # type: ignore[arg-type]
+                    t=t,
+                    is_linear_term=True,
+                    epsilon=epsilon,
+                    relative_epsilon=relative_epsilon,
+                    scale_cost=scale_cost,
+                    directed=directed,
+                    **kwargs,
+                )
+            if self.problem_kind == "quadratic":
+                problem_shape = x.shape if problem_shape is None else problem_shape
+                return _instantiate_geodesic_cost(
+                    arr=arr,
+                    problem_shape=problem_shape,
+                    t=t,
+                    is_linear_term=is_linear_term,
+                    epsilon=epsilon,
+                    relative_epsilon=relative_epsilon,
+                    scale_cost=scale_cost,
+                    directed=directed,
+                    **kwargs,
+                )
+
+            raise NotImplementedError(f"Invalid problem kind `{self.problem_kind}`.")
+        raise NotImplementedError(f"If the geometry is a graph, `cost` must be `geodesic`, found `{x.cost}`.")
+
     @property
     def solver(self) -> OTTSolver_t:
         """:mod:`ott` solver."""
@@ -160,6 +242,8 @@ class SinkhornSolver(OTTJaxSolver):
 
     def _prepare(
         self,
+        a: jnp.ndarray,
+        b: jnp.ndarray,
         xy: Optional[TaggedArray] = None,
         x: Optional[TaggedArray] = None,
         y: Optional[TaggedArray] = None,
@@ -170,24 +254,35 @@ class SinkhornSolver(OTTJaxSolver):
         scale_cost: Scale_t = 1.0,
         cost_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
         cost_matrix_rank: Optional[int] = None,
+        time_scales_heat_kernel: Optional[TimeScalesHeatKernel] = None,
         # problem
         **kwargs: Any,
     ) -> linear_problem.LinearProblem:
         del x, y
+        time_scales_heat_kernel = (
+            TimeScalesHeatKernel(None, None, None) if time_scales_heat_kernel is None else time_scales_heat_kernel
+        )
         if xy is None:
             raise ValueError(f"Unable to create geometry from `xy={xy}`.")
-
+        self._a = a
+        self._b = b
         geom = self._create_geometry(
             xy,
+            is_linear_term=True,
             epsilon=epsilon,
             relative_epsilon=relative_epsilon,
             batch_size=batch_size,
+            problem_shape=(len(self._a), len(self._b)),
             scale_cost=scale_cost,
+            t=time_scales_heat_kernel.xy,
             **cost_kwargs,
         )
         if cost_matrix_rank is not None:
             geom = geom.to_LRCGeometry(rank=cost_matrix_rank)
-        self._problem = linear_problem.LinearProblem(geom, **kwargs)
+        if isinstance(geom, geodesic.Geodesic):
+            a = jnp.concatenate((a, jnp.zeros_like(self._b)), axis=0)
+            b = jnp.concatenate((jnp.zeros_like(self._a), b), axis=0)
+        self._problem = linear_problem.LinearProblem(geom, a=a, b=b, **kwargs)
         return self._problem
 
     @property
@@ -201,7 +296,15 @@ class SinkhornSolver(OTTJaxSolver):
 
     @classmethod
     def _call_kwargs(cls) -> Tuple[Set[str], Set[str]]:
-        geom_kwargs = {"epsilon", "relative_epsilon", "batch_size", "scale_cost", "cost_kwargs", "cost_matrix_rank"}
+        geom_kwargs = {
+            "epsilon",
+            "relative_epsilon",
+            "batch_size",
+            "scale_cost",
+            "cost_kwargs",
+            "cost_matrix_rank",
+            "t",
+        }
         problem_kwargs = set(inspect.signature(linear_problem.LinearProblem).parameters.keys())
         problem_kwargs -= {"geom"}
         return geom_kwargs | problem_kwargs, {"epsilon"}
@@ -243,24 +346,30 @@ class GWSolver(OTTJaxSolver):
     ):
         super().__init__(jit=jit)
         if rank > -1:
-            linear_solver_kwargs = dict(linear_solver_kwargs)
-            linear_solver_kwargs.setdefault("gamma", 10)
-            linear_solver_kwargs.setdefault("gamma_rescale", True)
-            linear_ot_solver = sinkhorn_lr.LRSinkhorn(rank=rank, **linear_solver_kwargs)
+            kwargs.setdefault("gamma", 10)
+            kwargs.setdefault("gamma_rescale", True)
             initializer = "rank2" if initializer is None else initializer
+            self._solver = gromov_wasserstein_lr.LRGromovWasserstein(
+                rank=rank,
+                initializer=initializer,
+                kwargs_init=initializer_kwargs,
+                **kwargs,
+            )
         else:
             linear_ot_solver = sinkhorn.Sinkhorn(**linear_solver_kwargs)
             initializer = None
-        self._solver = gromov_wasserstein.GromovWasserstein(
-            rank=rank,
-            linear_ot_solver=linear_ot_solver,
-            quad_initializer=initializer,
-            kwargs_init=initializer_kwargs,
-            **kwargs,
-        )
+            self._solver = gromov_wasserstein.GromovWasserstein(
+                rank=rank,
+                linear_ot_solver=linear_ot_solver,
+                quad_initializer=initializer,
+                kwargs_init=initializer_kwargs,
+                **kwargs,
+            )
 
     def _prepare(
         self,
+        a: jnp.ndarray,
+        b: jnp.ndarray,
         xy: Optional[TaggedArray] = None,
         x: Optional[TaggedArray] = None,
         y: Optional[TaggedArray] = None,
@@ -271,32 +380,45 @@ class GWSolver(OTTJaxSolver):
         scale_cost: Scale_t = 1.0,
         cost_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
         cost_matrix_rank: Optional[int] = None,
+        time_scales_heat_kernel: Optional[TimeScalesHeatKernel] = None,
         # problem
         alpha: float = 0.5,
         **kwargs: Any,
     ) -> quadratic_problem.QuadraticProblem:
+        self._a = a
+        self._b = b
+        time_scales_heat_kernel = (
+            TimeScalesHeatKernel(None, None, None) if time_scales_heat_kernel is None else time_scales_heat_kernel
+        )
         if x is None or y is None:
             raise ValueError(f"Unable to create geometry from `x={x}`, `y={y}`.")
-        geom_kwargs: Any = {
+        geom_kwargs: dict[str, Any] = {
             "epsilon": epsilon,
             "relative_epsilon": relative_epsilon,
             "batch_size": batch_size,
             "scale_cost": scale_cost,
-            "cost_matrix_rank": cost_matrix_rank,
             **cost_kwargs,
         }
-        geom_xx = self._create_geometry(x, **geom_kwargs)
-        geom_yy = self._create_geometry(y, **geom_kwargs)
+        if cost_matrix_rank is not None:
+            geom_kwargs["cost_matrix_rank"] = cost_matrix_rank
+        geom_xx = self._create_geometry(x, t=time_scales_heat_kernel.x, is_linear_term=False, **geom_kwargs)
+        geom_yy = self._create_geometry(y, t=time_scales_heat_kernel.y, is_linear_term=False, **geom_kwargs)
         if alpha == 1.0 or xy is None:  # GW
             # arbitrary fused penalty; must be positive
             geom_xy, fused_penalty = None, 1.0
         else:  # FGW
             fused_penalty = alpha_to_fused_penalty(alpha)
-            geom_xy = self._create_geometry(xy, **geom_kwargs)
+            geom_xy = self._create_geometry(
+                xy,
+                t=time_scales_heat_kernel.xy,
+                problem_shape=(x.shape[0], y.shape[0]),
+                is_linear_term=True,
+                **geom_kwargs,
+            )
             check_shapes(geom_xx, geom_yy, geom_xy)
 
         self._problem = quadratic_problem.QuadraticProblem(
-            geom_xx, geom_yy, geom_xy, fused_penalty=fused_penalty, **kwargs
+            geom_xx, geom_yy, geom_xy, fused_penalty=fused_penalty, a=self._a, b=self._b, **kwargs
         )
         return self._problem