diff --git a/.agents/GUIDELINES.md b/.agents/GUIDELINES.md new file mode 100644 index 0000000..efd3b51 --- /dev/null +++ b/.agents/GUIDELINES.md @@ -0,0 +1,10 @@ +# Agent Notes + +- Prefer explanatory comments in numerical kernels, coordinate math, and callback boundaries. +- Target roughly one meaningful comment for every 5-10 lines in dense array code. +- Comment the reason for a layout transform, cached constant, broadcast shape, or numerical safeguard. +- Do not add filler comments that simply restate the next line. +- Write tests, follow red green TDD. +- Ensure all methods, functions, modules, and classes have a docstring. Private objects can have a single line, public objects should have a full numpy docstring with examples (following doctest). +- All lines must be covered by tests; delete unreachable edge case code or test it using public API and minimal monkey patching. +- Keep Markdown prose unwrapped. Do not hard-wrap paragraphs or list items in `.md` files; let editors soft-wrap them. Code blocks and formats that require line breaks are exceptions. diff --git a/.agents/README.md b/.agents/README.md deleted file mode 100644 index 958e217..0000000 --- a/.agents/README.md +++ /dev/null @@ -1,7 +0,0 @@ -# Agent Notes - -- Prefer explanatory comments in numerical kernels, coordinate math, and callback boundaries. -- Target roughly one meaningful comment for every 5-10 lines in dense array code. -- Comment the reason for a layout transform, cached constant, broadcast shape, or numerical safeguard. -- Do not add filler comments that simply restate the next line. -- When splitting modules, keep package-level re-exports stable unless the task explicitly narrows the public API. diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000..4d10d65 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,67 @@ +name: Documentation + +on: + push: + branches: + - main + pull_request: + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: pages-${{ github.ref }} + cancel-in-progress: true + +jobs: + build: + runs-on: ubuntu-latest + permissions: + contents: read + pages: write + steps: + - uses: actions/checkout@v4 + with: + fetch-tags: "true" + fetch-depth: "0" + + - if: github.event_name != 'pull_request' + uses: actions/configure-pages@v5 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install documentation dependencies + run: python -m pip install -e ".[docs]" + + - name: Generate API reference + run: python scripts/build_api_docs.py + + - name: Generate benchmark reference + run: python scripts/build_benchmark_docs.py + + - name: Build documentation + run: zensical build --clean + + - name: Upload Pages artifact + if: github.event_name != 'pull_request' + uses: actions/upload-pages-artifact@v4 + with: + path: site + + deploy: + if: github.event_name != 'pull_request' + needs: build + runs-on: ubuntu-latest + permissions: + pages: write + id-token: write + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + steps: + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 diff --git a/.gitignore b/.gitignore index 697dbf2..9428243 100644 --- a/.gitignore +++ b/.gitignore @@ -74,6 +74,9 @@ target/ # docs docs/api/* +docs/benchmarks/* +site/ +.benchmarks/ _autosummary .quarto/ docs/site_libs diff --git a/README.md b/README.md index 987e175..8f76153 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,8 @@ # dasjax -An experimental package for accelerating [DASCore](dascore.org) with [JAX](https://github.com/jax-ml/jax). +![dasjax logo](https://raw.githubusercontent.com/dasdae/dasjax/main/docs/static/dasjax_logo.png) + +An experimental package for accelerating [DASCore](https://dascore.org) with [JAX](https://github.com/jax-ml/jax). ## Installation @@ -10,11 +12,11 @@ python -m pip install -e ".[dev]" ## Usage -`dasjax`'s main feature is the ability to create compiled DAS pipelines that can run on CPU, GPU, or TPU. These also perform kernel fusions for increased efficiency. +`dasjax`'s main feature is the ability to create compiled DAS pipelines that can run on CPU, GPU, or TPU. These pipelines fuse adjacent JAX-backed operations where possible and cache metadata planning for repeated calls with the same static patch boundary. ### Compiled pipeline -Use `JaxPatchPipeline` when you want to compile a reusable sequence once and run it across many compatible patches. +Use `JaxPatchPipeline` when you want to build a reusable callable once and run it across many compatible patches. ```python import dascore as dc @@ -38,62 +40,60 @@ print(out.shape) ## Development -### Three-Tier Architecture +### Architecture -`dasjax` is organized as a small three-tier stack: +`dasjax` is organized around one core operation model: -1. Pipeline layer: - `src/dasjax/pipeline.py` records operation chains and compiles reusable patch transforms. This is the main user-facing API. -2. Operation layer: - `src/dasjax/operations/` defines the operation registry, execution policies, validation rules, eager patch implementations, and compiled leaf transforms. -3. Kernel layer: - `src/dasjax/kernels/` contains the array-level JAX and callback-backed kernels that actually do the numerical work, grouped by domain (`basic`, `signal`, `filters`, `spectral`). +1. Pipeline layer: `src/dasjax/pipeline.py` records operation chains, plans metadata boundaries, and compiles reusable patch transforms. This is the main user-facing API. +2. Operation layer: `src/dasjax/core.py` defines `PatchOperation`, `PatchBoundary`, `PatchPyTree`, and registry helpers. Registered operation classes live under `src/dasjax/operations/`, grouped by DASCore-style domains. +3. Kernel layer: `src/dasjax/kernels/` contains the array-level JAX kernels that actually do the numerical work, grouped by domain (`basic`, `signal`, `filters`, `spectral`). -This split keeps the package easier to extend: add or update numerical behavior in the kernel layer, describe how it plugs into compiled execution in the operation layer, and expose it through the pipeline layer. +Operation authors use `bind(boundary)` for Python-side metadata planning, `kernel(patch_tree)` for JAX-side data transforms, and `update_boundary(boundary)` for static metadata changes. -### Roadmap +### Operation Coverage -The table below tracks what is missing and roughly how much effort each addition requires. +`dasjax` currently registers 72 pipeline operations. Most operations use native JAX kernels; a smaller set of DASCore-compatible numeric transforms still use host callbacks where a fully static JAX kernel is not practical yet. The current operation set includes: -#### Near-term — straightforward pure-JAX array ops +- Elementwise math and masks: `abs`, `clip`, `real`, `imag`, `angle`, `conj`, `exp`, `log`, `log10`, `log2`, `is_finite`, `isinf`, `isnan`, `fillna`, `where`, and scalar arithmetic operations. +- Reductions and aggregation: `aggregate`, `all`, `any`, `max`, `mean`, `median`, `min`, `std`, and `sum`. +- Coordinate-aware array transforms: `flip`, `roll`, `pad`, `taper`, `taper_range`, `detrend`, `standardize`, `differentiate`, and `integrate`. +- Spectral and signal operations: `dft`, `idft`, `stft`, `istft`, `hilbert`, `envelope`, `phase_weighted_stack`, `whiten`, `fbe`, and `correlate_shift`. +- Filters, mutes, and DAS-domain operations: `pass_filter`, `gaussian_filter`, `hampel_filter`, `median_filter`, `notch_filter`, `savgol_filter`, `sobel_filter`, `slope_filter`, `wiener_filter`, `line_mute`, `slope_mute`, `correlate`, `decimate`, `interpolate`, `resample`, `dispersion_phase_shift`, `tau_p`, `velocity_to_strain_rate`, `velocity_to_strain_rate_edgeless`, and `radians_to_strain`. -Implemented in the current package: +Remaining DASCore patch methods are mostly metadata, selection, convenience, or data-dependent shape operations. `rolling` returns a roller object rather than a patch, and `dropna` has data-dependent output shape, so neither fits the current static compiled-pipeline model directly. -- `real`, `imag`, `angle`, `conj` -- `flip`, `roll`, `pad` -- `standardize`, `differentiate`, `integrate` -- `dft`, `idft` -- `hilbert`, `envelope` -- `taper`, `taper_range` -- `whiten` +## Performance Notes -#### Medium-term — moderate effort or shape-changing +- The intended fast path is to build a `JaxPatchPipeline`, call `.compile()` once, and reuse the returned callable. Patch-specific metadata binding and JIT segment creation happen lazily on the first call for a static boundary, then cached plans and segment runners are reused for subsequent calls with matching dims, dynamic coordinate values, coordinate units, and attrs. +- Equivalent pipeline definitions reuse cached compiled callables automatically. +- Callback-backed operations preserve DASCore compatibility but execute their operation body on the host, so they generally do not benefit as much from JAX fusion as native kernels. +- Benchmarks live under `benchmarks/` and compare compiled `dasjax` pipelines against equivalent DASCore operation chains. -These need either more work in the kernel layer or are shape-changing (segmented pipeline execution, same mechanism as `fbe`). +## Documentation -| Method | Implementation notes | -|---|---| -| `notch_filter` | SOS filter; same pattern as `pass_filter` | -| `savgol_filter` | polynomial fitting per frame; JAX-doable | -| `rolling` | rolling-window reductions (mean, std, …); needs strided views | -| `correlate` | cross-correlation via `jnp.fft` | -| `stft` / `istft` | expose the STFT kernel already used by `fbe` | -| `decimate` | anti-aliased downsampling; shape-changing | -| `aggregate` / `mean` / `std` / `sum` | axis reductions; shape-changing | +Documentation is built with Zensical. The public API reference is generated at build time from the installed `dasjax` package, so run the API generation script before building or serving the site. -## Performance Notes +```bash +uv run python scripts/build_api_docs.py +uv run --extra docs zensical build --clean +``` -- The intended fast path is to build a `JaxPatchPipeline`, call `.compile()` once, and reuse the returned callable across many patches of compatible shape and dtype. -- Equivalent pipeline definitions reuse cached compiled callables automatically. -- Benchmarks live under `benchmarks/` and compare compiled `dasjax` pipelines against equivalent DASCore operation chains. +For local preview, run: + +```bash +uv run python scripts/build_api_docs.py +uv run --extra docs zensical serve +``` + +Generated files under `docs/api/` and `site/` are ignored by version control. GitHub Pages builds the same generated API docs and static site on pushes to `main`, then deploys the `site/` artifact through the `github-pages` environment. ## Development Guidelines -- Add new JAX patch methods by defining an array kernel in `src/dasjax/kernels/` and one operation spec in the relevant `src/dasjax/operations/` family module. -- The operation spec is the single source of truth for pipeline support, validation, and shared parity test cases. +- Add new JAX patch methods by defining an array kernel in `src/dasjax/kernels/` and one `PatchOperation` subclass in the appropriate `src/dasjax/operations/` module. +- The `PatchOperation` subclass is the single source of truth for pipeline support, metadata binding, and boundary updates. - Every new patch method must be tested against a DASCore baseline across the shared mixed-patch fixture in `tests/conftest.py`. - Prefer comparing internal operation behavior and compiled pipeline outputs against the closest native DASCore method or operator. If DASCore has no direct method, compare against an equivalent `Patch.update(...)` baseline. - Method-equivalence assertions should check data closeness with `equal_nan=True` when needed and should also verify coordinate preservation. -- Compiled pipeline parity should come from the same declared operation cases rather than a separate hand-maintained test matrix. +- Compiled pipeline parity should compare `JaxPatchPipeline` output against DASCore baselines for each registered operation. - Install Git hooks locally with `prek install`. diff --git a/agents.md b/agents.md index ba777d7..ea04ae3 100644 --- a/agents.md +++ b/agents.md @@ -2,7 +2,7 @@ This file gives AI/code agents a practical checklist for contributing safely to dasjax. -Keep Markdown prose unwrapped. Do not hard-wrap paragraphs in `.md` files unless a specific format requires it. +Keep Markdown prose unwrapped. Do not hard-wrap paragraphs or list items in `.md` files; let editors soft-wrap them. Code blocks and formats that require line breaks are exceptions. ## Scope and priorities diff --git a/benchmarks/readme.md b/benchmarks/readme.md index 85efcd4..d131ac6 100644 --- a/benchmarks/readme.md +++ b/benchmarks/readme.md @@ -17,12 +17,14 @@ pytest benchmarks/test_pipeline_benchmarks.py ## Benchmark Structure -The first benchmark suite focuses on side-by-side comparisons between: +The benchmark suite focuses on side-by-side comparisons between: - compiled `dasjax` pipelines - equivalent DASCore-native operation chains +- individual compiled `dasjax` operations +- equivalent individual DASCore operations -Each comparison is exposed as a separate benchmark test per engine so CodSpeed output is easy to read. +Each comparison is exposed as a separate benchmark test per engine so CodSpeed output is easy to read. Pipeline benchmark groups use names like `scale_fbe`; individual operation benchmark groups use names like `operation_fbe`. To export benchmark results for ratio comparisons, use: diff --git a/benchmarks/results/current.json b/benchmarks/results/current.json new file mode 100644 index 0000000..043f592 --- /dev/null +++ b/benchmarks/results/current.json @@ -0,0 +1,521 @@ +{ + "command": "/home/derrick/Gits/dasjax/.venv/bin/python3 -m pytest benchmarks/test_pipeline_benchmarks.py --benchmark-json=.benchmarks/current.json", + "environment": { + "dascore": "0.1.16", + "dasjax": "0.0.2.post6+git.fdb52b38.dirty", + "jax": "0.9.2", + "machine": "x86_64", + "platform": "Linux-6.8.0-110-generic-x86_64-with-glibc2.39", + "processor": "x86_64", + "python": "3.14.0" + }, + "generated_at": "2026-04-25T14:27:13+00:00", + "rows": [ + { + "case": "large-f32", + "dascore_mean_s": 0.007935926129023082, + "dasjax_mean_s": 0.017582151127201292, + "group": "operation_abs", + "speedup": 0.45136263882668115 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.015053255730610825, + "dasjax_mean_s": 0.034838203482900884, + "group": "operation_abs", + "speedup": 0.43209047039406584 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.0016216116442719372, + "dasjax_mean_s": 0.002737910963199475, + "group": "operation_abs", + "speedup": 0.5922806351514623 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.0029411434471127837, + "dasjax_mean_s": 0.0039036083065406677, + "group": "operation_abs", + "speedup": 0.753442255511335 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.007861990409571157, + "dasjax_mean_s": 0.017857173290923872, + "group": "operation_add", + "speedup": 0.4402707125862474 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.014968957292976797, + "dasjax_mean_s": 0.03530251717828443, + "group": "operation_add", + "speedup": 0.42401954561428906 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.0013730891043464352, + "dasjax_mean_s": 0.00273294466455809, + "group": "operation_add", + "speedup": 0.5024211145411025 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.002872821541935874, + "dasjax_mean_s": 0.003951023883951072, + "group": "operation_add", + "speedup": 0.7271081183804482 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.007773525349543112, + "dasjax_mean_s": 0.017493574536813498, + "group": "operation_clip", + "speedup": 0.4443646055975865 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.014837235603408772, + "dasjax_mean_s": 0.03473483007420092, + "group": "operation_clip", + "speedup": 0.42715728194763897 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.0015795147206869437, + "dasjax_mean_s": 0.002797094222745197, + "group": "operation_clip", + "speedup": 0.5646984316233492 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.002816467924502995, + "dasjax_mean_s": 0.003984043106562846, + "group": "operation_clip", + "speedup": 0.7069371111631487 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.010819450778073565, + "dasjax_mean_s": 0.02165447399992055, + "group": "operation_detrend", + "speedup": 0.49964043357106075 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.019358336065045518, + "dasjax_mean_s": 0.03907591064346759, + "group": "operation_detrend", + "speedup": 0.4954033251245981 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.002359406132353501, + "dasjax_mean_s": 0.004217029320276005, + "group": "operation_detrend", + "speedup": 0.5594948370430295 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.004759153708495725, + "dasjax_mean_s": 0.005132301489444065, + "group": "operation_detrend", + "speedup": 0.9272942593657413 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.0329788664482827, + "dasjax_mean_s": 0.02926682580786851, + "group": "operation_differentiate", + "speedup": 1.1268344119305274 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.04246424634984578, + "dasjax_mean_s": 0.05601835818742984, + "group": "operation_differentiate", + "speedup": 0.7580416085699292 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.006672640640317695, + "dasjax_mean_s": 0.004160485464541186, + "group": "operation_differentiate", + "speedup": 1.6038129918219886 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.012786168552802919, + "dasjax_mean_s": 0.007639906760496704, + "group": "operation_differentiate", + "speedup": 1.673602696163223 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.20849201119999633, + "dasjax_mean_s": 0.067468552866679, + "group": "operation_fbe", + "speedup": 3.0902102141124366 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.21955316859966842, + "dasjax_mean_s": 0.11780996044439639, + "group": "operation_fbe", + "speedup": 1.8636214439889613 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.05749154618706598, + "dasjax_mean_s": 0.011731497455548379, + "group": "operation_fbe", + "speedup": 4.9006144701395735 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.05919101474978561, + "dasjax_mean_s": 0.025517563621228014, + "group": "operation_fbe", + "speedup": 2.3196185822593467 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.06394093606650131, + "dasjax_mean_s": 0.033948521290092744, + "group": "operation_integrate", + "speedup": 1.8834674865547474 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.07652714800077284, + "dasjax_mean_s": 0.06064815950069007, + "group": "operation_integrate", + "speedup": 1.2618214407628 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.01654486007932324, + "dasjax_mean_s": 0.005963402924941606, + "group": "operation_integrate", + "speedup": 2.7743991622845523 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.01953043829625703, + "dasjax_mean_s": 0.009111122707663354, + "group": "operation_integrate", + "speedup": 2.143581962717943 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.016861063919932348, + "dasjax_mean_s": 0.017541805865361172, + "group": "operation_mean", + "speedup": 0.961193166162382 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.02815186759403332, + "dasjax_mean_s": 0.029849777181935322, + "group": "operation_mean", + "speedup": 0.9431181821709023 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.0045291220615997, + "dasjax_mean_s": 0.0031456233550736638, + "group": "operation_mean", + "speedup": 1.4398170252311209 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.006446509749821416, + "dasjax_mean_s": 0.002789384320467339, + "group": "operation_mean", + "speedup": 2.31108696729942 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.022142910000597607, + "dasjax_mean_s": 0.02258763956830411, + "group": "operation_normalize", + "speedup": 0.9803109321643965 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.03985147412489217, + "dasjax_mean_s": 0.04604641334757933, + "group": "operation_normalize", + "speedup": 0.8654631539719515 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.003866081424726048, + "dasjax_mean_s": 0.003445870612513223, + "group": "operation_normalize", + "speedup": 1.1219461957413273 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.0076487489297601555, + "dasjax_mean_s": 0.006458191331293379, + "group": "operation_normalize", + "speedup": 1.184348455688807 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.0076031236100588704, + "dasjax_mean_s": 0.017752565218424107, + "group": "operation_pad", + "speedup": 0.4282830969221359 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.01408750785177054, + "dasjax_mean_s": 0.035332438148880436, + "group": "operation_pad", + "speedup": 0.39871315396944734 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.002024273651784072, + "dasjax_mean_s": 0.002737090742528947, + "group": "operation_pad", + "speedup": 0.739571260949038 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.0029554693321760857, + "dasjax_mean_s": 0.0038526478810573758, + "group": "operation_pad", + "speedup": 0.7671267718774611 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.13081185362625547, + "dasjax_mean_s": 0.2346832335999352, + "group": "operation_pass_filter", + "speedup": 0.5573975252499316 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.1514530785001019, + "dasjax_mean_s": 0.2844445848008036, + "group": "operation_pass_filter", + "speedup": 0.5324519663686493 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.032527580249864387, + "dasjax_mean_s": 0.052176482400136594, + "group": "operation_pass_filter", + "speedup": 0.6234145874459952 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.04164246772081242, + "dasjax_mean_s": 0.07034695579956557, + "group": "operation_pass_filter", + "speedup": 0.5919583476996682 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.007709398385319686, + "dasjax_mean_s": 0.01771389471709109, + "group": "operation_scale", + "speedup": 0.435217579671022 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.014517987178805924, + "dasjax_mean_s": 0.036403064703470725, + "group": "operation_scale", + "speedup": 0.3988122235604454 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.0014040707887125262, + "dasjax_mean_s": 0.0026716843637301637, + "group": "operation_scale", + "speedup": 0.5255376749490664 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.0027331539843771678, + "dasjax_mean_s": 0.004020802098549036, + "group": "operation_scale", + "speedup": 0.6797534216775966 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.03271859779325213, + "dasjax_mean_s": 0.02677430021631645, + "group": "operation_standardize", + "speedup": 1.2220150490922332 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.06161480506658942, + "dasjax_mean_s": 0.047110426545424904, + "group": "operation_standardize", + "speedup": 1.3078804329478757 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.006401470681010033, + "dasjax_mean_s": 0.005071113542843734, + "group": "operation_standardize", + "speedup": 1.2623402388699572 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.013299871999966949, + "dasjax_mean_s": 0.007162377103119821, + "group": "operation_standardize", + "speedup": 1.8569075334184415 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.009185174989084693, + "dasjax_mean_s": 0.01774878547171884, + "group": "operation_taper", + "speedup": 0.5175100574470561 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.015738343254567803, + "dasjax_mean_s": 0.0372611652959744, + "group": "operation_taper", + "speedup": 0.4223792554407345 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.002441625380864328, + "dasjax_mean_s": 0.0025713899945996347, + "group": "operation_taper", + "speedup": 0.9495352264697946 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.0031923615308187435, + "dasjax_mean_s": 0.004030060436244411, + "group": "operation_taper", + "speedup": 0.7921373838735 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.0497878854001101, + "dasjax_mean_s": 0.03216135399959187, + "group": "scale_add_detrend_normalize", + "speedup": 1.5480655883064474 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.089915300092054, + "dasjax_mean_s": 0.05392152589430273, + "group": "scale_add_detrend_normalize", + "speedup": 1.6675214323182632 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.010027643979318176, + "dasjax_mean_s": 0.005617196311436172, + "group": "scale_add_detrend_normalize", + "speedup": 1.7851688677681923 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.02164517268201135, + "dasjax_mean_s": 0.00887979291158391, + "group": "scale_add_detrend_normalize", + "speedup": 2.4375762923226154 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.217333702200267, + "dasjax_mean_s": 0.06826758373354096, + "group": "scale_fbe", + "speedup": 3.183556386711362 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.23468941240134883, + "dasjax_mean_s": 0.12220381800014163, + "group": "scale_fbe", + "speedup": 1.920475286632017 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.059396552500629696, + "dasjax_mean_s": 0.016094810917775995, + "group": "scale_fbe", + "speedup": 3.6904162965362253 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.062460162599988205, + "dasjax_mean_s": 0.025890661076958237, + "group": "scale_fbe", + "speedup": 2.4124591648830287 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.22368254179964425, + "dasjax_mean_s": 0.06857853256269664, + "group": "scale_fbe_normalize", + "speedup": 3.2616991562942337 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.23593561480010977, + "dasjax_mean_s": 0.11862296349954704, + "group": "scale_fbe_normalize", + "speedup": 1.9889539751804521 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.062028346937268, + "dasjax_mean_s": 0.015839042776019257, + "group": "scale_fbe_normalize", + "speedup": 3.9161676506853436 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.06444275512512831, + "dasjax_mean_s": 0.025173454000185124, + "group": "scale_fbe_normalize", + "speedup": 2.559948870133372 + }, + { + "case": "large-f32", + "dascore_mean_s": 0.14785859057051962, + "dasjax_mean_s": 0.23864778060087702, + "group": "scale_pass_filter_abs", + "speedup": 0.6195682616374445 + }, + { + "case": "large-f64", + "dascore_mean_s": 0.182844042000094, + "dasjax_mean_s": 0.28871820400017895, + "group": "scale_pass_filter_abs", + "speedup": 0.6332958554978427 + }, + { + "case": "medium-f32", + "dascore_mean_s": 0.03665519632020733, + "dasjax_mean_s": 0.05447101847312297, + "group": "scale_pass_filter_abs", + "speedup": 0.6729302544305041 + }, + { + "case": "medium-f64", + "dascore_mean_s": 0.04954743119978957, + "dasjax_mean_s": 0.07087507012533933, + "group": "scale_pass_filter_abs", + "speedup": 0.6990812300032606 + } + ], + "schema_version": 1, + "source": "local" +} diff --git a/benchmarks/test_pipeline_benchmarks.py b/benchmarks/test_pipeline_benchmarks.py index c8af98c..50764d5 100644 --- a/benchmarks/test_pipeline_benchmarks.py +++ b/benchmarks/test_pipeline_benchmarks.py @@ -5,6 +5,7 @@ from collections.abc import Callable import dascore as dc +import numpy as np import pytest from dasjax import JaxPatchPipeline @@ -81,6 +82,213 @@ def dascore_scale_add_detrend_normalize(example_patch) -> Callable[[], object]: ) +@pytest.fixture(scope="module") +def dasjax_scale(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax scale operation.""" + compiled = JaxPatchPipeline().scale(2.0).compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_scale(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore scale operation.""" + return lambda: _scale_patch(example_patch, 2.0) + + +@pytest.fixture(scope="module") +def dasjax_add(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax add operation.""" + compiled = JaxPatchPipeline().add(1.0).compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_add(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore add operation.""" + return lambda: _add_patch(example_patch, 1.0) + + +@pytest.fixture(scope="module") +def dasjax_abs(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax abs operation.""" + compiled = JaxPatchPipeline().abs().compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_abs(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore abs operation.""" + return lambda: example_patch.abs() + + +@pytest.fixture(scope="module") +def dasjax_clip(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax clip operation.""" + compiled = JaxPatchPipeline().clip(-0.25, 0.25).compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_clip(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore clip operation.""" + return lambda: example_patch.update(data=np.clip(example_patch.data, -0.25, 0.25)) + + +@pytest.fixture(scope="module") +def dasjax_detrend(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax detrend operation.""" + compiled = JaxPatchPipeline().detrend(dim="time", type="constant").compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_detrend(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore detrend operation.""" + return lambda: example_patch.detrend(dim="time", type="constant") + + +@pytest.fixture(scope="module") +def dasjax_normalize(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax normalize operation.""" + compiled = JaxPatchPipeline().normalize(dim="time").compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_normalize(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore normalize operation.""" + return lambda: example_patch.normalize(dim="time") + + +@pytest.fixture(scope="module") +def dasjax_standardize(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax standardize operation.""" + compiled = JaxPatchPipeline().standardize(dim="time").compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_standardize(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore standardize operation.""" + return lambda: example_patch.standardize(dim="time") + + +@pytest.fixture(scope="module") +def dasjax_differentiate(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax differentiate operation.""" + compiled = JaxPatchPipeline().differentiate(dim="time").compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_differentiate(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore differentiate operation.""" + return lambda: example_patch.differentiate(dim="time") + + +@pytest.fixture(scope="module") +def dasjax_integrate(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax integrate operation.""" + compiled = JaxPatchPipeline().integrate(dim="time").compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_integrate(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore integrate operation.""" + return lambda: example_patch.integrate(dim="time") + + +@pytest.fixture(scope="module") +def dasjax_taper(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax taper operation.""" + compiled = JaxPatchPipeline().taper(time=0.05, window_type="hann").compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_taper(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore taper operation.""" + return lambda: example_patch.taper(time=0.05, window_type="hann") + + +@pytest.fixture(scope="module") +def dasjax_pad(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax pad operation.""" + compiled = JaxPatchPipeline().pad(time=(16, 16), samples=True).compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_pad(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore pad operation.""" + return lambda: example_patch.pad(time=(16, 16), samples=True) + + +@pytest.fixture(scope="module") +def dasjax_mean(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax mean reduction.""" + compiled = JaxPatchPipeline().mean(dim="time").compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_mean(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore mean reduction.""" + return lambda: example_patch.mean(dim="time") + + +@pytest.fixture(scope="module") +def dasjax_pass_filter(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax pass_filter operation.""" + compiled = JaxPatchPipeline().pass_filter(time=(2.0, 10.0)).compile() + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_pass_filter(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore pass_filter operation.""" + return lambda: example_patch.pass_filter(time=(2.0, 10.0)) + + +@pytest.fixture(scope="module") +def dasjax_fbe(example_patch) -> Callable[[], object]: + """Return a warmed compiled dasjax fbe operation.""" + compiled = ( + JaxPatchPipeline() + .fbe(time=64, samples=True, overlap=32, fmin=2.0, fmax=10.0) + .compile() + ) + compiled(example_patch) + return lambda: compiled(example_patch) + + +@pytest.fixture(scope="module") +def dascore_fbe(example_patch) -> Callable[[], object]: + """Return the equivalent DASCore fbe baseline operation.""" + return lambda: _run_dascore_fbe( + example_patch, + time=64, + overlap=32, + samples=True, + fmin=2.0, + fmax=10.0, + ) + + @pytest.fixture(scope="module") def dasjax_scale_pass_filter_abs(example_patch) -> Callable[[], object]: """Return a warmed compiled dasjax filtering pipeline.""" @@ -203,3 +411,163 @@ def test_dascore_scale_fbe_normalize( ) -> None: """Benchmark DASCore fbe+normalize chain execution.""" benchmark(dascore_scale_fbe_normalize) + + +class TestIndividualOperationBenchmarks: + """Benchmarks for individual dasjax operations and DASCore baselines.""" + + @pytest.mark.benchmark(group="operation_scale") + def test_dasjax_compiled_operation_scale(self, benchmark, dasjax_scale) -> None: + """Benchmark warmed compiled dasjax scale execution.""" + benchmark(dasjax_scale) + + @pytest.mark.benchmark(group="operation_scale") + def test_dascore_operation_scale(self, benchmark, dascore_scale) -> None: + """Benchmark DASCore scale execution.""" + benchmark(dascore_scale) + + @pytest.mark.benchmark(group="operation_add") + def test_dasjax_compiled_operation_add(self, benchmark, dasjax_add) -> None: + """Benchmark warmed compiled dasjax add execution.""" + benchmark(dasjax_add) + + @pytest.mark.benchmark(group="operation_add") + def test_dascore_operation_add(self, benchmark, dascore_add) -> None: + """Benchmark DASCore add execution.""" + benchmark(dascore_add) + + @pytest.mark.benchmark(group="operation_abs") + def test_dasjax_compiled_operation_abs(self, benchmark, dasjax_abs) -> None: + """Benchmark warmed compiled dasjax abs execution.""" + benchmark(dasjax_abs) + + @pytest.mark.benchmark(group="operation_abs") + def test_dascore_operation_abs(self, benchmark, dascore_abs) -> None: + """Benchmark DASCore abs execution.""" + benchmark(dascore_abs) + + @pytest.mark.benchmark(group="operation_clip") + def test_dasjax_compiled_operation_clip(self, benchmark, dasjax_clip) -> None: + """Benchmark warmed compiled dasjax clip execution.""" + benchmark(dasjax_clip) + + @pytest.mark.benchmark(group="operation_clip") + def test_dascore_operation_clip(self, benchmark, dascore_clip) -> None: + """Benchmark DASCore-equivalent clip execution.""" + benchmark(dascore_clip) + + @pytest.mark.benchmark(group="operation_detrend") + def test_dasjax_compiled_operation_detrend(self, benchmark, dasjax_detrend) -> None: + """Benchmark warmed compiled dasjax detrend execution.""" + benchmark(dasjax_detrend) + + @pytest.mark.benchmark(group="operation_detrend") + def test_dascore_operation_detrend(self, benchmark, dascore_detrend) -> None: + """Benchmark DASCore detrend execution.""" + benchmark(dascore_detrend) + + @pytest.mark.benchmark(group="operation_normalize") + def test_dasjax_compiled_operation_normalize( + self, benchmark, dasjax_normalize + ) -> None: + """Benchmark warmed compiled dasjax normalize execution.""" + benchmark(dasjax_normalize) + + @pytest.mark.benchmark(group="operation_normalize") + def test_dascore_operation_normalize(self, benchmark, dascore_normalize) -> None: + """Benchmark DASCore normalize execution.""" + benchmark(dascore_normalize) + + @pytest.mark.benchmark(group="operation_standardize") + def test_dasjax_compiled_operation_standardize( + self, benchmark, dasjax_standardize + ) -> None: + """Benchmark warmed compiled dasjax standardize execution.""" + benchmark(dasjax_standardize) + + @pytest.mark.benchmark(group="operation_standardize") + def test_dascore_operation_standardize( + self, benchmark, dascore_standardize + ) -> None: + """Benchmark DASCore standardize execution.""" + benchmark(dascore_standardize) + + @pytest.mark.benchmark(group="operation_differentiate") + def test_dasjax_compiled_operation_differentiate( + self, benchmark, dasjax_differentiate + ) -> None: + """Benchmark warmed compiled dasjax differentiate execution.""" + benchmark(dasjax_differentiate) + + @pytest.mark.benchmark(group="operation_differentiate") + def test_dascore_operation_differentiate( + self, benchmark, dascore_differentiate + ) -> None: + """Benchmark DASCore differentiate execution.""" + benchmark(dascore_differentiate) + + @pytest.mark.benchmark(group="operation_integrate") + def test_dasjax_compiled_operation_integrate( + self, benchmark, dasjax_integrate + ) -> None: + """Benchmark warmed compiled dasjax integrate execution.""" + benchmark(dasjax_integrate) + + @pytest.mark.benchmark(group="operation_integrate") + def test_dascore_operation_integrate(self, benchmark, dascore_integrate) -> None: + """Benchmark DASCore integrate execution.""" + benchmark(dascore_integrate) + + @pytest.mark.benchmark(group="operation_taper") + def test_dasjax_compiled_operation_taper(self, benchmark, dasjax_taper) -> None: + """Benchmark warmed compiled dasjax taper execution.""" + benchmark(dasjax_taper) + + @pytest.mark.benchmark(group="operation_taper") + def test_dascore_operation_taper(self, benchmark, dascore_taper) -> None: + """Benchmark DASCore taper execution.""" + benchmark(dascore_taper) + + @pytest.mark.benchmark(group="operation_pad") + def test_dasjax_compiled_operation_pad(self, benchmark, dasjax_pad) -> None: + """Benchmark warmed compiled dasjax pad execution.""" + benchmark(dasjax_pad) + + @pytest.mark.benchmark(group="operation_pad") + def test_dascore_operation_pad(self, benchmark, dascore_pad) -> None: + """Benchmark DASCore pad execution.""" + benchmark(dascore_pad) + + @pytest.mark.benchmark(group="operation_mean") + def test_dasjax_compiled_operation_mean(self, benchmark, dasjax_mean) -> None: + """Benchmark warmed compiled dasjax mean reduction execution.""" + benchmark(dasjax_mean) + + @pytest.mark.benchmark(group="operation_mean") + def test_dascore_operation_mean(self, benchmark, dascore_mean) -> None: + """Benchmark DASCore mean reduction execution.""" + benchmark(dascore_mean) + + @pytest.mark.benchmark(group="operation_pass_filter") + def test_dasjax_compiled_operation_pass_filter( + self, benchmark, dasjax_pass_filter + ) -> None: + """Benchmark warmed compiled dasjax pass_filter execution.""" + benchmark(dasjax_pass_filter) + + @pytest.mark.benchmark(group="operation_pass_filter") + def test_dascore_operation_pass_filter( + self, benchmark, dascore_pass_filter + ) -> None: + """Benchmark DASCore pass_filter execution.""" + benchmark(dascore_pass_filter) + + @pytest.mark.benchmark(group="operation_fbe") + def test_dasjax_compiled_operation_fbe(self, benchmark, dasjax_fbe) -> None: + """Benchmark warmed compiled dasjax fbe execution.""" + benchmark(dasjax_fbe) + + @pytest.mark.benchmark(group="operation_fbe") + def test_dascore_operation_fbe(self, benchmark, dascore_fbe) -> None: + """Benchmark DASCore fbe-equivalent execution.""" + benchmark(dascore_fbe) diff --git a/docs/development/architecture.md b/docs/development/architecture.md new file mode 100644 index 0000000..7e3bca7 --- /dev/null +++ b/docs/development/architecture.md @@ -0,0 +1,17 @@ +# Architecture + +`dasjax` is organized around one core operation model. + +## Pipeline Layer + +`src/dasjax/pipeline.py` records operation chains, plans metadata boundaries, and compiles reusable patch transforms. This is the main user-facing API. Compiled callables cache both JIT segment runners and bound metadata plans for repeated calls with the same static boundary. + +## Operation Layer + +`src/dasjax/core.py` defines `PatchOperation`, `PatchBoundary`, `PatchPyTree`, and registry helpers. Registered operation classes live under `src/dasjax/operations/`, grouped by DASCore-style domains. + +Operation authors use `bind(boundary)` for Python-side metadata planning, `kernel(patch_tree)` for JAX-side data transforms, and `update_boundary(boundary)` for static metadata changes. + +## Kernel Layer + +`src/dasjax/kernels/` contains the array-level JAX kernels that do the numerical work, grouped by domain. diff --git a/docs/development/documentation.md b/docs/development/documentation.md new file mode 100644 index 0000000..46751ef --- /dev/null +++ b/docs/development/documentation.md @@ -0,0 +1,34 @@ +# Documentation + +Documentation is built with Zensical. The API reference page is generated from the public `dasjax` package API immediately before the site build. Benchmark documentation is generated from the checked-in benchmark snapshot. + +```bash +uv run python scripts/build_api_docs.py +uv run python scripts/build_benchmark_docs.py +uv run --extra docs zensical build --clean +``` + +Preview the site locally with: + +```bash +uv run python scripts/build_api_docs.py +uv run python scripts/build_benchmark_docs.py +uv run --extra docs zensical serve +``` + +Generated API reference pages are written under `docs/api/`, generated benchmark docs are written under `docs/benchmarks/`, and the static site is written to `site/`. These outputs are ignored by version control. + +Refresh the benchmark snapshot before building benchmark docs with: + +```bash +uv run dasjax-benchmark refresh +``` + +## Local Checks + +```bash +uv run python scripts/build_api_docs.py +uv run python scripts/build_benchmark_docs.py +uv run --extra docs zensical build --clean +uv run pytest +``` diff --git a/docs/development/index.md b/docs/development/index.md new file mode 100644 index 0000000..22f679e --- /dev/null +++ b/docs/development/index.md @@ -0,0 +1,9 @@ +# Development + +Use this section when you are changing `dasjax` internals or extending the operation surface. + +## Pages + +- [Architecture](architecture.md) explains the pipeline, operation, and kernel layers. +- [Operation Coverage](operation-coverage.md) summarizes the registered operation surface. +- [Documentation](documentation.md) describes local docs builds and checks. diff --git a/docs/development/operation-coverage.md b/docs/development/operation-coverage.md new file mode 100644 index 0000000..3a41fa0 --- /dev/null +++ b/docs/development/operation-coverage.md @@ -0,0 +1,7 @@ +# Operation Coverage + +The package currently registers 72 pipeline operations. Most are implemented as JAX-backed kernels with static metadata planning; heavier DASCore-compatible numeric transforms can use host callbacks when a full static JAX kernel is not yet practical. Callback-backed operations are compatibility paths and generally offer less fusion benefit than native kernels. + +Remaining DASCore patch methods are mostly metadata, selection, convenience, or data-dependent shape operations. In particular, `rolling` returns a roller object and `dropna` has data-dependent output shape, so they do not map directly to the current compiled `Patch -> Patch` pipeline model. + +The generated [API Reference](../api/index.md) lists the current operation registry. diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..612c7a5 --- /dev/null +++ b/docs/index.md @@ -0,0 +1 @@ +--8<-- "README.md" diff --git a/docs/static/dasjax_logo.png b/docs/static/dasjax_logo.png new file mode 100644 index 0000000..23d140c Binary files /dev/null and b/docs/static/dasjax_logo.png differ diff --git a/docs/static/favicon.ico b/docs/static/favicon.ico new file mode 100644 index 0000000..5b119e8 Binary files /dev/null and b/docs/static/favicon.ico differ diff --git a/docs/tutorial/compiled-pipelines.md b/docs/tutorial/compiled-pipelines.md new file mode 100644 index 0000000..aeb0889 --- /dev/null +++ b/docs/tutorial/compiled-pipelines.md @@ -0,0 +1,34 @@ +# Compiled Pipelines + +Use `JaxPatchPipeline` when you want to build a reusable callable once and run it across many compatible patches. + +```python +import dascore as dc +from dasjax import JaxPatchPipeline + +patch = dc.get_example_patch("example_event_1") + +pipeline = ( + JaxPatchPipeline() + .scale(2.0) + .add(1.0) + .detrend(dim="time", type="constant") + .normalize(dim="time") +) +compiled = pipeline.compile() + +out = patch.pipe(compiled) + +print(out.shape) +``` + +## Workflow + +1. Create a `JaxPatchPipeline`. +2. Add DASCore-style patch operations. +3. Call `.compile()` once. +4. Reuse the compiled callable with `patch.pipe(compiled)` or `compiled(patch)`. + +## Compatibility + +The compiled callable can be reused for patches with matching static metadata: dims, dynamic coordinate values, coordinate units, and attrs. When those metadata inputs change, `dasjax` plans and caches a new compatible execution path. diff --git a/docs/tutorial/index.md b/docs/tutorial/index.md new file mode 100644 index 0000000..b08c3d4 --- /dev/null +++ b/docs/tutorial/index.md @@ -0,0 +1,10 @@ +# Tutorial + +Use this section when you want to learn the user-facing workflow first. + +`dasjax` centers on `JaxPatchPipeline`: build a DASCore-style operation chain, compile it once, and reuse the returned callable for compatible patches. + +## Pages + +- [Compiled Pipelines](compiled-pipelines.md) shows the main workflow. +- [Performance](performance.md) explains what is cached and what still runs on the host. diff --git a/docs/tutorial/performance.md b/docs/tutorial/performance.md new file mode 100644 index 0000000..4921e45 --- /dev/null +++ b/docs/tutorial/performance.md @@ -0,0 +1,15 @@ +# Performance + +The intended fast path is to call `.compile()` once and reuse the returned callable. + +## Caching + +Patch-specific metadata binding and JIT segment creation happen lazily on the first call for a static patch boundary. Cached plans and segment runners are reused for later calls with matching dims, dynamic coordinate values, coordinate units, and attrs. + +Equivalent pipeline definitions reuse cached compiled callables automatically. + +## Native Kernels And Callbacks + +Native JAX-backed operations can be fused into compiled segments. Some heavier DASCore-compatible numeric transforms still use host callbacks; those operations are useful for compatibility, but they do not usually benefit as much from JAX fusion as native kernels. + +Benchmarks live under `benchmarks/` and compare compiled `dasjax` pipelines against equivalent DASCore operation chains. diff --git a/pyproject.toml b/pyproject.toml index 589869a..c50077b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,12 @@ dev = [ "prek>=0.3.8", "ruff>=0.14.5", ] +docs = [ + "zensical", +] + +[project.scripts] +dasjax-benchmark = "dasjax.benchmark_cli:main" [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/scripts/build_api_docs.py b/scripts/build_api_docs.py new file mode 100644 index 0000000..936ed6b --- /dev/null +++ b/scripts/build_api_docs.py @@ -0,0 +1,484 @@ +"""Generate Markdown API reference pages for the Zensical docs build.""" + +from __future__ import annotations + +import inspect +import os +import re +import shutil +import sys +from dataclasses import dataclass +from inspect import Signature +from pathlib import Path +from textwrap import dedent +from types import ModuleType +from typing import Any + + +ROOT = Path(__file__).resolve().parents[1] +SRC = ROOT / "src" +OUT_DIR = ROOT / "docs" / "api" + +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) + +import dasjax # noqa: E402 + + +@dataclass(frozen=True) +class ApiEntry: + """One generated API documentation page.""" + + name: str + display_name: str + obj: Any + kind: str + path: Path + owner: "ApiEntry | None" = None + + +def _doc(obj: Any) -> str: + """Return a cleaned docstring or fallback text.""" + return inspect.getdoc(obj) or "No public docstring is available." + + +def _summary(obj: Any) -> str: + """Return a compact one-line summary from a docstring.""" + lines = [line.strip() for line in _doc(obj).splitlines() if line.strip()] + return lines[0] if lines else "" + + +def _signature(obj: Any) -> str: + """Return a readable signature for functions, methods, and classes.""" + target = obj.__init__ if inspect.isclass(obj) else obj + try: + signature = inspect.signature(target) + except (TypeError, ValueError): + return "()" + params = [ + parameter + for name, parameter in signature.parameters.items() + if name not in {"self", "cls"} + ] + return str(signature.replace(parameters=params, return_annotation=Signature.empty)) + + +def _slug(value: str) -> str: + """Return a filesystem-safe, deterministic slug.""" + value = value.replace("__", "-") + value = re.sub(r"[^A-Za-z0-9_.-]+", "-", value) + return value.strip("-").replace(".", "-").lower() + + +def _table_cell(value: str) -> str: + """Escape Markdown table cell separators.""" + return value.replace("|", "\\|").replace("\n", " ") + + +def _page_link(from_page: Path, entry: ApiEntry) -> str: + """Return a relative Markdown link from one generated page to another.""" + return Path(os.path.relpath(entry.path, from_page.parent)).as_posix() + + +def _xref_map(entries: tuple[ApiEntry, ...]) -> dict[str, ApiEntry]: + """Return local reference aliases mapped to generated API pages.""" + refs: dict[str, ApiEntry] = {} + for entry in entries: + aliases = { + entry.name, + entry.display_name, + entry.name.removeprefix("dasjax."), + entry.display_name.rsplit(".", 1)[-1], + } + module_name = getattr(entry.obj, "__module__", None) + object_name = getattr(entry.obj, "__name__", None) + if module_name and object_name: + aliases.add(f"{module_name}.{object_name}") + aliases.add(object_name) + if entry.kind == "method" and entry.owner is not None: + method_name = entry.display_name.rsplit(".", 1)[-1] + aliases.add(f"{entry.owner.display_name}.{method_name}") + aliases.add(f"{entry.owner.name}.{method_name}") + for alias in aliases: + if alias and alias not in refs: + refs[alias] = entry + return refs + + +def _link_ref(token: str, *, page: Path, refs: dict[str, ApiEntry]) -> str | None: + """Return a Markdown link for a reference token if it can be resolved.""" + clean = token.strip() + suffix = "" + while clean.endswith("()"): + clean = clean[:-2] + suffix += "()" + entry = refs.get(clean) + if entry is None: + return None + return f"[`{token}`]({_page_link(page, entry)})" + + +def _link_docstring(text: str, *, page: Path, refs: dict[str, ApiEntry]) -> str: + """Convert common Python docstring references into Markdown links.""" + + def _replace_role(match: re.Match[str]) -> str: + label = match.group("label") or match.group("target") + target = match.group("target") + link = _link_ref(target, page=page, refs=refs) + if link is None: + return match.group(0) + href = link.rsplit("](", 1)[1][:-1] + return f"[`{label}`]({href})" + + def _replace_double_backtick(match: re.Match[str]) -> str: + token = match.group("token") + return _link_ref(token, page=page, refs=refs) or f"`{token}`" + + def _replace_backtick(match: re.Match[str]) -> str: + token = match.group("token") + return _link_ref(token, page=page, refs=refs) or match.group(0) + + text = re.sub( + r":(?:class|func|meth|mod|obj|py:class|py:func|py:meth|py:mod):" + r"`(?:(?P