Update tutorial structure (#560)

* Better tutorial hierarchy * Fix typos in tutorials * Fix more links
ott-jax · Jul 12, 2024 · 0bf7fda · 0bf7fda
1 parent c76e2ac
commit 0bf7fda
Show file tree

Hide file tree

Showing 32 changed files with 373 additions and 385 deletions.
diff --git a/docs/tutorials/Sinkhorn_Barycenters.ipynb → ...als/barycenter/Sinkhorn_Barycenters.ipynb b/docs/tutorials/Sinkhorn_Barycenters.ipynb → ...als/barycenter/Sinkhorn_Barycenters.ipynb
diff --git a/docs/tutorials/gmm_pair_demo.ipynb → .../tutorials/barycenter/gmm_pair_demo.ipynb b/docs/tutorials/gmm_pair_demo.ipynb → .../tutorials/barycenter/gmm_pair_demo.ipynb
diff --git a/docs/tutorials/barycenter/index.rst b/docs/tutorials/barycenter/index.rst
@@ -0,0 +1,8 @@
+Barycenter
+==========
+
+.. toctree::
+    :maxdepth: 1
+    :glob:
+
+    *
diff --git a/...orials/wasserstein_barycenters_gmms.ipynb → ...center/wasserstein_barycenters_gmms.ipynb b/...orials/wasserstein_barycenters_gmms.ipynb → ...center/wasserstein_barycenters_gmms.ipynb
diff --git a/docs/tutorials/basic_ot_between_datasets.ipynb b/docs/tutorials/basic_ot_between_datasets.ipynb
@@ -65,9 +65,7 @@
    "cell_type": "code",
    "execution_count": 3,
    "id": "06e772a6",
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "rngs = jax.random.split(jax.random.PRNGKey(0), 2)\n",
@@ -187,9 +185,7 @@
    "cell_type": "code",
    "execution_count": 8,
    "id": "4d1d9c65",
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -245,7 +241,7 @@
    "source": [
     "## Automatic differentiation using  {mod}`jax`\n",
     "\n",
-    "We finish this quick tour by illustrating one of the main features of {mod}`ott`: it can be seamlessly integrated into differentiable, end-to-end architectures built using {mod}`jax` (see also {doc}`Hessians`) for an example exploiting unrolling or implicit differentiation).\n",
+    "We finish this quick tour by illustrating one of the main features of {mod}`ott`: it can be seamlessly integrated into differentiable, end-to-end architectures built using {mod}`jax` (see also {doc}`linear/Hessians`) for an example exploiting unrolling or implicit differentiation).\n",
     "\n",
     "We provide a simple use-case where we differentiate the (regularized) OT transport cost w.r.t. `x`,\n",
     "by defining a wrapper that takes `x` and `y` as input, to output their regularized OT cost."
@@ -346,24 +342,17 @@
     "\n",
     "This tutorial gave you a glimpse of the most basic features of {mod}`ott` and how they integrate with {mod}`jax`.\n",
     "{mod}`ott` implements many more functionalities, that are described in the following tutorials:\n",
-    "- Seamless integration of other or even custom cost-functions in {doc}`point_clouds`,\n",
-    "- Better performance of {class}`~ott.solvers.linear.sinkhorn.Sinkhorn` solvers using various acceleration techniques in {doc}`One_Sinkhorn`\n",
-    "- Extensions of that approach to {doc}`gromov_wasserstein`, to compare distributions  defined on heterogeneous spaces (for which a `cost_fn` $c(x,y)$ cannot be easily defined).\n",
-    "- {doc}`LRSinkhorn` for faster solvers that constraint coupling matrices (see plot above) to have a low-rank factorization, and exploit low-rank properties of {class}{class}`~ott.geometry.geometry.Geometry` objects, both for the standard OT problem\n",
-    "and its GW variant in {doc}`GWLRSinkhorn`.\n",
-    "- Wasserstein barycenters, as in {doc}`wasserstein_barycenters_gmms` or {doc}`Sinkhorn_Barycenters`,\n",
-    "- Multimarginal generalizations in {doc}`mmsink`,\n",
-    "- Differentiable sorting in {doc}`soft_sort`,\n",
-    "- Neural solvers in {doc}`neural_dual`, to estimate maps in functional form.\n",
-    "- Visual interface to plot progress bars in {doc}`tracking_progress`."
+    "- Seamless integration of other or even custom cost-functions in {doc}`geometry/point_cloud`,\n",
+    "- Better performance of {class}`~ott.solvers.linear.sinkhorn.Sinkhorn` solvers using various acceleration techniques in {doc}`linear/One_Sinkhorn`\n",
+    "- Extensions of that approach to {doc}`quadratic/gromov_wasserstein`, to compare distributions  defined on heterogeneous spaces (for which a `cost_fn` $c(x,y)$ cannot be easily defined).\n",
+    "- {doc}`linear/LRSinkhorn` for faster solvers that constraint coupling matrices (see plot above) to have a low-rank factorization, and exploit low-rank properties of {class}{class}`~ott.geometry.geometry.Geometry` objects, both for the standard OT problem\n",
+    "and its GW variant in {doc}`quadratic/GWLRSinkhorn`.\n",
+    "- Wasserstein barycenters, as in {doc}`barycenter/wasserstein_barycenters_gmms` or {doc}`barycenter/Sinkhorn_Barycenters`,\n",
+    "- Multimarginal generalizations in {doc}`linear/mmsink`,\n",
+    "- Differentiable sorting in {doc}`misc/soft_sort`,\n",
+    "- Neural solvers in {doc}`neural/neural_dual`, to estimate maps in functional form.\n",
+    "- Visual interface to plot progress bars in {doc}`misc/tracking_progress`."
    ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "id": "df0ac96a",
-   "metadata": {},
-   "source": []
   }
  ],
  "metadata": {
@@ -382,7 +371,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.6"
+   "version": "3.10.6"
   },
   "vscode": {
    "interpreter": {

diff --git a/docs/tutorials/introduction_grid.ipynb → docs/tutorials/geometry/grid.ipynb b/docs/tutorials/introduction_grid.ipynb → docs/tutorials/geometry/grid.ipynb
diff --git a/docs/tutorials/geometry/index.rst b/docs/tutorials/geometry/index.rst
@@ -0,0 +1,8 @@
+Geometry
+========
+
+.. toctree::
+    :maxdepth: 1
+    :glob:
+
+    *
diff --git a/docs/tutorials/point_clouds.ipynb → docs/tutorials/geometry/point_cloud.ipynb b/docs/tutorials/point_clouds.ipynb → docs/tutorials/geometry/point_cloud.ipynb
@@ -89,7 +89,7 @@
     "\n",
     "A {class}`~ott.geometry.pointcloud.PointCloud` geometry holds two arrays of vectors (supporting the two measures of interest), along with a cost function (a {class}`~ott.geometry.costs.CostFn` object, set by default to {class}`~ott.geometry.costs.SqEuclidean`) and, possibly an {class}`~ott.geometry.epsilon_scheduler.Epsilon` regularization parameter.\n",
     "\n",
-    "This geometry object defines a {class}`~ott.problems.linear.linear_problem.LinearProblem` object, which contains all the data needed to instantiate a linear OT problem (see the {doc}`gromov_wasserstein` tutorial for *quadratic* OT problems).\n",
+    "This geometry object defines a {class}`~ott.problems.linear.linear_problem.LinearProblem` object, which contains all the data needed to instantiate a linear OT problem (see the {doc}`../quadratic/gromov_wasserstein` tutorial for *quadratic* OT problems).\n",
     "\n",
     "We can then call a {class}`~ott.solvers.linear.sinkhorn.Sinkhorn` solver to solve that problem, and compute the OT between these points clouds. Note that all weights are assumed to be uniform in this notebook, but non-uniform weights can be passed as `a=..., b=...` arguments when defining the {class}`~ott.problems.linear.linear_problem.LinearProblem` below."
    ]
@@ -229,7 +229,7 @@
     "\n",
     "OTT returns quantities that are differentiable. In the following example, we leverage the gradients to move $n$ points in a way that minimizes the overall regularized OT cost, given a ground cost function. \n",
     "\n",
-    "We start by defining a minimal optimization loop, that does fixed-length gradient descent, and records various `ot` objects along the way for plotting. By choosing various cost functions, we can then plot different types of gradient flows for the point cloud in $x$. See also the {doc}`sinkhorn_divergence_gradient_flow` tutorial."
+    "We start by defining a minimal optimization loop, that does fixed-length gradient descent, and records various `ot` objects along the way for plotting. By choosing various cost functions, we can then plot different types of gradient flows for the point cloud in $x$. See also the {doc}`../linear/sinkhorn_divergence_gradient_flow` tutorial."
    ]
   },
   {
@@ -367066,7 +367066,7 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "Python 3.8.16 ('ott')",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -367080,7 +367080,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.16"
+   "version": "3.10.6"
   },
   "vscode": {
    "interpreter": {
@@ -367089,5 +367089,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 1
+ "nbformat_minor": 4
 }
diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst
@@ -1,61 +1,12 @@
 Tutorials
 =========
 
-Geometry
---------
 .. toctree::
-    :maxdepth: 1
+    :maxdepth: 2
 
-    introduction_grid
-
-Linear Optimal Transport
-------------------------
-.. toctree::
-    :maxdepth: 1
-
-    point_clouds
-    One_Sinkhorn
-    OTT_&_POT
-    Hessians
-    LRSinkhorn
-    mmsink
-    sinkhorn_divergence_gradient_flow
-    sparse_monge_displacements
-
-Barycenters
-^^^^^^^^^^^
-.. toctree::
-    :maxdepth: 1
-
-    Sinkhorn_Barycenters
-    gmm_pair_demo
-    wasserstein_barycenters_gmms
-
-Miscellaneous
-^^^^^^^^^^^^^
-.. toctree::
-    :maxdepth: 1
-
-    tracking_progress
-    soft_sort
-    application_biology
-
-Quadratic Optimal Transport
----------------------------
-.. toctree::
-    :maxdepth: 1
-
-    gromov_wasserstein
-    GWLRSinkhorn
-    gromov_wasserstein_multiomics
-
-Neural Optimal Transport
-------------------------
-.. toctree::
-    :maxdepth: 1
-
-    neural_dual
-    icnn_inits
-    MetaOT
-    Monge_Gap
-    ENOT
+    geometry/index
+    linear/index
+    quadratic/index
+    neural/index
+    barycenter/index
+    misc/index
diff --git a/docs/tutorials/Hessians.ipynb → docs/tutorials/linear/Hessians.ipynb b/docs/tutorials/Hessians.ipynb → docs/tutorials/linear/Hessians.ipynb
diff --git a/docs/tutorials/LRSinkhorn.ipynb → docs/tutorials/linear/LRSinkhorn.ipynb b/docs/tutorials/LRSinkhorn.ipynb → docs/tutorials/linear/LRSinkhorn.ipynb
diff --git a/docs/tutorials/OTT_&_POT.ipynb → docs/tutorials/linear/OTT_&_POT.ipynb b/docs/tutorials/OTT_&_POT.ipynb → docs/tutorials/linear/OTT_&_POT.ipynb
diff --git a/docs/tutorials/One_Sinkhorn.ipynb → docs/tutorials/linear/One_Sinkhorn.ipynb b/docs/tutorials/One_Sinkhorn.ipynb → docs/tutorials/linear/One_Sinkhorn.ipynb
diff --git a/docs/tutorials/linear/index.rst b/docs/tutorials/linear/index.rst
@@ -0,0 +1,8 @@
+Linear OT
+=========
+
+.. toctree::
+    :maxdepth: 1
+    :glob:
+
+    *
diff --git a/docs/tutorials/linear/mmsink.ipynb b/docs/tutorials/linear/mmsink.ipynb
diff --git a/...s/sinkhorn_divergence_gradient_flow.ipynb → ...r/sinkhorn_divergence_gradient_flow.ipynb b/...s/sinkhorn_divergence_gradient_flow.ipynb → ...r/sinkhorn_divergence_gradient_flow.ipynb
@@ -17,7 +17,7 @@
     "\n",
     "The Sinkhorn divergence, defined in {cite}`genevay:18` as $\\mathrm{S}_\\varepsilon(\\alpha, \\beta) = \\mathrm{OT_\\varepsilon}(\\alpha, \\beta) - \\frac{1}{2}\\mathrm{OT_\\varepsilon}(\\alpha, \\alpha) - \\frac{1}{2}\\mathrm{OT_\\varepsilon}(\\beta, \\beta)$ removes this entropic bias.\n",
     "\n",
-    "In this tutorial we showcase the advantage of removing the entropic bias using gradient flows on 2-D distributions, as done in {cite}`feydy:19` and following the {doc}`point_clouds` tutorial."
+    "In this tutorial we showcase the advantage of removing the entropic bias using gradient flows on 2-D distributions, as done in {cite}`feydy:19` and following the {doc}`../geometry/point_cloud` tutorial."
    ]
   },
   {
@@ -27010,7 +27010,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -27024,7 +27024,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.10.6"
   }
  },
  "nbformat": 4,

diff --git a/...utorials/sparse_monge_displacements.ipynb → ...s/linear/sparse_monge_displacements.ipynb b/...utorials/sparse_monge_displacements.ipynb → ...s/linear/sparse_monge_displacements.ipynb
diff --git a/docs/tutorials/application_biology.ipynb → .../tutorials/misc/application_biology.ipynb b/docs/tutorials/application_biology.ipynb → .../tutorials/misc/application_biology.ipynb
diff --git a/docs/tutorials/misc/index.rst b/docs/tutorials/misc/index.rst
@@ -0,0 +1,8 @@
+Miscellaneous
+=============
+
+.. toctree::
+    :maxdepth: 1
+    :glob:
+
+    *
diff --git a/docs/tutorials/soft_sort.ipynb → docs/tutorials/misc/soft_sort.ipynb b/docs/tutorials/soft_sort.ipynb → docs/tutorials/misc/soft_sort.ipynb
diff --git a/docs/tutorials/tracking_progress.ipynb → docs/tutorials/misc/tracking_progress.ipynb b/docs/tutorials/tracking_progress.ipynb → docs/tutorials/misc/tracking_progress.ipynb
@@ -362,7 +362,7 @@
    "source": [
     "## Tracking progress of the Gromov-Wasserstein solver\n",
     "\n",
-    "We can track progress of the {class}`~ott.solvers.quadratic.gromov_wasserstein.GromovWasserstein` solver in the same way as with the Sinkhorn solvers. Let's define a small {class}`~ott.problems.quadratic.quadratic_problem.QuadraticProblem`, same as in the {doc}`gromov_wasserstein` notebook:"
+    "We can track progress of the {class}`~ott.solvers.quadratic.gromov_wasserstein.GromovWasserstein` solver in the same way as with the Sinkhorn solvers. Let's define a small {class}`~ott.problems.quadratic.quadratic_problem.QuadraticProblem`, same as in the {doc}`../quadratic/gromov_wasserstein` notebook:"
    ]
   },
   {

diff --git a/docs/tutorials/mmsink.ipynb b/docs/tutorials/mmsink.ipynb
diff --git a/docs/tutorials/ENOT.ipynb → docs/tutorials/neural/ENOT.ipynb b/docs/tutorials/ENOT.ipynb → docs/tutorials/neural/ENOT.ipynb
diff --git a/docs/tutorials/MetaOT.ipynb → docs/tutorials/neural/MetaOT.ipynb b/docs/tutorials/MetaOT.ipynb → docs/tutorials/neural/MetaOT.ipynb
diff --git a/docs/tutorials/Monge_Gap.ipynb → docs/tutorials/neural/Monge_Gap.ipynb b/docs/tutorials/Monge_Gap.ipynb → docs/tutorials/neural/Monge_Gap.ipynb
diff --git a/docs/tutorials/icnn_inits.ipynb → docs/tutorials/neural/icnn_inits.ipynb b/docs/tutorials/icnn_inits.ipynb → docs/tutorials/neural/icnn_inits.ipynb
diff --git a/docs/tutorials/neural/index.rst b/docs/tutorials/neural/index.rst
@@ -0,0 +1,8 @@
+Neural OT
+=========
+
+.. toctree::
+    :maxdepth: 1
+    :glob:
+
+    *
diff --git a/docs/tutorials/neural_dual.ipynb → docs/tutorials/neural/neural_dual.ipynb b/docs/tutorials/neural_dual.ipynb → docs/tutorials/neural/neural_dual.ipynb
diff --git a/docs/tutorials/GWLRSinkhorn.ipynb → docs/tutorials/quadratic/GWLRSinkhorn.ipynb b/docs/tutorials/GWLRSinkhorn.ipynb → docs/tutorials/quadratic/GWLRSinkhorn.ipynb
@@ -8,7 +8,7 @@
    "source": [
     "# Low-Rank GW\n",
     "\n",
-    "We provide in this tutorial a minimal example that shows the low-rank (LR) Gromov-Wasserstein solver in action. This quadratic OT solver is presented in {cite}`scetbon:22`, as a follow up to the (linear) LR Sinkhorn solver in {cite}`scetbon:21`, see {doc}`LRSinkhorn`."
+    "We provide in this tutorial a minimal example that shows the low-rank (LR) Gromov-Wasserstein solver in action. This quadratic OT solver is presented in {cite}`scetbon:22`, as a follow up to the (linear) LR Sinkhorn solver in {cite}`scetbon:21`, see {doc}`../linear/LRSinkhorn`."
    ]
   },
   {
@@ -276,5 +276,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 1
+ "nbformat_minor": 4
 }
diff --git a/docs/tutorials/gromov_wasserstein.ipynb → ...orials/quadratic/gromov_wasserstein.ipynb b/docs/tutorials/gromov_wasserstein.ipynb → ...orials/quadratic/gromov_wasserstein.ipynb
@@ -10,7 +10,7 @@
     "\n",
     "In this tutorial, we present the {class}`~ott.solvers.quadratic.gromov_wasserstein.GromovWasserstein` entropic solver, as presented in {cite}`peyre:16`. The GW problem {cite}`memoli:11` proposes to match two point clouds that may lie in two different spaces, each endowed with its own geometry, using a quadratic assignment problem {cite}:`koopmans:57`.\n",
     "\n",
-    "We illustrate this by aligning a 2-d to a 3-d {class}`point cloud <ott.geometry.pointcloud.PointCloud>`, see also {doc}`gromov_wasserstein_multiomics` for a more realistic application to single-cell omics."
+    "We illustrate this by aligning a 2-d to a 3-d {class}`point cloud <ott.geometry.pointcloud.PointCloud>`, see also {doc}`../quadratic/gromov_wasserstein_multiomics` for a more realistic application to single-cell omics."
    ]
   },
   {
@@ -20006,7 +20006,7 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "bilevel",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -20020,7 +20020,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.10.6"
   }
  },
  "nbformat": 4,

diff --git a/...rials/gromov_wasserstein_multiomics.ipynb → ...ratic/gromov_wasserstein_multiomics.ipynb b/...rials/gromov_wasserstein_multiomics.ipynb → ...ratic/gromov_wasserstein_multiomics.ipynb
diff --git a/docs/tutorials/quadratic/index.rst b/docs/tutorials/quadratic/index.rst
@@ -0,0 +1,8 @@
+Quadratic OT
+============
+
+.. toctree::
+    :maxdepth: 1
+    :glob:
+
+    *