Remove functional API (#222)

* Remove `sinkhorn function` * Fix `sinkhorn_divergence` test * Remove `gromov_wasserstein` function * Remove `make` functions * Fix `soft_sort` and Jacobian tests * Remove `Transport` interface * Fix Jacobian test * Fix `soft_sort` and tests * Clean up some tests * Fix wrong `value_and_grad` usage * Update notebooks, isort and pre-commit * [ci skip] Fix rendering in `Sinkhorn` * Handle TODOs, clean initializer tests * Add `sinkhorn.solve` utility * Re-add `gromov_wasserstein.solve`, polish docs * Remove redundant line from `pyproject.toml` * Polish quad docs * Add rank to `sinkhorn.solve` * Add `rank` to `sinkhorn.solve`
ott-jax · Jan 7, 2023 · 7b14b52 · 7b14b52
1 parent a6feea4
commit 7b14b52
Show file tree

Hide file tree

Showing 84 changed files with 328,619 additions and 320,937 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -11,12 +11,15 @@ repos:
   hooks:
   - id: yapf
     additional_dependencies: [toml]
-- repo: https://github.com/tomcatling/black-nb
-  rev: '0.7'
+- repo: https://github.com/nbQA-dev/nbQA
+  rev: 1.6.0
   hooks:
-  - id: black-nb
+  - id: nbqa-pyupgrade
+    args: [--py38-plus]
+  - id: nbqa-black
+  - id: nbqa-isort
 - repo: https://github.com/PyCQA/isort
-  rev: 5.10.1
+  rev: 5.11.4
   hooks:
   - id: isort
 - repo: https://github.com/asottile/yesqa
@@ -30,7 +33,7 @@ repos:
     - flake8-bugbear
     - flake8-blind-except
 - repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
-  rev: v2.4.0
+  rev: v2.5.0
   hooks:
   - id: pretty-format-yaml
     args: [--autofix, --indent, '2']
@@ -65,7 +68,7 @@ repos:
     - flake8-blind-except
     args: [--docstring-convention, google]
 - repo: https://github.com/asottile/pyupgrade
-  rev: v3.2.2
+  rev: v3.3.1
   hooks:
   - id: pyupgrade
     args: [--py3-plus, --py37-plus, --keep-runtime-typing]
diff --git a/README.md b/README.md
@@ -52,25 +52,31 @@ In the simple toy example below, we compute the optimal coupling matrix between
 ```python
 import jax
 import jax.numpy as jnp
-from ott.tools import transport
-# Samples two point clouds and their weights.
-rngs = jax.random.split(jax.random.PRNGKey(0),4)
+
+from ott.geometry import pointcloud
+from ott.problems.linear import linear_problem
+from ott.solvers.linear import sinkhorn
+
+# sample two point clouds and their weights.
+rngs = jax.random.split(jax.random.PRNGKey(0), 4)
 n, m, d = 12, 14, 2
 x = jax.random.normal(rngs[0], (n,d)) + 1
 y = jax.random.uniform(rngs[1], (m,d))
 a = jax.random.uniform(rngs[2], (n,))
 b = jax.random.uniform(rngs[3], (m,))
 a, b = a / jnp.sum(a), b / jnp.sum(b)
 # Computes the couplings using the Sinkhorn algorithm.
-ot = transport.solve(x, y, a=a, b=b)
-P = ot.matrix
+geom = pointcloud.PointCloud(x, y)
+prob = linear_problem.LinearProblem(geom, a, b)
+
+solver = sinkhorn.Sinkhorn()
+out = solver(prob)
 ```
 
-The call to `solve` above works out the optimal transport solution. The `ot` object contains a transport matrix
+The call to `solver(prob)` above works out the optimal transport solution. The `out` object contains a transport matrix
 (here of size $12\times 14$) that quantifies a `link strength` between each point of the first point cloud, to one or
-more points from the second, as illustrated in the plot below. In this toy example, most choices were arbitrary, and
-are reflected in the crude `solve` API. We provide far more flexibility to define custom cost functions, objectives,
-and solvers, as detailed in the [full documentation](https://ott-jax.readthedocs.io/en/latest/).
+more points from the second, as illustrated in the plot below. We provide more flexibility to define custom cost
+functions, objectives, and solvers, as detailed in the [full documentation](https://ott-jax.readthedocs.io/en/latest/).
 
 ![obtained coupling](https://raw.githubusercontent.com/ott-jax/ott/main/images/couplings.png)
 

diff --git a/docs/conf.py b/docs/conf.py
@@ -70,10 +70,6 @@
 source_suffix = ['.rst']
 
 autosummary_generate = True
-autosummary_filename_map = {
-    "ott.solvers.linear.sinkhorn.sinkhorn":
-        "ott.solvers.linear.sinkhorn.sinkhorn-function"
-}
 
 autodoc_typehints = 'description'
 

diff --git a/docs/notebooks/GWLRSinkhorn.ipynb b/docs/notebooks/GWLRSinkhorn.ipynb
diff --git a/docs/notebooks/Hessians.ipynb b/docs/notebooks/Hessians.ipynb
diff --git a/docs/notebooks/LRSinkhorn.ipynb b/docs/notebooks/LRSinkhorn.ipynb
diff --git a/docs/notebooks/MetaOT.ipynb b/docs/notebooks/MetaOT.ipynb
@@ -40,44 +40,39 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "id": "9fde1353",
+   "id": "368061d9",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Note: you may need to restart the kernel to use updated packages.\n",
-      "Note: you may need to restart the kernel to use updated packages.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "%pip install -q ott-jax\n",
-    "%pip install -q torchvision"
+    "import sys\n",
+    "\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    !pip install -q git+https://github.com/ott-jax/ott@main\n",
+    "    !pip install -q torchvision"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 2,
-   "id": "a0b4b464",
+   "id": "9fde1353",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from ott.geometry import pointcloud\n",
-    "from ott.initializers.linear import initializers as init_lib\n",
-    "from ott.problems.linear import linear_problem\n",
-    "from ott.solvers.linear import sinkhorn\n",
+    "from collections import namedtuple\n",
+    "\n",
+    "import torchvision\n",
     "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
-    "\n",
     "import numpy as np\n",
-    "from collections import namedtuple\n",
-    "import torchvision\n",
     "\n",
     "import matplotlib.pyplot as plt\n",
-    "from matplotlib import cm"
+    "from matplotlib import cm\n",
+    "\n",
+    "from ott.geometry import pointcloud\n",
+    "from ott.initializers.linear import initializers as init_lib\n",
+    "from ott.problems.linear import linear_problem\n",
+    "from ott.solvers.linear import sinkhorn"
    ]
   },
   {
@@ -296,7 +291,7 @@
     "We interpret the pair of MNIST digits as discrete measures\n",
     "$\\alpha = \\sum_{i=1}^{n_a} a_i \\delta_{x_i}$ and $\\beta = \\sum_{j=1}^{n_b} b_j \\delta_{y_j}$.\n",
     "The default Sinkhorn implementation in \n",
-    "[ott.solvers.linear.sinkhorn.sinkhorn](../solvers/_autosummary/ott.solvers.linear.sinkhorn.sinkhorn-function.html)\n",
+    "[ott.solvers.linear.sinkhorn.Sinkhorn](../solvers/_autosummary/ott.solvers.linear.sinkhorn.sinkhorn.html)\n",
     "can easily compute their optimal coupling and associated\n",
     "dual potentials $f$ and $g$ from scratch.\n",
     "The optimal coupling between the measures can be used\n",
@@ -332,7 +327,11 @@
    ],
    "source": [
     "a, b = demo_batch.a[0], demo_batch.b[0]\n",
-    "base_sink_out = sinkhorn.sinkhorn(geom, a=a, b=b)\n",
+    "prob = linear_problem.LinearProblem(geom, a=a, b=b)\n",
+    "\n",
+    "solver = sinkhorn.Sinkhorn()\n",
+    "\n",
+    "base_sink_out = solver(prob)\n",
     "interpolate(\n",
     "    key, base_sink_out.f, base_sink_out.g, a, b, title=\"Sinkhorn interpolation\"\n",
     ")"
@@ -455,7 +454,7 @@
     }
    ],
    "source": [
-    "num_train_iterations = 50000\n",
+    "num_train_iterations = 50_000\n",
     "\n",
     "for train_iter in range(num_train_iterations):\n",
     "    key, step_key = jax.random.split(key)\n",
@@ -608,15 +607,16 @@
     "\n",
     "\n",
     "def get_sinkhorn_potentials(a, b):\n",
-    "    base_sink_out = sinkhorn.sinkhorn(geom, a=a, b=b)\n",
+    "    prob = linear_problem.LinearProblem(geom, a=a, b=b)\n",
+    "    base_sink_out = sinkhorn.Sinkhorn()(prob)\n",
     "    return base_sink_out.f, base_sink_out.g\n",
     "\n",
     "\n",
     "plot_demo_initializations(get_sinkhorn_potentials, title=\"Ground-truth\")\n",
     "\n",
     "\n",
     "def get_meta_ot_potentials(a, b):\n",
-    "    ot_problem = linear_problem.LinearProblem(geom, a, b)\n",
+    "    ot_problem = linear_problem.LinearProblem(geom, a=a, b=b)\n",
     "    f = meta_initializer.init_dual_a(ot_problem, lse_mode=True)\n",
     "    g = geom.update_potential(f, jnp.zeros_like(b), jnp.log(b), 0, axis=0)\n",
     "    return f, g\n",
@@ -687,21 +687,17 @@
     "    }\n",
     "\n",
     "    ot_problem = linear_problem.LinearProblem(geom, a=a, b=b)\n",
-    "    base_sink_out = sinkhorn.sinkhorn(\n",
-    "        geom, a=a, b=b, init_dual_a=None, **sink_kwargs\n",
-    "    )\n",
+    "    solver = sinkhorn.Sinkhorn(**sink_kwargs)\n",
+    "\n",
+    "    base_sink_out = solver((None, None))\n",
     "\n",
     "    init_dual_a = meta_initializer.init_dual_a(ot_problem, lse_mode=True)\n",
-    "    meta_sink_out = sinkhorn.sinkhorn(\n",
-    "        geom, a=a, b=b, init_dual_a=init_dual_a, **sink_kwargs\n",
-    "    )\n",
+    "    meta_sink_out = solver((init_dual_a, None))\n",
     "\n",
     "    init_dual_a = init_lib.GaussianInitializer().init_dual_a(\n",
     "        ot_problem, lse_mode=True\n",
     "    )\n",
-    "    gaus_sink_out = sinkhorn.sinkhorn(\n",
-    "        geom, a=a, b=b, init_dual_a=init_dual_a, **sink_kwargs\n",
-    "    )\n",
+    "    gaus_sink_out = solver((init_dual_a, None))\n",
     "\n",
     "    error_log[\"base\"].append(base_sink_out.errors)\n",
     "    error_log[\"meta_ot\"].append(meta_sink_out.errors)\n",
@@ -773,7 +769,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.6"
+   "version": "3.10.8"
   }
  },
  "nbformat": 4,

diff --git a/docs/notebooks/OTT_&_POT.ipynb b/docs/notebooks/OTT_&_POT.ipynb
@@ -28,10 +28,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {
-    "id": "IO2KLVZ1KWvq"
-   },
+   "execution_count": null,
+   "metadata": {},
    "outputs": [],
    "source": [
     "import sys\n",
@@ -43,38 +41,43 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "id": "02VJX2uXYHDX"
-   },
+   "metadata": {},
    "source": [
     "... and import them, along with their numerical environments, `jax` and `numpy`."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {
-    "id": "ysURew0UKhHE"
-   },
+   "execution_count": null,
+   "metadata": {},
    "outputs": [],
    "source": [
-    "# import JAX and OTT\n",
+    "import timeit\n",
+    "\n",
+    "import ot\n",
+    "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
-    "import ott\n",
-    "from ott.geometry import pointcloud\n",
-    "from ott.solvers.linear import sinkhorn\n",
-    "\n",
-    "# import OT, from POT\n",
     "import numpy as np\n",
-    "import ot\n",
     "\n",
-    "# misc\n",
     "import matplotlib.pyplot as plt\n",
-    "\n",
-    "plt.rc(\"font\", size=20)\n",
     "import mpl_toolkits.axes_grid1\n",
-    "import timeit"
+    "\n",
+    "import ott\n",
+    "from ott.geometry import pointcloud\n",
+    "from ott.problems.linear import linear_problem\n",
+    "from ott.solvers.linear import sinkhorn"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "id": "ysURew0UKhHE"
+   },
+   "outputs": [],
+   "source": [
+    "plt.rc(\"font\", size=20)"
    ]
   },
   {
@@ -162,14 +165,14 @@
     "\n",
     "@jax.jit\n",
     "def solve_ott(a, b, x, y, 𝜀, threshold):\n",
-    "    out = sinkhorn.sinkhorn(\n",
-    "        pointcloud.PointCloud(x, y, epsilon=𝜀),\n",
-    "        a,\n",
-    "        b,\n",
-    "        threshold=threshold,\n",
-    "        lse_mode=True,\n",
-    "        max_iterations=1000,\n",
+    "    geom = pointcloud.PointCloud(x, y, epsilon=𝜀)\n",
+    "    prob = linear_problem.LinearProblem(geom, a=a, b=b)\n",
+    "\n",
+    "    solver = sinkhorn.Sinkhorn(\n",
+    "        threshold=threshold, lse_mode=True, max_iterations=1000\n",
     "    )\n",
+    "    out = solver(prob)\n",
+    "\n",
     "    f, g = out.f, out.g\n",
     "    f, g = f - np.mean(f), g + np.mean(\n",
     "        f\n",
@@ -322,8 +325,8 @@
     "    reg_ot[name] = np.ones((len(n_range), len(𝜀_range))) * np.nan\n",
     "    for i, n in enumerate(n_range):\n",
     "        for j, 𝜀 in enumerate(𝜀_range):\n",
-    "            exec, out = run_simulation(rng, n, 𝜀, threshold, solver_spec)\n",
-    "            exec_time[name][i, j] = exec\n",
+    "            t, out = run_simulation(rng, n, 𝜀, threshold, solver_spec)\n",
+    "            exec_time[name][i, j] = t\n",
     "            reg_ot[name][i, j] = out[-1]"
    ]
   },
@@ -380,9 +383,7 @@
     "    )\n",
     "    p[0].set_linestyle(\"dotted\")\n",
     "    p[1].set_linestyle(\"solid\")\n",
-    "    list_legend += [\n",
-    "        name + r\"  $\\varepsilon $=\" + \"{:.2g}\".format(𝜀) for 𝜀 in 𝜀_range\n",
-    "    ]\n",
+    "    list_legend += [name + r\"  $\\varepsilon $=\" + f\"{𝜀:.2g}\" for 𝜀 in 𝜀_range]\n",
     "\n",
     "plt.xticks(ticks=np.arange(len(n_range)), labels=n_range)\n",
     "plt.legend(list_legend)\n",
@@ -509,9 +510,9 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "ott",
    "language": "python",
-   "name": "python3"
+   "name": "ott"
   },
   "language_info": {
    "codemirror_mode": {