Feature/src layout (#188)

* Switch to `src/` layout

* Switch to `pyproject.toml`

* Prune `examples/` from `MANIFEST.in`

* Update pre-commits

* Update `pyproject.toml`

* Remove unused packages

* Add `scikit-learn` to tests reqs.

* Fix `coverage`

* Run tests on schedule

* Remove `Lint as: python3` artifact

* [ci skip] Update URLs

.codecov.yml

@@ -1,5 +1,5 @@
 codecov:
-  require_ci_to_pass: no
+  require_ci_to_pass: yes
   strict_yaml_branch: main
 
 coverage:

.github/workflows/notebook_tests.yml

@@ -31,8 +31,8 @@ jobs:
     - name: Print versions
       run: |
         python -VV
-        python -c "import jax; print('jax', jax.__version__)"
-        python -c "import jaxlib; print('jaxlib', jaxlib.__version__)"
+        python -c "import jax; print('jax==', jax.__version__)"
+        python -c "import jaxlib; print('jaxlib==', jaxlib.__version__)"
 
     - name: Intall Jupyter kernel
       run: |

.github/workflows/tests.yml

@@ -1,6 +1,8 @@
 name: tests
 
 on:
+  schedule:
+  - cron: 00 00 * * 1
   push:
     branches: [main]
   pull_request:
@@ -39,18 +41,18 @@ jobs:
     - name: Print versions
       run: |
         python -VV
-        python -c "import jax; print('jax', jax.__version__)"
-        python -c "import jaxlib; print('jaxlib', jaxlib.__version__)"
+        python -c "import jax; print('jax==', jax.__version__)"
+        python -c "import jaxlib; print('jaxlib==', jaxlib.__version__)"
 
     - name: Run fast tests
       if: ${{ matrix.test_mark == 'fast' }}
       run: |
-        python -m pytest --cov=ott --cov-append --cov-report=xml --cov-report=term-missing --cov-config=setup.cfg --memray -m fast -n auto
+        python -m pytest --cov=ott --cov-append --cov-report=xml --cov-report=term-missing --cov-config=pyproject.toml --memray -m fast -n auto
 
     - name: Run all tests
       if: ${{ matrix.test_mark == 'all' }}
       run: |
-        python -m pytest --cov=ott --cov-append --cov-report=xml --cov-report=term-missing --cov-config=setup.cfg --memray
+        python -m pytest --cov=ott --cov-append --cov-report=xml --cov-report=term-missing --cov-config=pyproject.toml --memray
 
     - name: Upload coverage
       uses: codecov/codecov-action@v3

.pre-commit-config.yaml

@@ -20,7 +20,7 @@ repos:
   hooks:
   - id: isort
 - repo: https://github.com/asottile/yesqa
-  rev: v1.3.0
+  rev: v1.4.0
   hooks:
   - id: yesqa
     additional_dependencies:
@@ -35,7 +35,7 @@ repos:
   - id: pretty-format-yaml
     args: [--autofix, --indent, '2']
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.3.0
+  rev: v4.4.0
   hooks:
   - id: detect-private-key
   - id: check-ast
@@ -45,7 +45,7 @@ repos:
   - id: trailing-whitespace
   - id: check-case-conflict
 - repo: https://github.com/myint/autoflake
-  rev: v1.4
+  rev: v2.0.0
   hooks:
   - id: autoflake
     args:
@@ -54,7 +54,7 @@ repos:
     - --remove-unused-variable
     - --ignore-init-module-imports
 - repo: https://github.com/PyCQA/flake8
-  rev: 4.0.1
+  rev: 6.0.0
   hooks:
   - id: flake8
     additional_dependencies:
@@ -65,7 +65,7 @@ repos:
     - flake8-blind-except
     args: [--docstring-convention, google]
 - repo: https://github.com/asottile/pyupgrade
-  rev: v2.37.1
+  rev: v3.2.2
   hooks:
   - id: pyupgrade
     args: [--py3-plus, --py37-plus, --keep-runtime-typing]

MANIFEST.in

@@ -1,4 +1,5 @@
 prune docs
-prune .github
+prune examples
 prune images
 prune tests
+prune .github

pyproject.toml

@@ -1,19 +1,135 @@
 [build-system]
-requires = [
-    "setuptools>=45",
-    "setuptools-scm[toml]>=6.2",
-]
+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.7"
+dynamic = ["version"]
+readme = {file = "README.md", content-type = "text/markdown"}
+license = {file = "LICENSE"}
+authors = [
+    {name = "OTT team", email = "[email protected]"}
+]
+dependencies = [
+    'importlib-metadata>=1.0; python_version<"3.8"',
+    "jax>=0.1.67",
+    "jaxlib>=0.1.47",
+    # https://github.com/google/jax/discussions/9951#discussioncomment-3017784
+    "numpy>=1.18.4, !=1.23.0",
+    "matplotlib>=3.0.0",
+    "flax>=0.5.2",
+    "optax>=0.1.1",
+    'typing_extensions; python_version<"3.8"',
+    "scipy>=1.7.0",
+]
+keywords = [
+    "optimal transport",
+    "gromov wasserstein",
+    "sinkhorn",
+    "low-rank sinkhorn",
+    "sinkhorn divergences",
+    "wasserstein",
+    "wasserstein barycenter",
+    "jax",
+    "autodiff",
+    "implicit differentiation",
+]
+classifiers = [
+    "Typing :: Typed",
+    "Development Status :: 5 - Production/Stable",
+    "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.7",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+]
+
+[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]
+dev = [
+    "pre-commit",
+]
+test = [
+    "pytest",
+    "pytest-xdist",
+    "pytest-cov",
+    "coverage[toml]",
+    "testbook",
+    "chex",
+    "networkx>=2.5",
+    "scikit-learn>=1.0"
+]
+experimental = [
+    "scikit-sparse>=0.4.6",
+]
+docs = [
+    "sphinx>=4.0",
+    "nbsphinx>=0.8.0",
+    "recommonmark>=0.7.1",
+    "ipython>=7.20.0",
+    "sphinx_autodoc_typehints>=1.12.0",
+    "sphinx-book-theme>=0.3.3",
+    "sphinxcontrib-bibtex",
+]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+packages = {find = {where = ["src"], namespaces = false}}
+
 [tool.setuptools_scm]
 
+[tool.black]
+line-length = 80
+target-version = ["py38"]
+include = '\.ipynb$'
+
 [tool.isort]
 profile = "black"
 include_trailing_comma = true
 multi_line_output = 3
 skip_glob = ["docs/*"]
 
-[tool.black]
-line-length = 80
-target-version = ['py38']
-include = '\.ipynb$'
+[tool.pytest.ini_options]
+minversion = "6.0"
+addopts = '-v -m "not notebook"'
+testpaths = [
+    "tests",
+]
+markers = [
+    "fast: Mark tests as fast.",
+    "notebook: Mark tests as notebook related.",
+]
+
+[tool.coverage.run]
+branch = true
+parallel = true
+source = ["src/"]
+omit = ["*/__init__.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"

ott/geometry/costs.py → src/ott/geometry/costs.py

@@ -11,8 +11,6 @@
 # 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.
-
-# Lint as: python3
 """Several cost/norm functions for relevant vector types."""
 import abc
 import functools

ott/geometry/epsilon_scheduler.py → src/ott/geometry/epsilon_scheduler.py

@@ -11,8 +11,6 @@
 # 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.
-
-# Lint as: python3
 """A class to define a scheduler for the entropic regularization epsilon."""
 from typing import Any, Optional
 

ott/geometry/geometry.py → src/ott/geometry/geometry.py

@@ -11,8 +11,6 @@
 # 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.
-
-# Lint as: python3
 """A class describing operations used to instantiate and use a geometry."""
 import functools
 from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Union

ott/geometry/grid.py → src/ott/geometry/grid.py

@@ -11,8 +11,6 @@
 # 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.
-
-# Lint as: python3
 """Implements a geometry class for points supported on a cartesian product."""
 import itertools
 from typing import Any, List, NoReturn, Optional, Sequence, Tuple

ott/geometry/low_rank.py → src/ott/geometry/low_rank.py

@@ -11,8 +11,6 @@
 # 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.
-
-# Lint as: python3
 """A class describing low-rank geometries."""
 from typing import Any, Callable, Optional, Tuple, Union
 

ott/geometry/pointcloud.py → src/ott/geometry/pointcloud.py

@@ -11,8 +11,6 @@
 # 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.
-
-# Lint as: python3
 """A geometry defined using 2 point clouds and a cost function between them."""
 import math
 from typing import Any, Callable, Optional, Tuple, Union