Complete pyQSC_JAX refactor with pyQSC parity and plasma–coil field jets

.github/workflows/benchmarks.yml

@@ -0,0 +1,30 @@
+name: benchmarks
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "29 8 * * 0"
+
+jobs:
+  verify:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    env:
+      JAX_ENABLE_X64: "true"
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.14"
+          cache: pip
+      - run: python -m pip install -e '.[dev]'
+      - run: pytest -m numerics
+      - run: python benchmarks/benchmark_core.py
+        env:
+          PYQSC_BENCHMARK_OUTPUT: benchmarks/results/ci.json
+      - name: Record environment
+        run: python -m pip freeze
+      - uses: actions/upload-artifact@v4
+        with:
+          name: benchmark-report
+          path: benchmarks/results/ci.json

.github/workflows/compatibility.yml

@@ -0,0 +1,88 @@
+name: compatibility
+
+on:
+  pull_request:
+  workflow_dispatch:
+  schedule:
+    - cron: "43 6 * * 3"
+
+jobs:
+  frozen-contracts:
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    env:
+      JAX_ENABLE_X64: "true"
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.14"
+          cache: pip
+      - run: python -m pip install -e '.[dev]'
+      - run: pytest tests/compatibility tests/regression
+
+  upstream-head-smoke:
+    if: github.event_name == 'schedule' || github.event_name == 'workflow_dispatch'
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    continue-on-error: true
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.14"
+          cache: pip
+      - run: python -m pip install -e '.[dev]'
+      - run: python -m pip install 'git+https://github.com/landreman/pyQSC.git'
+      - run: >-
+          python -c "from qsc import Qsc;
+          from pyqsc_jax.near_axis import near_axis;
+          kwargs = dict(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3, etabar=-0.9, nphi=31);
+          assert abs(float(Qsc(**kwargs).iota) - float(near_axis(**kwargs).iota)) < 1e-10"
+
+  essos-field-jet:
+    runs-on: ubuntu-latest
+    timeout-minutes: 25
+    env:
+      JAX_ENABLE_X64: "true"
+      MPLBACKEND: Agg
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/checkout@v4
+        with:
+          repository: uwplasma/ESSOS
+          ref: refactor/pyqsc-external-field-jet
+          path: essos-integration
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.14"
+          cache: pip
+      - run: python -m pip install -e '.[dev,plot]'
+      - run: python -m pip install -e ./essos-integration
+      - run: python -m pip install -e .
+      - run: >-
+          pytest
+          essos-integration/tests/test_field_jet.py
+          essos-integration/tests/test_field_jet_examples.py
+
+  vmex-implicit:
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    env:
+      JAX_ENABLE_X64: "true"
+      PYQSC_RUN_VMEX: "1"
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/checkout@v4
+        with:
+          repository: uwplasma/vmex
+          ref: main
+          path: vmex-integration
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+      - run: python -m pip install -e '.[dev]'
+      - run: python -m pip install -e ./vmex-integration
+      - run: python -m pip install -e .
+      - run: pytest tests/integration/test_vmex_live.py

.github/workflows/docs.yml

@@ -0,0 +1,30 @@
+name: docs
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+  workflow_dispatch:
+  schedule:
+    - cron: "17 7 * * 1"
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    env:
+      JAX_ENABLE_X64: "true"
+      MPLBACKEND: Agg
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.14"
+          cache: pip
+      - run: python -m pip install -e '.[docs,plot]'
+      - run: python -m sphinx -W --keep-going -b html docs docs/_build/html
+      - run: python -m sphinx -W --keep-going -b doctest docs docs/_build/doctest
+      - run: python examples/01_first_order_qa.py
+      - name: Check external links
+        if: github.event_name == 'schedule'
+        run: python -m sphinx -W --keep-going -b linkcheck docs docs/_build/linkcheck

.github/workflows/publish.yml

@@ -0,0 +1,87 @@
+name: publish
+
+on:
+  release:
+    types: [published]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  verify-and-build:
+    runs-on: ubuntu-latest
+    timeout-minutes: 45
+    env:
+      JAX_ENABLE_X64: "true"
+      MPLBACKEND: Agg
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - name: Require production release commit on main
+        if: github.event_name == 'release'
+        run: >-
+          git fetch origin main:refs/remotes/origin/main &&
+          git merge-base --is-ancestor "$GITHUB_SHA" origin/main
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.14"
+          cache: pip
+      - run: python -m pip install -e '.[dev,docs,plot]'
+      - run: ruff check src tests examples benchmarks
+      - run: ruff format --check src tests examples benchmarks
+      - run: pytest --cov=pyqsc_jax --cov-branch
+      - run: python -m sphinx -W --keep-going -b html docs docs/_build/html
+      - run: python -m sphinx -W --keep-going -b doctest docs docs/_build/doctest
+      - run: python -m build
+      - run: python -m twine check dist/*
+      - run: python -m venv "${{ runner.temp }}/wheel-env"
+      - run: "${{ runner.temp }}/wheel-env/bin/python -m pip install dist/*.whl"
+      - run: >-
+          JAX_ENABLE_X64=true
+          ${{ runner.temp }}/wheel-env/bin/python -c
+          "import pyqsc_jax as q; s=q.solve_configuration('qa', nphi=31);
+          assert s.root_report.converged and s.linear_report.converged"
+      - run: python -m venv "${{ runner.temp }}/sdist-env"
+      - run: "${{ runner.temp }}/sdist-env/bin/python -m pip install dist/*.tar.gz"
+      - run: >-
+          JAX_ENABLE_X64=true
+          ${{ runner.temp }}/sdist-env/bin/python -c
+          "from pyqsc_jax.near_axis import near_axis;
+          s=near_axis(rc=[1.0,0.045],zs=[0.0,-0.045],nfp=3,etabar=-0.9);
+          assert abs(float(s.iota)-0.41830690943386617)<1e-10"
+      - uses: actions/upload-artifact@v4
+        with:
+          name: distributions
+          path: dist/
+
+  publish-testpypi:
+    if: github.event_name == 'workflow_dispatch'
+    needs: verify-and-build
+    runs-on: ubuntu-latest
+    environment: testpypi
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: distributions
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://test.pypi.org/legacy/
+
+  publish-pypi:
+    if: github.event_name == 'release'
+    needs: verify-and-build
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: distributions
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1

.github/workflows/tests.yml

@@ -0,0 +1,81 @@
+name: tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.14"
+          cache: pip
+      - run: python -m pip install ruff
+      - run: ruff check src tests examples benchmarks
+      - run: ruff format --check src tests examples benchmarks
+
+  test:
+    permissions:
+      contents: read
+      id-token: write
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - os: ubuntu-latest
+            python: "3.12"
+          - os: ubuntu-latest
+            python: "3.14"
+          - os: macos-latest
+            python: "3.14"
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 30
+    env:
+      JAX_ENABLE_X64: "true"
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+          cache: pip
+      - run: python -m pip install -e '.[dev,plot]'
+      - run: python -c "import jax, pyqsc_jax; print(jax.__version__, pyqsc_jax.__version__, pyqsc_jax.__file__)"
+      - run: pytest --cov=pyqsc_jax --cov-branch --cov-report=term --cov-report=xml
+      - uses: codecov/codecov-action@v5
+        if: matrix.os == 'ubuntu-latest' && matrix.python == '3.14'
+        with:
+          use_oidc: true
+          fail_ci_if_error: true
+
+  package:
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    env:
+      JAX_ENABLE_X64: "true"
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.14"
+          cache: pip
+      - run: python -m pip install build twine
+      - run: python -m build
+      - run: python -m twine check dist/*
+      - run: python -m pip install --force-reinstall dist/*.whl
+      - run: >-
+          python -c "from pyqsc_jax.near_axis import near_axis;
+          q = near_axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3, etabar=-0.9);
+          assert abs(float(q.iota) - 0.41830690943386617) < 1e-10"

.gitignore

@@ -0,0 +1,13 @@
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.coverage
+coverage.xml
+htmlcov/
+.audit/
+docs/_build/
+build/
+dist/
+*.egg-info/
+examples/output/
+benchmarks/results/

.readthedocs.yaml

@@ -0,0 +1,17 @@
+version: 2
+
+build:
+  os: ubuntu-24.04
+  tools:
+    python: "3.12"
+
+sphinx:
+  configuration: docs/conf.py
+  fail_on_warning: true
+
+python:
+  install:
+    - method: pip
+      path: .
+      extra_requirements:
+        - docs