dasjax 0.0.2__tar.gz

dasjax-0.0.2/LICENSE

@@ -0,0 +1,167 @@
+dasjax is licensed under LGPL version 3 license:
+
+                   GNU LESSER GENERAL PUBLIC LICENSE
+                       Version 3, 29 June 2007
+
+ Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
+
+
+  This version of the GNU Lesser General Public License incorporates
+the terms and conditions of version 3 of the GNU General Public
+License, supplemented by the additional permissions listed below.
+
+  0. Additional Definitions.
+
+  As used herein, "this License" refers to version 3 of the GNU Lesser
+General Public License, and the "GNU GPL" refers to version 3 of the GNU
+General Public License.
+
+  "The Library" refers to a covered work governed by this License,
+other than an Application or a Combined Work as defined below.
+
+  An "Application" is any work that makes use of an interface provided
+by the Library, but which is not otherwise based on the Library.
+Defining a subclass of a class defined by the Library is deemed a mode
+of using an interface provided by the Library.
+
+  A "Combined Work" is a work produced by combining or linking an
+Application with the Library.  The particular version of the Library
+with which the Combined Work was made is also called the "Linked
+Version".
+
+  The "Minimal Corresponding Source" for a Combined Work means the
+Corresponding Source for the Combined Work, excluding any source code
+for portions of the Combined Work that, considered in isolation, are
+based on the Application, and not on the Linked Version.
+
+  The "Corresponding Application Code" for a Combined Work means the
+object code and/or source code for the Application, including any data
+and utility programs needed for reproducing the Combined Work from the
+Application, but excluding the System Libraries of the Combined Work.
+
+  1. Exception to Section 3 of the GNU GPL.
+
+  You may convey a covered work under sections 3 and 4 of this License
+without being bound by section 3 of the GNU GPL.
+
+  2. Conveying Modified Versions.
+
+  If you modify a copy of the Library, and, in your modifications, a
+facility refers to a function or data to be supplied by an Application
+that uses the facility (other than as an argument passed when the
+facility is invoked), then you may convey a copy of the modified
+version:
+
+   a) under this License, provided that you make a good faith effort to
+   ensure that, in the event an Application does not supply the
+   function or data, the facility still operates, and performs
+   whatever part of its purpose remains meaningful, or
+
+   b) under the GNU GPL, with none of the additional permissions of
+   this License applicable to that copy.
+
+  3. Object Code Incorporating Material from Library Header Files.
+
+  The object code form of an Application may incorporate material from
+a header file that is part of the Library.  You may convey such object
+code under terms of your choice, provided that, if the incorporated
+material is not limited to numerical parameters, data structure
+layouts and accessors, or small macros, inline functions and templates
+(ten or fewer lines in length), you do both of the following:
+
+   a) Give prominent notice with each copy of the object code that the
+   Library is used in it and that the Library and its use are
+   covered by this License.
+
+   b) Accompany the object code with a copy of the GNU GPL and this license
+   document.
+
+  4. Combined Works.
+
+  You may convey a Combined Work under terms of your choice that,
+taken together, effectively do not restrict modification of the
+portions of the Library contained in the Combined Work and reverse
+engineering for debugging such modifications, if you also do each of
+the following:
+
+   a) Give prominent notice with each copy of the Combined Work that
+   the Library is used in it and that the Library and its use are
+   covered by this License.
+
+   b) Accompany the Combined Work with a copy of the GNU GPL and this license
+   document.
+
+   c) For a Combined Work that displays copyright notices during
+   execution, include the copyright notice for the Library among
+   these notices, as well as a reference directing the user to the
+   copies of the GNU GPL and this license document.
+
+   d) Do one of the following:
+
+       0) Convey the Minimal Corresponding Source under the terms of this
+       License, and the Corresponding Application Code in a form
+       suitable for, and under terms that permit, the user to
+       recombine or relink the Application with a modified version of
+       the Linked Version to produce a modified Combined Work, in the
+       manner specified by section 6 of the GNU GPL for conveying
+       Corresponding Source.
+
+       1) Use a suitable shared library mechanism for linking with the
+       Library.  A suitable mechanism is one that (a) uses at run time
+       a copy of the Library already present on the user's computer
+       system, and (b) will operate properly with a modified version
+       of the Library that is interface-compatible with the Linked
+       Version.
+
+   e) Provide Installation Information, but only if you would otherwise
+   be required to provide such information under section 6 of the
+   GNU GPL, and only to the extent that such information is
+   necessary to install and execute a modified version of the
+   Combined Work produced by recombining or relinking the
+   Application with a modified version of the Linked Version. (If
+   you use option 4d0, the Installation Information must accompany
+   the Minimal Corresponding Source and Corresponding Application
+   Code. If you use option 4d1, you must provide the Installation
+   Information in the manner specified by section 6 of the GNU GPL
+   for conveying Corresponding Source.)
+
+  5. Combined Libraries.
+
+  You may place library facilities that are a work based on the
+Library side by side in a single library together with other library
+facilities that are not Applications and are not covered by this
+License, and convey such a combined library under terms of your
+choice, if you do both of the following:
+
+   a) Accompany the combined library with a copy of the same work based
+   on the Library, uncombined with any other library facilities,
+   conveyed under the terms of this License.
+
+   b) Give prominent notice with the combined library that part of it
+   is a work based on the Library, and explaining where to find the
+   accompanying uncombined form of the same work.
+
+  6. Revised Versions of the GNU Lesser General Public License.
+
+  The Free Software Foundation may publish revised and/or new versions
+of the GNU Lesser General Public License from time to time. Such new
+versions will be similar in spirit to the present version, but may
+differ in detail to address new problems or concerns.
+
+  Each version is given a distinguishing version number. If the
+Library as you received it specifies that a certain numbered version
+of the GNU Lesser General Public License "or any later version"
+applies to it, you have the option of following the terms and
+conditions either of that published version or of any later version
+published by the Free Software Foundation. If the Library as you
+received it does not specify a version number of the GNU Lesser
+General Public License, you may choose any version of the GNU Lesser
+General Public License ever published by the Free Software Foundation.
+
+  If the Library as you received it specifies that a proxy can decide
+whether future versions of the GNU Lesser General Public License shall
+apply, that proxy's public statement of acceptance of any version is
+permanent authorization for you to choose that version for the
+Library.

dasjax-0.0.2/PKG-INFO

@@ -0,0 +1,126 @@
+Metadata-Version: 2.4
+Name: dasjax
+Version: 0.0.2
+Summary: A dascore compatibility layer for JAX.
+License-Expression: LGPL-3.0-or-later
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: dascore>=0.1.15
+Requires-Dist: jax>=0.4.33
+Requires-Dist: scipy
+Requires-Dist: pandas
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest-benchmark; extra == "test"
+Requires-Dist: pytest-codspeed; extra == "test"
+Provides-Extra: dev
+Requires-Dist: prek>=0.3.8; extra == "dev"
+Requires-Dist: ruff>=0.14.5; extra == "dev"
+Dynamic: license-file
+
+# dasjax
+
+An experimental package for accelerating [DASCore](dascore.org) with [JAX](https://github.com/jax-ml/jax).
+
+## Installation
+
+```bash
+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.
+
+### Compiled pipeline
+
+Use `JaxPatchPipeline` when you want to compile a reusable sequence 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)
+```
+
+## Development
+
+### Three-Tier Architecture
+
+`dasjax` is organized as a small three-tier stack:
+
+1. Pipeline layer:
+   `src/dasjax/pipeline.py` records operation chains and compiles reusable patch transforms. This is the main user-facing API.
+2. Operation and pipeline layer:
+   `src/dasjax/operations.py` defines the operation registry, validation rules, internal eager implementations, and compiled behavior.
+3. Kernel layer:
+   `src/dasjax/kernels.py` contains the array-level JAX and callback-backed kernels that actually do the numerical work.
+
+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.
+
+
+### Roadmap
+
+The table below tracks what is missing and roughly how much effort each addition requires.
+
+#### Near-term — straightforward pure-JAX array ops
+
+No new infrastructure needed; each maps directly to one or two `jnp` calls.
+
+| Method | Implementation notes |
+|---|---|
+| `real`, `imag`, `angle`, `conj` | `jnp` one-liners for complex patches |
+| `flip` | `jnp.flip` along a named dim |
+| `roll` | `jnp.roll` circular shift along a dim |
+| `pad` | `jnp.pad` with DASCore coordinate extension |
+| `standardize` | zero-mean + unit-std (compare `normalize`) |
+| `differentiate` | `jnp.diff` finite differences along a dim |
+| `integrate` | `jnp.cumsum` / trapezoid along a dim |
+| `dft` / `idft` | `jnp.fft.rfft` / `irfft` with coord reconstruction |
+| `hilbert` / `envelope` | hilbert via FFT; `envelope = abs(hilbert(data))` |
+| `taper` / `taper_range` | hann / cosine windows broadcast along axis |
+| `whiten` | spectral divide-by-amplitude via FFT |
+
+#### Medium-term — moderate effort or shape-changing
+
+These need either more work in kernels.py or are shape-changing (segmented pipeline execution, same mechanism as `fbe`).
+
+| 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 |
+
+## Performance Notes
+
+- 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.
+
+## Development Guidelines
+
+- Add new JAX patch methods by defining an array kernel in `src/dasjax/kernels.py` and one operation spec in `src/dasjax/operations.py`.
+- The operation spec is the single source of truth for pipeline support, validation, and shared parity test cases.
+- 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.
+- Install Git hooks locally with `prek install`.

dasjax-0.0.2/README.md

@@ -0,0 +1,105 @@
+# dasjax
+
+An experimental package for accelerating [DASCore](dascore.org) with [JAX](https://github.com/jax-ml/jax).
+
+## Installation
+
+```bash
+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.
+
+### Compiled pipeline
+
+Use `JaxPatchPipeline` when you want to compile a reusable sequence 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)
+```
+
+## Development
+
+### Three-Tier Architecture
+
+`dasjax` is organized as a small three-tier stack:
+
+1. Pipeline layer:
+   `src/dasjax/pipeline.py` records operation chains and compiles reusable patch transforms. This is the main user-facing API.
+2. Operation and pipeline layer:
+   `src/dasjax/operations.py` defines the operation registry, validation rules, internal eager implementations, and compiled behavior.
+3. Kernel layer:
+   `src/dasjax/kernels.py` contains the array-level JAX and callback-backed kernels that actually do the numerical work.
+
+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.
+
+
+### Roadmap
+
+The table below tracks what is missing and roughly how much effort each addition requires.
+
+#### Near-term — straightforward pure-JAX array ops
+
+No new infrastructure needed; each maps directly to one or two `jnp` calls.
+
+| Method | Implementation notes |
+|---|---|
+| `real`, `imag`, `angle`, `conj` | `jnp` one-liners for complex patches |
+| `flip` | `jnp.flip` along a named dim |
+| `roll` | `jnp.roll` circular shift along a dim |
+| `pad` | `jnp.pad` with DASCore coordinate extension |
+| `standardize` | zero-mean + unit-std (compare `normalize`) |
+| `differentiate` | `jnp.diff` finite differences along a dim |
+| `integrate` | `jnp.cumsum` / trapezoid along a dim |
+| `dft` / `idft` | `jnp.fft.rfft` / `irfft` with coord reconstruction |
+| `hilbert` / `envelope` | hilbert via FFT; `envelope = abs(hilbert(data))` |
+| `taper` / `taper_range` | hann / cosine windows broadcast along axis |
+| `whiten` | spectral divide-by-amplitude via FFT |
+
+#### Medium-term — moderate effort or shape-changing
+
+These need either more work in kernels.py or are shape-changing (segmented pipeline execution, same mechanism as `fbe`).
+
+| 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 |
+
+## Performance Notes
+
+- 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.
+
+## Development Guidelines
+
+- Add new JAX patch methods by defining an array kernel in `src/dasjax/kernels.py` and one operation spec in `src/dasjax/operations.py`.
+- The operation spec is the single source of truth for pipeline support, validation, and shared parity test cases.
+- 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.
+- Install Git hooks locally with `prek install`.

dasjax-0.0.2/pyproject.toml

@@ -0,0 +1,47 @@
+[build-system]
+requires = ["setuptools>=80", "setuptools-git-versioning"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools-git-versioning]
+enabled = true
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[project]
+name = "dasjax"
+dynamic = ["version"]
+description = "A dascore compatibility layer for JAX."
+readme = "README.md"
+requires-python = ">=3.12"
+license = "LGPL-3.0-or-later"
+dependencies = [
+  "dascore>=0.1.15",
+  "jax>=0.4.33",
+  "scipy",
+  "pandas",
+]
+
+[project.optional-dependencies]
+test = [
+  "pytest",
+  "pytest-benchmark",
+  "pytest-codspeed",
+]
+dev = [
+  "prek>=0.3.8",
+  "ruff>=0.14.5",
+]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+markers = [
+  "benchmark: performance benchmark tests",
+]
+
+[tool.ruff]
+line-length = 88
+target-version = "py312"

dasjax-0.0.2/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

dasjax-0.0.2/src/dasjax/__init__.py

@@ -0,0 +1,22 @@
+"""DASCore extensions for JAX-backed patch processing."""
+
+# Ruff would normally require imports before executable statements, but this
+# flag must be set before importing package modules that initialize JAX objects.
+# ruff: noqa: E402
+
+from importlib.metadata import PackageNotFoundError, version
+
+import jax
+
+jax.config.update("jax_enable_x64", True)
+
+from . import pytree as pytree
+from .operations import JaxPatchOperation, list_operations
+from .pipeline import JaxPatchPipeline
+
+__all__ = ["JaxPatchOperation", "JaxPatchPipeline", "list_operations"]
+
+try:
+    __version__ = version("dasjax")
+except PackageNotFoundError:
+    __version__ = "0+unknown"