Enable D* checks (#319)

* Re-enable D* flake8

* Fix `D*` errors

* Do not ignore D411

* Remove ignored `F401`

* Enable `B008`, add `optax` to `intersphinx_mapping`

* Remove old `.flake8` config

* Add `SortingInitializer` to docs

* Enable `E501`

* Fix docs in `{costs,geometry}.py`

* Update `pointcloud.py`

* Fix `initializers.py`

* Fix `matrix_square_root.py`

* Fix `unbalanced_functions.py`

* Fix more docstrings, add missing licenses

* Fix typo in `FreeWassersteinBarycenter`

* Fix `sinkhorn_lr.py`

* Fix `models.py`

* Fix `neuraldual.py`

* Fix `gromov_wasserstein.py`

* Fix `plot.py`

* Fix `tests/`

* Fix `utils.py`

* Fix `k_means.py`

* Fix `gaussian.py`

* Fix `fit_gmm_pair.py`

* Fix typo

* Make `create_sample_generators` private

* Unable more flake8's in docs

* Add more testing flake8

* Add `SIM`

* Add `RET`

* Undo adding docs

* Add more type hints

* Use double quotes always

* Add explicit quotes normalization

* Udpate comment

* Remove commented flake8s in `pyproject.toml`

* Fix typo in typing

* Skip graph factor cache test

* Fix float in test instead of an integer

* Improve `CostFn` docstring

* Update `pairwise` docs

* Remove `[d,]` from docs

* [ci skip] Mention `gamma > 0`

* Remove module docs, add missing copyright

* [ci skip] Fix docs rendering

docs/conf.py

@@ -26,9 +26,8 @@
 import logging
 from datetime import datetime
 
-from sphinx.util import logging as sphinx_logging
-
 import ott
+from sphinx.util import logging as sphinx_logging
 
 # -- Project information -----------------------------------------------------
 needs_sphinx = "4.0"
@@ -47,17 +46,17 @@
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = [
-    'sphinx.ext.autodoc',
-    'sphinx.ext.autosummary',
-    'sphinx.ext.intersphinx',
-    'sphinx.ext.mathjax',
-    'sphinx.ext.napoleon',
-    'sphinx.ext.viewcode',
-    'sphinxcontrib.bibtex',
-    'sphinx_copybutton',
-    'myst_nb',
-    'IPython.sphinxext.ipython_console_highlighting',
-    'sphinx_autodoc_typehints',
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "sphinx.ext.intersphinx",
+    "sphinx.ext.mathjax",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.viewcode",
+    "sphinxcontrib.bibtex",
+    "sphinx_copybutton",
+    "myst_nb",
+    "IPython.sphinxext.ipython_console_highlighting",
+    "sphinx_autodoc_typehints",
 ]
 
 intersphinx_mapping = {
@@ -68,27 +67,28 @@
     "scipy": ("https://docs.scipy.org/doc/scipy/reference/", None),
     "pot": ("https://pythonot.github.io/", None),
     "jaxopt": ("https://jaxopt.github.io/stable", None),
+    "optax": ("https://optax.readthedocs.io/en/latest/", None),
     "matplotlib": ("https://matplotlib.org/stable/", None),
 }
 
-master_doc = 'index'
+master_doc = "index"
 source_suffix = {
-    '.rst': 'restructuredtext',
-    '.ipynb': 'myst-nb',
+    ".rst": "restructuredtext",
+    ".ipynb": "myst-nb",
 }
 todo_include_todos = False
 
 autosummary_generate = True
-autodoc_typehints = 'description'
+autodoc_typehints = "description"
 
 # myst-nb
 myst_heading_anchors = 2
 nb_execution_mode = "off"
 nb_mime_priority_overrides = [("spelling", "text/plain", 0)]
 myst_enable_extensions = [
-    'amsmath',
-    'colon_fence',
-    'dollarmath',
+    "amsmath",
+    "colon_fence",
+    "dollarmath",
 ]
 
 # bibliography
@@ -108,37 +108,37 @@
 ]
 
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
+templates_path = ["_templates"]
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = ['_build', '**.ipynb_checkpoints']
+exclude_patterns = ["_build", "**.ipynb_checkpoints"]
 
 # -- Options for HTML output -------------------------------------------------
 
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 
-html_theme = 'sphinx_book_theme'
-html_logo = '_static/images/logoOTT.png'
-html_favicon = '_static/images/logoOTT.ico'
+html_theme = "sphinx_book_theme"
+html_logo = "_static/images/logoOTT.png"
+html_favicon = "_static/images/logoOTT.ico"
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ['_static']
+html_static_path = ["_static"]
 html_theme_options = {
-    'repository_url': 'https://github.com/ott-jax/ott',
-    'repository_branch': 'main',
-    'path_to_docs': 'docs/',
-    'use_repository_button': True,
-    'use_fullscreen_button': False,
-    'logo_only': True,
-    'launch_buttons': {
-        'colab_url': 'https://colab.research.google.com',
-        'binderhub_url': 'https://mybinder.org',
-        'notebook_interface': 'jupyterlab',
+    "repository_url": "https://github.com/ott-jax/ott",
+    "repository_branch": "main",
+    "path_to_docs": "docs/",
+    "use_repository_button": True,
+    "use_fullscreen_button": False,
+    "logo_only": True,
+    "launch_buttons": {
+        "colab_url": "https://colab.research.google.com",
+        "binderhub_url": "https://mybinder.org",
+        "notebook_interface": "jupyterlab",
     },
 }
 

docs/initializers/linear.rst

@@ -12,7 +12,7 @@ Sinkhorn Initializers
 
     initializers.DefaultInitializer
     initializers.GaussianInitializer
-    initializers.SinkhornInitializer
+    initializers.SortingInitializer
     initializers.SubsampleInitializer
 
 Low-rank Sinkhorn Initializers

pyproject.toml

@@ -250,66 +250,43 @@ exclude = [
     "dist"
 ]
 ignore = [
-    # line too long -> we accept long comment lines; black gets rid of long code lines
-    "E501",
     # Do not assign a lambda expression, use a def -> lambda expression assignments are convenient
     "E731",
     # allow I, O, l as variable names -> I is the identity matrix, i, j, k, l is reasonable indexing notation
     "E741",
     # Missing docstring in public package
     "D104",
-    # ... imported but unused
-    "F401",
     # Missing docstring in public module
     "D100",
     # Missing docstring in __init__
     "D107",
-    # Do not perform function calls in argument defaults.
-    "B008",
     # Missing docstring in magic method
     "D105",
-    # Missing blank line before section
-    "D411",
-    ## Flake8 rules not supported by ruff:
-    # line break before a binary operator -> black does not adhere to PEP8
-    # "W503",
-    # line break occured after a binary operator -> black does not adhere to PEP8
-    # "W504",
-    # whitespace before : -> black does not adhere to PEP8
-    # "E203",
-    # whitespace before : -> black does not adhere to PEP8
-    # "E203",
-    # missing whitespace after ,', ';', or ':' -> black does not adhere to PEP8
-    # "E231",
-    # continuation line over-indented for hanging indent -> black does not adhere to PEP8
-    # "E126",
-    # inline comment should start with '#' -> Scanpy allows them for specific explanations
-    # "E266",
-    # format string does contain unindexed parameters
-    # "P101",
-    # indentation is not a multiple of 4
-    # "E111",
-    # "E114",
 ]
 line-length = 80
 select = [
+    "D", # flake8-docstrings
     "I", # isort
     "E", # pycodestyle
     "F", # pyflakes
     "W", # pycodestyle
-    # below are not autofixed
+    "Q", # flake8-quotes
+    "SIM", # flake8-simplify
+    "NPY",  # NumPy-specific rules
+    "PT",  # flake8-pytest-style
+    "B", # flake8-bugbear
     "UP", # pyupgrade
     "C4", # flake8-comprehensions
-    "B", # flake8-bugbear
     "BLE", # flake8-blind-except
+    "T20",  # flake8-print
+    "RET", # flake8-raise
 ]
-unfixable = ["B", "UP", "C4", "BLE"]
+unfixable = ["B", "UP", "C4", "BLE", "T20", "RET"]
 target-version = "py38"
 [tool.ruff.per-file-ignores]
-    "tests/*" = ["D", "E", "F", "W", "I", "C408"]
+    "tests/*" = ["D", "PT004"]  # TODO(michalk8): remove `self.initialize` in `tests/`
     "*/__init__.py" = ["F401"]
-    "examples/*" = ["D101", "D102", "D103"]
-    "docs/*" = ["E", "F", "W", "I", "D101", "D102"]
+    "docs/*" = ["D"]
     "src/ott/types.py" = ["D102"]
 [tool.ruff.pydocstyle]
 convention = "google"
@@ -319,3 +296,7 @@ keep-runtime-typing = true
 [tool.ruff.flake8-tidy-imports]
 # Disallow all relative imports.
 ban-relative-imports = "parents"
+[tool.ruff.flake8-bugbear]
+extend-immutable-calls = ["jax.random.PRNGKey"]
+[tool.ruff.flake8-quotes]
+inline-quotes = "double"

src/ott/_version.py

@@ -1,3 +1,16 @@
+# Copyright OTT-JAX
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from importlib.metadata import PackageNotFoundError, version
 
 try: