Refactor regularized TI costs (#553)

* Update docs Add new costs Update tests Final polishing * Decrease tolerance in `test_l2_moreau_envelope` * Fix not passing `kwargs` * Expose initial estimate in `h_transform` * Test for `x_init` * Change variable name in `h_transform` * Better `h_transform` parameterization * Add more references * Add a missing reference * Update docs of p-norm * Improve the docs of `RegTICost.h_transform` * Rename `is_squared` -> `is_factor` * Better `h_legendre` docs * Expose solver in `TICost.h_transform`
ott-jax · Jul 3, 2024 · d5d8a47 · d5d8a47
1 parent c6fb25c
commit d5d8a47
Show file tree

Hide file tree

Showing 11 changed files with 1,241 additions and 557 deletions.
diff --git a/docs/conf.py b/docs/conf.py
@@ -95,6 +95,17 @@
     "dollarmath",
 ]
 
+# mathjax
+mathjax3_config = {
+    "tex": {
+        "macros": {
+            "prox": r"\mathrm{prox}",
+            "min": r"\mathrm{min}",
+            "argmin": r"\mathrm{argmin}",
+        }
+    }
+}
+
 # bibliography
 bibtex_bibfiles = ["references.bib"]
 bibtex_reference_style = "author_year"

diff --git a/docs/geometry.rst b/docs/geometry.rst
@@ -60,18 +60,29 @@ Cost Functions
     costs.SqPNorm
     costs.PNormP
     costs.SqEuclidean
+    costs.RegTICost
     costs.Euclidean
     costs.Cosine
     costs.Arccos
     costs.Bures
     costs.UnbalancedBures
-    costs.ElasticL1
-    costs.ElasticL2
-    costs.ElasticSTVS
-    costs.ElasticSqKOverlap
     costs.SoftDTW
     distrib_costs.UnivariateWasserstein
 
+Regularizers
+------------
+.. autosummary::
+    :toctree: _autosummary
+
+    regularizers.ProximalOperator
+    regularizers.Regularization
+    regularizers.Orthogonal
+    regularizers.Quadratic
+    regularizers.L1
+    regularizers.L2
+    regularizers.STVS
+    regularizers.SqKOverlap
+
 Utilities
 ---------
 .. autosummary::

diff --git a/docs/references.bib b/docs/references.bib
@@ -740,6 +740,17 @@ @inproceedings{argyriou:12
   year      = {2012},
 }
 
+@article{mcdonald:16,
+  author  = {McDonald, Andrew M. and Pontil, Massimiliano and Stamos, Dimitris},
+  url     = {http://jmlr.org/papers/v17/15-151.html},
+  journal = {Journal of Machine Learning Research},
+  number  = {155},
+  pages   = {1--38},
+  title   = {New Perspectives on k-Support and Cluster Norms},
+  volume  = {17},
+  year    = {2016},
+}
+
 @article{zou:05,
   author    = {Zou, Hui and Hastie, Trevor},
   publisher = {[Royal Statistical Society, Wiley]},
@@ -819,14 +830,14 @@ @misc{scetbon:23
   year        = {2023},
 }
 
-@misc{klein:23,
+@misc{klein:24,
   author      = {Klein, Michal and Pooladian, Aram-Alexandre and Ablin, Pierre and Ndiaye, Eugène and Niles-Weed, Jonathan and Cuturi, Marco},
-  doi         = {10.48550/arXiv.2306.11895},
+  url         = {https://arxiv.org/abs/2306.11895},
   eprint      = {2306.11895},
   eprintclass = {stat.ML},
   eprinttype  = {arXiv},
-  title       = {Learning Costs for Structured Monge Displacements},
-  year        = {2023},
+  title       = {Learning Elastic Costs to Shape Monge Displacements},
+  year        = {2024},
 }
 
 @inproceedings{scetbon:20,
@@ -925,12 +936,12 @@ @inproceedings{cho:09
 }
 
 @misc{buzun:24,
-  title         = {ENOT: Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport},
-  author        = {Nazar Buzun and Maksim Bobrin and Dmitry V. Dylov},
-  year          = {2024},
-  eprint        = {2403.03777},
-  archivePrefix = {arXiv},
-  primaryClass  = {cs.LG}
+  author      = {Buzun, Nazar and Bobrin, Maksim and Dylov, Dmitry V.},
+  eprint      = {2403.03777},
+  eprintclass = {cs.LG},
+  eprinttype  = {arXiv},
+  title       = {ENOT: Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport},
+  year        = {2024},
 }
 
 @misc{vesseron:24,
@@ -941,3 +952,34 @@ @misc{vesseron:24
   title       = {On a Neural Implementation of Brenier's Polar Factorization},
   year        = {2024},
 }
+
+@book{bauschke:17,
+  author    = {Bauschke, Heinz H. and Combettes, Patrick Louis},
+  publisher = {Springer},
+  url       = {https://hal.sorbonne-universite.fr/hal-01517477},
+  title     = {Convex analysis and monotone operator theory in Hilbert spaces},
+  year      = {2017},
+}
+
+@article{combettes:07,
+  author   = {Combettes, Patrick L. and Pesquet, Jean-Christophe},
+  doi      = {10.1109/JSTSP.2007.910264},
+  journal  = {IEEE Journal of Selected Topics in Signal Processing},
+  keywords = {Signal processing algorithms;Hilbert space;Signal processing;Signal analysis;Convergence;Image denoising;Noise reduction;Projection algorithms;Mathematical model;Convex optimization;denoising;Douglas–Rachford;frame;nondifferentiable optimization;Poisson noise;proximal algorithm;wavelets},
+  number   = {4},
+  pages    = {564--574},
+  title    = {A Douglas–Rachford Splitting Approach to Nonsmooth Convex Variational Signal Recovery},
+  volume   = {1},
+  year     = {2007},
+}
+
+@article{parikh:14,
+  author    = {Parikh, Neal and Boyd, Stephen and others},
+  publisher = {Now Publishers, Inc.},
+  journal   = {Foundations and trends{®} in Optimization},
+  number    = {3},
+  pages     = {127--239},
+  title     = {Proximal algorithms},
+  volume    = {1},
+  year      = {2014},
+}
diff --git a/docs/tutorials/sparse_monge_displacements.ipynb b/docs/tutorials/sparse_monge_displacements.ipynb
diff --git a/src/ott/geometry/__init__.py b/src/ott/geometry/__init__.py
@@ -20,5 +20,6 @@
     graph,
     grid,
     pointcloud,
+    regularizers,
     segment,
 )