refjax 0.1.1__py3-none-any.whl

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/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/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/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.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,9 @@
+refjax/__init__.py,sha256=UlkcAIZuqzpcWqixVnKhdeuJzpsleSbCZKjdvFfLSqY,204
+refjax/model.py,sha256=UyKwaDqt6orONRcpcjcZVSCzuVKcRhE5aLrp3-jluxw,4372
+refjax/smear.py,sha256=pvnlmm9Mmy4TbfJYfegguSgjkqdzuIbGCqiy-h4pPyk,1254
+refjax/tools.py,sha256=sBCyHxieRQ-nb65nD8wJVaII5YyVmsJwNNe7M9zI6Go,670
+refjax-0.1.1.dist-info/licenses/LICENSE,sha256=yQu-_ub0ZmzUBIAXpzP79YqemMrHyItf4lJJ_pktDzI,1069
+refjax-0.1.1.dist-info/METADATA,sha256=F_nfJRIHUaOwy5gjVWzPjUH_cRyUELgjLen8ofdgjpc,9931
+refjax-0.1.1.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+refjax-0.1.1.dist-info/top_level.txt,sha256=0TA82h8w_ZeRvpC8adSyM5u-rxyNtPBxOf8CrjBt4oE,7
+refjax-0.1.1.dist-info/RECORD,,

refjax-0.1.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

refjax-0.1.1.dist-info/licenses/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.dist-info/top_level.txt

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