reflectosim 0.1.0__py3-none-any.whl

reflectosim/__init__.py

@@ -0,0 +1,32 @@
+"""X-ray reflectivity forward physics and realistic synthetic data for machine learning."""
+
+from importlib.metadata import version
+
+from reflectosim import materials
+from reflectosim.forward import Instrument, forward
+from reflectosim.priors import InstrumentPrior, OxidePrior, StackPrior, SubstratePrior
+from reflectosim.simulate import Batch, simulate
+from reflectosim.structure import Layer, Stack
+from reflectosim.units import CU_KALPHA1, CU_KALPHA2, fwhm_to_sigma, q_to_tth, tth_to_q
+
+__version__ = version("reflectosim")
+
+__all__ = [
+    "CU_KALPHA1",
+    "CU_KALPHA2",
+    "Batch",
+    "Instrument",
+    "InstrumentPrior",
+    "Layer",
+    "OxidePrior",
+    "Stack",
+    "StackPrior",
+    "SubstratePrior",
+    "__version__",
+    "forward",
+    "fwhm_to_sigma",
+    "materials",
+    "q_to_tth",
+    "simulate",
+    "tth_to_q",
+]

reflectosim/cli.py

@@ -0,0 +1,76 @@
+"""Command line: reflectosim sim | describe."""
+
+import argparse
+import json
+import sys
+
+import numpy as np
+
+from reflectosim import __version__, dataset, priors
+from reflectosim.grid import q_grid
+from reflectosim.simulate import simulate
+
+
+def _seed(text: str):
+    parts = [int(p) for p in text.split(",")]
+    return parts[0] if len(parts) == 1 else tuple(parts)
+
+
+def _sim(args) -> int:
+    stack, inst = priors.load(args.prior)
+    q = q_grid(args.grid, inst.wavelength)
+    batch = simulate(stack, inst, q, n=args.n, seed=_seed(args.seed))
+    out = dataset.write(args.out, batch, stack, inst)
+    print(f"wrote {out}: {args.n} curves x {q.size} points, oxide {stack.oxide}, "
+          f"resolution nodes {batch.nodes or 'per point'}")  # fmt: skip
+    return 0
+
+
+def _describe(args) -> int:
+    arrays, params, meta = dataset.read(args.dataset)
+    n, nq = arrays["counts"].shape
+    print(
+        f"{args.dataset}: {n} curves x {nq} points, q {arrays['q'][0]:.4g}-{arrays['q'][-1]:.4g} 1/A"
+    )
+    width = max(len(k) for k in params)
+    for k, v in params.items():
+        ok = np.isfinite(v)
+        if ok.any():
+            print(f"  {k:<{width}}  {v[ok].min():.4g} .. {v[ok].max():.4g}  ({ok.sum()} rows)")
+    print(
+        json.dumps(
+            {
+                k: meta[k]
+                for k in ("reflectosim", "reflectofit", "seed", "resolution_nodes")
+                if k in meta
+            }
+        )
+    )
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    p = argparse.ArgumentParser(prog="reflectosim", description="synthetic X-ray reflectivity data")
+    p.add_argument("--version", action="version", version=f"reflectosim {__version__}")
+    sub = p.add_subparsers(dest="command", required=True)
+
+    s = sub.add_parser("sim", help="simulate a dataset from a prior TOML")
+    s.add_argument(
+        "--prior", required=True, help="prior TOML: [stack], [oxide], [substrate], [instrument]"
+    )
+    s.add_argument("--n", type=int, required=True, help="number of curves")
+    s.add_argument("--grid", default="tth:0.05:8:500", help="tth:START:STOP:N or q:START:STOP:N")
+    s.add_argument("--seed", default="0", help="an int or comma-separated ints, e.g. 20261024,1")
+    s.add_argument("-o", "--out", required=True, help="output .h5 (needs h5py) or .npz")
+    s.set_defaults(run=_sim)
+
+    d = sub.add_parser("describe", help="summarise a dataset file")
+    d.add_argument("dataset")
+    d.set_defaults(run=_describe)
+
+    args = p.parse_args(argv)
+    return args.run(args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

reflectosim/dataset.py

@@ -0,0 +1,84 @@
+"""Dataset files: HDF5 (with the `hdf5` extra) or compressed npz, chosen by the file suffix.
+
+Arrays: `q` (Q,), `counts` and `reflectivity` (n, Q), `counts_scale` (n,), and one (n,) array per
+parameter under `params/<name>` (HDF5) or `params.<name>` (npz). The metadata is one JSON string with
+the priors, the seed, the resolution node count, the package versions and the conventions, enough
+to regenerate the file.
+"""
+
+import json
+from importlib.metadata import version
+from pathlib import Path
+
+import numpy as np
+
+from reflectosim import priors
+from reflectosim.simulate import Batch
+
+CONVENTIONS = (
+    "ambient first, roughness of the interface above, Im SLD > 0 "
+    "absorbs, SLD 1e-6 A^-2, q A^-1, angles in degrees; reflectivity is I0 * smeared R + background "
+    "before noise; counts are Poisson(counts_scale * reflectivity)"
+)
+
+
+def metadata(batch: Batch, stack: priors.StackPrior, instrument: priors.InstrumentPrior) -> dict:
+    seed = list(batch.seed) if isinstance(batch.seed, tuple | list) else batch.seed
+    return {
+        "reflectosim": version("reflectosim"),
+        "stack_prior": priors.to_dict(stack),
+        "instrument_prior": priors.to_dict(instrument),
+        "seed": seed,
+        "n": int(batch.counts.shape[0]),
+        "resolution_nodes": batch.nodes,
+        "backend": "reflectosim numba float64",
+        "conventions": CONVENTIONS,
+    }
+
+
+def write(
+    path: str | Path, batch: Batch, stack: priors.StackPrior, instrument: priors.InstrumentPrior
+) -> Path:
+    path = Path(path)
+    meta = json.dumps(metadata(batch, stack, instrument))
+    arrays = {
+        "q": batch.q,
+        "counts": batch.counts.astype(np.float32),
+        "reflectivity": batch.reflectivity.astype(np.float32),
+        "counts_scale": batch.counts_scale,
+    }
+    if path.suffix in (".h5", ".hdf5"):
+        import h5py
+
+        with h5py.File(path, "w") as f:
+            for k, v in arrays.items():
+                f.create_dataset(k, data=v, compression="gzip" if v.ndim == 2 else None)
+            g = f.create_group("params")
+            for k, v in batch.params.items():
+                g.create_dataset(k, data=v)
+            f.attrs["metadata"] = meta
+    elif path.suffix == ".npz":
+        np.savez_compressed(
+            path, **arrays, **{f"params.{k}": v for k, v in batch.params.items()}, metadata=meta
+        )
+    else:
+        raise ValueError(f"unknown dataset suffix {path.suffix!r}; use .h5 or .npz")
+    return path
+
+
+def read(path: str | Path) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray], dict]:
+    """(arrays, params, metadata) of a dataset file."""
+    path = Path(path)
+    if path.suffix in (".h5", ".hdf5"):
+        import h5py
+
+        with h5py.File(path, "r") as f:
+            arrays = {k: f[k][()] for k in ("q", "counts", "reflectivity", "counts_scale")}
+            params = {k: f["params"][k][()] for k in f["params"]}
+            meta = json.loads(f.attrs["metadata"])
+        return arrays, params, meta
+    with np.load(path) as z:
+        arrays = {k: z[k] for k in ("q", "counts", "reflectivity", "counts_scale")}
+        params = {k.removeprefix("params."): z[k] for k in z.files if k.startswith("params.")}
+        meta = json.loads(str(z["metadata"]))
+    return arrays, params, meta

reflectosim/forward.py

@@ -0,0 +1,305 @@
+"""Observed reflectivity: the layer kernel plus the instrument.
+
+The nominal 2θ plus the alignment offset is the true angle; the resolution is a Gaussian in q around
+the q of that angle. At each quadrature node: R for Kα₁ (and Kα₂ at the same angle, mixed by intensity
+ratio) times the footprint factor at that angle. The weighted sum over nodes is scaled by I0 and the
+background is added (DESIGN.md section 2). Every instrument field may be a scalar or a (B,) array, one value per curve.
+"""
+
+import warnings
+from dataclasses import dataclass, fields
+from typing import Any
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+from scipy.special import erf
+
+from reflectosim import resolution
+from reflectosim.kernel import parratt, smear
+from reflectosim.structure import as_batch
+from reflectosim.units import CU_KALPHA1, CU_KALPHA2
+
+MAX_TTH = 179.0
+
+
+@dataclass(frozen=True)
+class Instrument:
+    """Instrument model. Angles in degrees, lengths in mm, resolution as Gaussian σ in q."""
+
+    wavelength: float = CU_KALPHA1
+    i0: Any = 1.0
+    background: Any = 0.0
+    tth_offset: Any = 0.0
+    sigma_q: Any = 0.0
+    """Constant resolution term, Gaussian σ [Å⁻¹]."""
+    dq_over_q: Any = 0.0
+    """Relative resolution term: σ = dq_over_q · q. Combined in quadrature with sigma_q."""
+    kalpha2_ratio: Any = 0.0
+    """Kα₂/Kα₁ intensity ratio at the detector; 0 for a monochromatic beam."""
+    kalpha2_wavelength: float = CU_KALPHA2
+    sample_length: Any = None
+    """Sample length along the beam [mm]; with beam_sigma enables the footprint factor."""
+    beam_sigma: Any = None
+    """Gaussian σ of the beam profile at the sample [mm] (FWHM / 2.3548), beam centred."""
+    nodes: int | None = None
+    """None: each point gets the node count its resolution width, the fringes and any nearby
+    critical edge need (`reflectosim.resolution`). An int: that many nodes for every point, faster
+    and of the accuracy the count allows; also the way to vmap over resolution parameters."""
+
+    def replace(self, **changes) -> "Instrument":
+        values = {f.name: getattr(self, f.name) for f in fields(self)}
+        return Instrument(**(values | changes))
+
+
+class _NumpyOps:
+    asarray = staticmethod(np.asarray)
+    sqrt = staticmethod(np.sqrt)
+    sin = staticmethod(np.sin)
+    arcsin = staticmethod(np.arcsin)
+    clip = staticmethod(np.clip)
+    erf = staticmethod(erf)
+    exp = staticmethod(np.exp)
+    cat = staticmethod(np.concatenate)
+
+    @staticmethod
+    def column(x, b):
+        return np.broadcast_to(np.asarray(x, dtype=float), (b,)).reshape(b, 1, 1)
+
+    @staticmethod
+    def take_rows(values, idx):
+        """values (B, Q) gathered at idx (B, M)."""
+        return np.take_along_axis(values, idx, axis=1)
+
+    @staticmethod
+    def sum_rows(values, idx, width):
+        """(B, width) sums of values (B, M) by idx (B, M)."""
+        b = values.shape[0]
+        flat = (idx + width * np.arange(b)[:, None]).ravel()
+        return np.bincount(flat, values.ravel(), minlength=b * width).reshape(b, width)
+
+
+def _numpy(x) -> np.ndarray:
+    """Values of an array or tensor as numpy, without its graph (also under torch.func.jacfwd)."""
+    if not hasattr(x, "detach"):
+        return np.asarray(x)
+    x = x.detach()
+    try:
+        return x.cpu().numpy()
+    except RuntimeError:  # a functorch-wrapped tensor has no storage; its values are still readable
+        return np.array(x.tolist())
+
+
+def _active(x) -> bool:
+    """Whether an instrument term is switched on. A tensor counts as on: its value may be traced
+    (torch.func.jacfwd, vmap), so no branch may depend on it."""
+    if x is None:
+        return False
+    if hasattr(x, "detach"):
+        return True
+    return bool(np.any(np.asarray(x) != 0))
+
+
+ROWS_PER_BLOCK = 256
+"""Curves per block of the resolution integral (bounds the flat node arrays)."""
+NODES_PER_CALL = 2_000_000
+"""Padded resolution nodes per kernel call (bounds the memory of a block)."""
+
+
+def _at_angles(ops, kernel, x, packed, inst: Instrument):
+    """R at nominal-true q values x (B, M): Kα₁ (and Kα₂ at the same angle) times the footprint."""
+    b = x.shape[0]
+    wl1 = inst.wavelength
+    # a point below q = 0 is the mirror image of a grazing angle; keep it at grazing incidence
+    sin_t = ops.clip(x * (wl1 / (4.0 * np.pi)), 1e-12, 1.0)
+    r = kernel((4.0 * np.pi / wl1) * sin_t, *packed)
+    if _active(inst.kalpha2_ratio):
+        a = ops.column(inst.kalpha2_ratio, b)[..., 0]
+        r2 = kernel((4.0 * np.pi / inst.kalpha2_wavelength) * sin_t, *packed)
+        r = (r + a * r2) / (1.0 + a)
+    if inst.sample_length is not None and inst.beam_sigma is not None:
+        length = ops.column(inst.sample_length, b)[..., 0]
+        beam = ops.column(inst.beam_sigma, b)[..., 0]
+        r = r * ops.erf(length * sin_t / (2.0 * np.sqrt(2.0) * beam))
+    return r
+
+
+def _point_smear(ops, kernel, q_true, sigma, packed, inst: Instrument):
+    """Resolution with a node count per point (`reflectosim.resolution`) for one block of curves:
+    all nodes of a curve in one flat list, one kernel call, summed back per point."""
+    b, nq = q_true.shape[0], q_true.shape[1]
+    qt, sg = _numpy(q_true)[..., 0], _numpy(sigma)[..., 0]
+    ka2 = _active(inst.kalpha2_ratio)
+    counts, u_all, w_all, err, capped = resolution.layout(
+        qt, sg, tuple(_numpy(a) for a in packed), inst.wavelength,
+        inst.kalpha2_wavelength if ka2 else None,
+    )  # fmt: skip
+    loose = capped | (err > resolution.REFINE_TOL)
+    if loose.any():
+        warnings.warn(
+            f"resolution: {int(loose.sum())} points above the refinement tolerance "
+            f"({int(capped.sum())} stopped at MAX_DEPTH); estimated relative error up to "
+            f"{err[loose].max():.1e} (resolution_report gives it per point)",
+            RuntimeWarning,
+            stacklevel=4,
+        )
+    if ops is _NumpyOps:
+        # one fused numba loop over the points and their nodes (no padding, no temporaries)
+        footprint = inst.sample_length is not None and inst.beam_sigma is not None
+        nan = np.full(b, np.nan)
+        return smear(
+            qt,
+            sg,
+            counts,
+            u_all,
+            w_all,
+            *packed,
+            inst.wavelength,
+            inst.kalpha2_wavelength if ka2 else 0.0,
+            ops.column(inst.kalpha2_ratio, b).ravel(),
+            ops.column(inst.sample_length, b).ravel() if footprint else nan,
+            ops.column(inst.beam_sigma, b).ravel() if footprint else nan,
+        )
+    # per curve: owner point of each node, its u and weight
+    ends = np.cumsum(counts.reshape(b, nq).sum(axis=1))
+    layouts = [
+        (np.repeat(np.arange(nq), c), x, y)
+        for c, x, y in zip(
+            counts.reshape(b, nq),
+            np.split(u_all, ends[:-1]),
+            np.split(w_all, ends[:-1]),
+            strict=True,
+        )
+    ]
+    sizes = [o.size for o, _, _ in layouts]
+    out = []
+    start = 0
+    while start < b:
+        # curves in one kernel call while the padded node array stays within NODES_PER_CALL
+        stop, m = start + 1, sizes[start]
+        while stop < b and max(m, sizes[stop]) * (stop + 1 - start) <= NODES_PER_CALL:
+            m = max(m, sizes[stop])
+            stop += 1
+        rows = slice(start, stop)
+        part = layouts[rows]
+        # padding nodes sit on point 0 with weight 0, so a curve's result does not depend on its batch
+        owner = np.stack([np.pad(o, (0, m - o.size)) for o, _, _ in part])
+        u = np.stack([np.pad(x, (0, m - x.size)) for _, x, _ in part])
+        w = np.stack([np.pad(x, (0, m - x.size)) for _, _, x in part])
+        nodes = ops.take_rows(q_true[rows, :, 0], owner) + ops.take_rows(
+            sigma[rows, :, 0], owner
+        ) * ops.asarray(u)
+        r = _at_angles(ops, kernel, nodes, tuple(a[rows] for a in packed), _rows(inst, rows, b))
+        out.append(ops.sum_rows(r * ops.asarray(w), owner, nq))
+        start = stop
+    return ops.cat(out) if len(out) > 1 else out[0]
+
+
+def _observed(ops, kernel, q, b, packed, inst: Instrument):
+    """Shared pipeline. q (Q,) nominal; returns (B, Q).
+
+    The offset comes first: the true angle is the nominal 2θ plus the offset, and the resolution
+    is a Gaussian in q around the q of that true angle. Kα₁, Kα₂ and the footprint are evaluated at
+    the angle of each point of the resolution integral.
+    """
+    q = q.reshape(1, -1, 1)
+    wl1 = inst.wavelength
+    to_deg = 180.0 / np.pi
+    theta = ops.arcsin(ops.clip(q * (wl1 / (4.0 * np.pi)), 0.0, 1.0)) * to_deg
+    theta = ops.clip(theta + 0.5 * ops.column(inst.tth_offset, b), 1e-9, 0.5 * MAX_TTH)
+    q_true = (4.0 * np.pi / wl1) * ops.sin(theta / to_deg)  # (B, Q, 1)
+    resolved = _active(inst.sigma_q) or _active(inst.dq_over_q)
+    if not resolved or inst.nodes == 1:
+        r = _at_angles(ops, kernel, q_true[..., 0], packed, inst)
+    else:
+        sigma2 = ops.column(inst.sigma_q, b) ** 2 + (ops.column(inst.dq_over_q, b) * q_true) ** 2
+        # the floor keeps d sqrt / d sigma finite (zero) for a curve without resolution
+        sigma = ops.sqrt(ops.clip(sigma2, 1e-300, None))
+        if inst.nodes is None:
+            blocks = []
+            for start in range(0, b, ROWS_PER_BLOCK):
+                rows = slice(start, min(start + ROWS_PER_BLOCK, b))
+                part = tuple(a[rows] for a in packed)
+                blocks.append(
+                    _point_smear(ops, kernel, q_true[rows], sigma[rows], part, _rows(inst, rows, b))
+                )
+            r = ops.cat(blocks) if len(blocks) > 1 else blocks[0]
+        else:
+            u = np.linspace(-resolution.SPAN, resolution.SPAN, inst.nodes)
+            w = np.exp(-0.5 * u**2)
+            nodes = q_true + sigma * ops.asarray(u).reshape(1, 1, -1)
+            shape = nodes.shape
+            r = _at_angles(ops, kernel, nodes.reshape(b, -1), packed, inst).reshape(shape)
+            r = (r * ops.asarray(w / w.sum()).reshape(1, 1, -1)).sum(-1)
+    return ops.column(inst.i0, b)[..., 0] * r + ops.column(inst.background, b)[..., 0]
+
+
+def _rows(inst: Instrument, rows: slice, total: int) -> Instrument:
+    """The instrument for a block of curves (per-curve fields sliced)."""
+    changes = {}
+    for f in fields(inst):
+        v = getattr(inst, f.name)
+        per_curve = v is not None and np.ndim(v) == 1 and len(v) == total
+        if f.name not in ("wavelength", "kalpha2_wavelength", "nodes") and per_curve:
+            changes[f.name] = v[rows]
+    return inst.replace(**changes) if changes else inst
+
+
+def forward(
+    q: ArrayLike, model, instrument: Instrument | None = None, backend: str = "numba"
+) -> NDArray[np.float64]:
+    """Observed reflectivity for one model (returns (Q,)) or a batch (returns (B, Q)).
+
+    `model` is a Stack, a packed (sld, thickness, roughness) tuple of (N,) or (B, N) arrays, or a
+    sequence of either. Without an instrument this is the bare |r|² of the stack.
+    """
+    if backend == "torch":
+        from reflectosim.torch_backend import forward as torch_forward
+
+        return torch_forward(q, model, instrument)
+    if backend != "numba":
+        raise ValueError(f"unknown backend {backend!r}; use 'numba' or 'torch'")
+    packed, single = as_batch(model)
+    q = np.asarray(q, dtype=float)
+    if q.ndim != 1:
+        raise ValueError("q must be one-dimensional; all curves of a batch share the q grid")
+    b = packed[0].shape[0]
+    if instrument is None:
+        r = parratt(np.broadcast_to(q, (b, q.size)), *packed)
+    else:
+        r = _observed(_NumpyOps, parratt, q, b, packed, instrument)
+    return r[0] if single else r
+
+
+def resolution_report(q: ArrayLike, model, instrument: Instrument) -> dict[str, NDArray]:
+    """For the default resolution rule: nodes per point, the estimated relative error of the refined
+    range between the critical q, and whether refinement stopped at `resolution.MAX_DEPTH`, each
+    (B, Q) (or (Q,) for one model). What `forward` integrates with, without evaluating the curve."""
+    packed, single = as_batch(model)
+    q = np.asarray(q, dtype=float)
+    b = packed[0].shape[0]
+    inst = instrument
+    wl1 = inst.wavelength
+
+    def col(v):
+        return np.broadcast_to(np.asarray(v, dtype=float), (b,))[:, None]
+
+    theta = np.degrees(np.arcsin(np.clip(q * wl1 / (4.0 * np.pi), 0.0, 1.0)))[None, :]
+    theta = np.clip(theta + 0.5 * col(inst.tth_offset), 1e-9, 0.5 * MAX_TTH)
+    q_true = 4.0 * np.pi / wl1 * np.sin(np.radians(theta))
+    sigma = np.sqrt(
+        np.clip(col(inst.sigma_q) ** 2 + (col(inst.dq_over_q) * q_true) ** 2, 1e-300, None)
+    )
+    ka2 = _active(inst.kalpha2_ratio)
+    counts, _, _, err, capped = resolution.layout(
+        q_true,
+        sigma,
+        packed,
+        wl1,
+        inst.kalpha2_wavelength if ka2 else None,
+    )
+    out = {
+        "nodes": counts.reshape(b, -1),
+        "error": err.reshape(b, -1),
+        "capped": capped.reshape(b, -1),
+    }
+    return {k: v[0] for k, v in out.items()} if single else out

reflectosim/grid.py

@@ -0,0 +1,16 @@
+"""The q grid of the command lines."""
+
+import numpy as np
+
+from reflectosim.units import tth_to_q
+
+
+def q_grid(spec: str, wavelength: float) -> np.ndarray:
+    """'tth:START:STOP:N' (degrees) or 'q:START:STOP:N' (Å⁻¹) as q in Å⁻¹."""
+    axis, start, stop, n = spec.split(":")
+    values = np.linspace(float(start), float(stop), int(n))
+    if axis == "tth":
+        return tth_to_q(values, wavelength)
+    if axis == "q":
+        return values
+    raise SystemExit(f"grid axis must be tth or q, got {axis!r}")

reflectosim/kernel.py

@@ -0,0 +1,148 @@
+"""Float64 Parratt recursion compiled with numba.
+
+k_z² = (q/2)² − 4π(ρ − ρ_ambient) with ρ in Å⁻², Névot-Croce roughness. The SLD is conjugated so
+that the principal square root gives Im k_z ≥ 0 (decay into absorbers) with the exp(+2i k_z d)
+recursion; callers pass Im SLD > 0 for absorption.
+"""
+
+import math
+
+import numba
+import numpy as np
+from numpy.typing import NDArray
+
+
+@numba.njit(cache=True)
+def _curve_points(q, sld, thickness, roughness, start, stop, i, out):
+    n = sld.shape[1]
+    kz = np.empty(n, np.complex128)
+    rho = np.empty(n, np.complex128)
+    for j in range(n):
+        rho[j] = 4.0e-6 * np.pi * np.conj(sld[i, j] - sld[i, 0])
+    for iq in range(start, stop):
+        h = 0.5 * q[i, iq]
+        for j in range(n):
+            kz[j] = np.sqrt(h * h - rho[j])
+        x = 0j
+        for j in range(n - 2, -1, -1):
+            k1 = kz[j]
+            k2 = kz[j + 1]
+            s = roughness[i, j + 1]
+            r = (k1 - k2) / (k1 + k2) * np.exp(-2.0 * k1 * k2 * s * s)
+            if j + 1 < n - 1:
+                p = np.exp(2j * k2 * thickness[i, j + 1])
+                xp = x * p
+            else:
+                xp = 0j
+            x = (r + xp) / (1.0 + r * xp)
+        out[i, iq] = x.real * x.real + x.imag * x.imag
+
+
+@numba.njit(parallel=True, cache=True)
+def _parratt(q, sld, thickness, roughness, out):
+    b, nq = q.shape
+    chunk = 64
+    per_curve = (nq + chunk - 1) // chunk
+    for t in numba.prange(b * per_curve):
+        i = t // per_curve
+        start = (t % per_curve) * chunk
+        _curve_points(q, sld, thickness, roughness, start, min(start + chunk, nq), i, out)
+
+
+def parratt(
+    q: NDArray[np.float64],
+    sld: NDArray[np.complex128],
+    thickness: NDArray[np.float64],
+    roughness: NDArray[np.float64],
+) -> NDArray[np.float64]:
+    """|r|² for q (B, Q) and media arrays (B, N), ambient first. Returns (B, Q)."""
+    q = np.ascontiguousarray(q, dtype=np.float64)
+    out = np.empty(q.shape)
+    _parratt(
+        q,
+        np.ascontiguousarray(sld, dtype=np.complex128),
+        np.ascontiguousarray(thickness, dtype=np.float64),
+        np.ascontiguousarray(roughness, dtype=np.float64),
+        out,
+    )
+    return out
+
+
+@numba.njit(cache=True)
+def _reflectivity(qv, rho, thickness, roughness, i):
+    """|r|² at one q for curve i (rho = 4π · conj(ρ − ρ_ambient) [Å⁻²], ambient first)."""
+    n = rho.shape[1]
+    h = 0.5 * qv
+    k2 = np.sqrt(h * h - rho[i, n - 1])
+    x = 0j
+    for j in range(n - 2, -1, -1):
+        k1 = np.sqrt(h * h - rho[i, j])
+        s = roughness[i, j + 1]
+        r = (k1 - k2) / (k1 + k2) * np.exp(-2.0 * k1 * k2 * s * s)
+        xp = x * np.exp(2j * k2 * thickness[i, j + 1]) if j + 1 < n - 1 else 0j
+        x = (r + xp) / (1.0 + r * xp)
+        k2 = k1
+    return x.real * x.real + x.imag * x.imag
+
+
+@numba.njit(parallel=True, cache=True)
+def _smear(
+    q_true, sigma, starts, u, w, rho, thickness, roughness, wl1, wl2, ratio, length, beam, out
+):
+    b, nq = q_true.shape
+    for t in numba.prange(b * nq):
+        i, iq = t // nq, t % nq
+        acc = 0.0
+        for m in range(starts[t], starts[t + 1]):
+            sin_t = min(
+                max((q_true[i, iq] + sigma[i, iq] * u[m]) * wl1 / (4.0 * np.pi), 1e-12), 1.0
+            )
+            r = _reflectivity(4.0 * np.pi / wl1 * sin_t, rho, thickness, roughness, i)
+            if wl2 > 0:
+                r2 = _reflectivity(4.0 * np.pi / wl2 * sin_t, rho, thickness, roughness, i)
+                r = (r + ratio[i] * r2) / (1.0 + ratio[i])
+            if length[i] == length[i]:
+                r *= math.erf(length[i] * sin_t / (2.0 * np.sqrt(2.0) * beam[i]))
+            acc += w[m] * r
+        out[i, iq] = acc
+
+
+def smear(
+    q_true: NDArray[np.float64],
+    sigma: NDArray[np.float64],
+    counts: NDArray[np.int64],
+    u: NDArray[np.float64],
+    w: NDArray[np.float64],
+    sld: NDArray[np.complex128],
+    thickness: NDArray[np.float64],
+    roughness: NDArray[np.float64],
+    wl1: float,
+    wl2: float,
+    ratio: NDArray[np.float64],
+    length: NDArray[np.float64],
+    beam: NDArray[np.float64],
+) -> NDArray[np.float64]:
+    """Σ w · R over each point's nodes q_true + sigma · u, for (B, Q) points whose nodes are listed
+    point after point (counts (B·Q,) nodes each): Kα₁ (and Kα₂ at the same angle when wl2 > 0, mixed
+    by ratio), times the footprint factor where length is not NaN. Returns (B, Q)."""
+    sld = np.ascontiguousarray(sld, dtype=np.complex128)
+    rho = 4.0e-6 * np.pi * np.conj(sld - sld[:, :1])
+    starts = np.concatenate([[0], np.cumsum(counts)])
+    out = np.empty(q_true.shape)
+    _smear(
+        np.ascontiguousarray(q_true, dtype=np.float64),
+        np.ascontiguousarray(sigma, dtype=np.float64),
+        starts,
+        np.ascontiguousarray(u, dtype=np.float64),
+        np.ascontiguousarray(w, dtype=np.float64),
+        rho,
+        np.ascontiguousarray(thickness, dtype=np.float64),
+        np.ascontiguousarray(roughness, dtype=np.float64),
+        float(wl1),
+        float(wl2),
+        np.ascontiguousarray(ratio, dtype=np.float64),
+        np.ascontiguousarray(length, dtype=np.float64),
+        np.ascontiguousarray(beam, dtype=np.float64),
+        out,
+    )
+    return out

reflectosim/materials.py

@@ -0,0 +1,15 @@
+"""X-ray SLD from composition and mass density (DECISIONS #8: computed, never copied).
+
+Uses the X-ray scattering factor tables of `periodictable`; at Cu Kα₁ Si at 2.33 g/cm³ gives
+20.071 + 0.457i and SiO₂ at 2.20 g/cm³ gives 18.865 + 0.244i (1e-6 Å⁻²).
+"""
+
+from periodictable import xsf
+
+from reflectosim.units import CU_KALPHA1
+
+
+def sld(formula: str, density: float, wavelength: float = CU_KALPHA1) -> complex:
+    """Complex SLD [1e-6 Å⁻²] of a compound at a mass density [g/cm³]; Im > 0 is absorption."""
+    re, im = xsf.xray_sld(formula, density=density, wavelength=wavelength)
+    return complex(re, im)