refjax 0.1.1__tar.gz

refjax-0.1.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Max Champneys
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

refjax-0.1.1/PKG-INFO

@@ -0,0 +1,186 @@
+Metadata-Version: 2.4
+Name: refjax
+Version: 0.1.1
+Summary: Lightweight `jit` and `grad` compatible reflectometry in jax.
+Author: Stephanie L. Burg
+Author-email: "Max D. Champneys" <max.champneys@sheffield.ac.uk>
+Maintainer-email: Max Champneys <max.champneys@sheffield.ac.uk>
+License-Expression: MIT
+Project-URL: Repository, https://github.com/MDCHAMP/refjax
+Project-URL: Issues, https://github.com/MDCHAMP/refjax/issues
+Project-URL: Paper, https://doi.org/10.1038/s43246-026-01341-4
+Keywords: Reflectometry,Neutron reflectometry,jax,autograd
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=2.0
+Requires-Dist: jax>=0.4.26
+Requires-Dist: jaxlib>=0.4.26
+Provides-Extra: test
+Requires-Dist: pytest>=7.0; extra == "test"
+Requires-Dist: pandas>=2.2.2; extra == "test"
+Requires-Dist: refnx>=0.1.45; extra == "test"
+Dynamic: license-file
+
+# Refjax - Lightweight `jit` and `grad` compatible reflectometry in jax
+
+This package provides lightweight, low-level access to neutron reflectometry and spectroscopic ellipsometry kernels in jax.
+
+The provided kernels are all `jit` and `grad` compatible and easy to integrate with your favourite optimisation or uncertainty quantification library for fast parameterisation of stratified reflectometry models. Several examples for popular libraries are available in the [`examples`](https://github.com/MDCHAMP/refjax/tree/main/examples) folder.
+
+
+## Installation
+
+From PyPI:
+
+```bash
+pip install refjax
+# or
+uv add refjax
+```
+
+Or from GitHub, pinned to the 0.1.1 release by its tag:
+
+```bash
+pip install "refjax @ git+https://github.com/MDCHAMP/refjax@v0.1.1"
+# or
+uv add git+https://github.com/MDCHAMP/refjax --tag v0.1.1
+```
+
+refjax needs Python >= 3.11, `jax` and `jaxlib` >= 0.4.26 and `numpy` >= 2.0. The test suite runs in CI at these minimum versions and at the latest releases.
+
+### 64-bit precision is on by default
+
+`import refjax` switches on jax's 64-bit mode (`jax.config.update("jax_enable_x64", True)`), so refjax computes in float64 and complex128. This setting is global: it applies to all jax code in the same Python process, not only to refjax. To work in 32-bit, switch it off after importing refjax, before creating any arrays: `jax.config.update("jax_enable_x64", False)`. refjax then computes in float32 and complex64; its tests run only in 64-bit.
+
+## Neutron reflectometry
+
+### Design pattern
+
+Reflectometry models are created by creating reflectivity kernels with `refjax.kernel`:
+
+```python
+kernel = refjax.kernel(model, smear=refjax.smear.no_smear, smear_params=None, jit=True)
+```
+
+Currently implemented models:
+- `refjax.model.abeles` Abeles matrix formalism for stratified media
+
+Currently implemented smearing kernels:
+- `refjax.smear.no_smear` No smearing kernel (default)
+- `refjax.smear.constant_gaussian` Constant gaussian smearing at every Q value (the third argument is the FWHM as a percentage of Q)
+- `refjax.smear.pointwise_gaussian` Pointwise Gaussian approximation to the true smearing kernel (requires FWHM resolution measurements), by 17th-order Gauss-Hermite quadrature
+
+Once a kernel has been created, it is called with the model parameters and Q (and FWHM values for smeared kernels).
+
+```python
+R = kernel(params, Q, FWHM) # fast because jit
+```
+
+### Parameters of `abeles`
+
+For a slab model with `n` layers between a fronting and a backing medium, `params` is a dict:
+
+| key     | shape     | meaning                                                                    |
+|---------|-----------|----------------------------------------------------------------------------|
+| `thick` | `(n,)`    | layer thicknesses, in the length unit of 1/Q (Å for Q in Å⁻¹)                |
+| `SLD`   | `(n + 2,)` | scattering length densities in 10⁻⁶ Å⁻², fronting first, backing last     |
+| `rough` | `(n + 1,)` | interfacial roughnesses, one per interface                                 |
+| `scale` | `()`      | multiplies the reflectivity                                                |
+| `bkg`   | `()`      | constant background added to the reflectivity                              |
+
+### Example
+
+```python
+import jax.numpy as jnp
+import refjax
+
+params = {
+    "thick": jnp.array([255.0, 17.0]),
+    "SLD": jnp.array([0.0, 6.4, 3.47, 2.07]),
+    "rough": jnp.array([5.2, 5.0, 2.07]),
+    "scale": 1.0,
+    "bkg": 2e-7,
+}
+Q = jnp.linspace(0.005, 0.3, 200)  # Å⁻¹
+FWHM = 0.05 * Q  # resolution FWHM at each Q
+kernel = refjax.kernel(refjax.model.abeles, refjax.smear.pointwise_gaussian)
+R = kernel(params, Q, FWHM)
+```
+
+[`examples/01_basic_usage.py`](https://github.com/MDCHAMP/refjax/blob/main/examples/01_basic_usage.py) walks through this model against the example data in [`examples/.example_NR_data.csv`](https://github.com/MDCHAMP/refjax/blob/main/examples/.example_NR_data.csv); the other scripts in [`examples`](https://github.com/MDCHAMP/refjax/tree/main/examples) show parameter identification with optax and numpyro.
+
+## Spectroscopic ellipsometry
+
+`refjax.model.spec_ellips(params, lam)` computes the ellipsometric angles Psi and Delta of a stratified medium (transfer-matrix method) at the wavelengths `lam`. It is `jit` and `grad` compatible and is called directly (not through `refjax.kernel`).
+
+For `n` layers between an ambient medium and a substrate, `params` is a dict:
+
+| key        | shape                | meaning                                                                         |
+|------------|----------------------|---------------------------------------------------------------------------------|
+| `nbar`     | `(n + 2, lam.shape)` | complex refractive index `n + ik` (`k >= 0`) of each medium, ambient first, substrate last |
+| `thick`    | `(n,)`               | layer thicknesses, in the same length unit as `lam`                             |
+| `aoi_rads` | `()`                 | angle of incidence in the ambient, in radians                                   |
+
+It returns `(psi, delta)` in degrees, each of shape `lam.shape`.
+
+Conventions:
+- `tan(psi) exp(i delta) = rp / rs` in Woollam's convention, as used by J.A. Woollam's instruments and CompleteEASE and set out in H. Fujiwara, *Spectroscopic Ellipsometry: Principles and Applications* (Wiley, 2007), [doi:10.1002/9780470060193](https://doi.org/10.1002/9780470060193): the complex index written n - ik, so refjax's n + ik enters conjugated. `delta` is wrapped to `[0, 360)`; compare it to data modulo 360.
+- The complex index is `n + ik` with `k >= 0`. This is the complex conjugate of Woollam's `n - ik`, hence the minus sign on the phase of `rp / rs` in the code. A Woollam-convention index `N = n - ik` is passed as `N.conj()`.
+- `n cos(theta)` in each medium follows from Snell's law with the principal square root, which gives `Im >= 0`: the decaying wave for `n + ik`.
+
+```python
+import jax.numpy as jnp
+import numpy as np
+import refjax
+
+lam = jnp.linspace(370.0, 1000.0, 50)  # nm
+# ambient, film, SiO2, Si substrate (n + ik), each constant in wavelength here
+nbar = jnp.array([1.0, 1.58, 1.46, 3.9 + 0.02j])[:, None] * jnp.ones_like(lam)
+params = {"nbar": nbar, "thick": jnp.array([17.0, 2.0]), "aoi_rads": np.radians(75.0)}
+psi, delta = refjax.model.spec_ellips(params, lam)
+```
+
+[`examples/06_spectroscopic_ellipsometry.py`](https://github.com/MDCHAMP/refjax/blob/main/examples/06_spectroscopic_ellipsometry.py) models and fits a thick polymer film on Si against a measured spectrum in [`examples/.example_ellips_data.csv`](https://github.com/MDCHAMP/refjax/blob/main/examples/.example_ellips_data.csv) (which carries the substrate's indices). [`test/test_ellips.py`](https://github.com/MDCHAMP/refjax/blob/main/test/test_ellips.py) checks `spec_ellips` against an independent numpy implementation and against that spectrum.
+
+### Changes in 0.1.1
+
+The first release on PyPI. Over the 0.1.0 tag: the smearing quadrature is computed in float64 (it was float32, an error of about 2e-8 relative in smeared reflectivity), importing refjax no longer starts jax's backend, `abeles` and `spec_ellips` have docstrings, and the packaging metadata is updated.
+
+### Changes in 0.1.0
+
+`spec_ellips` now follows the conventions above. `delta` is wrapped to `[0, 360)`; before 0.1.0 it was unwrapped along the wavelength axis, so it could differ from the 0.1.0 value by multiples of 360. `n cos(theta)` now comes from the principal square root rather than `arcsin` of Snell's law. Code written against the earlier `spec_ellips` should compare `delta` modulo 360.
+
+## Development
+
+```bash
+uv run --extra test pytest
+```
+
+## Citing
+
+If you use refjax in published work, please cite:
+
+> Max D. Champneys, Andrew J. Parnell and Stephanie L. Burg, "Towards real-time surrogate-free Bayesian inversion for neutron reflectometry", *Communications Materials* (2026). [doi:10.1038/s43246-026-01341-4](https://doi.org/10.1038/s43246-026-01341-4)
+
+```bibtex
+@article{champneys2026refjax,
+  author  = {Champneys, Max D. and Parnell, Andrew J. and Burg, Stephanie L.},
+  title   = {Towards real-time surrogate-free {B}ayesian inversion for neutron reflectometry},
+  journal = {Communications Materials},
+  year    = {2026},
+  doi     = {10.1038/s43246-026-01341-4},
+  url     = {https://doi.org/10.1038/s43246-026-01341-4}
+}
+```
+
+Please also state the refjax version you used (e.g. 0.1.1). Citation metadata is in [`CITATION.cff`](https://github.com/MDCHAMP/refjax/blob/main/CITATION.cff); GitHub's "Cite this repository" button reads it.
+
+Having trouble? Feature request? Raise an issue on the [GitHub issues page](https://github.com/MDCHAMP/refjax/issues).

refjax-0.1.1/README.md

@@ -0,0 +1,155 @@
+# Refjax - Lightweight `jit` and `grad` compatible reflectometry in jax
+
+This package provides lightweight, low-level access to neutron reflectometry and spectroscopic ellipsometry kernels in jax.
+
+The provided kernels are all `jit` and `grad` compatible and easy to integrate with your favourite optimisation or uncertainty quantification library for fast parameterisation of stratified reflectometry models. Several examples for popular libraries are available in the [`examples`](https://github.com/MDCHAMP/refjax/tree/main/examples) folder.
+
+
+## Installation
+
+From PyPI:
+
+```bash
+pip install refjax
+# or
+uv add refjax
+```
+
+Or from GitHub, pinned to the 0.1.1 release by its tag:
+
+```bash
+pip install "refjax @ git+https://github.com/MDCHAMP/refjax@v0.1.1"
+# or
+uv add git+https://github.com/MDCHAMP/refjax --tag v0.1.1
+```
+
+refjax needs Python >= 3.11, `jax` and `jaxlib` >= 0.4.26 and `numpy` >= 2.0. The test suite runs in CI at these minimum versions and at the latest releases.
+
+### 64-bit precision is on by default
+
+`import refjax` switches on jax's 64-bit mode (`jax.config.update("jax_enable_x64", True)`), so refjax computes in float64 and complex128. This setting is global: it applies to all jax code in the same Python process, not only to refjax. To work in 32-bit, switch it off after importing refjax, before creating any arrays: `jax.config.update("jax_enable_x64", False)`. refjax then computes in float32 and complex64; its tests run only in 64-bit.
+
+## Neutron reflectometry
+
+### Design pattern
+
+Reflectometry models are created by creating reflectivity kernels with `refjax.kernel`:
+
+```python
+kernel = refjax.kernel(model, smear=refjax.smear.no_smear, smear_params=None, jit=True)
+```
+
+Currently implemented models:
+- `refjax.model.abeles` Abeles matrix formalism for stratified media
+
+Currently implemented smearing kernels:
+- `refjax.smear.no_smear` No smearing kernel (default)
+- `refjax.smear.constant_gaussian` Constant gaussian smearing at every Q value (the third argument is the FWHM as a percentage of Q)
+- `refjax.smear.pointwise_gaussian` Pointwise Gaussian approximation to the true smearing kernel (requires FWHM resolution measurements), by 17th-order Gauss-Hermite quadrature
+
+Once a kernel has been created, it is called with the model parameters and Q (and FWHM values for smeared kernels).
+
+```python
+R = kernel(params, Q, FWHM) # fast because jit
+```
+
+### Parameters of `abeles`
+
+For a slab model with `n` layers between a fronting and a backing medium, `params` is a dict:
+
+| key     | shape     | meaning                                                                    |
+|---------|-----------|----------------------------------------------------------------------------|
+| `thick` | `(n,)`    | layer thicknesses, in the length unit of 1/Q (Å for Q in Å⁻¹)                |
+| `SLD`   | `(n + 2,)` | scattering length densities in 10⁻⁶ Å⁻², fronting first, backing last     |
+| `rough` | `(n + 1,)` | interfacial roughnesses, one per interface                                 |
+| `scale` | `()`      | multiplies the reflectivity                                                |
+| `bkg`   | `()`      | constant background added to the reflectivity                              |
+
+### Example
+
+```python
+import jax.numpy as jnp
+import refjax
+
+params = {
+    "thick": jnp.array([255.0, 17.0]),
+    "SLD": jnp.array([0.0, 6.4, 3.47, 2.07]),
+    "rough": jnp.array([5.2, 5.0, 2.07]),
+    "scale": 1.0,
+    "bkg": 2e-7,
+}
+Q = jnp.linspace(0.005, 0.3, 200)  # Å⁻¹
+FWHM = 0.05 * Q  # resolution FWHM at each Q
+kernel = refjax.kernel(refjax.model.abeles, refjax.smear.pointwise_gaussian)
+R = kernel(params, Q, FWHM)
+```
+
+[`examples/01_basic_usage.py`](https://github.com/MDCHAMP/refjax/blob/main/examples/01_basic_usage.py) walks through this model against the example data in [`examples/.example_NR_data.csv`](https://github.com/MDCHAMP/refjax/blob/main/examples/.example_NR_data.csv); the other scripts in [`examples`](https://github.com/MDCHAMP/refjax/tree/main/examples) show parameter identification with optax and numpyro.
+
+## Spectroscopic ellipsometry
+
+`refjax.model.spec_ellips(params, lam)` computes the ellipsometric angles Psi and Delta of a stratified medium (transfer-matrix method) at the wavelengths `lam`. It is `jit` and `grad` compatible and is called directly (not through `refjax.kernel`).
+
+For `n` layers between an ambient medium and a substrate, `params` is a dict:
+
+| key        | shape                | meaning                                                                         |
+|------------|----------------------|---------------------------------------------------------------------------------|
+| `nbar`     | `(n + 2, lam.shape)` | complex refractive index `n + ik` (`k >= 0`) of each medium, ambient first, substrate last |
+| `thick`    | `(n,)`               | layer thicknesses, in the same length unit as `lam`                             |
+| `aoi_rads` | `()`                 | angle of incidence in the ambient, in radians                                   |
+
+It returns `(psi, delta)` in degrees, each of shape `lam.shape`.
+
+Conventions:
+- `tan(psi) exp(i delta) = rp / rs` in Woollam's convention, as used by J.A. Woollam's instruments and CompleteEASE and set out in H. Fujiwara, *Spectroscopic Ellipsometry: Principles and Applications* (Wiley, 2007), [doi:10.1002/9780470060193](https://doi.org/10.1002/9780470060193): the complex index written n - ik, so refjax's n + ik enters conjugated. `delta` is wrapped to `[0, 360)`; compare it to data modulo 360.
+- The complex index is `n + ik` with `k >= 0`. This is the complex conjugate of Woollam's `n - ik`, hence the minus sign on the phase of `rp / rs` in the code. A Woollam-convention index `N = n - ik` is passed as `N.conj()`.
+- `n cos(theta)` in each medium follows from Snell's law with the principal square root, which gives `Im >= 0`: the decaying wave for `n + ik`.
+
+```python
+import jax.numpy as jnp
+import numpy as np
+import refjax
+
+lam = jnp.linspace(370.0, 1000.0, 50)  # nm
+# ambient, film, SiO2, Si substrate (n + ik), each constant in wavelength here
+nbar = jnp.array([1.0, 1.58, 1.46, 3.9 + 0.02j])[:, None] * jnp.ones_like(lam)
+params = {"nbar": nbar, "thick": jnp.array([17.0, 2.0]), "aoi_rads": np.radians(75.0)}
+psi, delta = refjax.model.spec_ellips(params, lam)
+```
+
+[`examples/06_spectroscopic_ellipsometry.py`](https://github.com/MDCHAMP/refjax/blob/main/examples/06_spectroscopic_ellipsometry.py) models and fits a thick polymer film on Si against a measured spectrum in [`examples/.example_ellips_data.csv`](https://github.com/MDCHAMP/refjax/blob/main/examples/.example_ellips_data.csv) (which carries the substrate's indices). [`test/test_ellips.py`](https://github.com/MDCHAMP/refjax/blob/main/test/test_ellips.py) checks `spec_ellips` against an independent numpy implementation and against that spectrum.
+
+### Changes in 0.1.1
+
+The first release on PyPI. Over the 0.1.0 tag: the smearing quadrature is computed in float64 (it was float32, an error of about 2e-8 relative in smeared reflectivity), importing refjax no longer starts jax's backend, `abeles` and `spec_ellips` have docstrings, and the packaging metadata is updated.
+
+### Changes in 0.1.0
+
+`spec_ellips` now follows the conventions above. `delta` is wrapped to `[0, 360)`; before 0.1.0 it was unwrapped along the wavelength axis, so it could differ from the 0.1.0 value by multiples of 360. `n cos(theta)` now comes from the principal square root rather than `arcsin` of Snell's law. Code written against the earlier `spec_ellips` should compare `delta` modulo 360.
+
+## Development
+
+```bash
+uv run --extra test pytest
+```
+
+## Citing
+
+If you use refjax in published work, please cite:
+
+> Max D. Champneys, Andrew J. Parnell and Stephanie L. Burg, "Towards real-time surrogate-free Bayesian inversion for neutron reflectometry", *Communications Materials* (2026). [doi:10.1038/s43246-026-01341-4](https://doi.org/10.1038/s43246-026-01341-4)
+
+```bibtex
+@article{champneys2026refjax,
+  author  = {Champneys, Max D. and Parnell, Andrew J. and Burg, Stephanie L.},
+  title   = {Towards real-time surrogate-free {B}ayesian inversion for neutron reflectometry},
+  journal = {Communications Materials},
+  year    = {2026},
+  doi     = {10.1038/s43246-026-01341-4},
+  url     = {https://doi.org/10.1038/s43246-026-01341-4}
+}
+```
+
+Please also state the refjax version you used (e.g. 0.1.1). Citation metadata is in [`CITATION.cff`](https://github.com/MDCHAMP/refjax/blob/main/CITATION.cff); GitHub's "Cite this repository" button reads it.
+
+Having trouble? Feature request? Raise an issue on the [GitHub issues page](https://github.com/MDCHAMP/refjax/issues).

refjax-0.1.1/pyproject.toml

@@ -0,0 +1,45 @@
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "refjax"
+dependencies = ["numpy>=2.0", "jax>=0.4.26", "jaxlib>=0.4.26"]
+requires-python = ">=3.11"
+authors = [
+    { name = "Max D. Champneys", email = "max.champneys@sheffield.ac.uk" },
+    { name = "Stephanie L. Burg" },
+]
+maintainers = [
+    { name = "Max Champneys", email = "max.champneys@sheffield.ac.uk" },
+]
+description = "Lightweight `jit` and `grad` compatible reflectometry in jax."
+license = "MIT"
+license-files = ["LICENSE"]
+keywords = ["Reflectometry", "Neutron reflectometry", "jax", "autograd"]
+classifiers = [
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "Topic :: Scientific/Engineering :: Physics",
+]
+dynamic = ["version", "readme"]
+
+[project.optional-dependencies]
+test = ["pytest>=7.0", "pandas>=2.2.2", "refnx>=0.1.45"]
+
+[project.urls]
+Repository = "https://github.com/MDCHAMP/refjax"
+Issues = "https://github.com/MDCHAMP/refjax/issues"
+Paper = "https://doi.org/10.1038/s43246-026-01341-4"
+
+
+[tool.setuptools.dynamic]
+readme = {file = ["README.md"], content-type = "text/markdown"}
+version = {attr = "refjax.__version__"}
+
+[tool.setuptools.packages.find]
+where = ["src"]

refjax-0.1.1/setup.cfg

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

refjax-0.1.1/src/refjax/__init__.py

@@ -0,0 +1,9 @@
+__version__ = "0.1.1"
+
+from jax import config
+from refjax import model, smear
+from refjax.tools import kernel
+
+__all__ = ["__version__", "kernel", "model", "smear"]
+
+config.update("jax_enable_x64", True)

refjax-0.1.1/src/refjax/model.py

@@ -0,0 +1,112 @@
+import jax.numpy as jnp
+from functools import reduce
+
+# %% Abeles Matrix formalism
+
+
+def abeles(params, Q):
+    """Reflectivity of a stratified medium by the Abeles matrix formalism.
+
+    Parameters
+    ----------
+    params : dict
+        Slab model with n layers between a fronting and a backing medium:
+
+        - "thick": (n,) layer thicknesses, in the length unit of 1/Q (Å for Q in Å⁻¹)
+        - "rough": (n + 1,) interfacial roughnesses, one per interface
+        - "SLD": (n + 2,) scattering length densities in 10⁻⁶ Å⁻², fronting first, backing last
+        - "scale": () multiplies the reflectivity
+        - "bkg": () constant background added to the reflectivity
+    Q : array, shape (m,)
+        Momentum transfer.
+
+    Returns
+    -------
+    array, shape (m,)
+        The reflectivity at each Q.
+    """
+    # SLD profile differences
+    d_rho = (params["SLD"] - params["SLD"][0]).astype(complex) * 1e-6
+    # compute wavevectors for the nth layer
+    kn = jnp.sqrt(((Q[:, None] / 2) ** 2) - (4 * jnp.pi * d_rho))
+    # Fresnel reflection
+    a, b = kn[:, :-1], kn[:, 1:]
+    r = (a - b) / (a + b) * jnp.exp(-2 * a * b * params["rough"] ** 2)
+    # Compute and multiply characteristic matricies elementwise (n_layers + 1, Q.shape)
+    C00 = jnp.exp(1j * a * jnp.array([0, *params["thick"]])).T
+    C11 = C00**-1
+    C10 = r.T * C00
+    C01 = r.T * C11
+
+    # For some reason manually performing the 2x2 matrix multiplications
+    # is a huge speedup (~2x) in jit-ed calls compared to matmul (???)
+    # Could be something to do with the very wide broadcast size?
+    # The batch dim size is often >> 1k and idk if this is not efficient
+    # Perhaps there is an einsum that works better? - there isn't 08/04/24
+    # Tested it and seemingly not. Bizzare behavior. Absolute rabbithole.
+    # Perhaps constructing the large array (n_layers+1, Q.shape, 2,2) is expensive somehow? - Nope 08/04/24.
+    # Hard to imagine why this is - something to do with optimisations in jaxprs at a guess.
+
+    M00, M10, M01, M11 = C00[0], C10[0], C01[0], C11[0]
+    for c00, c10, c01, c11 in zip(C00[1:], C10[1:], C01[1:], C11[1:]):
+        a = M00 * c00 + M10 * c01
+        b = M00 * c10 + M10 * c11
+        c = M01 * c00 + M11 * c01
+        d = M01 * c10 + M11 * c11
+        M00, M10, M01, M11 = a, b, c, d
+    R = M01 / M00
+    return params["scale"] * (R * R.conj()).real + params["bkg"]
+
+
+# %% Spectroscopic ellipsometry
+
+
+def spec_ellips(params, lam):
+    """Ellipsometric angles Psi and Delta of a stratified medium by the transfer-matrix method.
+
+    Parameters
+    ----------
+    params : dict
+        n layers between an ambient medium and a substrate:
+
+        - "nbar": (n + 2, lam.shape) complex refractive index n + ik (k >= 0) of each medium,
+          ambient first, substrate last
+        - "thick": (n,) layer thicknesses, in the same length unit as lam
+        - "aoi_rads": () angle of incidence in the ambient, in radians
+    lam : array
+        Wavelengths.
+
+    Returns
+    -------
+    psi, delta : arrays, each of shape lam.shape
+        In degrees: tan(psi) exp(i delta) = rp / rs in Woollam's convention, delta in [0, 360).
+        Compare delta to data modulo 360.
+    """
+    nbar = params["nbar"]
+    # n cos(theta) in each medium (Snell); principal sqrt gives Im >= 0, the decaying wave for n + ik
+    ncos = jnp.sqrt(nbar**2 - (nbar[0] * jnp.sin(params["aoi_rads"])) ** 2)
+    cos = ncos / nbar
+    ncos22 = ncos[1:]
+    ncos11 = ncos[:-1]
+    ncos21 = nbar[1:] * cos[:-1]
+    ncos12 = nbar[:-1] * cos[1:]
+    rs = (ncos11 - ncos22) / (ncos11 + ncos22)
+    ts = 2 * ncos11 / (ncos11 + ncos22)
+    rp = (ncos21 - ncos12) / (ncos12 + ncos21)
+    tp = 2 * ncos11 / (ncos12 + ncos21)
+    dz = 1j * params["thick"][:, None] * 2 * jnp.pi * ncos[1:-1] / lam
+    one = jnp.ones_like(rs)
+    edr = jnp.exp(dz)
+    nedr = jnp.exp(-dz)
+
+    def TMM(r, t):
+        ms = jnp.array([[nedr, nedr * r[1:]], [edr * r[1:], edr]]) / t[1:]
+        ms0 = jnp.array([[one[0], r[0]], [r[0], one[0]]]) / t[0]
+        Ms = reduce(jnp.matmul, ms.transpose(2, 3, 0, 1), ms0.T)
+        return Ms[:, 1, 0] / Ms[:, 0, 0]
+
+    rho = TMM(rp, tp) / TMM(rs, ts)
+    # n + ik is the conjugate of Woollam's n - ik, hence the minus sign on delta
+    _psi = jnp.degrees(jnp.arctan(jnp.abs(rho)))
+    _delta = jnp.degrees(-jnp.angle(rho)) % 360
+    return _psi, _delta

refjax-0.1.1/src/refjax/smear.py

@@ -0,0 +1,34 @@
+import jax
+import jax.numpy as jnp
+from numpy.polynomial.hermite import hermgauss
+
+_ORDER = 17
+herm_root, herm_weight = hermgauss(_ORDER)
+
+def no_smear(params, Q, FWHM, model, **kwargs):
+    return model(params, Q)
+
+def constant_gaussian(params, Q, const_smear, model, **kwargs):
+    '''
+    FWHM = const_smear% * Q
+    Gaussian kernel smearing.
+    There are a few options here can solve as a convolution or just do gauss-hermite? 
+    For now treat as pointwise.
+    '''
+    return pointwise_gaussian(params, Q, Q*const_smear/100, model, **kwargs)
+
+def pointwise_gaussian(params, Q, FWHMs, model, **kwargs):
+    '''
+    Pointwise gaussian smearing kernel. 
+    Compute the expectation integral using a 17th order Gauss-Hermite quadrature. 
+    Note that this differs from the approach in refnx which uses a finite width interval so numerics may differ.'''
+    Q_eval = 0.5 * FWHMs / jnp.sqrt(jnp.log(2.0)) * herm_root[:, None] + Q
+    R = jax.vmap(model, in_axes=(None, 0))(params, Q_eval)
+    return herm_weight @ R / jnp.sqrt(jnp.pi)
+
+def _custom_kernel(params, Q, kernel, model, **kwargs):
+    '''TODO
+    Manually compute the expectation interval of R(q) w.r.t k(q)
+    Can use simpsons or a Gauss-quad maybe?
+    '''
+    raise NotImplementedError

refjax-0.1.1/src/refjax/tools.py

@@ -0,0 +1,25 @@
+import jax
+from refjax.smear import no_smear
+
+
+def kernel(model, smear=no_smear, smear_params=None, jit=True):
+    """Build a reflection kernel from a reflection model and a smearing kernel"""
+    if smear_params is None:
+        smear_params = {}
+
+    def kern(params, q, dq=None):
+        return smear(params, q, dq, model, **smear_params)    
+    if jit:
+        kern = jax.jit(kern)
+    return kern
+
+
+def _from_refnx(refnx_model):
+    """TODO Pull out required parameters from a refnx model"""
+    raise NotImplementedError
+
+def _validata_params(params):
+    """TODO Validata params manually because don't want to do this in a jit"""
+    raise NotImplementedError
+
+

refjax-0.1.1/src/refjax.egg-info/PKG-INFO

@@ -0,0 +1,186 @@
+Metadata-Version: 2.4
+Name: refjax
+Version: 0.1.1
+Summary: Lightweight `jit` and `grad` compatible reflectometry in jax.
+Author: Stephanie L. Burg
+Author-email: "Max D. Champneys" <max.champneys@sheffield.ac.uk>
+Maintainer-email: Max Champneys <max.champneys@sheffield.ac.uk>
+License-Expression: MIT
+Project-URL: Repository, https://github.com/MDCHAMP/refjax
+Project-URL: Issues, https://github.com/MDCHAMP/refjax/issues
+Project-URL: Paper, https://doi.org/10.1038/s43246-026-01341-4
+Keywords: Reflectometry,Neutron reflectometry,jax,autograd
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=2.0
+Requires-Dist: jax>=0.4.26
+Requires-Dist: jaxlib>=0.4.26
+Provides-Extra: test
+Requires-Dist: pytest>=7.0; extra == "test"
+Requires-Dist: pandas>=2.2.2; extra == "test"
+Requires-Dist: refnx>=0.1.45; extra == "test"
+Dynamic: license-file
+
+# Refjax - Lightweight `jit` and `grad` compatible reflectometry in jax
+
+This package provides lightweight, low-level access to neutron reflectometry and spectroscopic ellipsometry kernels in jax.
+
+The provided kernels are all `jit` and `grad` compatible and easy to integrate with your favourite optimisation or uncertainty quantification library for fast parameterisation of stratified reflectometry models. Several examples for popular libraries are available in the [`examples`](https://github.com/MDCHAMP/refjax/tree/main/examples) folder.
+
+
+## Installation
+
+From PyPI:
+
+```bash
+pip install refjax
+# or
+uv add refjax
+```
+
+Or from GitHub, pinned to the 0.1.1 release by its tag:
+
+```bash
+pip install "refjax @ git+https://github.com/MDCHAMP/refjax@v0.1.1"
+# or
+uv add git+https://github.com/MDCHAMP/refjax --tag v0.1.1
+```
+
+refjax needs Python >= 3.11, `jax` and `jaxlib` >= 0.4.26 and `numpy` >= 2.0. The test suite runs in CI at these minimum versions and at the latest releases.
+
+### 64-bit precision is on by default
+
+`import refjax` switches on jax's 64-bit mode (`jax.config.update("jax_enable_x64", True)`), so refjax computes in float64 and complex128. This setting is global: it applies to all jax code in the same Python process, not only to refjax. To work in 32-bit, switch it off after importing refjax, before creating any arrays: `jax.config.update("jax_enable_x64", False)`. refjax then computes in float32 and complex64; its tests run only in 64-bit.
+
+## Neutron reflectometry
+
+### Design pattern
+
+Reflectometry models are created by creating reflectivity kernels with `refjax.kernel`:
+
+```python
+kernel = refjax.kernel(model, smear=refjax.smear.no_smear, smear_params=None, jit=True)
+```
+
+Currently implemented models:
+- `refjax.model.abeles` Abeles matrix formalism for stratified media
+
+Currently implemented smearing kernels:
+- `refjax.smear.no_smear` No smearing kernel (default)
+- `refjax.smear.constant_gaussian` Constant gaussian smearing at every Q value (the third argument is the FWHM as a percentage of Q)
+- `refjax.smear.pointwise_gaussian` Pointwise Gaussian approximation to the true smearing kernel (requires FWHM resolution measurements), by 17th-order Gauss-Hermite quadrature
+
+Once a kernel has been created, it is called with the model parameters and Q (and FWHM values for smeared kernels).
+
+```python
+R = kernel(params, Q, FWHM) # fast because jit
+```
+
+### Parameters of `abeles`
+
+For a slab model with `n` layers between a fronting and a backing medium, `params` is a dict:
+
+| key     | shape     | meaning                                                                    |
+|---------|-----------|----------------------------------------------------------------------------|
+| `thick` | `(n,)`    | layer thicknesses, in the length unit of 1/Q (Å for Q in Å⁻¹)                |
+| `SLD`   | `(n + 2,)` | scattering length densities in 10⁻⁶ Å⁻², fronting first, backing last     |
+| `rough` | `(n + 1,)` | interfacial roughnesses, one per interface                                 |
+| `scale` | `()`      | multiplies the reflectivity                                                |
+| `bkg`   | `()`      | constant background added to the reflectivity                              |
+
+### Example
+
+```python
+import jax.numpy as jnp
+import refjax
+
+params = {
+    "thick": jnp.array([255.0, 17.0]),
+    "SLD": jnp.array([0.0, 6.4, 3.47, 2.07]),
+    "rough": jnp.array([5.2, 5.0, 2.07]),
+    "scale": 1.0,
+    "bkg": 2e-7,
+}
+Q = jnp.linspace(0.005, 0.3, 200)  # Å⁻¹
+FWHM = 0.05 * Q  # resolution FWHM at each Q
+kernel = refjax.kernel(refjax.model.abeles, refjax.smear.pointwise_gaussian)
+R = kernel(params, Q, FWHM)
+```
+
+[`examples/01_basic_usage.py`](https://github.com/MDCHAMP/refjax/blob/main/examples/01_basic_usage.py) walks through this model against the example data in [`examples/.example_NR_data.csv`](https://github.com/MDCHAMP/refjax/blob/main/examples/.example_NR_data.csv); the other scripts in [`examples`](https://github.com/MDCHAMP/refjax/tree/main/examples) show parameter identification with optax and numpyro.
+
+## Spectroscopic ellipsometry
+
+`refjax.model.spec_ellips(params, lam)` computes the ellipsometric angles Psi and Delta of a stratified medium (transfer-matrix method) at the wavelengths `lam`. It is `jit` and `grad` compatible and is called directly (not through `refjax.kernel`).
+
+For `n` layers between an ambient medium and a substrate, `params` is a dict:
+
+| key        | shape                | meaning                                                                         |
+|------------|----------------------|---------------------------------------------------------------------------------|
+| `nbar`     | `(n + 2, lam.shape)` | complex refractive index `n + ik` (`k >= 0`) of each medium, ambient first, substrate last |
+| `thick`    | `(n,)`               | layer thicknesses, in the same length unit as `lam`                             |
+| `aoi_rads` | `()`                 | angle of incidence in the ambient, in radians                                   |
+
+It returns `(psi, delta)` in degrees, each of shape `lam.shape`.
+
+Conventions:
+- `tan(psi) exp(i delta) = rp / rs` in Woollam's convention, as used by J.A. Woollam's instruments and CompleteEASE and set out in H. Fujiwara, *Spectroscopic Ellipsometry: Principles and Applications* (Wiley, 2007), [doi:10.1002/9780470060193](https://doi.org/10.1002/9780470060193): the complex index written n - ik, so refjax's n + ik enters conjugated. `delta` is wrapped to `[0, 360)`; compare it to data modulo 360.
+- The complex index is `n + ik` with `k >= 0`. This is the complex conjugate of Woollam's `n - ik`, hence the minus sign on the phase of `rp / rs` in the code. A Woollam-convention index `N = n - ik` is passed as `N.conj()`.
+- `n cos(theta)` in each medium follows from Snell's law with the principal square root, which gives `Im >= 0`: the decaying wave for `n + ik`.
+
+```python
+import jax.numpy as jnp
+import numpy as np
+import refjax
+
+lam = jnp.linspace(370.0, 1000.0, 50)  # nm
+# ambient, film, SiO2, Si substrate (n + ik), each constant in wavelength here
+nbar = jnp.array([1.0, 1.58, 1.46, 3.9 + 0.02j])[:, None] * jnp.ones_like(lam)
+params = {"nbar": nbar, "thick": jnp.array([17.0, 2.0]), "aoi_rads": np.radians(75.0)}
+psi, delta = refjax.model.spec_ellips(params, lam)
+```
+
+[`examples/06_spectroscopic_ellipsometry.py`](https://github.com/MDCHAMP/refjax/blob/main/examples/06_spectroscopic_ellipsometry.py) models and fits a thick polymer film on Si against a measured spectrum in [`examples/.example_ellips_data.csv`](https://github.com/MDCHAMP/refjax/blob/main/examples/.example_ellips_data.csv) (which carries the substrate's indices). [`test/test_ellips.py`](https://github.com/MDCHAMP/refjax/blob/main/test/test_ellips.py) checks `spec_ellips` against an independent numpy implementation and against that spectrum.
+
+### Changes in 0.1.1
+
+The first release on PyPI. Over the 0.1.0 tag: the smearing quadrature is computed in float64 (it was float32, an error of about 2e-8 relative in smeared reflectivity), importing refjax no longer starts jax's backend, `abeles` and `spec_ellips` have docstrings, and the packaging metadata is updated.
+
+### Changes in 0.1.0
+
+`spec_ellips` now follows the conventions above. `delta` is wrapped to `[0, 360)`; before 0.1.0 it was unwrapped along the wavelength axis, so it could differ from the 0.1.0 value by multiples of 360. `n cos(theta)` now comes from the principal square root rather than `arcsin` of Snell's law. Code written against the earlier `spec_ellips` should compare `delta` modulo 360.
+
+## Development
+
+```bash
+uv run --extra test pytest
+```
+
+## Citing
+
+If you use refjax in published work, please cite:
+
+> Max D. Champneys, Andrew J. Parnell and Stephanie L. Burg, "Towards real-time surrogate-free Bayesian inversion for neutron reflectometry", *Communications Materials* (2026). [doi:10.1038/s43246-026-01341-4](https://doi.org/10.1038/s43246-026-01341-4)
+
+```bibtex
+@article{champneys2026refjax,
+  author  = {Champneys, Max D. and Parnell, Andrew J. and Burg, Stephanie L.},
+  title   = {Towards real-time surrogate-free {B}ayesian inversion for neutron reflectometry},
+  journal = {Communications Materials},
+  year    = {2026},
+  doi     = {10.1038/s43246-026-01341-4},
+  url     = {https://doi.org/10.1038/s43246-026-01341-4}
+}
+```
+
+Please also state the refjax version you used (e.g. 0.1.1). Citation metadata is in [`CITATION.cff`](https://github.com/MDCHAMP/refjax/blob/main/CITATION.cff); GitHub's "Cite this repository" button reads it.
+
+Having trouble? Feature request? Raise an issue on the [GitHub issues page](https://github.com/MDCHAMP/refjax/issues).

refjax-0.1.1/src/refjax.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+LICENSE
+README.md
+pyproject.toml
+src/refjax/__init__.py
+src/refjax/model.py
+src/refjax/smear.py
+src/refjax/tools.py
+src/refjax.egg-info/PKG-INFO
+src/refjax.egg-info/SOURCES.txt
+src/refjax.egg-info/dependency_links.txt
+src/refjax.egg-info/requires.txt
+src/refjax.egg-info/top_level.txt
+test/test_ellips.py
+test/test_refjax.py

refjax-0.1.1/src/refjax.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

refjax-0.1.1/src/refjax.egg-info/requires.txt

@@ -0,0 +1,8 @@
+numpy>=2.0
+jax>=0.4.26
+jaxlib>=0.4.26
+
+[test]
+pytest>=7.0
+pandas>=2.2.2
+refnx>=0.1.45

refjax-0.1.1/src/refjax.egg-info/top_level.txt

@@ -0,0 +1 @@
+refjax

refjax-0.1.1/test/test_ellips.py

@@ -0,0 +1,90 @@
+from pathlib import Path
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pandas as pd
+import pytest
+import refjax as rj
+
+
+def reference(N, d, lam, aoi):
+    # Independent numpy Psi, Delta: Woollam/Fujiwara convention (N = n - ik), Airy recursion
+    ncos = np.sqrt(N**2 - (N[0] * np.sin(aoi)) ** 2)
+    ncos = np.where(ncos.imag > 0, -ncos, ncos)
+    cos = ncos / N
+    out = []
+    for pol in "sp":
+        def fres(i, j):
+            if pol == "s":
+                return (ncos[i] - ncos[j]) / (ncos[i] + ncos[j])
+            return (N[j] * cos[i] - N[i] * cos[j]) / (N[j] * cos[i] + N[i] * cos[j])
+
+        r = fres(-2, -1)
+        for j in range(len(d) - 1, -1, -1):
+            e = np.exp(-4j * np.pi * d[j] * ncos[j + 1] / lam)
+            rij = fres(j, j + 1)
+            r = (rij + r * e) / (1 + rij * r * e)
+        out.append(r)
+    rho = out[1] / out[0]
+    return np.degrees(np.arctan(np.abs(rho))), np.degrees(np.angle(rho)) % 360
+
+
+lam = np.linspace(370, 1000, 50)
+aoi = np.radians(75.1)
+nSi = 3.9 - 0.02j
+
+cases = {
+    "bare Si": ([1.0, nSi], []),
+    "SiO2/Si": ([1.0, 1.46, nSi], [2.0]),
+    "PS/SiO2/Si": ([1.0, 1.58, 1.46, nSi], [17.0, 2.0]),
+    "thick PS/SiO2/Si": ([1.0, 1.58, 1.46, nSi], [600.0, 2.0]),
+    "absorbing film": ([1.0, 2.0 - 0.3j, nSi], [30.0]),
+}
+
+
+def params(n, d):
+    N = np.array(n, complex)[:, None] * np.ones_like(lam)
+    # refjax takes n + ik
+    return {"nbar": jnp.array(N.conj()), "thick": jnp.array(d, float), "aoi_rads": aoi}, N
+
+
+@pytest.mark.parametrize("name", cases)
+def test_spec_ellips(name):
+    p, N = params(*cases[name])
+    psi_ref, delta_ref = reference(N, np.array(cases[name][1], float), lam, aoi)
+    psi, delta = rj.model.spec_ellips(p, jnp.array(lam))
+    assert np.allclose(psi, psi_ref, atol=1e-9)
+    assert np.allclose(delta, delta_ref, atol=1e-9)
+    assert np.all((delta >= 0) & (delta < 360))
+
+
+def test_ellips_can_jit_and_grad():
+    p, _ = params(*cases["PS/SiO2/Si"])
+
+    def f(thick):
+        return jnp.sum(rj.model.spec_ellips({**p, "thick": thick}, jnp.array(lam))[0])
+
+    g = jax.jit(jax.grad(f))(p["thick"])
+    assert np.all(np.isfinite(g))
+
+
+# %% Measured data: one spectrum of a thick polymer film on Si (native oxide), 450-1000 nm, with the substrate's indices
+# at the measurement temperature (Franta's Si, Malitson's SiO2). The parameters are a fit of a Cauchy film to it,
+# whose reduced chi2 against the stated errors was 12.860.
+
+data = pd.read_csv(Path(__file__).parent.parent / "examples" / ".example_ellips_data.csv")
+film = {"A": 1.5592380152192518, "B": 0.009669302459697682, "thick_nm": 269.12932822155117}
+oxide_nm = 1.7976495299127038
+aoi_deg = 75.11362116  
+
+
+def test_spec_ellips_vs_data():
+    lam = jnp.array(data.wavelength_nm)
+    n_film = film["A"] + film["B"] / (lam / 1000) ** 2
+    nbar = jnp.stack([jnp.ones_like(lam), n_film, jnp.array(data.n_sio2), jnp.array(data.n_si) + 1j * jnp.array(data.k_si)])
+    p = {"nbar": nbar.astype(complex), "thick": jnp.array([film["thick_nm"], oxide_nm]), "aoi_rads": jnp.radians(aoi_deg)}
+    psi, delta = rj.model.spec_ellips(p, lam)
+    r = np.concatenate([(psi - data.psi) / data.dpsi, ((delta - data.delta + 180) % 360 - 180) / data.ddelta])
+    chi2_red = np.sum(r**2) / (r.size - 3)
+    assert abs(chi2_red - 12.860) < 1e-2  # the fit's own value

refjax-0.1.1/test/test_refjax.py

@@ -0,0 +1,70 @@
+from pathlib import Path
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pandas as pd
+import refjax as rj
+from refnx.reflect.reflect_model import reflectivity
+
+
+# %% Fixture some data
+
+#data
+data = pd.read_csv(Path(__file__).parent.parent / "examples" / ".example_NR_data.csv")
+
+#Parameters
+struct = jnp.array(
+    [
+        [0.0, 0.0, 0.0, 0.0],
+        [254.9, 6.37, 0.0, 6.8],
+        [17.0, 3.47, 0.0, 5.2],
+        [0.0, 2.07, 0.0, 2.07],
+    ]
+)  # blame refnx for this
+
+scale = 1.0
+bkg = 2e-7
+
+
+F_refnx = reflectivity(np.array(data.Q), np.array(struct), scale,bkg, np.array(data.FWHM))
+
+params = {
+    "thick": jnp.array([254.9, 17.0]),
+    "SLD": jnp.array([0.0, 6.37, 3.47, 2.07]),
+    "rough": jnp.array([6.8, 5.2, 2.07]),
+    "scale": scale,
+    "bkg": bkg,
+}
+
+
+# %%
+
+def test_fit():
+    kernel = rj.kernel(rj.model.abeles, rj.smear.pointwise_gaussian, jit=0)
+    F = kernel(params, jnp.array(data.Q), jnp.array(data.FWHM))
+    
+    # Vs refnx
+    # tolerance because the smearing quadrature methods are not identical
+    assert ((F-F_refnx)**2).sum()/F.shape[0] < 1e-6 
+    
+    # Vs data
+    assert ((F-data.R)**2).sum()/F.shape[0] < 1e-1 # low MSE
+
+
+def test_can_jit():
+    kernel = rj.kernel(rj.model.abeles, rj.smear.pointwise_gaussian)
+    kernel(params, jnp.array(data.Q), jnp.array(data.FWHM))
+    
+    
+def test_can_grad():
+    kernel = jax.jit(rj.kernel(rj.model.abeles, rj.smear.pointwise_gaussian))
+    dk_dtheta = jax.jacfwd(kernel)
+    dk_dtheta(params, jnp.array(data.Q), jnp.array(data.FWHM))
+
+
+def test_quadrature_nodes_float64():
+    # made at import, before refjax switches on x64: they must not be float32 jax arrays
+    assert rj.smear.herm_root.dtype == np.float64
+    assert rj.smear.herm_weight.dtype == np.float64
+