Improve timing scripts

.gitattributes

@@ -1,2 +1,3 @@
 docs/** linguist-vendored
-notebooks/** linguist-vendored
+notebooks/** linguist-vendored
+tests/** linguist-vendored

.github/workflows/CI.yml

@@ -0,0 +1,66 @@
+name: CI
+
+on:
+  pull_request:
+    branches: [main, ripple-dev]
+  push:
+    branches: [main, ripple-dev]
+  release:
+    types: [published]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+      - name: Set up Python ${{ matrix.python-version }}
+        run: uv python install ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: uv sync --python ${{ matrix.python-version }} --group test
+      - name: Run tests
+        run: uv run pytest --cov=ripplegw --cov-report=term-missing tests/unit tests/integration --ignore=tests/cross_validation --ignore=tests/benchmarks
+      - name: Upload coverage to Coveralls
+        if: matrix.python-version == '3.12'
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          COVERALLS_REPO_TOKEN: ${{ secrets.COVERALLS_REPO_TOKEN }}
+        run: uv run coveralls
+
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+      - name: Install dependencies
+        run: uv sync --python 3.12 --group test
+      - name: Run pre-commit
+        run: uv run pre-commit run --all-files
+
+  publish:
+    needs: [test, pre-commit]
+    runs-on: ubuntu-latest
+    if: github.event_name == 'release'
+    environment:
+      name: pypi
+      url: https://pypi.org/project/ripplegw/
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+      - name: Build package
+        run: uv build
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.PYPI_API_TOKEN }}

.gitignore

@@ -1,10 +1,32 @@
+# Python
+__pycache__/
+*.py[cod]
+
+*.egg-info/
+
+# Testing / coverage
+.pytest_cache/
+.coverage
+.coverage.*
+htmlcov/
+
+# Virtual environments
+.venv/
+venv/
+
+# IDE / Editor
+.vscode/
+*.swp
+
+# OS
 .DS_Store
-.ipynb_checkpoints/
+
+# Python environment
 .python-version
-*__pycache__
-*.egg-info
-docs/build/*
-figures/*.pdf
-.vscode/*
 
-build/*
+# Documentation
+docs/_build/
+/site
+
+# Jupyter
+.ipynb_checkpoints/

.pre-commit-config.yaml

@@ -0,0 +1,26 @@
+ci:
+  autoupdate_branch: ripple-dev
+
+repos:
+  - repo: https://github.com/astral-sh/uv-pre-commit
+    rev: 0.10.8
+    hooks:
+      - id: uv-lock
+        args: ["--upgrade"]
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.15.4
+    hooks:
+      - id: ruff-check
+        args: ["--fix"]
+        files: src/
+      - id: ruff-format
+        files: src/
+
+  # Temporarily disabled - re-enable once type errors are fixed
+  # - repo: https://github.com/RobertCraigie/pyright-python
+  #   rev: v1.1.408
+  #   hooks:
+  #   - id: pyright
+  #     files: src/
+  #     additional_dependencies: [jax, jaxtyping, pytest, typing_extensions, numpy]

LICENSE

@@ -1,6 +1,7 @@
 MIT License
 
-Copyright (c) [2022] [Adam Coogan, Thomas Edwards]
+Copyright (c) 2022 Adam Coogan, Thomas Edwards
+Copyright (c) 2025 GW-JAX-Team
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

README.md

@@ -1,100 +1,128 @@
-# Ripple :ocean:
+# Ripple 🌊
 
-A small `jax` package for differentiable and fast gravitational wave data analysis. 
+**A small `jax` package for differentiable and fast gravitational wave data analysis**
 
-# Getting Started
+[![doc](https://badgen.net/badge/Read/the%20doc/blue)](https://ripplegw.readthedocs.io/) [![license](https://badgen.net/badge/License/MIT/blue)](https://github.com/GW-JAX-Team/ripple/blob/main/LICENSE) [![coverage](https://badgen.net/coveralls/c/github/GW-JAX-Team/ripple/main)](https://coveralls.io/github/GW-JAX-Team/ripple?branch=main) [![pre-commit.ci status](https://results.pre-commit.ci/badge/github/GW-JAX-Team/ripple/main.svg)](https://results.pre-commit.ci/latest/github/GW-JAX-Team/ripple/main)
 
-Note that Thibeau Wouters and Kaze Wong will soon become the main developers of ripple so please contact them if you have comments/questions.
+ripple is a JAX-based package for differentiable and hardware-accelerated gravitational wave data analysis.It is maintained by the GW-JAX-Team organization and was originally developed by Thomas Edwards and Adam Coogan, with significant contributions from Kaze Wong and the community.
 
-## Documentation
+See the accompanying paper, [Edwards et al. (2024)](https://journals.aps.org/prd/abstract/10.1103/PhysRevD.110.064028), for more details. For questions or comments, please open an issue on the [GitHub repository](https://github.com/GW-JAX-Team/ripple).
 
-You can find the full documentation at [Read the Docs](https://ripplegw.readthedocs.io/).
+## Installation
 
-### Installation
+The simplest way to install ripple is through pip:
 
-Both waveforms have been tested extensively and match `lalsuite` implementations to machine precision across all the parameter space. 
+```bash
+pip install ripplegw
+```
 
-Ripple can be installed using 
+This will install the latest stable release and its dependencies.
+ripple is built on [JAX](https://github.com/google/jax).
+By default, this installs the CPU version of JAX from [PyPI](https://pypi.org).
+If you have a GPU and want to leverage hardware acceleration, install the CUDA-enabled version:
 
-```
-pip3 install ripplegw
+```bash
+pip install ripplegw[cuda]
 ```
 
-Note that by default we do not include enable float64 in `jax`` since we want allow users to use float32 to improve performance.
-If you require float64, please include the following code at the start of the script:
+If you want to install the latest version of ripple, you can clone this repo and install it locally:
 
+```bash
+git clone https://github.com/GW-JAX-Team/ripple.git
+cd ripple
+pip install -e .
 ```
+
+**Note:** By default, ripple uses float32 precision for improved performance. If you require float64 precision, add the following at the start of your script:
+
+```python
 from jax import config
 config.update("jax_enable_x64", True)
 ```
 
-See https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html for other common `jax` gotchas.
+See [JAX - The Sharp Bits](https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html#jax-the-sharp-bits) for other common `jax` gotchas.
+
+## Supported Waveforms
 
-### Supported waveforms
+All waveforms have been extensively tested and match `lalsuite` implementations to machine precision across the full parameter space.
 
-- IMRPhenomXAS (aligned spin)
-- IMRPhenomD (aligned spin)
-- IMRPhenomPv2 (Still finalizing sampling checks)
-- TaylorF2 with tidal effects
-- IMRPhenomD_NRTidalv2, verified for the low spin regime (chi1, chi2 < 0.05), further testing is required for higher spins
+- **IMRPhenomXAS** (aligned spin)
+- **IMRPhenomD** (aligned spin)
+- **IMRPhenomPv2** (finalizing sampling validation)
+- **TaylorF2** with tidal effects
+- **IMRPhenomD_NRTidalv2** (verified for low spin: $\chi_1$ and $\chi_2$ < 0.05; higher spins require further testing)
+- **IMRPhenomXAS_NRTidalv3**
 
-### Generating a waveform and its derivative
+## Usage
 
-Generating a waveform is incredibly easy. Below is an example of calling the PhenomXAS waveform model
-to get the h_+ and h_x polarizations of the waveform model
+### Generating a Waveform
 
-We start with some basic imports:
+Generating waveforms with ripple is straightforward. Below is an example using the IMRPhenomXAS model to compute the $h_+$ and $h_\times$ polarizations.
+
+Start with the basic imports:
 
 ```python
 import jax.numpy as jnp
-
-from ripple.waveforms import IMRPhenomXAS
-from ripple import ms_to_Mc_eta
+from ripplegw import IMRPhenomXAS
+from ripplegw.conversions import ms_to_Mc_eta
 ```
 
 And now we can just set the parameters and call the waveform!
 
 ```python
-# Get a frequency domain waveform
-# source parameters
-
-m1_msun = 20.0 # In solar masses
-m2_msun = 19.0
-chi1 = 0.5 # Dimensionless spin
-chi2 = -0.5
-tc = 0.0 # Time of coalescence in seconds
-phic = 0.0 # Time of coalescence
-dist_mpc = 440 # Distance to source in Mpc
-inclination = 0.0 # Inclination Angle
-
-# The PhenomD waveform model is parameterized with the chirp mass and symmetric mass ratio
+# Define source parameters
+m1_msun = 20.0           # Primary mass (solar masses)
+m2_msun = 19.0           # Secondary mass (solar masses)
+chi1 = 0.5               # Primary dimensionless spin
+chi2 = -0.5              # Secondary dimensionless spin
+tc = 0.0                 # Time of coalescence (seconds)
+phic = 0.0               # Phase at coalescence (radians)
+dist_mpc = 440           # Luminosity distance (Mpc)
+inclination = 0.0        # Inclination angle (radians)
+
+# Convert to chirp mass and symmetric mass ratio
 Mc, eta = ms_to_Mc_eta(jnp.array([m1_msun, m2_msun]))
 
-# These are the parametrs that go into the waveform generator
-# Note that JAX does not give index errors, so if you pass in the
-# the wrong array it will behave strangely
+# Construct parameter array
+# Note: JAX does not raise index errors, so ensure the array is correctly ordered
 theta_ripple = jnp.array([Mc, eta, chi1, chi2, dist_mpc, tc, phic, inclination])
 
-# Now we need to generate the frequency grid
-f_l = 24
-f_u = 512
-del_f = 0.01
+# Generate frequency grid
+f_l = 24                 # Lower frequency bound (Hz)
+f_u = 512                # Upper frequency bound (Hz)
+del_f = 0.01             # Frequency resolution (Hz)
 fs = jnp.arange(f_l, f_u, del_f)
-f_ref = f_l
+f_ref = f_l              # Reference frequency
 
-# And finally lets generate the waveform!
+# Generate the waveform
 hp_ripple, hc_ripple = IMRPhenomXAS.gen_IMRPhenomXAS_hphc(fs, theta_ripple, f_ref)
 
-# Note that we have not internally jitted the functions since this would
-# introduce an annoying overhead each time the user evaluated the function with a different length frequency array
-# We therefore recommend that the user jit the function themselves to accelerate evaluations. For example:
+# For better performance, we recommend JIT-compiling the waveform function.
+# This avoids recompilation overhead when the frequency array length changes:
 
 import jax
 
 @jax.jit
 def waveform(theta):
-    return IMRPhenomXAS.gen_IMRPhenomXAS_hphc(fs, theta)
+    return IMRPhenomXAS.gen_IMRPhenomXAS_hphc(fs, theta, f_ref)
 ```
 
-
-
+## Attribution
+
+If you use ripple in your research, please cite the accompanying paper:
+
+```bibtex
+@article{Edwards:2023sak,
+    author = "Edwards, Thomas D. P. and Wong, Kaze W. K. and Lam, Kelvin K. H. and Coogan, Adam and Foreman-Mackey, Daniel and Isi, Maximiliano and Zimmerman, Aaron",
+    title = "{Differentiable and hardware-accelerated waveforms for gravitational wave data analysis}",
+    eprint = "2302.05329",
+    archivePrefix = "arXiv",
+    primaryClass = "astro-ph.IM",
+    doi = "10.1103/PhysRevD.110.064028",
+    journal = "Phys. Rev. D",
+    volume = "110",
+    number = "6",
+    pages = "064028",
+    year = "2024"
+}
+```

docs/requirements.txt

@@ -6,10 +6,10 @@ equinox>=0.10.6
 optax>=0.1.5
 evosax>=0.1.4
 tqdm
-mkdocs==1.4.3               # Main documentation generator.
-mkdocs-material==9.1.18     # Theme
-pymdown-extensions==10.1    # Markdown extensions e.g. to handle LaTeX.
-mkdocstrings[python]==0.22.0        # Autogenerate documentation from docstrings.
-mkdocs-jupyter==0.24.2       # Turn Jupyter Lab notebooks into webpages.
-mkdocs-gen-files==0.5.0
-mkdocs-literate-nav==0.6.0
+mkdocs==1.6.1               # Main documentation generator.
+mkdocs-material==9.7.1     # Theme
+pymdown-extensions==10.20.1    # Markdown extensions e.g. to handle LaTeX.
+mkdocstrings[python]==1.0.3        # Autogenerate documentation from docstrings.
+mkdocs-jupyter==0.25.1       # Turn Jupyter Lab notebooks into webpages.
+mkdocs-gen-files==0.6.0
+mkdocs-literate-nav==0.6.2