NaNIsaac docs

Appearance

Python API ​

Signatures and docstrings of the public functions and classes in reflectolearn at docs/site (dddcc54), read from the source without importing it. The guides on this site say how the pieces fit together; this page is for looking up an argument.

reflectolearn.benchmark ​

Benchmark: the full predict pipeline on two sets of synthetic curves with known truth.

"synthetic": random films made by an independent float64 Parratt recursion written here in numpy,
not by the package's torch engine, so the benchmark does not grade the pipeline with its own forward
model (the "inverse crime"). Each curve mimics a lab scan: 2θ 0.2–6° in 0.02° steps, a Gaussian beam
footprint, Gaussian resolution with constant dq/q, a background and Poisson counts. The stack is a
film on a thin SiO2 layer on Si, so the native oxide the pipeline treats as a nuisance is present.

"synthetic_fixed": a few films with the arbitrary parameters in FIXED_CASES, simulated by the
training forward model (AbelesMatrix, same stack layout as the generator) with the training
augmentor's footprint, resolution smearing, background and Poisson noise at stated values. These
curves come from the model's own simulator, so they check the pipeline in its training domain, not
its transfer to independent physics.

parratt ​

python
def parratt(q: np.ndarray, layers: list[tuple[complex, float, float]]) -> np.ndarray
|r|^2 for [(sld, thickness, roughness of its top interface)], ambient first, substrate last.

SLD in 1e-6 Å^-2 with positive imaginary part for absorption; Névot–Croce roughness. The SLD is
conjugated so that the principal square root gives Im k > 0 (decay into absorbers) in every layer.

smeared ​

python
def smeared(q: np.ndarray, layers, dq_over_q: float, n: int=17) -> np.ndarray
Reflectivity convolved with a Gaussian of sigma = dq_over_q · q (Gauss–Hermite-like grid, ±3σ).

SyntheticCase ​

python
@dataclass
class SyntheticCase:
    name: str
    thickness: float
    roughness: float
    sld: float
    isld: float
    oxide_thickness: float
    oxide_roughness: float
    substrate_roughness: float
    sample_length_mm: float
    dq_over_q: float
    i0: float
    background: float

make_synthetic ​

python
def make_synthetic(
    n: int,
    seed: int = 2026,
    wavelength: float = 1.5406,
    tth: np.ndarray | None = None,
) -> list[tuple[SyntheticCase, np.ndarray, np.ndarray]]
n synthetic curves: list of (case, 2θ [deg], counts).

FixedCase ​

python
@dataclass
class FixedCase:
    name: str
    thickness: float
    roughness: float
    sld: float
    oxide_thickness: float
    oxide_sld: float
    oxide_roughness: float
    sub_sld: float
    sub_roughness: float
    beam_width: float
    sample_length: float
    res_sigma: float
    bkg: float
    poisson_scale: float
A film with stated parameters plus the instrument effects applied to its curve.

Units as in the config: Å, 1e-6 Å^-2; beam_width and sample_length as in
augmentation.beam_width_range / sample_len_range; res_sigma in Å^-1; bkg and poisson_scale
as log10 values (augmentation.bkg / poisson_log_scale). oxide_thickness 0 means no oxide.

make_fixed ​

python
def make_fixed(
    cfg,
    seed: int = 2026,
    cases=FIXED_CASES,
    tth_range: tuple[float, float] = (0.2, 8.0),
) -> list[tuple[FixedCase, np.ndarray, np.ndarray]]
Synthetic curves of `cases` from the training forward model: list of (case, q [Å^-1], intensity).

Same media layout as data.generator.simulate_reflectivity, then the XRRAugmentor footprint and
smearing, background and Poisson noise at the case's values, cropped to `tth_range` [deg].
The seed only drives the Poisson draw.

score ​

python
def score(kind: str, name: str, res: dict, truth: dict, seconds: float) -> dict
One benchmark row from a predict_curve result and the known truth.

summarize ​

python
def summarize(rows: list[dict]) -> dict
Per-set error statistics of the rows from `score`.

run_benchmark ​

python
def run_benchmark(
    cfg,
    checkpoint,
    *,
    stats_path=None,
    n_synthetic: int = 12,
    seed: int = 2026,
    alias_check: bool = False,
    mcmc_steps: int = 0,
    out_dir: str | Path = 'benchmark_out',
    device=None,
) -> dict
Run predict on both synthetic sets and score them; writes JSON and CSV.

reflectolearn.cli ​

ReflectoLearn CLI (typer): train / evaluate / config.

Run via:  uv run -m reflectolearn   (or)  reflectolearn  (installed script)

config ​

python
@app.command()
def config(
    out: Path = typer.Option('config.yaml', '--out', '-o', help='Output YAML path.'),
)
Write a default configuration YAML.

train ​

python
@app.command()
def train(
    config_path: Path | None = typer.Option(None, '--config', '-c', help='Config YAML.'),
    save_dir: Path | None = typer.Option(None, '--save-dir', '-s', help="Output dir. Defaults to config's save_root/experiment_name."),
    override: list[str] = typer.Option([], '--set', help='Override e.g. training.lr=1e-3'),
    resume: Path | None = typer.Option(None, '--resume', help='Checkpoint to resume.'),
    seed: int = typer.Option(0, '--seed', help='Random seed.'),
)
Train the XRR model on streaming synthetic data.

tune ​

python
@app.command()
def tune(
    config_path: Path | None = typer.Option(None, '--config', '-c', help='Base config YAML. Defaults to cwd/config.yaml.'),
    n_trials: int = typer.Option(30, '--trials', '-n', help='Number of trials to search.'),
    epochs: int = typer.Option(12, '--epochs', help='Proxy training epochs per trial (reduced).'),
    steps: int = typer.Option(400, '--steps', help='Steps per epoch per trial (reduced).'),
    val_steps: int = typer.Option(80, '--val-steps', help='Validation steps per trial (reduced).'),
    seed: int = typer.Option(0, '--seed', help='Sampler/data seed (for fair comparison).'),
    out: Path = typer.Option('config_best.yaml', '--out', '-o', help='Path to save the best config.'),
    storage: str | None = typer.Option(None, '--storage', help="Optuna storage URL. Defaults to auto-created <out>_study.db SQLite (resumable). Use ':memory:' to run in-memory."),
    study_name: str = typer.Option('reflectolearn', '--study-name', help='Optuna study name.'),
    param: list[str] = typer.Option([], '--param', '-p', help='Search only these parameters (repeatable). Defaults to searching all; the rest are fixed at base-config values. e.g. -p model.fourier_features'),
    override: list[str] = typer.Option([], '--set', help='Override fixed base values e.g. simulation.q_points=2000'),
)
Hyperparameter search (Optuna): find the best config via reduced training and save it as YAML.

evaluate ​

python
@app.command()
def evaluate(
    checkpoint: Path | None = typer.Argument(None, help="Checkpoint path. Defaults to config's save_dir/best.pt."),
    config_path: Path | None = typer.Option(None, '--config', '-c', help='Config YAML. Defaults to cwd/config.yaml or config.yaml next to the checkpoint.'),
    stats: Path | None = typer.Option(None, '--stats', help='stats.pt. Defaults to stats.pt next to the checkpoint.'),
    n_samples: int = typer.Option(2000, '--n-samples', help='Number of synthetic samples for evaluation.'),
    batch_size: int = typer.Option(200, '--batch-size', help='Inference batch size.'),
    refine: bool = typer.Option(False, '--refine/--no-refine', help='Physics-refine AI predictions against the measured curve and compare before/after.'),
    trust_margin: float = typer.Option(0.4, '--trust-margin', help='Trust-region half-width ratio about the AI estimate during refine.'),
)
Full checkpoint evaluation: metrics (MAE/MAPE/RMSE/R2) + CSV + visualization report.

predict ​

python
@app.command()
def predict(
    data: Path = typer.Argument(..., help='Measured curve file (2 columns: 2theta[deg] or q, intensity).'),
    checkpoint: Path | None = typer.Option(None, '--checkpoint', '-k', help="Checkpoint. Defaults to config's save_dir/best.pt."),
    config_path: Path | None = typer.Option(None, '--config', '-c', help='Config YAML. Defaults to cwd or config.yaml next to the checkpoint.'),
    stats: Path | None = typer.Option(None, '--stats', help='stats.pt. Defaults to stats.pt next to the checkpoint.'),
    x_axis: str = typer.Option('tth', '--x-axis', help="Input x-axis: 'tth' (2theta deg) or 'q' (Å⁻¹)."),
    refine: bool = typer.Option(True, '--refine/--no-refine', help='Physics-refine AI predictions against the measured curve.'),
    trust_margin: float = typer.Option(0.4, '--trust-margin', help='Refine trust-region half-width ratio.'),
    alias_check: bool = typer.Option(False, '--alias-check/--no-alias-check', help='Experimental: also refine from d/2, 2d and the FFT thickness; pick by FOM plus FFT prior.'),
    mcmc_steps: int = typer.Option(0, '--mcmc-steps', help="emcee steps for the posterior after refine (0 = off; needs the 'mcmc' extra)."),
    out_dir: Path | None = typer.Option(None, '--out-dir', '-o', help='Output folder. Defaults to next to the input file.'),
)
Infer one measured XRR curve, optionally refine the physics, and save the prediction + overlay plot.

benchmark ​

python
@app.command()
def benchmark(
    checkpoint: Path = typer.Option(..., '--checkpoint', '-k', help='Checkpoint to benchmark.'),
    config_path: Path | None = typer.Option(None, '--config', '-c', help='Config YAML. Defaults to config.yaml next to the checkpoint.'),
    stats: Path | None = typer.Option(None, '--stats', help='stats.pt. Defaults to next to the checkpoint.'),
    n_synthetic: int = typer.Option(12, '--n', help='Synthetic curves (independent numpy Parratt, known truth).'),
    seed: int = typer.Option(2026, '--seed', help='Seed of the synthetic sets.'),
    alias_check: bool = typer.Option(False, '--alias-check/--no-alias-check', help='Refine from d/2, d, 2d and the FFT thickness (experimental).'),
    mcmc_steps: int = typer.Option(0, '--mcmc-steps', help='emcee steps per curve (0 = off).'),
    out_dir: Path = typer.Option(Path('benchmark_out'), '--out-dir', '-o', help='Where benchmark.json/.csv go.'),
)
Score predict on synthetic curves with known truth: random films and a few fixed example films.

stitch ​

python
@app.command()
def stitch(
    scans: list[Path] = typer.Argument(..., help='Scan files (x, intensity, error) in increasing 2theta order.'),
    out: Path = typer.Option(..., '--out', '-o', help='Output text file.'),
)
Stitch attenuator-segmented scans into one curve (last scan is the intensity reference).

main ​

python
def main()

reflectolearn.config ​

ReflectoLearn Configuration (Pydantic, Streaming + 1-Layer)
===========================================================
Single-film 3-parameter (thickness/roughness/sld) prediction, for streaming
Abeles data generation.

ParamRanges ​

python
class ParamRanges(BaseModel):
    thickness: Range = [10.0, 1200.0]
    roughness: Range = [0.0, 15.0]
    sld: Range = [5.0, 100.0]
Range [min, max] of the film parameters to predict.

SubstrateConfig ​

python
class SubstrateConfig(BaseModel):
    sld: Range = [18.0, 22.0]
    roughness: Range = [1.0, 6.0]
Substrate: not predicted, used only in simulation. Set min==max to fix it.

SimulationConfig ​

python
class SimulationConfig(BaseModel):
    wavelength: float = 1.54
    q_min_tth: float = 0.05
    q_max_tth: float = 15.0
    q_points: int = 2000
    n_samples: int = 50000

ModelConfig ​

python
class ModelConfig(BaseModel):
    use_fourier: bool = True
    fourier_scale: float = 15.0
    fourier_features: int = 32
    n_channels: int = 64
    depth: int = 6
    mlp_hidden: int = 512
    dropout: float = 0.1
1:1 correspondence with XRRPhysicsModel arguments.

TrainingConfig ​

python
class TrainingConfig(BaseModel):
    batch_size: int = 64
    lr: float = 0.0002
    weight_decay: float = 0.0001
    epochs: int = 150
    steps_per_epoch: int = 3000
    val_steps: int = 300
    patience: int = 50

AugmentationConfig ​

python
class AugmentationConfig(BaseModel):
    prob: float = 0.9
    i0: Range = [0.8, 1.2]
    bkg: Range = [-8.0, -5.0]
    poisson_log_scale: Range = [5.0, 8.0]
    intensity_noise_scale: float = 0.25
    res_sigma_range: Range = [0.0001, 0.006]
    do_footprint: bool = True
    beam_width_range: Range = [0.005, 0.015]
    sample_len_range: Range = [5.0, 30.0]
    tth_offset: Range = [-0.02, 0.02]
    oxide_prob: float = 0.7
    oxide_thickness: Range = [0.0, 40.0]
    oxide_sld: Range = [5.0, 25.0]
    oxide_roughness: Range = [0.0, 8.0]
    kalpha2_ratio: float = 0.5
    kalpha2_wavelength: float = 1.54439
Batch torch augmentation parameters (read by XRRAugmentor).

Config ​

python
class Config(BaseModel):
    experiment_name: str = 'exp'
    save_root: Path = Path('runs')
    device: torch.device = Field(default_factory=lambda: torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
    param_ranges: ParamRanges = Field(default_factory=ParamRanges)
    substrate: SubstrateConfig = Field(default_factory=SubstrateConfig)
    simulation: SimulationConfig = Field(default_factory=SimulationConfig)
    model: ModelConfig = Field(default_factory=ModelConfig)
    training: TrainingConfig = Field(default_factory=TrainingConfig)
    augmentation: AugmentationConfig = Field(default_factory=AugmentationConfig)

Config.save_dir ​

python
@property
def save_dir(self) -> Path
Path for training artifacts = save_root / experiment_name.

Config.q_grid ​

python
@property
def q_grid(self) -> torch.Tensor
@q_grid.setter
def q_grid(self, value: torch.Tensor)

Config.save_yaml ​

python
def save_yaml(self, path: str | Path)

Config.load_yaml ​

python
@classmethod
def load_yaml(cls, path: str | Path) -> Self

reflectolearn.data.augmentations ​

Realistic XRR augmentations: instrument effects applied to linear R curves.

Contract: forward(q, R) takes linear-scale reflectivity R of shape (B, Q) and
the q grid (Q,), and returns linear-scale R of shape (B, Q). Randomness is
drawn fresh here, since the streaming generator calls this once per batch.

XRRAugmentor ​

python
class XRRAugmentor:
    def __init__(self, config: Config, q: torch.Tensor)
Apply realistic instrument effects to reflectivity curves.

All ops operate on linear-scale R of shape (B, Q) -> (B, Q). Effects:
  - footprint attenuation at low angles (erf geometric correction)
  - resolution (Gaussian q-smearing)
  - i0  : per-sample intensity scaling   (R * i0)
  - bkg : additive background            (+ 10**bkg)
  - Poisson counting noise
  - low/high angle masking (sets points to 0 -> masked downstream)
Each effect is gated independently by `config.augmentation.prob`.

XRRAugmentor.forward ​

python
def forward(self, q: torch.Tensor, R: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]
Apply instrument effects.

R: (B, Q) linear -> returns (R, mask):
  R    : (B, Q) linear, instrument effects applied (still > 0)
  mask : (B, Q) in {0, 1}; 0 marks angles that were NOT measured
         (low/high-angle cutoff). The generator uses this to fill the
         masked region with the log-floor (-15) consistently, exactly
         like the offline exp07 pipeline.

reflectolearn.data.generator ​

On-the-fly streaming XRR data generator using the Abeles physics engine.

param_stats ​

python
def param_stats(config: Config) -> dict[str, torch.Tensor]
Analytic mean/std for the uniform parameter ranges.

For U(lo, hi): mean=(lo+hi)/2, std=(hi-lo)/sqrt(12).
Returns tensors of shape (3,) ordered [thickness, roughness, sld].

simulate_reflectivity ​

python
def simulate_reflectivity(
    config: Config,
    q: Tensor,
    thickness: Tensor,
    roughness: Tensor,
    sld: Tensor,
    sub_sld: Tensor,
    sub_rough: Tensor,
    device,
) -> Tensor
Map film parameters to reflectivity R(q) with forward-model domain randomization.

For robustness on real measurements, the following are applied per-sample
at simulation time as nuisance effects (not prediction targets):
  * tth_offset : q-axis zero-point (alignment) error
  * oxide      : native oxide / surface layer (nuisance layer between ambient and film)
  * Kalpha2    : Cu K-alpha2 doublet (weighted sum)

Each effect is controlled by the range/probability in config.augmentation and
is disabled when set to 0. All input film parameters are (B,) tensors; the
return is (B, Q) linear R.

generate_synthetic_batch ​

python
def generate_synthetic_batch(
    config: Config,
    batch_size: int,
    augmentor: XRRAugmentor | None = None,
    device=None,
) -> tuple[torch.Tensor, torch.Tensor]
Generate one batch of synthetic XRR curves and normalized targets.

Returns:
    x: (B, 2, Q) model input [log10-normalized R, valid-q mask]
    y: (B, 3) normalized targets [thickness, roughness, sld]

XRRStreamingDataset ​

python
class XRRStreamingDataset(IterableDataset):
    def __init__(self, config: Config, augmentor: XRRAugmentor | None=None, steps: int=200)
Infinite streaming dataset that generates XRR batches on-the-fly.

get_xrr_dataloader ​

python
def get_xrr_dataloader(
    config: Config,
    augmentor: XRRAugmentor | None = None,
    steps: int = 200,
) -> DataLoader
DataLoader over streaming XRR data (batching happens inside the generator).

reflectolearn.evaluate ​

ReflectoLearn Evaluator
=======================
Analyze the performance of a trained XRRPhysicsModel on a synthetic test set and
generate a report.

Ideas are borrowed from the PINN (XRefine) evaluator but simplified for this
project's structure:
  - No ParamSet -> target is a normalized (B, 3) tensor [thickness, roughness, sld]
  - Model output is a single (B, 3) (no tuple/fixed_values)
  - Re-simulate with Abeles using the predicted parameters to assess curve quality
    (best/worst)

Outputs (all saved in the checkpoint's parent directory):
  - eval_metrics_full.csv       per-sample true/pred/error
  - eval_summary.csv            per-parameter MAE/MAPE/RMSE/R2
  - eval_1_best_worst.png       best/worst curves + parameter table
  - eval_2_learning_curves.png  train/val loss, lr
  - eval_3_error_vs_value.png   error (bias) vs physical value scatter
  - eval_4_parity.png           true vs pred parity (R2)

XRREvaluator ​

python
class XRREvaluator:
    def __init__(
        self,
        checkpoint_path: str | Path,
        config: Config,
        stats_path: str | Path | None = None,
        n_samples: int = 2000,
        batch_size: int = 200,
        device: torch.device | None = None,
        refine_enabled: bool = False,
        refine_config: RefineConfig | None = None,
    )

XRREvaluator.setup ​

python
def setup(self)

XRREvaluator.run_inference ​

python
@torch.no_grad()
def run_inference(self)

XRREvaluator.calculate_metrics ​

python
def calculate_metrics(self)

XRREvaluator.plot_best_worst ​

python
def plot_best_worst(self, n: int=3)

XRREvaluator.plot_learning_curves ​

python
def plot_learning_curves(self)

XRREvaluator.plot_error_vs_value ​

python
def plot_error_vs_value(self)

XRREvaluator.plot_parity ​

python
def plot_parity(self)

XRREvaluator.report ​

python
def report(self)

XRREvaluator.evaluate_all ​

python
def evaluate_all(self)

reflectolearn.genx_fit ​

Classical GenX fit (differential evolution via bumps) seeded with a CNN estimate.

Optional: needs the ``genx`` extra (``uv sync --extra genx``). Model: film + oxide
layer on a fixed substrate, with I0, background and footprint (sample length) fitted
in two stages (oxide + I0, then everything).

genx_fit ​

python
def genx_fit(
    q: np.ndarray,
    intensity: np.ndarray,
    thickness: float,
    roughness: float,
    sld: float,
    *,
    wavelength: float = 1.54,
    sub_sld: float = 20.07,
    sub_rough: float = 3.0,
    beam_width: float = 0.1,
    sample_len: float = 10.0,
    film_margin: float = 0.1,
    pop_size: int = 40,
    steps_scale: float = 1.0,
) -> dict
Fit a measured curve with GenX, starting from (thickness, roughness, sld).

Film parameters are bounded to ±``film_margin`` of the guess. ``intensity`` is
normalized to its max. Returns fitted parameters plus ``r_fit`` and ``fom`` (logR2).

reflectolearn.model ​

FourierFeatureMapping ​

python
class FourierFeatureMapping(nn.Module):
    def __init__(self, input_dim=1, mapping_size=64, scale=10.0)
Random Fourier Features.

Maps the coordinate (q) into a high-dimensional frequency space to help the
model learn high-frequency patterns.

FourierFeatureMapping.forward ​

python
def forward(self, x)

XRRPhysicsModel ​

python
class XRRPhysicsModel(nn.Module):
    def __init__(
        self,
        q_len: int,
        input_channels: int = 2,
        output_dim: int = 3,
        n_channels: int = 64,
        depth: int = 5,
        mlp_hidden: int = 256,
        dropout: float = 0.1,
        use_fourier: bool = True,
        fourier_scale: float = 10.0,
        fourier_features: int = 32,
    )

XRRPhysicsModel.forward ​

python
def forward(self, x)

reflectolearn.physics.abeles ​

Abeles matrix method.

Physics engine for X-ray reflectivity calculation, optimized for batch
processing and gradient descent.

abeles_core ​

python
def abeles_core(q: Tensor, thickness: Tensor, roughness: Tensor, sld: Tensor) -> Tensor
Core Abeles matrix calculation.

Length convention (important):
    There are N media and N-1 interfaces.
    thickness / roughness are per-interface, so they are 1 shorter than sld.

Args:
    q:         (B, Q, 1) complex  - scattering vector (k_z = q/2)
    sld:       (B, 1, N) complex  - relative SLD per medium (ambient=0 reference)
    thickness: (B, 1, N-1) real   - thickness of each finite layer/interface
    roughness: (B, 1, N-1) real   - roughness of each interface (Nevot-Croce)

Returns:
    (B, Q) real - reflectivity R(q) = |M_10 / M_00|^2

AbelesMatrix ​

python
class AbelesMatrix(nn.Module):
    def __init__(self, device: torch.device | None=None)

AbelesMatrix.forward ​

python
def forward(
    self,
    q: Tensor,
    thickness: Tensor,
    roughness: Tensor,
    sld: Tensor,
    sample_length: Tensor | None = None,
    beam_width: Tensor | None = None,
    wavelength: float = 1.5406,
) -> Tensor
Pre/post-processing layer wrapping the physics kernel (abeles_core).

Responsibilities:
  1) Input preparation: dtype casting, relative-SLD conversion, unit scaling, dimension alignment
  2) Call abeles_core (pure physics calculation)
  3) Footprint correction (optional)
As an nn.Module, its operations are recorded on the autograd graph and are differentiable.

Args:
    q:         (B, Q) real - scattering vector [1/Angstrom]
    thickness: (B, N-1)    - thickness of the upper medium at each interface [Angstrom];
                             the first entry is the ambient's and has no effect
    roughness: (B, N-1)    - rms roughness of each interface, top first [Angstrom]
    sld:       (B, N)      - SLD per medium, ambient first [1e-6/Angstrom^2],
                             real or complex; computed in complex64
    sample_length, beam_width: for footprint correction (optional)
    wavelength: X-ray wavelength [Angstrom], defaults to Cu K-alpha
Returns:
    (B, Q) real - reflectivity

abeles ​

python
def abeles(*args, **kwargs)
Functional interface wrapper.

reflectolearn.physics.utils ​

tth2q ​

python
def tth2q(tth: float, wavelen: float=1.54) -> float
def tth2q(tth: NDArray[np.float64], wavelen: float=1.54) -> NDArray[np.float64]
Convert 2θ (in degrees) and wavelength (in Å) to scattering vector q (in 1/Å).
Formula: q = (4 * pi / lambda) * sin(theta)

apply_tth_offset ​

python
def apply_tth_offset(q: Tensor, offset_deg: Tensor, wavelength: float) -> Tensor
Apply a q-axis zero-point (alignment) offset: add offset[deg] to the
measured 2θ, then convert back to q.

Differentiable (torch). q: (Q,) or (B, Q); offset_deg: (B, 1) -> returns (B, Q).

reflectolearn.predict ​

Inference on real measured XRR curves plus physics refinement.

Training uses synthetic data only, but inference targets real measured files.
The model only works correctly if the measured input goes through exactly the
same preprocessing as the synthetic training input:

  measured I(2theta) -> convert to q -> interpolate onto fixed q_grid
            -> per-curve max normalize -> log10 -> clamp(<=0)
            -> unmeasured region: mask=0 + log-floor

There are no ground-truth parameters, so evaluation centers on the refinement
FOM (LogMSE): how well the measured curve is reproduced when the AI prediction
is used as the initial guess for physics fitting. The substrate is unknown for
real measurements, so it is fixed to the midpoint of the config.substrate range.

load_measured_curve ​

python
def load_measured_curve(
    path: str | Path,
    *,
    x_axis: str,
    wavelength: float,
) -> tuple[np.ndarray, np.ndarray]
Read a 2-column text file (2theta or q, intensity) into (q[Å⁻¹], intensity).

Comment/header lines (non-numeric) are skipped automatically. Both whitespace-
and comma-separated formats are supported.

preprocess_measured ​

python
def preprocess_measured(
    q_meas: np.ndarray,
    intensity: np.ndarray,
    cfg: Config,
    device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
Convert a measured curve into model input (1, 2, Q).

Returns:
    x      : (1, 2, Q)  [log10-normalized R, mask]  — model input
    r_norm : (1, Q)     reflectivity-like normalized linear values (for refine's log_r_obs)
    mask   : (1, Q)     {0,1} valid-q mask

load_predictor ​

python
def load_predictor(
    cfg: Config,
    checkpoint: str | Path,
    stats_path: str | Path | None,
    device: torch.device,
) -> tuple[XRRPhysicsModel, torch.Tensor, torch.Tensor]
Model in eval mode plus the target de-normalisation (mean, std).

predict_curve ​

python
def predict_curve(
    q_meas: np.ndarray,
    intensity: np.ndarray,
    cfg: Config,
    predictor: tuple[XRRPhysicsModel, torch.Tensor, torch.Tensor],
    *,
    refine_enabled: bool = True,
    refine_config: RefineConfig | None = None,
    alias_check: bool = False,
    mcmc_steps: int = 0,
    device: torch.device | None = None,
    name: str = 'curve',
) -> tuple[dict, dict]
AI prediction, refinement (from the AI thickness and, with alias_check, from half and double
of it) and, with mcmc_steps > 0, posterior sampling for one measured curve.

Returns (result, curves): result holds the parameters, FOMs and posterior; curves holds the q grid,
observed, AI and refined curves and the mask, for plotting.

predict_file ​

python
def predict_file(
    path: str | Path,
    cfg: Config,
    checkpoint: str | Path,
    *,
    stats_path: str | Path | None = None,
    x_axis: str = 'tth',
    refine_enabled: bool = True,
    refine_config: RefineConfig | None = None,
    alias_check: bool = False,
    mcmc_steps: int = 0,
    device: torch.device | None = None,
    out_dir: str | Path | None = None,
) -> dict
Infer a single measured file, optionally physics-refine, then save results/plots/summary.

reflectolearn.refine.losses ​

Curve-fitting loss functions for physics refinement.

All losses follow the same signature over the (B, Q) tensor convention:

    r_sim:     (B, Q) simulated reflectivity (linear, > 0)
    log_r_obs: (B, Q) log10 of observed reflectivity
    mask:      (B, Q) valid-q mask (1 = valid, 0 = ignore; float or bool accepted)

Returns a scalar loss over the batch (per-sample masked mean -> batch mean).

log_mse_loss ​

python
def log_mse_loss(r_sim: Tensor, log_r_obs: Tensor, mask: Tensor) -> Tensor
MSE in log10 space. Sensitive to dynamic-range (fringe) shape such as roughness/attenuation.

linear_mse_loss ​

python
def linear_mse_loss(r_sim: Tensor, log_r_obs: Tensor, mask: Tensor) -> Tensor
MSE in linear space. Sensitive to absolute intensity (scale) near the critical angle.

correlation_loss ​

python
def correlation_loss(r_sim: Tensor, log_r_obs: Tensor, mask: Tensor) -> Tensor
1 - Pearson correlation. Invariant to scale/attenuation, so strong at matching fringe frequency (period).

log_gradient_loss ​

python
def log_gradient_loss(r_sim: Tensor, log_r_obs: Tensor, mask: Tensor) -> Tensor
Value MSE + gradient (Sobolev) MSE. Strongly penalizes solutions that flatten the curve (vanishing slope).

poisson_deviance_loss ​

python
def poisson_deviance_loss(counts_scale: float) -> LossFn
Poisson deviance per point for counts = counts_scale · R (counts_scale: counts at R = 1).

Unlike the log losses this is the likelihood of counting data: each point is weighted by its
counts, and points with zero counts inside the mask contribute too. Built per curve set because
it needs the count scale, so it is not in LOSS_PRESETS; a stage selects it with loss="poisson".

It is the likelihood only if counts_scale · 10**log_r_obs are the raw per-point counts on the
measured points. Values interpolated onto another grid (as `predict.preprocess_measured` does
for the network) or averaged over equal q are not Poisson counts, and the deviance of such
values is only an approximation. Pass refine the measured q grid and the raw counts instead.

hybrid_linear_log_loss ​

python
def hybrid_linear_log_loss(r_sim: Tensor, log_r_obs: Tensor, mask: Tensor) -> Tensor
Mix of linear (intensity/critical angle) + log (fringe/dynamic range), linear:log = 20:1.

reflectolearn.refine.mcmc ​

Posterior sampling after refinement: emcee on the same forward model refine fits with.

refine gives one best fit and no uncertainty. Sampling the posterior around it gives the spread of
each parameter and exposes degeneracies (thickness against SLD, roughness against resolution).

Likelihood: Gaussian in log10 R over the valid q points with an unknown spread exp(lnsigma) that
is fitted with the rest. Without it the curve's own scatter has to be assumed, and when the model
does not describe the data to within that scatter the posterior comes out absurdly narrow. The
measured curve is interpolated onto the fixed q grid, so grid points are not independent; the
log-likelihood is scaled to the number of measured points (``n_independent``), or the intervals
would shrink by about sqrt(grid points / measured points).

Priors are flat within the same bounds refine used (trust region for the film, global ranges for
the nuisances). All walkers are evaluated in one batch on the refine device.

The kept half of the chain is checked against its integrated autocorrelation time tau: estimates are
usable when it spans at least 50 tau. The check is necessary, not sufficient: on a short chain tau
itself is underestimated, so a flag set near the boundary can be optimistic; "tau_margin" (kept steps
/ max tau) shows how far from the boundary the chain is. With refine's twelve parameters free a 2000-step chain is often
several times shorter than that (ReflectoLearn#29); the result reports tau and a warning is logged.

sample_posterior ​

python
def sample_posterior(
    info: dict,
    substrate: Tensor,
    q_grid: Tensor,
    log_r_obs: Tensor,
    mask: Tensor,
    *,
    wavelength: float,
    config: RefineConfig | None = None,
    n_independent: int | None = None,
    steps: int = 2000,
    walkers_per_dim: int = 4,
    seed: int = 0,
    device: torch.device | None = None,
) -> dict
Sample the posterior of every parameter refine fitted, for one curve.

Args:
    info: refine's (or refine_multistart's) info for a single curve (batch of 1).
    n_independent: measured points inside the q grid; defaults to the number of valid grid points.
    steps: emcee steps; the first half is discarded as burn-in.

Returns:
    {"params": {key: {"median", "lo", "hi"}}, "lnsigma": {...}, "acceptance", "steps", "walkers",
     "samples": (N, K) array of the kept samples, "keys": list, "log_prob_start", "log_prob_max",
     "best": the highest-probability point visited, "tau": integrated autocorrelation time of
     each key and of lnsigma over the kept half, "tau_margin": kept steps / max(tau),
     "converged": tau_margin >= 50}.
    lo/hi are the 16th/84th percentiles (68 % interval).

reflectolearn.refine.pipeline ​

Gradient-based physics refinement of AI parameter predictions.

Refines CNN predictions (initial guess) by directly fitting a physics-simulated
curve to the measured curve. The ``AbelesMatrix`` forward model is fully
differentiable via autograd, so the refinement runs as **Adam gradient descent**.
Core elements:

  * Multi-stage strategy: scale/alignment -> film structure -> fine tuning
  * Per-stage loss choice (linear / correlation / log, see
    :mod:`reflectolearn.refine.losses`)
  * Trust-region bounds that narrow the search around the AI estimate
  * Valid-q masking + FOM (LogMSE) to track improvement

To fit the measured curve accurately, **nuisance parameters** (values that affect
the curve shape but are not of interest) are fitted simultaneously. This prevents
forward-model mismatch from leaking into the parameters of interest (thickness etc.):

  * i0  : intensity scale,  bkg : background (log10)
  * tth_offset : q-axis zero (alignment) offset
  * dq         : resolution (Gaussian q-smearing) — disentangles from roughness
  * dq_rel     : q-proportional resolution term, sigma(q) = sqrt(dq² + (dq_rel·q)²); off by default
  * footprint  : low-angle beam footprint attenuation (beam > sample) —
                 disentangles from low-angle amplitude
  * oxide      : native oxide / surface layer (nuisance layer between ambient and film)

Bound constraints map a latent variable ``z`` through a sigmoid onto [lo, hi],
keeping parameters within their physical range throughout gradient descent.

RefineStage ​

python
@dataclass
class RefineStage:
    name: str
    optimize: tuple[str, ...]
    loss: str = 'log_mse'
    steps: int = 200
    lr: float = 0.03
    q_min: float | None = None
    q_max: float | None = None
One refinement stage: which parameters to optimize, with which loss, for how many steps.

likelihood_stages ​

python
def likelihood_stages(budget: int=10) -> list[RefineStage]
Stages that fit counting data to its Poisson maximum likelihood (ReflectoLearn#25).

Changes to the default strategy, each found necessary in the R2b cause separation (removing any
one of them stops refine short of the optimum on curves from its own forward model):
  * every stage runs `budget` times as many steps;
  * no correlation stage: its loss ignores the level of the curve and accepts other fringe orders;
  * the fine stage minimizes the Poisson deviance instead of the log10 MSE, whose optimum is not
    the maximum-likelihood point;
  * the fine stage is followed by two stages over the same keys at lr 1e-3 and 1e-4, so that
    Adam settles instead of jittering around the optimum.
About 30 times the default number of steps.

RefineConfig ​

python
@dataclass
class RefineConfig:
    trust_margin: float = 0.4
    fit_i0: bool = True
    fit_bkg: bool = True
    fit_tth_offset: bool = True
    tth_offset_last_stage: bool = False
    fit_dq: bool = True
    fit_dq_rel: bool = False
    fit_kalpha2: bool = True
    fit_footprint: bool = True
    fit_oxide: bool = True
    nuisance_reg: float = 0.2
    i0_bounds: tuple[float, float] = (0.5, 2.0)
    bkg_bounds: tuple[float, float] = (-10.0, -4.0)
    tth_offset_bounds: tuple[float, float] = (-0.08, 0.08)
    tth_offset_bound_tol: float = 0.01
    offset_off_margin: float = 0.05
    dq_bounds: tuple[float, float] = (0.0, 0.012)
    dq_rel_bounds: tuple[float, float] = (0.0, 0.05)
    kalpha2_ratio_bounds: tuple[float, float] = (0.0, 0.6)
    kalpha2_wavelength: float = 1.54439
    footprint_bounds: tuple[float, float] = (150.0, 8000.0)
    oxide_thickness_bounds: tuple[float, float] = (0.0, 60.0)
    oxide_sld_bounds: tuple[float, float] = (15.0, 22.0)
    oxide_roughness_bounds: tuple[float, float] = (0.0, 12.0)
    bkg_init: float = -7.0
    dq_init: float = 0.002
    dq_rel_init: float = 0.01
    kalpha2_ratio_init: float = 0.5
    footprint_init: float = 2000.0
    oxide_thickness_init: float = 5.0
    oxide_sld_init: float = 18.0
    oxide_roughness_init: float = 3.0
    counts_scale: float | None = None
    stages: list[RefineStage] = field(default_factory=_default_stages)

RefineConfig.likelihood ​

python
@classmethod
def likelihood(cls, counts_scale: float, **overrides) -> RefineConfig
The configuration that reaches the Poisson maximum-likelihood fit (ReflectoLearn#25):
likelihood_stages() with the count scale of the data. Other fields keep their defaults
unless overridden.

Use it on raw counts on the measured q points (refine's q_grid = the measured q,
log_r_obs = log10(counts / counts_scale)), not on a curve interpolated onto the model grid:
see `poisson_deviance_loss`.

make_forward ​

python
def make_forward(
    q_grid: Tensor,
    substrate: Tensor,
    *,
    wavelength: float,
    config: RefineConfig,
    device: torch.device,
) -> Callable[[dict[str, Tensor], set[str]], Tensor]
Build the forward model refine fits with: film (+oxide) on a fixed substrate, plus the active
nuisance effects. The returned function takes a dict of (W,) parameter tensors for any batch size W
(refine's batch, or every MCMC walker at once) and returns (W, Q) reflectivity.

refine ​

python
def refine(
    p_init: Tensor,
    substrate: Tensor,
    q_grid: Tensor,
    log_r_obs: Tensor,
    mask: Tensor,
    param_ranges,
    *,
    wavelength: float,
    config: RefineConfig | None = None,
    device: torch.device | None = None,
    verbose: bool = False,
) -> tuple[Tensor, dict]
Refine AI-estimated film parameters to the measured curve via gradient descent
(with simultaneous nuisance fitting).

Returns:
    p_refined: (B, 3) refined [thickness, roughness, sld]
    info: {"fom_before", "fom_after" (scalar LogMSE), fitted nuisances (B,)}
          fom_before is measured with film parameters fixed at AI-init and only
          i0/bkg (global scale/offset) aligned, so the gap to fom_after reflects
          film-parameter improvement rather than scale matching.

per_curve_log_mse ​

python
def per_curve_log_mse(r: Tensor, log_r_obs: Tensor, mask: Tensor) -> Tensor
(B,) log10 MSE over each curve's valid points (the refine FOM, per curve).

fft_thickness ​

python
def fft_thickness(
    q_grid: Tensor,
    log_r_obs: Tensor,
    mask: Tensor,
    d_range: tuple[float, float],
    sld: Tensor | None = None,
    d_min: float = 60.0,
) -> Tensor
(B,) film thickness from the Kiessig period, independent of the network: the strongest peak of
the real-space transform of the detrended log R (the Patterson-type peak at z = d).

Fringes are periodic in the film's own wavevector sqrt(q^2 - q_c^2), not in q; transforming over q
overestimates d by several percent near the edge. With ``sld`` (B,) [1e-6/Å^2] the transform runs
over the film wavevector, q_c = 4 sqrt(pi sld 1e-6); a rough SLD is enough.

Fresnel decay is removed with a quintic fit over the valid region above the critical edge, the
rest is Hann-windowed and transformed; the estimate is the highest local maximum with z in
[max(d_min, lo), hi]. d_min keeps out the native oxide and the leakage tail at small z. NaN where
there are fewer than 40 valid points or no such maximum (films thinner than d_min).

refine_multistart ​

python
def refine_multistart(
    p_init: Tensor,
    substrate: Tensor,
    q_grid: Tensor,
    log_r_obs: Tensor,
    mask: Tensor,
    param_ranges,
    *,
    wavelength: float,
    config: RefineConfig | None = None,
    device: torch.device | None = None,
    verbose: bool = False,
    thickness_factors: tuple[float, ...] = (0.5, 1.0, 2.0),
    fom_tolerance: float = 0.25,
    fft_weight: float = 1.5,
) -> tuple[Tensor, dict]
refine from several thickness starts and keep the one the fringe period supports.

A single refine cannot correct a guess off by about a factor of two: the trust region keeps
thickness within ±trust_margin of its start. So refine also runs from multiples of the network's
thickness and from the FFT estimate of the Kiessig period (fft_thickness), all in one batch.

The lowest log-MSE alone is not a safe choice: with the oxide, resolution and Kα2 nuisances free,
a film two or four times as thick often fits as well or better. The FFT estimate alone is not
safe either: now and then it locks onto a wrong peak. So each start is scored by
ln FOM + fft_weight · |ln(d / d_fft)|, a soft prior on the fringe period, and the lowest wins.
On the benchmark the right start wins for weights between about 0.9 and 2.5. Without an estimate,
the network's own start wins when its FOM is within ``fom_tolerance`` of the best, else the best
FOM. info["alias"] lists every start's label, refined thickness and FOM, the FFT estimate and the
choice.

reflectolearn.report.figures ​

Figures for the data slides.

Slide 2 ("Domain Randomization"): a 3x3 panel grid that takes one clean simulated
XRR curve and overlays each randomized effect (oxide layer, 2theta offset, Cu Kalpha2
doublet, footprint, resolution, intensity scale, background, Poisson noise) plus a
final "all combined" realistic curve.

All effects reuse the production code paths (AbelesMatrix, apply_tth_offset,
XRRAugmentor) so the figure stays faithful to what the generator actually does.
Everything is deterministic (fixed film params + fixed seed) for reproducible slides.

FilmCase ​

python
@dataclass
class FilmCase:
    thickness: float = 150.0
    roughness: float = 5.0
    sld: float = 40.0
    sub_sld: float = 20.0
    sub_rough: float = 3.0
    oxide_thickness: float = 30.0
    oxide_sld: float = 18.0
    oxide_roughness: float = 3.0
Representative single film on a substrate (units match config).

domain_randomization_figure ​

python
def domain_randomization_figure(
    cfg: Config | None = None,
    film: FilmCase | None = None,
    out: str | Path = 'domain_randomization.png',
    *,
    seed: int = 0,
) -> Path
Render the slide-2 domain-randomization panel grid and save it as PNG.

sim_to_real_figure ​

python
def sim_to_real_figure(
    cfg: Config | None = None,
    film: FilmCase | None = None,
    out: str | Path = 'sim_to_real.png',
    *,
    seed: int = 0,
) -> Path
Single-axes figure: clean simulation -> + systematic distortions -> + stochastic noise.

Visualizes the sim-to-real gap as two stacked layers: deterministic (systematic)
distortions reshape the curve, and Poisson (stochastic) noise adds randomness.

model_architecture_figure ​

python
def model_architecture_figure(
    cfg: Config | None = None,
    out: str | Path = 'model_architecture.png',
) -> Path
Left-to-right block diagram of XRRPhysicsModel (curve -> film parameters).

Stages and shapes are derived from the config so the diagram stays accurate:
Input (2,Q) -> Fourier features -> 7x Conv block -> Global Avg Pool -> MLP -> 3 params.

write_synthetic_curve ​

python
def write_synthetic_curve(cfg: Config, path: str | Path, case: str='fixed_mid') -> Path
Write the benchmark's synthetic `case` as a 2-column (q, intensity) text file.

refine_figure ​

python
def refine_figure(
    cfg: Config | None = None,
    data_path: str | Path | None = DEFAULT_REFINE_DATA,
    out: str | Path = 'refine.png',
    *,
    checkpoint: str | Path | None = None,
    x_axis: str = 'tth',
) -> Path
Physics-refinement figure: run the full inference + refine pipeline on one curve and
produce the data / AI-predict / refined overlay (with FOM).

Without `data_path` the curve is the benchmark's synthetic "fixed_mid" film, written to
synthetic_fixed_mid.txt next to `out`. Mirrors the CLI 'predict' resolution: checkpoint
defaults to config's save_dir/best.pt, and config/stats next to that checkpoint are
auto-discovered. The plot is rendered by predict_file; this wrapper just resolves paths and
copies the PNG to `out`.

refine_process_figure ​

python
def refine_process_figure(out: str | Path='refine_process.png') -> Path
Block diagram of the refinement gradient-descent loop.

AI init -> [Abeles forward -> LogMSE vs measured -> backprop -> Adam update] -> theta*.
Schematic only (no computation), for the 'refinement principle' slide.

reflectolearn.stitch ​

Stitch XRR scans measured in consecutive 2θ ranges with different attenuators.

The last scan is the intensity reference. Going backwards, each earlier scan is
scaled so its last point equals the first point of the following scan, and that
duplicate boundary point is dropped. Adjacent scans must share one x value.

stitch_scans ​

python
def stitch_scans(scans: list[np.ndarray]) -> np.ndarray
Combine scans of shape (N, 2+) [x, intensity, (error, ...)] in increasing-x order.

Every column after x is rescaled together.

read_scan ​

python
def read_scan(path: str | Path) -> np.ndarray
Read a whitespace-separated scan file (x, intensity, error), no header.

reflectolearn.train ​

Trainer ​

python
class Trainer:
    def __init__(
        self,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader,
        save_dir: str | Path,
        lr: float = 0.0002,
        weight_decay: float = 0.0001,
        patience: int = 20,
        verbose: bool = True,
    )

Trainer.train ​

python
def train(
    self,
    epochs: int,
    resume_from: str | Path | None = None,
    on_epoch_end: Callable[[int, float, float], None] | None = None,
)
Training loop.

on_epoch_end(epoch, train_loss, val_loss): hook called after validation
on every epoch. During hyperparameter search it can report the
intermediate val_loss and raise an exception to prune an unpromising trial.

reflectolearn.tune ​

Hyperparameter search (Optuna) for the XRR inverse-problem model.

Search targets (all independent of `q_len`/physical ranges, so they are safe to
vary freely because the encoder folds the length with ``AdaptiveAvgPool1d(1)``):

  * training : lr, weight_decay, batch_size
  * model    : n_channels, depth, mlp_hidden, dropout, fourier_scale, fourier_features

The objective is the **normalized validation loss** (the MSE used by Trainer).
Target normalization is analytic and not a search target, so trials are directly
comparable.

Each trial runs a reduced "proxy" training over fewer epochs/steps and returns
the best val_loss; ``MedianPruner`` watches the val_loss reported each epoch and
prunes unpromising trials early. Data generation uses the same seed per trial to
reduce noise.

Checkpoints are written to a per-trial temporary directory and deleted on exit,
so they do not clutter the disk. Finally the best hyperparameters are applied to
the base config and saved as YAML, and the full trial history is saved as CSV —
just run ``reflectolearn train`` with the saved YAML.

suggest_overrides ​

python
def suggest_overrides(
    trial: optuna.Trial,
    cfg: Config,
    only: list[str] | None = None,
) -> None
Apply the search space directly to cfg (in-place). Key names match dotted paths.

only: list of dotted-keys to search. None searches everything. If given, only
      those keys are searched and the rest stay fixed at the base config values
      (e.g. ["model.fourier_features"]).

objective ​

python
def objective(
    trial: optuna.Trial,
    base_cfg: Config,
    *,
    epochs: int,
    steps: int,
    val_steps: int,
    seed: int,
    search_params: list[str] | None = None,
) -> float
One trial: run reduced training and return the best validation loss; TrialPruned if unpromising.
python
def run_search(
    base_cfg: Config,
    *,
    n_trials: int = 30,
    epochs: int = 12,
    steps: int = 400,
    val_steps: int = 80,
    seed: int = 0,
    out: str | Path = 'config_best.yaml',
    storage: str | None = None,
    study_name: str = 'reflectolearn',
    search_params: list[str] | None = None,
) -> optuna.Study
Run the Optuna search and save the best config (YAML) + trial history (CSV).

When storage is unset, results are written to ``<out>_study.db`` SQLite, which
supports interrupt/resume by default (re-running with the same study_name picks
up existing trials and continues searching). To run in-memory, pass
``storage=":memory:"`` explicitly.

Each project section is built from that project’s docs/ directory.

Isaac Yong