pyproject.toml

[build-system]
requires = ["setuptools>=61", "setuptools-scm[toml]>=6.2"]
build-backend = "setuptools.build_meta"

[project]
name = "ott-jax"
description = "Optimal Transport Tools in JAX."
requires-python = ">=3.9"
dynamic = ["version"]
readme = {file = "README.md", content-type = "text/markdown"}
license = {file = "LICENSE"}
authors = [
    {name = "OTT team", email = "optimal.transport.tools@gmail.com"}
]
dependencies = [
    "jax>=0.4.0",
    "jaxopt>=0.8",
    "lineax>=0.0.5",
    "numpy>=1.20.0",
]
keywords = [
    "optimal transport",
    "gromov wasserstein",
    "sinkhorn",
    "low-rank sinkhorn",
    "sinkhorn divergences",
    "wasserstein",
    "wasserstein barycenter",
    "jax",
    "autodiff",
    "implicit differentiation",
]
classifiers = [
    "Typing :: Typed",
    "Development Status :: 4 - Beta",
    "License :: OSI Approved :: Apache Software License",
    "Topic :: Scientific/Engineering :: Mathematics",
    "Natural Language :: English",
    "Intended Audience :: Developers",
    "Intended Audience :: Science/Research",
    "Operating System :: POSIX :: Linux",
    "Operating System :: MacOS :: MacOS X",
    "Operating System :: Microsoft :: Windows",
    "Programming Language :: Python :: 3",
    "Programming Language :: Python :: 3.9",
    "Programming Language :: Python :: 3.10",
    "Programming Language :: Python :: 3.11",
    "Programming Language :: Python :: 3.12",
]

[project.urls]
"Source Code" = "https://github.com/ott-jax/ott"
Documentation = "https://ott-jax.readthedocs.io"
"Issue Tracker" = "https://github.com/ott-jax/ott/issues"
Changelog = "https://github.com/ott-jax/ott/releases"

[project.optional-dependencies]
neural = [
    "flax>=0.6.6",
    "optax>=0.1.1",
    "diffrax>=0.4.1",
]
dev = [
    "pre-commit>=2.16.0",
    "tox>=4",
]
test = [
    "pytest",
    "pytest-xdist",
    "pytest-cov",
    "pytest-memray",
    "coverage[toml]",
    "chex",
    "networkx>=2.5",
    "scikit-learn>=1.0",
    "tqdm",
    "tslearn>=0.5",
    "matplotlib",
]
docs = [
    "sphinx>=4.0",
    "sphinx-book-theme>=1.0.1",
    "sphinx_autodoc_typehints>=1.12.0",
    "sphinx-copybutton>=0.5.1",
    "sphinxcontrib-bibtex>=2.5.0",
    "sphinxcontrib-spelling>=7.7.0",
    "myst-nb>=0.17.1",
]

[tool.setuptools]
package-dir = {"" = "src"}
include-package-data = true

[tool.setuptools_scm]

[tool.black]
line-length = 80
target-version = ["py39"]
include = '\.ipynb$'

[tool.isort]
profile = "black"
line_length = 80
include_trailing_comma = true
multi_line_output = 3
sections = ["FUTURE", "STDLIB", "THIRDPARTY", "TEST", "NUMERIC", "NEURAL", "PLOTTING", "FIRSTPARTY", "LOCALFOLDER"]
# also contains what we import in notebooks/tests
known_neural = ["flax", "optax", "diffrax", "orbax"]
known_numeric = ["numpy", "scipy", "jax", "flax", "optax", "jaxopt", "ot", "torch", "torchvision", "pandas", "sklearn", "tslearn"]
known_test = ["_pytest", "pytest"]
known_plotting = ["IPython", "matplotlib", "mpl_toolkits", "seaborn"]

[tool.pytest.ini_options]
minversion = "6.0"
testpaths = [
    "tests",
]
markers = [
    "cpu: Mark tests as CPU only.",
    "fast: Mark tests as fast.",
]
filterwarnings = [
    "ignore:\\n*.*scipy.sparse array",
    "ignore:jax.random.KeyArray is deprecated:DeprecationWarning",
    "ignore:.*jax.config:DeprecationWarning",
    "ignore:jax.core.Shape is deprecated:DeprecationWarning:chex",
]

[tool.coverage.run]
branch = true
source = ["src/"]
omit = [
    "*/__init__.py",
    "*/_version.py",
    "*/types.py",
]

[tool.coverage.report]
exclude_lines = [
    '\#.*pragma:\s*no.?cover',
    "^if __name__ == .__main__.:$",
    '^\s*raise AssertionError\b',
    '^\s*raise NotImplementedError\b',
    '^\s*return NotImplemented\b',
]
precision = 2
show_missing = true
skip_empty = true
sort = "Miss"

[tool.yapf]
based_on_style = "yapf"
column_limit = 80
indent_width = 2
split_before_named_assigns = true
spaces_around_power_operator = true
dedent_closing_brackets = true
coalesce_brackets = true

[tool.rstcheck]
ignore_directives = [
    "include",
    "toctree",
    "module",
    "currentmodule",
    "autosummary",
    "automodule",
    "autoclass",
    "bibliography"
]
ignore_roles = [
    "class",
    "doc",
    "mod",
    "cite",
]

[tool.doc8]
max_line_length = 80
# Parser "myst_parser.sphinx_" not found. No module named 'myst_parser'.
ignore_path = ["docs/**/_autosummary", "docs/contributing.rst"]

[tool.tox]
legacy_tox_ini = """
[tox]
min_version = 4.0
env_list = lint-code,py{3.9,3.10,3.11,3.12},py3.9-jax-default
skip_missing_interpreters = true

[testenv]
extras =
    test
    # https://github.com/google/flax/issues/3329
    py{3.9,3.10,3.11,3.12},py3.9-jax-default: neural
pass_env = CUDA_*,PYTEST_*,CI
commands_pre =
    gpu: python -I -m pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
    jax-latest: python -I -m pip install 'git+https://github.com/google/jax@main'
commands =
    python -m pytest {tty:--color=yes} {posargs: \
      --cov={env_site_packages_dir}{/}ott --cov-config={tox_root}{/}pyproject.toml \
      --no-cov-on-fail --cov-report=xml --cov-report=term-missing:skip-covered}

[testenv:lint-code]
description = Lint the code.
deps = pre-commit>=3.0.0
skip_install = true
commands =
    pre-commit run --all-files --show-diff-on-failure

[testenv:lint-docs]
description = Lint the documentation.
deps =
extras = docs,neural
ignore_errors = true
allowlist_externals = make
pass_env = PYENCHANT_LIBRARY_PATH
set_env = SPHINXOPTS = -W -q --keep-going
changedir = {tox_root}{/}docs
commands =
    make linkcheck {posargs}
    make spelling {posargs}

[testenv:build-docs]
description = Build the documentation.
use_develop = true
deps =
extras = docs,neural
allowlist_externals = make
changedir = {tox_root}{/}docs
commands =
    make html {posargs}
commands_post =
    python -c 'import pathlib; print("Documentation is under:", pathlib.Path("{tox_root}") / "docs" / "_build" / "html" / "index.html")'

[testenv:clean-docs]
description = Remove the documentation.
deps =
skip_install = true
changedir = {tox_root}{/}docs
allowlist_externals = make
commands =
    make clean

[testenv:build-package]
description = Build the package.
deps =
    build
    twine
commands =
    python -m build --sdist --wheel --outdir {tox_root}{/}dist{/} {posargs:}
    twine check {tox_root}{/}dist{/}*
commands_post =
    python -c 'import pathlib; print(f"Package is under:", pathlib.Path("{tox_root}") / "dist")'

[testenv:format-references]
description = Format references.bib.
skip_install = true
allowlist_externals = biber
commands = biber --tool --output_file={tox_root}{/}docs{/}references.bib --nolog \
    --output_align --output_indent=2 --output_fieldcase=lower \
    --output_legacy_dates --output-field-replace=journaltitle:journal,thesis:phdthesis,institution:school \
    {tox_root}{/}docs{/}references.bib
"""

[tool.ruff]
exclude = [
    ".git",
    "__pycache__",
    "build",
    "docs/_build",
    "dist"
]
line-length = 80
target-version = "py38"  # TODO(michalk8): use py39 and fix the type hints

[tool.ruff.lint]
ignore = [
    # 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",
    # Missing docstring in public module
    "D100",
    # Missing docstring in __init__
    "D107",
    # Missing docstring in magic method
    "D105",
]
select = [
    "D", # flake8-docstrings
    "E", # pycodestyle
    "F", # pyflakes
    "W", # pycodestyle
    "Q", # flake8-quotes
    "SIM", # flake8-simplify
    "TID",  # flake-8-tidy-imports
    "NPY",  # NumPy-specific rules
    "PT",  # flake8-pytest-style
    "B", # flake8-bugbear
    "UP", # pyupgrade
    "C4", # flake8-comprehensions
    "BLE", # flake8-blind-except
    "T20",  # flake8-print
    "RET", # flake8-raise
]
unfixable = ["B", "UP", "C4", "BLE", "T20", "RET"]

[tool.ruff.lint.per-file-ignores]
# TODO(michalk8): PO004 - remove `self.initialize`
"tests/*" = ["D", "PT004", "E402"]
"*/__init__.py" = ["F401"]
"docs/*" = ["D"]
"src/ott/types.py" = ["D102"]
[tool.ruff.lint.pydocstyle]
convention = "google"
[tool.ruff.lint.pyupgrade]
# Preserve types, even if a file imports `from __future__ import annotations`.
keep-runtime-typing = true
[tool.ruff.lint.flake8-tidy-imports]
# Disallow all relative imports.
ban-relative-imports = "parents"
[tool.ruff.lint.flake8-quotes]
inline-quotes = "double"