Feature/move taggedarray

src/moscot/_docs/_docs.py

@@ -323,7 +323,10 @@
 """
 _scale_cost = """\
 scale_cost
-    Method to scale cost matrices. If `None` no scaling is applied.
+    How to rescale the cost matrix. Implemented scalings are
+    'median', 'mean', 'max_cost', 'max_norm' and 'max_bound'.
+    Alternatively, a float factor can be given to rescale the cost such
+    that ``cost_matrix /= scale_cost``.
 """
 _cost_lin = """\
 cost

src/moscot/backends/ott/_output.py

@@ -77,7 +77,7 @@ def select_values(last_k: Optional[int] = None) -> Tuple[str, jnp.ndarray, jnp.n
         kind, values, xs = select_values(last_k)
 
         ax.plot(xs, values, **kwargs)
-        ax.set_xlabel("iteration")
+        ax.set_xlabel("iteration (logged)")
         ax.set_ylabel(kind)
         if title is None:
             title = "converged" if self.converged else "not converged"  # type: ignore[attr-defined]

src/moscot/problems/generic/_mixins.py

@@ -103,7 +103,7 @@ def push(
         **kwargs: Any,
     ) -> Optional[ApplyOutput_t[K]]:
         """
-        Push distribution of cells through time.
+        Push distribution of cells from source to target.
 
         Parameters
         ----------
@@ -160,7 +160,7 @@ def pull(
         **kwargs: Any,
     ) -> Optional[ApplyOutput_t[K]]:
         """
-        Pull distribution of cells through time.
+        Pull distribution of cells from target to source.
 
         Parameters
         ----------

src/moscot/problems/time/_mixins.py

@@ -450,7 +450,7 @@ def _get_data(
         else:
             raise ValueError(f"No data found for `{target}` time point.")
 
-        return (  # type: ignore[return-value]
+        return (  # type:ignore[return-value]
             source_data,
             growth_rates_source,
             intermediate_data,

src/moscot/solvers/_tagged_array.py

@@ -16,6 +16,7 @@
 
 
 def get_cost_function(cost: str, *, backend: Literal["ott"] = "ott", **kwargs: Any) -> Callable[..., Any]:
+    """Get backend-dependent cost function."""
     if backend == "ott":
         from moscot.backends.ott._solver import OTTCost