diff --git a/.github/ISSUE_TEMPLATE/bug-report.yml b/.github/ISSUE_TEMPLATE/bug-report.yml index 01bed0d..bc71300 100644 --- a/.github/ISSUE_TEMPLATE/bug-report.yml +++ b/.github/ISSUE_TEMPLATE/bug-report.yml @@ -48,10 +48,11 @@ body: attributes: label: "Runtime Environment" description: > - Please provide a description of the environment in which the error - occurred. + Please provide the versions of fiber, Python, and JAX (jax and jaxlib), + the operating system, and whether JAX is running on CPU or GPU. placeholder: > - Raspberry Pi 4 running Ubuntu 22.04 natively. + fiber 0.1.0, Python 3.14, jax 0.4.38 / jaxlib 0.4.38 on CUDA GPU, + Ubuntu 24.04. validations: required: true diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..1035b83 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,40 @@ +name: Continuous Integration + +on: + push: + branches: + - main + pull_request: + workflow_dispatch: + +jobs: + lint: + name: Lint + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v7 + + - name: Install uv + uses: astral-sh/setup-uv@v7 + + - name: Run Ruff linter + run: uvx ruff@0.16.5 check --output-format=github + + - name: Run Ruff formatter + run: uvx ruff@0.16.5 format --check + + lock: + name: Lockfile + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v7 + + - name: Install uv + uses: astral-sh/setup-uv@v7 + with: + python-version: "3.14" + + - name: Check that uv.lock is up to date + run: uv lock --check diff --git a/.github/workflows/format.yml b/.github/workflows/format.yml new file mode 100644 index 0000000..3d85e85 --- /dev/null +++ b/.github/workflows/format.yml @@ -0,0 +1,31 @@ +name: Formatting (pre-commit) + +on: + pull_request: + push: + branches: + - main + workflow_dispatch: + +jobs: + pre-commit: + name: Format + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v7 + + - name: Setup Python + uses: actions/setup-python@v7 + with: + python-version: "3.14" + + - name: Run pre-commit + uses: pre-commit/action@v3.0.1 + id: precommit + + - name: Upload pre-commit changes + if: failure() && steps.precommit.outcome == 'failure' + uses: rhaschke/upload-git-patch-action@main + with: + name: pre-commit diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..79af7ca --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,47 @@ +name: Release + +on: + push: + tags: + - "v*" + +permissions: + contents: write + +jobs: + release: + name: Build and publish release + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v7 + with: + fetch-depth: 0 + + - name: Setup Python + uses: actions/setup-python@v7 + with: + python-version: "3.14" + + - name: Install uv + uses: astral-sh/setup-uv@v7 + + - name: Verify tag matches project version + run: | + version=$(python -c 'import tomllib, pathlib; print(tomllib.loads(pathlib.Path("pyproject.toml").read_text())["project"]["version"])') + if [ "v${version}" != "${GITHUB_REF_NAME}" ]; then + echo "::error::Tag ${GITHUB_REF_NAME} does not match pyproject.toml version v${version}" + exit 1 + fi + + - name: Build distributions + run: uv build + + - name: Create GitHub release + env: + GH_TOKEN: ${{ github.token }} + run: | + prerelease="" + case "${GITHUB_REF_NAME}" in *-*) prerelease="--prerelease" ;; esac + gh release create "${GITHUB_REF_NAME}" dist/* \ + --generate-notes --verify-tag ${prerelease} diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 55cd696..4575618 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,14 +1,16 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.14.14 + rev: v0.16.5 hooks: - - id: ruff + - id: ruff-check args: ["--fix", "--exit-non-zero-on-fix"] + - id: ruff-format - repo: https://github.com/codespell-project/codespell rev: v2.4.1 hooks: - id: codespell + args: ["-L", "groupt"] - repo: https://github.com/pre-commit/pre-commit-hooks rev: v6.0.0 diff --git a/README.md b/README.md index 48e3bd8..eb192c7 100644 --- a/README.md +++ b/README.md @@ -71,7 +71,7 @@ the typed wrapper classes: ```python import fiber.numpy as fnp -w_hat = fnp.skew3(w_vec) # 3x3 skew-symmetric matrix +w_hat = fnp.skew3(w_vec) # 3x3 skew-symmetric matrix R_mat = fnp.so3.expm(w_hat) # 3x3 rotation matrix ``` diff --git a/fiber/solvers/_euler.py b/fiber/solvers/_euler.py index 5cdffcb..1ff9396 100644 --- a/fiber/solvers/_euler.py +++ b/fiber/solvers/_euler.py @@ -24,7 +24,7 @@ import equinox as eqx from diffrax import RESULTS, AbstractItoSolver, AbstractTerm -from .._custom_types import VF, Args, BoolScalarLike, DenseInfo, RealScalarLike +from .._custom_types import VF, Args, BoolScalarLike, Control, DenseInfo, RealScalarLike from .._groups._element import AbstractTangentVector from .._local_interpolation import LocalLeftBundleInterpolation as LocalInterpolation from .._operations import rplus @@ -35,7 +35,7 @@ class LieEuler(AbstractItoSolver): - term_structure: ClassVar = AbstractTerm + term_structure: ClassVar = AbstractTerm[AbstractTangentVector, Control] interpolation_cls: ClassVar[Callable[..., LocalInterpolation]] = LocalInterpolation def order(self, terms): diff --git a/fiber/solvers/_euler_heun.py b/fiber/solvers/_euler_heun.py index 43d153e..69d791b 100644 --- a/fiber/solvers/_euler_heun.py +++ b/fiber/solvers/_euler_heun.py @@ -19,22 +19,29 @@ # THE SOFTWARE. from collections.abc import Callable -from typing import ClassVar +from typing import ClassVar, cast import equinox as eqx from diffrax import RESULTS, AbstractStratonovichSolver, AbstractTerm, MultiTerm +from diffrax._term import WrapTerm -from .._custom_types import VF, Args, BoolScalarLike, DenseInfo, RealScalarLike -from .._groups import AbstractTangentVector +from .._custom_types import VF, Args, BoolScalarLike, Control, DenseInfo, RealScalarLike +from .._groups import AbstractCotangentVector, AbstractTangentVector from .._local_interpolation import LocalLeftBundleInterpolation as LocalInterpolation from .._operations import rplus from ._term import SharpTerm type _ErrorEstimate = None type _SolverState = None -type _Terms = MultiTerm[tuple[SharpTerm, AbstractTerm]] type _V = AbstractTangentVector +_Terms = MultiTerm[ + tuple[ + SharpTerm[AbstractCotangentVector], + AbstractTerm[AbstractCotangentVector, Control], + ] +] + class EulerHeun(AbstractStratonovichSolver): term_structure: ClassVar = _Terms @@ -72,20 +79,20 @@ def step( del solver_state, made_jump drift, diffusion = terms.terms + dual_metric = cast(SharpTerm, cast(WrapTerm, drift).term).dual_metric + dt = drift.contr(t0, t1) dw = diffusion.contr(t0, t1) f0 = drift.vf_prod(t0, y0, args, dt) - h0 = diffusion.prod(diffusion.vf(t0, y0, args), dw) - h0 = drift.term.dual_metric(y0, h0) # type: ignore + h0 = diffusion.vf_prod(t0, y0, args, dw) - y_prime = y0 + h0 + y_prime = y0 + dual_metric(y0, h0) # type: ignore[reportOperatorIssue] h_prime = diffusion.vf_prod(t0, y_prime, args, dw) - h_prime = drift.term.dual_metric(y_prime, h_prime) # type: ignore - vf = f0 + 0.5 * (h0 + h_prime) - y1 = y0 + vf - y1 = eqx.tree_at(lambda w: w.point.value, y1, rplus(y0.point, f0).value) + k = f0 + 0.5 * (h0 + h_prime) + y1 = y0 + dual_metric(y0, k) # type: ignore[reportOperatorIssue] + y1 = eqx.tree_at(lambda w: w.point.value, y1, rplus(y0.point, y0 * dt).value) dense_info = {"y0": y0, "y1": y1} return y1, None, dense_info, None, RESULTS.successful @@ -94,7 +101,8 @@ def func( self, terms: _Terms, t0: RealScalarLike, - y0: AbstractTangentVector, + y0: _V, args: Args, ) -> VF: - return terms.vf(t0, y0, args) + drift, diffusion = terms.terms + return drift.vf(t0, y0, args), diffusion.vf(t0, y0, args) diff --git a/fiber/solvers/_variational_integrator.py b/fiber/solvers/_variational_integrator.py index a372beb..8e08287 100644 --- a/fiber/solvers/_variational_integrator.py +++ b/fiber/solvers/_variational_integrator.py @@ -40,10 +40,16 @@ type _ErrorEstimate = None type _SolverState = None -type _Terms = MultiTerm[tuple[ImplicitVariationalTerm, VariationalDiffusionTerm]] type _V = AbstractTangentVector type _CV = AbstractCotangentVector +_Terms = MultiTerm[ + tuple[ + ImplicitVariationalTerm[AbstractCotangentVector], + VariationalDiffusionTerm[AbstractCotangentVector], + ] +] + def _implicit_relation(v: Array, solver_args: Args) -> Array: implicit_step, y_prime, t1, vector_cls, integrator_args, dt = solver_args diff --git a/pyproject.toml b/pyproject.toml index 7bd2189..467989b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,3 +1,7 @@ +[build-system] +requires = ["flit_core >=3.8.0,<4"] +build-backend = "flit_core.buildapi" + [project] name = "fiber" version = "0.1.0" diff --git a/uv.lock b/uv.lock index 84d4478..5fa164b 100644 --- a/uv.lock +++ b/uv.lock @@ -47,7 +47,7 @@ wheels = [ [[package]] name = "fiber" version = "0.1.0" -source = { virtual = "." } +source = { editable = "." } dependencies = [ { name = "diffrax" }, { name = "equinox" },