Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .flake8
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
[flake8]
max-line-length = 88
select = E, F
ignore = E402, E721, E731, E741, F722
per-file-ignores = **/__init__.py:F401
8 changes: 5 additions & 3 deletions .github/workflows/lint.yml
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ jobs:
max-parallel: 4
matrix:
os: [ubuntu-latest]
python-version: ["3.10"]
python-version: ["3.12"]
steps:
- uses: actions/checkout@v3
- name: Set up python ${{ matrix.python-version }}
Expand All @@ -28,6 +28,8 @@ jobs:
- name: Run styles & lint check
run: |
pip install .[testing]
ruff format --check jumpax/ tests/ examples/
ruff check jumpax/ tests/ examples/
pip install "black[jupyter]" flake8 isort
black --check jumpax/ tests/ examples/
isort --check jumpax/ tests/ examples/
flake8 jumpax/ tests/ examples/
pyright jumpax/
2 changes: 2 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -206,3 +206,5 @@ cython_debug/
marimo/_static/
marimo/_lsp/
__marimo__/

.vscode
22 changes: 22 additions & 0 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
repos:
- repo: https://github.com/PyCQA/isort
rev: 8.0.1
hooks:
- id: isort
types_or: [ python, pyi, jupyter ]
- repo: https://github.com/psf/black
rev: 26.3.1
hooks:
- id: black
types_or: [ python, pyi, jupyter ]
additional_dependencies: ["black[jupyter]"]
- repo: https://github.com/PyCQA/flake8
rev: 7.3.0
hooks:
- id: flake8
types_or: [ python, pyi, jupyter ]
- repo: https://github.com/RobertCraigie/pyright-python
rev: v1.1.408
hooks:
- id: pyright
additional_dependencies: ["equinox", "pytest", "jax", "jaxtyping"]
1 change: 0 additions & 1 deletion jumpax/_aggregators/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
from .._custom_meta import AbstractStrictModule
from .._custom_types import Args, Info, JumpState, RealScalarLike, SolverState, U


_SolverState = TypeVar("_SolverState", bound=SolverState)


Expand Down
1 change: 0 additions & 1 deletion jumpax/_custom_meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
import equinox as eqx
import ihoop


try:

class _StrictEqxMeta(ihoop.strict._StrictMeta, eqx._module._module._ModuleMeta):
Expand Down
1 change: 0 additions & 1 deletion jumpax/_custom_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
Shaped,
)


if TYPE_CHECKING:
RealScalarLike = bool | int | float | Array | np.ndarray
else:
Expand Down
9 changes: 4 additions & 5 deletions jumpax/_jumps.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
from ._custom_meta import AbstractStrictModule
from ._custom_types import Args, JumpState, Rate, RealScalarLike, U


_Rate = TypeVar("_Rate", bound=Rate)
_JumpState = TypeVar("_JumpState", bound=JumpState)

Expand Down Expand Up @@ -197,10 +196,10 @@ def leap_delta(self, t: RealScalarLike, u: U, args: Args) -> Int[Array, "1 S"]:


def _massaction_comb(
n: Float[Array, "*shape"],
k: Int[Array, "*shape"],
lgamma_k_plus_1: Float[Array, "*shape"],
) -> Float[Array, "*shape"]:
n: Float[Array, " *shape"],
k: Int[Array, " *shape"],
lgamma_k_plus_1: Float[Array, " *shape"],
) -> Float[Array, " *shape"]:
r"""
Compute the binomial coefficient $\binom{n}{k}$ using precomputed log-factorials.

Expand Down
34 changes: 8 additions & 26 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,7 @@ examples = [
testing = [
"pytest>=7.4.0",
"nbmake>=1.5.0",
"ruff==0.11.0",
"pyright==1.1.399",
"pyright==1.1.408",
]
docs = [
"mkdocs>=1.5.0",
Expand All @@ -54,31 +53,14 @@ docs = [
[tool.setuptools]
packages = { find = { include = ["jumpax", "jumpax.*"] } }

[tool.ruff]
extend-include = ["*.ipynb"]
src = []
[tool.black]
target-version = ["py312"]

[tool.ruff.lint]
# I001: isort import sorting
# F401: Unused imports
fixable = ["I001", "F401"]

# E402: Module level import not at top of file
# E731: Do not assign a lambda expression, use a def
# E741: Ambiguous variable name
# F722: Syntax error in forward annotation
ignore = ["E402", "E731", "E741", "F722"]

# Error codes to check for during linting
# E: pycodestyle errors
# F: pyflakes errors
# I001: isort import sorting
select = ["E", "F", "I001"]

[tool.ruff.lint.isort]
combine-as-imports = true
lines-after-imports = 2
order-by-type = false
[tool.isort]
profile = "black"
combine_as_imports = true
lines_after_imports = -1
order_by_type = false

[tool.pyright]
reportIncompatibleMethodOverride = true
Expand Down
3 changes: 2 additions & 1 deletion tests/test_cme.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
import unittest

import jax.numpy as jnp
import jumpax as jx
from jax.scipy.special import gammaln

import jumpax as jx


def poisson_pmf(k, lam):
"""Poisson probability mass function."""
Expand Down
1 change: 1 addition & 0 deletions tests/test_readme.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import jax.numpy as jnp
import jax.random as jr

import jumpax as jpx


Expand Down
Loading