add ruff and remove redundant pre-commits

.pre-commit-config.yaml

@@ -18,20 +18,6 @@ repos:
     args: [--py38-plus]
   - id: nbqa-black
   - id: nbqa-isort
-- repo: https://github.com/PyCQA/isort
-  rev: 5.12.0
-  hooks:
-  - id: isort
-- repo: https://github.com/asottile/yesqa
-  rev: v1.4.0
-  hooks:
-  - id: yesqa
-    additional_dependencies:
-    - flake8-tidy-imports
-    - flake8-docstrings
-    - flake8-comprehensions
-    - flake8-bugbear
-    - flake8-blind-except
 - repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
   rev: v2.6.0
   hooks:
@@ -48,31 +34,12 @@ repos:
     args: [--fix=lf]
   - id: trailing-whitespace
   - id: check-case-conflict
-- repo: https://github.com/myint/autoflake
-  rev: v2.0.1
-  hooks:
-  - id: autoflake
-    args:
-    - --in-place
-    - --remove-all-unused-imports
-    - --remove-unused-variable
-    - --ignore-init-module-imports
-- repo: https://github.com/PyCQA/flake8
-  rev: 6.0.0
-  hooks:
-  - id: flake8
-    additional_dependencies:
-    - flake8-tidy-imports
-    - flake8-docstrings
-    - flake8-comprehensions
-    - flake8-bugbear
-    - flake8-blind-except
-    args: [--docstring-convention, google]
-- repo: https://github.com/asottile/pyupgrade
-  rev: v3.3.1
+- repo: https://github.com/charliermarsh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.0.246
   hooks:
-  - id: pyupgrade
-    args: [--py38-plus, --keep-runtime-typing]
+  - id: ruff
+    args: [--fix, --exit-non-zero-on-fix]
 - repo: https://github.com/rstcheck/rstcheck
   rev: v6.1.1
   hooks:

examples/fairness/losses.py

@@ -17,7 +17,6 @@
 
 import jax
 import jax.numpy as jnp
-
 import ott
 
 

examples/fairness/main.py

@@ -15,12 +15,10 @@
 
 from typing import Sequence
 
+import jax
 from absl import app, flags, logging
 from clu import platform
 from ml_collections import config_flags
-
-import jax
-
 from ott.examples.fairness import train
 
 FLAGS = flags.FLAGS

examples/fairness/train.py

@@ -17,15 +17,13 @@
 import functools
 from typing import Any
 
-import ml_collections
-
 import flax
 import jax
 import jax.numpy as jnp
+import ml_collections
 from flax import jax_utils
 from flax.metrics import tensorboard
 from flax.training import checkpoints, common_utils
-
 from ott.examples.fairness import data, losses, models
 
 

examples/soft_error/data.py

@@ -13,10 +13,9 @@
 # limitations under the License.
 """Data loading and data augmentation."""
 
+import jax
 import tensorflow as tf
 import tensorflow_datasets as tfds
-
-import jax
 from flax import jax_utils
 
 

examples/soft_error/losses.py

@@ -18,7 +18,6 @@
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
-
 from ott.tools import soft_sort
 
 

examples/soft_error/main.py

@@ -15,12 +15,10 @@
 
 from typing import Sequence
 
+import jax
 from absl import app, flags, logging
 from clu import platform
 from ml_collections import config_flags
-
-import jax
-
 from ott.examples.soft_error import train
 
 FLAGS = flags.FLAGS

examples/soft_error/train.py

@@ -17,16 +17,14 @@
 import functools
 from typing import Any
 
-import ml_collections
-import tensorflow_datasets as tfds
-
 import flax
 import jax
 import jax.numpy as jnp
+import ml_collections
+import tensorflow_datasets as tfds
 from flax import jax_utils
 from flax.metrics import tensorboard
 from flax.training import checkpoints, common_utils
-
 from ott.examples.soft_error import data, losses
 from ott.examples.soft_error import model as model_lib
 

pyproject.toml

@@ -233,3 +233,82 @@ legacy_tox_ini = """
     commands_post =
         python -c 'import pathlib; print(f"Package is under:", pathlib.Path(f"{tox_root}") / "dist")'
 """
+
+[tool.ruff]
+exclude = [
+    ".git",
+    "__pycache__",
+    "build",
+    "docs/_build",
+    "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 = [
+    "I", # isort
+    "E", # pycodestyle
+    "F", # pyflakes
+    "W", # pycodestyle
+    # below are not autofixed
+    "UP", # pyupgrade
+    "C4", # flake8-comprehensions
+    "B", # flake8-bugbear
+    "BLE", # flake8-blind-except
+]
+unfixable = ["B", "UP", "C4", "BLE"]
+target-version = "py38"
+[tool.ruff.per-file-ignores]
+    "tests/*" = ["D", "E", "F", "W", "I", "C408"]
+    "*/__init__.py" = ["F401"]
+    "examples/*" = ["D101", "D102", "D103"]
+    "docs/*" = ["E", "F", "W", "I", "D101", "D102"]
+    "src/ott/types.py" = ["D102"]
+[tool.ruff.pydocstyle]
+convention = "google"
+[tool.ruff.pyupgrade]
+# Preserve types, even if a file imports `from __future__ import annotations`.
+keep-runtime-typing = true
+[tool.ruff.flake8-tidy-imports]
+# Disallow all relative imports.
+ban-relative-imports = "parents"

src/ott/geometry/__init__.py

@@ -11,4 +11,12 @@
 # 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 . import costs, epsilon_scheduler, geometry, graph, grid, pointcloud, segment
+from . import (
+    costs,
+    epsilon_scheduler,
+    geometry,
+    graph,
+    grid,
+    pointcloud,
+    segment,
+)

src/ott/problems/linear/potentials.py

@@ -26,10 +26,9 @@
 import jax.numpy as jnp
 import jax.scipy as jsp
 import jax.tree_util as jtu
-import numpy as np
-
 import matplotlib
 import matplotlib.pyplot as plt
+import numpy as np
 
 from ott.problems.linear import linear_problem
 

src/ott/solvers/linear/sinkhorn_lr.py

@@ -12,7 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """A Jax implementation of the Low-Rank Sinkhorn algorithm."""
-from typing import Any, Literal, Mapping, NamedTuple, NoReturn, Optional, Tuple, Union
+from typing import (
+    Any,
+    Literal,
+    Mapping,
+    NamedTuple,
+    NoReturn,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import jax
 import jax.numpy as jnp

src/ott/solvers/nn/conjugate_solvers.py

@@ -16,9 +16,8 @@
 import abc
 from typing import Callable, Literal, NamedTuple, Optional
 
-from jaxopt import LBFGS
-
 import jax.numpy as jnp
+from jaxopt import LBFGS
 
 from ott import utils
 

src/ott/solvers/nn/neuraldual.py

@@ -14,7 +14,16 @@
 """A Jax implementation of the neural-based Kantorovich dual."""
 
 import warnings
-from typing import Callable, Dict, Iterable, List, Literal, Optional, Tuple, Union
+from typing import (
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import jax
 import jax.numpy as jnp

src/ott/tools/gaussian_mixture/fit_gmm_pair.py

@@ -85,7 +85,11 @@
 import jax.numpy as jnp
 import optax
 
-from ott.tools.gaussian_mixture import fit_gmm, gaussian_mixture, gaussian_mixture_pair
+from ott.tools.gaussian_mixture import (
+    fit_gmm,
+    gaussian_mixture,
+    gaussian_mixture_pair,
+)
 
 __all__ = ["get_fit_model_em_fn"]
 

src/ott/tools/gaussian_mixture/gaussian_mixture.py

@@ -18,7 +18,12 @@
 import jax
 import jax.numpy as jnp
 
-from ott.tools.gaussian_mixture import gaussian, linalg, probabilities, scale_tril
+from ott.tools.gaussian_mixture import (
+    gaussian,
+    linalg,
+    probabilities,
+    scale_tril,
+)
 
 __all__ = ["GaussianMixture"]
 

src/ott/tools/plot.py

@@ -16,10 +16,9 @@
 from typing import List, Optional, Sequence, Tuple, Union
 
 import jax.numpy as jnp
+import matplotlib.pyplot as plt
 import numpy as np
 import scipy
-
-import matplotlib.pyplot as plt
 from matplotlib import animation
 
 from ott import utils