Preprocess

The composable Plan → Plan preprocess pipeline: build the initial plan from an experiment, sharpen it with swappable steps (beam selection, orientation/thickness determination, rocking-curve integration, mosaicity, convergence/coverage sweeps), then hand the final plan to a terminal (run_inference or engine.refine).

Spine

The Plan: the invariant geometry the differentiable refinement is conditioned on.

Plan is the spine of the preprocess pipeline – one immutable value bundling the shared StructureFactorGrid and the per-rotation OrientationPlans. Each Plan -> Plan step returns a sharpened copy (via dataclasses.replace()); refine consumes the final Plan. The dependency points preprocess -> engine (this module imports the engine’s geometry plans), never the reverse – the engine stays unaware of Plan and remains a pure consumer of grid + orientations.

class diffBloch.preprocess.plan.CandidatePlan(orientation: NDArray[float64], energy: float, u0: float, thickness: NDArray[float64], beam_hkl: NDArray[int64], pattern: PatternBatch)

Bases: object

The pre-build candidate phase for one rotation: source only, no built geometry.

from_experiment lays down a CandidatePlan per rotation (the difference-safe candidate beam_hkl + orientation source), select_beams prunes beam_hkl (cheap, source-only), and build_orientation_plans() then builds it into an OrientationPlan – the one place the structure-factor gather is built, over the already-pruned beam set. A CandidatePlan has no beam_plans, so it is unsolvable by construction (the engine consumes only built OrientationPlans); building the expensive gather over the full candidate pool is thereby avoided entirely.

orientation: NDArray[float64]

energy: float

u0: float

thickness: NDArray[float64]

beam_hkl: NDArray[int64]

pattern: PatternBatch

classmethod seed(beam_hkl: NDArray[int64], pattern: PatternBatch, *, energy: float, thickness: Tensor | NDArray[float64] | Sequence[float], u0: float = 0.0, orientation: Tensor | NDArray[float64] | None = None) → CandidatePlan

Assemble a candidate seed for one orientation (plain-numpy source; builds no gather).

class diffBloch.preprocess.plan.Plan(structure_factor_grid: StructureFactorGrid, orientations: tuple[CandidatePlan | OrientationPlanLike, ...], provenance: tuple[StepRecord, ...] = <factory>)

Bases: object

Shared structure_factor_grid plus per-rotation orientations (the refinement spine).

structure_factor_grid fixes the Fgb support and metric; orientations is one OrientationPlan per rotation (each already coupled to the grid at build time). Immutable: preprocess steps return dataclasses.replace() copies rather than mutating in place.

provenance is the ordered tuple of StepRecords that produced this plan – pipeline() appends one per step as it runs. A freshly built plan (from from_experiment) has empty provenance; the recipe identity a checkpoint locks against is this tuple. Steps do not touch it (their replace preserves it); the combinator owns it.

structure_factor_grid: StructureFactorGrid

orientations: tuple[CandidatePlan | OrientationPlanLike, ...]

provenance: tuple[StepRecord, ...]

diffBloch.preprocess.plan.coupling_stats(op: CandidatePlan | OrientationPlan | CoupledOrientationPlan) → dict[str, int]

One rotation’s solve-geometry shape: the (unions, tilts-per-union, beams-per-union) cost.

Phase-robust (a plan is summarised after every pipeline step, from the pre-build candidate on): a CoupledOrientationPlan reports its real coupling (n_coupling_segments unions, per-union cover widths and union_beam_index beam counts); a built OrientationPlan is one implicit union spanning all its tilts; a pre-build CandidatePlan knows only its beam-pool size (no tilts/segments yet). These are exactly the (B, T, N) drivers of the segmented Bloch solve the refinement loop repeats.

diffBloch.preprocess.plan.require_built_plans(plan: Plan) → tuple[OrientationPlan | CoupledOrientationPlan, ...]

Parse the orientation phase to built (OrientationPlan / CoupledOrientationPlan).

Plan.orientations is a phase-union (CandidatePlan before the build, built plans after) rather than a phase-indexed Plan[P], as the recipe is a runtime pipeline list of homogeneous Plan -> Plan steps: the phase cannot ride in the type across that list (nor across read_plan, which reconstructs a Plan from .npz bytes). This is the parse that re-establishes the phase at that erased boundary – parse, don’t validate: it returns the narrowed type once, and the terminals (inference, the fits, couple_beams, checkpoint serialize) then hold built geometry without re-checking. The phase-indexed alternative (Plan[P]) does not survive contact with Python – there is no phase-changing composition over a homogeneous step list, and disallow_any_generics rejects the erased reconstruction.

A CandidatePlan has no beam_plans and is unsolvable, so this raises unless build_orientation_plans has run (the default recipe runs it right after select_beams). Returns the plan’s own orientations tuple (identity preserved) once narrowed.

diffBloch.preprocess.plan.require_candidate_plans(plan: Plan) → tuple[CandidatePlan, ...]

Parse the orientation phase to the pre-build candidate (CandidatePlan).

The candidate-phase counterpart of require_built_plans() (which documents why the phase is parsed at a runtime boundary rather than carried in the type). select_beams and build_orientation_plans operate on the candidate phase from_experiment lays down; this raises with a clear error if the plan is already built (holds OrientationPlans), since build_orientation_plans is the single step that builds them and runs once, right after select_beams.

diffBloch.preprocess.plan.require_orientation_plans(plan: Plan) → tuple[OrientationPlan, ...]

Narrow a plan’s orientations to plain OrientationPlans (reject segmented ones).

The tilt-independent-only plan-shaping steps (select_beams, integrate_rocking_curve) transform the OrientationPlan, which carries one shared beam set. couple_beams replaces each orientation with a CoupledOrientationPlan (a per-tilt-chunk beam set) that those steps cannot consume, so they must all precede couple_beams in a pipeline. This helper enforces that ordering with a clear error and narrows the element type for the caller.

The fitting steps (optimize_orientation, optimize_thickness) are deliberately not narrowed: they are plan-agnostic (they rebuild via OrientationPlan.with_orientation() / replace(thickness=...), both defined on the segmented plan too), so they iterate plan.orientations directly and run either before or after couple_beams.

diffBloch.preprocess.plan.summarize_plan(plan: Plan) → dict[str, float]

Plan-level shape as numeric measurements (the observability summary of a settled/mid Plan).

Emitted after every pipeline step (and once for the seed, as PlanSeeded), so consecutive summaries are the per-stage survival counts: how many solve beams and scored reflections each filtering step left behind. The names are scoped because the sets are independent – SOLVE (n_solve_beams_*, the beams that couple dynamically) is not SCORED (n_matched_hkl, the reflections that enter the R-factor) is not the structure-factor support (n_grid_hkl). n_observed_hkl/n_matched_hkl are deduplicated distinct (h, k, l) counts across every rotation (unique_hkl_count()), not a sum of each rotation’s own count – a reflection re-observed (or matched) in more than one rotation is counted once, not once per rotation.

n_matched_hkl is absent, not zero, before build_orientation_plans runs: a CandidatePlan has no alignment, and reporting 0 there would make “not built yet” indistinguishable from “matched nothing”.

diffBloch.preprocess.plan.unique_hkl_count(hkl_batches: Iterable[Tensor]) → int

Count of distinct (h, k, l) triples across hkl_batches (each (M, 3)).

A reflection recorded in multiple rotations (rotation electron diffraction frames overlap in angle, so the same reciprocal-lattice point is routinely re-observed in several consecutive frames) is counted once here, not once per rotation it appears in – a plain sum(len(batch) ...) double-counts exactly those overlaps.

from_experiment: the boundary constructor from parsed records + config to refinement inputs.

This is the initial total construction of the preprocess pipeline (it is not Plan -> Plan – there is no Plan yet). It assembles two separable products from the same records/config:

• a Plan pair (train / validation) – the invariant geometry the preprocess steps then sharpen and refine consumes;

• a RefinementSetup – the structure-side static + refinable inputs the RefinementEngine needs (ASU expansion, constraint spec, initial parameters, atomic numbers).

The structure side lives here so the structure/experimental split mirrors the two parsed records.

class diffBloch.preprocess.experiment.DatasetSetup(plan: Plan, integration: IntegrationGeometry, mosaicity: MosaicSmoothed | None, energy: float, n_rotations: int, ignored_rotations: tuple[int, ...])

Bases: object

One dataset’s seeded geometry: the candidate plan plus its own measurement context.

The per-dataset product of setup_datasets(). plan holds one CandidatePlan per non-ignored rotation, with rotation_index file-local (the rotation’s original position within this dataset’s own PETS file, so ignoring a rotation leaves a gap rather than renumbering later frames); pool() maps these onto the pooled global index space. integration is this file’s own rocking-curve semiangle – pooled datasets may differ (each file’s recipe runs with its own geometry). energy is the snapped beam energy (snap_to_standard_energy() over the PETS wavelength); pool guards that pooled datasets agree, since the engine solves one experiment at one energy. n_rotations is the full pre-ignore rotation count (the pooled offset arithmetic and the train/val mask both run over original counts), and ignored_rotations the sorted file-local ignore slice (part of this dataset’s checkpoint-lock identity).

plan: Plan

integration: IntegrationGeometry

mosaicity: MosaicSmoothed | None

energy: float

n_rotations: int

ignored_rotations: tuple[int, ...]

class diffBloch.preprocess.experiment.ExperimentSetup(plans: PlanSplit, refinement: RefinementSetup, integration: IntegrationGeometry, mosaicity: MosaicSmoothed | None)

Bases: object

The full product of from_experiment: the geometry plans + structure refinement.

Two separable concerns from the same records/config: plans (the Plan -> Plan geometry spine) and refinement (the static structure context the engine is built from). Kept distinct so only the Plan pair flows through the preprocess pipeline.

plans: PlanSplit

refinement: RefinementSetup

integration: IntegrationGeometry

mosaicity: MosaicSmoothed | None

class diffBloch.preprocess.experiment.PlanSplit(train: Plan, validation: Plan)

Bases: object

A train / validation Plan pair sharing one StructureFactorGrid.

from_experiment splits the rotations into the two plans (validation = every 10th rotation by default); both reference the same grid object, so the shared Fgb support cannot diverge.

The split is currently dormant: nothing downstream distinguishes train from validation (inference and the engine take a single Plan), and whole-experiment work uses combined (all rotations). A whole-rotation holdout is a weak cross-validation guard for over-determined physics refinement anyway – the principled analog holds out reflections (R_free), not orientations. It is retained because it becomes informative for learned modes, where a learned theta -> thickness can overfit per rotation.

train: Plan

validation: Plan

property combined: Plan

One Plan over all rotations (train + validation) on the shared grid.

For whole-experiment evaluation (e.g. inference over every rotation), where the train/val split is irrelevant. train orientations come first, then validation.

class diffBloch.preprocess.experiment.RefinementSetup(asu_plan: AsuExpansionPlan, spec: ConstraintSpec, params: RefinableParams, numbers: Tensor, cell_parameters: NDArray[float64] | None = None)

Bases: object

The structure-side inputs a RefinementEngine is built from.

Kept separate from the geometry Plan (which carries the grid + orientations): the Plan flows through the Plan -> Plan preprocess steps, while this static structure context is handed to the engine at refinement time. params are the initial refinable parameters seeded from the CIF (positions at their CIF values, ADPs inverted from the CIF ADPs); spec freezes the constraint metadata (fixed positions, occupancies, ADP kinds, the reciprocal frame ADPs map through).

asu_plan: AsuExpansionPlan

spec: ConstraintSpec

params: RefinableParams

numbers: Tensor

cell_parameters: NDArray[float64] | None = None

classmethod from_structure(structure: StructureRecord, *, cell_parameters: NDArray[float64] | None = None, isotropic_displacements_only: bool = False) → RefinementSetup

Assemble the structure-side refinement inputs from a parsed StructureRecord.

Positions and symmetry are read directly off the CIF’s fractional coordinates and symmetry operators – both are metric-independent. ADPs are mapped to the reciprocal U* frame by diffBloch.params.constrain(), using cell_parameters as the metric when given (PETS’s authoritative cell from setup_datasets()); otherwise the structure’s own CIF cell, for direct construction with no PETS record in scope. Per-rotation thickness lives elsewhere – on each per-rotation plan (seeded by from_experiment, fitted by optimize_thickness) – because it varies per rotation, not per structure.

Special-position degrees of freedom are constrained by the site-symmetry projector built natively from the structure’s symmetry operators (symmetry_constraints()): an atom special position is held on its site under refinement, so it is neither over-parameterized nor free to drift off. A general-position atom gets the identity projector (unconstrained).

isotropic_displacements_only (inputs.isotropic_displacements_only) forces every atom onto Uiso via a derived ADP record (_force_isotropic_adp()), never by mutating structure itself.

diffBloch.preprocess.experiment.from_experiment(structure: StructureRecord, experimental_data: ExperimentalRecord, config: ExperimentConfig) → ExperimentSetup

Construct the geometry Plan pair + structure RefinementSetup from parsed inputs.

The single-dataset public boundary, kept for API users and the inference/e2e paths: it is setup_datasets() over one record, with the train/validation split applied on top. Rotations split into train / validation plans sharing the grid; split membership is defined on the original PETS order (ignoring a rotation must not renumber later frames and silently move them between train and validation). The app’s preprocess spine does not come through here – it runs setup_datasets() per dataset and applies the split after pool().

diffBloch.preprocess.experiment.resolve_dataset_mosaicity(enabled: bool, record: ExperimentalRecord, rocking: RockingCurve) → MosaicSmoothed | None

Resolve PETS apparent mosaicity to the applied tilt reduction, or None when disabled.

diffBloch.preprocess.experiment.resolve_dataset_orientations(record: ExperimentalRecord) → NDArray[float64]

This dataset’s per-rotation orientations, in the frame the rest of the pipeline assumes.

Composes the two concerns kept separate in preprocess.orientation: the as-collected derivation (orientation_matrices()) and the goniometer-axis correction (rotation_axis_correction()), which brings the rotation axis onto x so the left-multiplied rocking tilts and klar_beam_mask’s (g_y, g_z) lever arm are measured about the right axis.

The azimuth is read from the PETS file and there is no override: a file that does not record it cannot be processed. A missing value is an error rather than an assumed zero, because absent, unparsed, and genuinely-zero are three different situations, and silently treating the first two as the third is how a wrong integration axis goes unnoticed.

diffBloch.preprocess.experiment.seed_beam_hkl(grid: StructureFactorGrid, *, g_max: float) → NDArray[int64]

Difference-safe seed beams: the grid reflections within g_max (includes 000).

The orientation-independent candidate pool {hkl in grid : |g| <= g_max} that from_experiment lays down. Selecting from the shared grid keeps every beam difference inside the Fgb support as long as 2 * g_max <= grid.g_max (the caller’s responsibility).

diffBloch.preprocess.experiment.setup_datasets(structure: StructureRecord, records: Sequence[ExperimentalRecord], config: ExperimentConfig) → tuple[RefinementSetup, tuple[DatasetSetup, ...]]

Seed one candidate Plan per dataset + the structure side.

The initial total construction of the preprocess pipeline (not Plan -> Plan – there is no Plan yet), generalized over one or more PETS files (inputs.multi_dataset). The structure-side products are dataset-independent and built once, shared by every dataset’s plan: the structure-factor grid is derived from the solve cutoff (from_cell_for_beam_cutoff() sizes it to 2x the cutoff so it spans every coupled g - h difference – the same grid object rides on every per-dataset plan, so their Fgb support cannot diverge), as are the difference-safe seed beams and the RefinementSetup.

Per dataset: the beam energy is derived from that file’s PETS wavelength and snapped onto the nearest standard TEM voltage when close (PETS records wavelength to only 4-5 significant figures, so the exact inverse lands a few hundred eV off 100/200/300 kV rather than on it), and the mean-inner-potential u0 follows the energy – computed once per distinct snapped energy, since it depends on nothing else that varies between datasets. One CandidatePlan per rotation carries its crystal orientation matrix (native PETS derivation, no side-car file) and the observed pattern for that zone axis.

blochwave.ignore_orientations indexes the pooled rotation space (files concatenated in records order, original pre-ignore counts); it is validated against the pooled total and translated to each dataset’s file-local slice here. A dataset whose every rotation is ignored raises – its recipe would have nothing to fit.

This is intentionally a module-level function rather than a classmethod: it is the single documented public boundary of the preprocess pipeline (records + config -> setups), and it returns a composite of products rather than constructing one domain object. The per-object constructors it delegates to follow the classmethod idiom (StructureFactorGrid.from_cell_for_beam_cutoff, CandidatePlan.seed, RefinementSetup.from_structure).

diffBloch.preprocess.experiment.validation_mask(n_rotations: int, split: DataSplitConfig) → NDArray[bool]

Boolean per-rotation validation mask from the split policy.

train_test=False holds out nothing (every rotation trains). Otherwise every round(1 / val_frac)-th rotation (1-based count -> 0-based indices) is held out for validation, e.g. val_frac=0.2 holds out every 5th rotation.

The preprocess composition combinators: sequence and fixpoint over Plan -> Plan steps.

A step is a pure Plan -> Plan transformer (it fits something – numerics, orientation, thickness – and returns a sharpened Plan). pipeline() chains steps left to right; iterate_until() drives one step to a fixpoint (for convergence testing or alternating fits). Both return a Plan -> Plan step, so they nest – a fixpoint of steps is itself a step. refine is deliberately not expressible here: it is the terminal Plan -> Result estimator, not a Plan -> Plan transform.

Provenance. A step is self-describing: it carries a StepRecord (its name + serialized params). As pipeline() applies each step it stamps that record onto the resulting Plan’s provenance tuple, so the final Plan records the ordered recipe that produced it – each step appends its record as it runs. This is what lets a checkpoint bind its identity to the recipe, not just the inputs. A step with no record (a bare closure, a nested composite) stamps OPAQUE – a plan whose provenance contains it can never be reused (safe: a miss, never a false hit).

type diffBloch.preprocess.pipeline.ConvergenceCheck = Callable[[Plan, Plan], bool]

class diffBloch.preprocess.pipeline.Fork(predicate: Callable[[StructureFactorGrid], bool], when_true: tuple[PlanStep, ...], when_false: tuple[PlanStep, ...])

Bases: object

The choice combinator: run one of two step lists, chosen by a predicate on the grid.

The one rule that makes it checkpointable: the predicate reads only the :class:`~diffBloch.engine.plan.StructureFactorGrid`, invariant across every preprocess step (steps replace orientations; nothing resizes the grid). So the branch is a deterministic function of the experiment’s fixed inputs – knowable before running – rather than of the mutating Plan. That keeps the fork’s shape static, so resolve_recipe() can splice the chosen branch inline into a flat, fork-free recipe before the checkpoint lock ever looks at it. A predicate over the Plan would make the shape depend on intermediate results and is deliberately unrepresentable here.

Branches are step lists, not pre-composed pipeline([...]) closures, so each branch step’s StepRecord survives into the resolved recipe (a composed closure would collapse to one OPAQUE). __call__() lets a Fork also run ad hoc inside a raw pipeline – it produces the right Plan but records OPAQUE (a non-Step in the stamping loop), a safe miss; checkpointable identity comes only from resolve_recipe().

predicate: Callable[[StructureFactorGrid], bool]

when_true: tuple[PlanStep, ...]

when_false: tuple[PlanStep, ...]

resolve(grid: StructureFactorGrid) → tuple[PlanStep, ...]

The branch this fork takes for grid (the invariant discriminant).

type diffBloch.preprocess.pipeline.PlanStep = Callable[[Plan], Plan]

type diffBloch.preprocess.pipeline.StateInitializer = Callable[[Plan], State]

type diffBloch.preprocess.pipeline.StatefulPlanStep = Callable[[Plan, State], tuple[Plan, State]]

class diffBloch.preprocess.pipeline.Step(record: StepRecord, run: PlanStep)

Bases: object

A self-describing Plan -> Plan step: its provenance record + the run transform.

Callable, so a Step satisfies PlanStep structurally and every existing caller (pipeline, run_inference(prepare=...)) treats it as before; pipeline() also reads record to stamp provenance.

record: StepRecord

run: PlanStep

class diffBloch.preprocess.pipeline.StepRecord(name: str, params: dict[str, Any] | None = None)

Bases: object

A step’s provenance entry: its name and canonical serialized params (or None).

Two records compare equal iff the step and its params are identical, so a recipe’s provenance is a stable, comparable identity. params is the spec_to_params() form (JSON-able, with __type__ tags), so the record round-trips through the lock and the .npz __meta__.

name: str

params: dict[str, Any] | None = None

diffBloch.preprocess.pipeline.as_step(name: str, spec: Any, run: PlanStep) → Step

Wrap a step’s run closure with its StepRecord (name + serialized spec).

diffBloch.preprocess.pipeline.fork(predicate: Callable[[StructureFactorGrid], bool], *, when_true: Sequence[PlanStep], when_false: Sequence[PlanStep]) → Fork

Build a Fork choosing between two step lists by a predicate on the grid.

predicate receives the shared StructureFactorGrid (e.g. a cell-volume / grid-size test routing a large cell to a coarse-precision branch); when_true / when_false are the branch step lists (kept as lists so their records survive resolution).

predicate must be a pure function of the grid – no external or mutable state. The grid argument is only half the contract: the type stops it reading the mutating Plan, but a predicate that closed over a global flag would desync the pre-run resolution (for the lock) from the runtime branch just as badly. Purity over an input that is itself pipeline-invariant is what makes the branch deterministic per experiment.

diffBloch.preprocess.pipeline.iterate_until(step: PlanStep, *, until: ConvergenceCheck, max_iterations: int = 50) → PlanStep

Drive step to a fixpoint: re-apply it until until(previous, current) holds.

Returns a Plan -> Plan step that applies step repeatedly, checking until against the (previous, just-produced) Plan pair after each application, and returns the first Plan that satisfies it. Raises RuntimeError if max_iterations is reached without convergence – silent non-convergence is never returned. max_iterations must be >= 1.

Provenance: the fixpoint stamps a single OPAQUE record – the number of iterations is input-dependent, so a per-iteration log would not be a stable recipe identity. A plan produced through iterate_until is therefore not checkpoint-reusable (a safe miss).

diffBloch.preprocess.pipeline.pipeline(steps: Sequence[PlanStep], *, logger: Logger = NULL_LOGGER) → PlanStep

Compose steps left to right, stamping each step’s record onto the plan’s provenance.

After applying each step, appends its StepRecord (or OPAQUE for a bare closure) to the plan’s provenance, so the composed result records the ordered recipe. An empty list yields the identity (provenance unchanged).

logger (default the null sink) receives one PlanSeeded for the incoming plan and then a PlanStepCompleted after each step – the step’s name as the event channel, its ordinal as the step, and summarize_plan() of the resulting plan – so a fresh preprocess run streams the plan’s shape as it evolves. The baseline is what makes the stream survival counts rather than absolute ones: each stage’s beam and reflection totals are only legible as a filter’s effect against what entered it. Reusing a checkpoint bypasses this runner, so those fire only on a fresh run (the boundary CouplingSummary covers the reuse case). Emission is alongside the provenance tell; the null default keeps the pure composition path unchanged.

diffBloch.preprocess.pipeline.resolve_recipe(steps: Sequence[PlanStep], grid: StructureFactorGrid) → tuple[PlanStep, ...]

Compile every Fork away against grid -> a flat, fork-free step list.

Splices each fork’s chosen branch inline (recursively, so nested forks flatten too). Because the grid is invariant across the pipeline, resolving against the base grid here yields exactly the recipe that will run – which is what lets the checkpoint lock key on a flat step_records list with no knowledge of forks.

diffBloch.preprocess.pipeline.spec_to_params(spec: Any) → dict[str, Any] | None

Serialize a frozen-dataclass value-type to a canonical, deterministic dict for provenance.

Recurses into nested dataclasses (the specs nest: TrialCoupling holds a policy + ScoredHklSelection holds a BeamSelection holds an IntegrationGeometry), tagging each with __type__ = its class name so a fieldless discriminated-union arm (e.g. TiltIndependent, whose asdict is {}) is distinguishable from any other empty spec. Non-dataclass leaves (int/float/str/bool/None, Literals-as-str) pass through; tuples/lists recurse elementwise. None returns None (a paramless step).

diffBloch.preprocess.pipeline.stateful_pipeline(steps: Sequence[StatefulPlanStep[State]]) → StatefulPlanStep[State]

Compose state-threading plan phases left to right.

A StatefulPlanStep has shape (Plan, State) -> (Plan, State): it can transform the plan while also carrying live driver state that should not become part of the public Plan. This helper is the explicit, immutable-state counterpart to pipeline(): it folds the phases left to right, threading (plan, state) through each one without mutating either in place.

The returned value is still stateful. Use stateful_plan_step() to adapt it back to the ordinary Plan -> Plan preprocess boundary.

diffBloch.preprocess.pipeline.stateful_plan_step(init_state: StateInitializer[State], step: StatefulPlanStep[State]) → PlanStep

Adapt a state-threading driver to the ordinary Plan -> Plan pipeline shape.

init_state derives the driver’s initial state from the incoming plan; step runs the stateful computation; the final state is intentionally discarded. This formalizes drivers such as numerical convergence, whose public product is a settled Plan but whose internals must thread transient scalar choices between phases.

Provenance note: this returns a bare closure. Wrap it with as_step() if the driver has a stable, serializable recipe identity; otherwise it will stamp OPAQUE like any custom closure, which is the safe checkpoint behavior for input-dependent loops.

diffBloch.preprocess.pipeline.step_records(steps: Sequence[PlanStep]) → tuple[StepRecord, ...]

The recipe pipeline(steps) will stamp – one record per step, in order.

The checkpoint driver reads this before running to compare against a lock (does the intended recipe match / extend the snapshot’s?). A bare step contributes OPAQUE, so a recipe containing one can be detected and refused (never checkpointed) up front.

Native crystal-orientation derivation for the preprocess pipeline.

Reconstructs per-rotation crystal orientation matrices from the experiment’s goniometer geometry – the UB matrix and per-rotation tilt angles recorded in the PETS data – with no side-car orientation file. The orientations are first-class inputs to the Plan; optimize_orientation refines them in-Plan – it must not re-orthonormalise them: U carries PETS’s own small UB-vs-cell-parameters fit residual, so a polar/SVD projection would silently drop it.

Convention:

orientation = R_z(omega) . R_x(alpha) . R_y(beta) @ U,    U = UB @ B^-1

This is the as-collected convention – what PETS recorded. The goniometer axis is additionally brought onto x by R_z(-rotation_axis_position) (rotation_axis_correction()), composed at the dataset boundary in resolve_dataset_orientations() rather than here, so this module’s derivation stays a pure function of the PETS-recorded geometry.

where B is the Busing-Levy reciprocal matrix built from this dataset’s own PETS cell parameters (the same cell UB was fit against, so U is close to a pure rotation – see diffBloch.preprocess.experiment._resolve_authoritative_cell for how a combined experiment’s shared cell is chosen and cross-checked) and the goniometer rotations are active, in degrees. Geometry then uses orientation_basis() = reciprocal_cell(cell @ orientation.T) (NOT reciprocal_basis @ orientation.T), because orientation is not guaranteed exactly orthonormal even so.

Reference: W. R. Busing & H. A. Levy, Acta Cryst. 22, 457 (1967) (the UB-matrix formalism).

diffBloch.preprocess.orientation.busing_levy_matrix(cell_parameters: FloatArray) → FloatArray

Busing-Levy reciprocal B matrix from (a, b, c, alpha, beta, gamma), angles in degrees.

Rows follow the standard a*-along-x setting. The cell volume is computed exactly from the parameters; we deliberately do not consume a rounded _cell_volume field if the source file carries one (for the quartz anchor that rounding shifts orientations by ~1e-6, negligible).

diffBloch.preprocess.orientation.goniometer_rotation(alpha: float, beta: float, omega: float) → FloatArray

Active goniometer rotation R_z(omega) . R_x(alpha) . R_y(beta), angles in degrees.

diffBloch.preprocess.orientation.hexagonal_tilt(azimuth: float, polar: float) → FloatArray

Palatinus hexagonal-search tilt R_z(azimuth) . R_x(polar) . R_z(-azimuth), in degrees.

A tilt of magnitude polar about the in-plane axis at azimuth – the delta rotation optimize_orientation right-multiplies onto an orientation (orientation @ tilt). Being a true true rotation (det = 1) it preserves whatever small fit residual U already carries exactly, so the re-orthonormalisation trap is dodged by construction.

Reference: L. Palatinus et al., Acta Cryst. A69, 171-188 (2013), the hexagonal modified-simplex search.

diffBloch.preprocess.orientation.orientation_basis(cell: FloatArray, orientation: FloatArray) → FloatArray

Lab-frame reciprocal basis for one orientation: reciprocal_cell(cell @ orientation.T).

The single home for the orientation convention. orientation is not guaranteed exactly orthonormal – it carries PETS’s own UB-vs-cell-parameters fit residual (see preprocess.orientation) – so this is NOT reciprocal_basis @ orientation.T. cell rows are the real-space basis; returns rows a*, b*, c* in inverse Angstrom. orientation = I reproduces reciprocal_cell(cell).

diffBloch.preprocess.orientation.orientation_matrices(ub_matrix: FloatArray, cell_parameters: FloatArray, alphas: FloatArray, betas: FloatArray, omegas: FloatArray) → FloatArray

Per-rotation as-collected orientation matrices R_goni @ U, shape (R, 3, 3).

alphas/betas/omegas are the per-rotation goniometer angles (degrees), one entry per PETS zone axis, in the same order as the record’s zone_axis_ids.

As-collected means exactly what PETS recorded: this deliberately applies no goniometer-axis correction, so the result is a pure function of UB, the cell, and the angles. Bringing the rotation axis onto x is a separate concern with a separate input, and lives in resolve_dataset_orientations().

diffBloch.preprocess.orientation.rocking_curve_tilts(semiangle: float, sampling: int, *, geometry: str = 'continuous_rotation') → FloatArray

Rocking-curve integration tilts as (N, 3, 3) rotation matrices, N = sampling.

For continuous_rotation, sampling tilts span linspace(-semiangle, +semiangle, sampling) degrees about x, the goniometer axis in the PETS coordinate frame. sampling = 1 is the identity so that unit-sample rocking integration composes off. For precession, samples lie at the fixed cone semi-angle and uniformly spaced azimuths over [0, 360) using R_z(phi) R_x(semiangle) R_z(-phi). This is the convention used by the legacy preprocessing path.

In both modes these matrices left-multiply the already-PETS-rotated nominal orientation (R_tilt @ orientation). Callers unpack a validated RockingCurve into these raw arguments (the value-type owns the invariants), matching hexagonal_tilt()’s raw-float style.

diffBloch.preprocess.orientation.rotation_axis_correction(rotation_axis_position: float) → FloatArray

R_z(-rotation_axis_position), degrees: brings the true goniometer axis onto x.

Two independent parts of the pipeline assume the goniometer/rotation axis lies along x in PETS’s own coordinate frame: rocking_curve_tilts(), whose tilts are left-multiplied onto an orientation (R_tilt @ orientation), and klar_beam_mask(), whose lever arm is the lab-frame (g_y, g_z) – the distance from the x rock axis. Both hold only when PETS’s rotation axis position is zero.

A nonzero value means the real axis sits at that azimuth instead, so every per-rotation orientation must be pre-rotated back by its negative before any x-axis tilt is composed onto it, or the whole rocking-curve integration runs about the wrong axis. This is the pure matrix; the composition and the record-vs-config resolution live in resolve_dataset_orientations().

diffBloch.preprocess.orientation.u_matrix(ub_matrix: FloatArray, cell_parameters: FloatArray) → FloatArray

Crystal U matrix U = UB @ B^-1, from this dataset’s own PETS UB and cell parameters.

Close to a pure rotation (B is built from the same cell PETS fit UB against), but not guaranteed to be exactly orthonormal – PETS’s own UB-vs-cell-parameters fit carries a small residual. Not re-orthonormalised (see the module docstring).

Assemble a RefinementEngine and score orientations against data (computes wR2).

Bridges the two products diffBloch.preprocess.from_experiment() returns – the geometry Plan and the structure-side RefinementSetup – into a runnable RefinementEngine, then exposes the per-orientation scaling-optimised wR2 the orientation refinement (optimize_orientation) minimises.

build_engine is the general Plan + RefinementSetup -> engine assembly (the same engine refine will consume); score_orientations is the thin convenience that computes the orientation-invariant F_gb once and scores every orientation of a Plan. The forward simulation inside is deterministic and depends only on its inputs (same inputs always give the same result), so it does not change any shared state – it is ordinary computation reading captured read-only context, not a side effect.

diffBloch.preprocess.scoring.active_structure_factor_indices(orientations: tuple[OrientationPlan | CoupledOrientationPlan, ...], gpts: tuple[int, int, int]) → Tensor

Return grid rows referenced by the settled plans’ structure-factor gathers.

diffBloch.preprocess.scoring.build_engine(plan: Plan, refinement: RefinementSetup, *, loss: LossFn = wr2_loss, scores: ScoresFn = wr2_scores, method: SolverMethod = 'matrix_exp', max_batch: int | None = None, absorption: Absorption = NO_ABSORPTION, compact_structure_factors: bool = True, profile: bool = False, checkpoint_activations: bool = True) → RefinementEngine

Wire a geometry plan and a structure refinement into a runnable engine (no compute).

Pure assembly, not a forward pass. Plan and RefinementSetup are kept deliberately separate – the Plan (shared grid + per-rotation orientations) flows through the Plan -> Plan preprocess steps, while refinement (constraint spec, ASU-expansion plan, ASU atomic numbers) is static structure context. build_engine is the single place that rejoins them when a simulation is actually needed; both score_orientations here and refine later go through it. loss/scores are the matching scalar/per-thickness forms of one objective (see diffBloch.engine.losses): loss is the per-orientation term refine minimises, scores is what RefinementEngine.score_orientation() / score_orientation_per_thickness() search over (optimize_orientation/optimize_thickness). Both default to wR2 (wr2_loss()/wr2_scores()) after matching calculated total intensity to observed – calc and obs are on different scales, so the raw metric would be flat/gradient-free. Callers should pass both from the same ExperimentConfig.loss_metrics (to_loss() / to_scores()) so the search and the gradient objective agree.

max_batch (default None) caps the matrix_exp propagator block; None lets each solve pick a memory-safe block from its beam count, bounding peak memory while matching the unbounded solve to machine precision (a pin is only needed for a specific device budget). Execution-only, like method.

compact_structure_factors computes only support-grid rows referenced by the settled solve gathers and scatters them into the unchanged grid-shaped interface. It changes neither the solve nor its gradients; callers whose beam sets change dynamically may disable it or extend the support lazily.

profile logs per-phase wall time (structure factors, each rotation’s solve) on the built engine; see RefinementEngine. Execution-only and off by default – it forces a CUDA sync around every measured block.

checkpoint_activations (default True) trades peak memory for one extra forward recompute per per-orientation/per-segment solve on the refinement backward pass; disabling it removes that recompute at the cost of retaining every solve’s intermediates until backward. Execution-only – gradients are identical either way. See RefinementEngine.

diffBloch.preprocess.scoring.score_orientations(plan: Plan, refinement: RefinementSetup, *, method: SolverMethod = 'matrix_exp') → tuple[Tensor, ...]

Scaling-optimised wR2 for every orientation in plan at the seeded refinement.params.

Computes the orientation-invariant F_gb once and reuses it across orientations. This is the objective surface optimize_orientation searches over per rotation; here it evaluates the current (seed) orientations, returning one scalar score per rotation.

Tilt-segment-union beam coupling: per-tilt-chunk beam sets across a rocking curve.

Rocking-curve integration solves the crystal at many slightly-tilted sub-orientations. A single beam set for the whole curve either over-couples (every beam any tilt needs, slow) or drops beams a tilt needs (a reflection drifts through the Ewald sphere as the crystal rocks). This coupling policy instead partitions the tilts into contiguous chunks and couples, within each chunk, the union of the excited-beam sets at the chunk’s two boundary tilts.

This module is the pure geometry of that partition: given a UnionCoupling policy and the rotation’s tilt geometry, it returns the ordered Segment list (each a beam set + the disjoint tilt indices it covers). It computes the excitation mask (|Sg| < sg_max and |g| < g_max) from which the union sets are built, reusing excitation_errors() and orientation_basis(). It does not build or solve anything – an engine step turns the segments into per-chunk plans and reassembles their curves before reduction.

class diffBloch.preprocess.coupling.Segment(union_hkl: NDArray[int64], covered_tilt_indices: tuple[int, ...])

Bases: object

One tilt chunk’s coupling: the beam set it solves and the tilt indices it covers.

union_hkl (n_beams, 3) is the union of the excited beams at the chunk’s boundary tilts (includes (0, 0, 0), always excited). covered_tilt_indices are the global rocking-curve tilt indices this segment is responsible for – contiguous and, across a rotation’s segments, disjoint and covering every tilt exactly once. The segment solves its beam set at each covered tilt; the per-tilt intensities are later scattered back onto each reflection’s full rocking curve.

union_hkl: NDArray[int64]

covered_tilt_indices: tuple[int, ...]

diffBloch.preprocess.coupling.build_coupling_segments(policy: UnionCoupling | PerTiltCoupling, candidate_beam_hkl: NDArray[int64], *, cell: NDArray[float64], orientation: NDArray[float64], tilts: NDArray[float64], energy: float, u0: float) → tuple[Segment, ...]

Partition the rocking curve into boundary-union coupled segments (pure geometry).

candidate_beam_hkl (G, 3) is the beam candidate pool to select from (the shared StructureFactorGrid structure_factor_hkl – radius 2 * g_max, so the |g| < g_max mask filters it to each tilt’s excited set). cell (3, 3) is the real-space basis, orientation (3, 3) the rotation’s crystal orientation, tilts (B, 3, 3) the rocking-curve tilt matrices (each left-multiplying orientation). energy (eV) and u0 (mean-inner-potential correction) set the Ewald geometry.

Each boundary tilt’s excited mask is |Sg| < sg_max and |g| < g_max with g = candidate_beam_hkl @ orientation_basis(cell, tilt @ orientation). Segment i couples the union of the masks at boundary tilts i and i + 1 and covers the half-open tilt range between them; the final segment includes the end, so the covers tile 0 .. B - 1 exactly.

Steps

select_beams: prune each orientation’s beams to its active set (Klar et al. 2023 filter).

A Plan -> Plan step that re-picks every OrientationPlan’s active beam_hkl using the relative-/minimum-excitation-error criterion of the SI of Klar et al. (2023), then rebuilds its BeamPlan + AlignmentPlan against the shared grid (pattern unchanged). This is the per-orientation selection that replaces the orientation-independent g_max seed laid down by from_experiment.

sg_max is the excitation-error span a reflection sweeps during the actual integration, so its transverse lever arm is set by the tilt geometry (BeamSelection.geometry), which must match the integrator’s (RockingCurve). For continuous_rotation the crystal rocks about the goniometer axis (x in the PETS frame; rocking_curve_tilts builds R_x), so the swept excitation error has amplitude |(g_y, g_z)| – the distance from the rock axis – and a reflection on that axis (g_y = g_z = 0) never sweeps and is correctly dropped. For precession (an isotropic cone about the -z beam) the lever arm is instead |(g_x, g_y)|, the distance from the beam. The frame convention has the beam along -z and the rock axis along x.

diffBloch.preprocess.steps.beams.build_orientation_plans(rocking: RockingCurve | None = None, mosaicity: MosaicSmoothed | None = None, *, coupling: UnionCoupling | PerTiltCoupling | None = None, scoring_selection: BeamSelection | None = None, workers: int = 1) → PlanStep

Build each candidate’s final tilted Bloch geometry and intensity reduction.

The single build boundary of the preprocess pipeline: it materialises each orientation’s structure-factor gather (the dominant cost) over its beam set via OrientationPlan.build, and the rebuilt AlignmentPlan re-bridges the simulator output to the observed pattern. A custom pipeline may compose it after select_beams(); the default coupled path instead derives the SOLVE beams directly from g_max/sg_max and the explicit sub-tilts. The engine consumes only these built plans; a CandidatePlan has no beam_plans and is unsolvable by construction.

When rocking is supplied, the builder directly creates its complete sub-tilt geometry instead of first building a temporary central-orientation plan and rebuilding it later. mosaicity selects the reduction applied to those sub-tilt intensities and therefore requires rocking. When coupling is supplied, each segment’s beam set is selected from the full support grid by |g| < g_max and |Sg| < sg_max at its boundary tilts, then the ordinary alignment intersects the resulting simulator HKLs with the PETS experimental data. scoring_selection optionally applies the former Klar rsg/dsg/semiangle filter to the candidate scoring pool before that intersection; it does not alter the coupled SOLVE beams. workers fans independent rotation builds over threads while preserving input order. It is execution-only and therefore intentionally absent from the step’s provenance record. Omitting coupling preserves the simple builder used by focused APIs/tests.

diffBloch.preprocess.steps.beams.klar_beam_mask(g: NDArray[float64], *, energy: float, u0: float = 0.0, rsg: float, dsg: float, semiangle: float, geometry: Literal['continuous_rotation', 'precession'] = 'continuous_rotation') → NDArray[bool]

Boolean keep-mask for reflections g (N, 3) under the Klar (2023) rsg/dsg filter.

Each reflection’s excitation error |Sg| (Spence & Zuo, via excitation_errors(); beam along -z) is compared against sg_max, the excitation-error span it sweeps during integration: sg_max = |g_lever| * deg2rad(semiangle). The lever arm depends on geometry – for continuous_rotation the rock is about the goniometer x axis, so g_lever = (g_y, g_z) (distance from the rock axis); for precession (cone about the beam) it is g_lever = (g_x, g_y) (distance from the -z beam). A reflection is kept when both |Sg| / sg_max < rsg (relative excitation error small) and sg_max - |Sg| > dsg (a minimum absolute margin). Reflections with sg_max = 0 (on the rock axis, resp. optic axis) fail the relative test and are dropped – they never sweep through the Ewald sphere; the 000-beam retention required by the Bloch system is handled by the caller.

diffBloch.preprocess.steps.beams.select_beams(selection: BeamSelection) → PlanStep

Return a Plan -> Plan step that prunes each candidate to its Klar active beam set.

A source-level prune on the CandidatePlan phase: for every orientation the candidate beam_hkl is re-selected by klar_beam_mask() against that orientation’s lab-frame g (derived from its stored orientation and the grid cell), keeping only the active set. No geometry is built here – the structure-factor gather is built later, over the pruned set, by build_orientation_plans(). selection is a pre-validated BeamSelection (rsg relative excitation-error cutoff, dsg minimum margin, integration_semiangle in degrees); invalid cutoffs are unrepresentable, so this step never re-validates. The observed pattern is untouched.

The 000 transmitted beam is retained whenever present (the from_experiment seed always includes it): BeamPlan anchors psi0 on hkl == 000, and 000 has g = 0 so its sg_max = 0 would otherwise reject it. Beams stay within the seed radius (blochwave.g_max), so the Fgb difference support remains valid once build_orientation_plans runs.

optimize_orientation: per-rotation crystal-orientation refinement.

A Plan -> Plan step that sharpens each orientation by minimising the scaling-optimised wR2 of the full dynamical simulation against the observed pattern – the objective exposed by score_orientation(). A local scipy.optimize.minimize(method="Nelder-Mead") simplex search over the three goniometer-correction angles directly, seeded from a fixed initial simplex of edge length search.step_size around (alpha, beta, omega) = (0, 0, 0) (see _refine_one()).

The captured refinement is read-only context the step never mutates; the simulation inside is deterministic and depends only on its inputs, so it is ordinary computation, not a side effect.

Without coupling=, the SCORED reflection set is held fixed at each orientation’s seed selection across the search (with_orientation() only recomputes the orientation-dependent bases). The rocking-curve tilt set carried by the Plan is threaded through every trial unchanged, so each candidate is scored under the same integration as the seed – the optimize/eval consistency invariant. Ordering integrate_rocking_curve before this step therefore couples the search to the integrated model; with rocking off the tilt set is a single identity, identical to a static search.

Coupling (opt-in, ``coupling=TrialCoupling(…)``) re-derives the excitation-selected SOLVE beams and re-selects the SCORED set at every trial’s own orientation (coupling.scored’s Klar rsg/dsg window + resolution cap, reapplied to that trial’s fresh lab-frame geometry) – mirroring the reference implementation’s per-trial filter_hkls, rather than pinning scoring to the seed’s selection. A trial’s matched-reflection count can therefore differ from the seed’s and from other trials’; the default wR2 formula itself has no reflection-count term (diffBloch.core.losses.w_rbragg(): plain sum/sum), so the search can in principle prefer a trial that improves partly by matching a different, easier subset – search.penalize_fewer_reflections (NelderMeadSearch, off by default) guards against exactly this by dividing the comparison score by the matched count (_comparable_score()); the reported score stays the plain metric regardless. The seed is rebuilt through the same builder, and the last accepted trial is already the coupled-at-optimized-orientation plan – no separate couple_beams step is needed. Atomic F_gb values are cached lazily by support-grid row: each trial computes only previously unseen beam differences, while every segment’s structure matrix remains a cheap gather-index into that cache.

With coupling=None (the default) each trial is current.with_orientation(...), defined on both the tilt-independent OrientationPlan and the CoupledOrientationPlan, so an already-segmented plan is optimized under its frozen union.

diffBloch.preprocess.steps.optimize_orientation.optimize_orientation(refinement: RefinementSetup, search: NelderMeadSearch, *, method: SolverMethod = 'matrix_exp', coupling: TrialCoupling | None = None, validate: bool = True, workers: int = 1, device: Device | None = None, max_batch: int | None = None, logger: Logger = NULL_LOGGER, absorption: Absorption = NO_ABSORPTION, scores: ScoresFn = wr2_scores, residual: str = 'wr2') → PlanStep

Return a Plan -> Plan step refining each orientation by orientation search.

residual (default "wr2") is the display name for scores – pass cfg.loss_metrics.residual alongside scores=cfg.loss_metrics.to_scores() so OrientationOptimized reports the score under its real name.

scores (default wr2_scores()) is the per-thickness metric score_orientation() searches – pass cfg.loss_metrics.to_scores() to search the same residual the gradient refinement stage minimises (to_scores()). Execution-only like method: it changes what the search optimizes for, not the recipe’s own identity (the resolved ExperimentConfig.loss_metrics already rides in dataset_config_digest()).

refinement (constraint spec, ASU expansion, atomic numbers, seeded params) is captured read-only and rejoined to the geometry Plan via build_engine(); the orientation-invariant F_gb is computed once and reused across every orientation and trial. search is a pre-validated NelderMeadSearch (invalid bounds are unrepresentable, so this function never re-validates). method configures the engine’s solver (score_orientation scores with scores, a scaling-optimised wR2 by default).

coupling (default None) opts the optimization into per-trial re-coupling: a TrialCoupling re-derives the solve union and re-selects the scored set at every trial orientation (see the module docstring for the non-stationary-objective nuance). None keeps the tilt-independent search (one fixed beam set across the search).

validate (default True) forwards to the per-trial coupled gather rebuild (build_structure_factor_gather()). False skips its O(N^2) integrity checks – the dominant per-trial cost over a large coupled union – for the large-cell fast path. It is sound only because the coupled coverage guard (above) proves the grid spans the beam-difference support; without that guarantee a skipped check would let a gather silently read zeros. Inert unless coupling is set (the tilt-independent path rebuilds no gather in the search), and, like the coverage guard, it does not enter the recipe identity: the checks are pure, so False yields identical gather indices when coverage holds.

device (default None = CPU) places the search’s forward solve on the given accelerator: the seed params are moved there and engine.fgb is computed on-device, so every per-trial score_orientation co-locates onto the param-derived fgb.device at the use site (the CPU trial rebuilds are cheap numpy; only their tensors reach the device). Kept out of the recipe identity like workers/logger – but unlike those it is not bit-exact: the solve shifts ~1e-11 cross-device, and because the greedy search accepts on a threshold, that shift can flip a near-tie into a full-radius orientation difference (a well-conditioned optimization stays; a knife-edge one legitimately diverges). Safe regardless: reproducibility is anchored at the checkpoint boundary, so a committed CPU checkpoint is reused (not recomputed) on GPU, cannot restale – only a fresh GPU-computed checkpoint would differ from a CPU one.

workers (default 1, sequential) fans the per-rotation searches over a thread pool. Rotations are independent, the engine and F_gb are read-only shared context, results keep plan order, and each rotation’s gather cache is thread-local – so the results are identical to a sequential run. Threads (not processes) suffice because torch’s CPU linalg releases the GIL.

max_batch (default None) caps the matrix_exp propagator block; None lets each solve derive a memory-safe block from its beam count. Execution-only and matches the unbounded solve to machine precision (memory only), like device – raise it to fill a larger GPU. See build_engine().

logger receives an OrientationOptimized per rotation as its search completes (the optimization is the run’s long phase, so this is the progress stream); the default NULL_LOGGER discards them. With workers > 1 events arrive in completion order.

The greedy search restarts at the same radius on every accepting (improving) pass, so the radius schedule alone does not bound the pass count. Mirroring iterate_until(), search.max_iterations caps the total passes per orientation and a RuntimeError is raised if it is reached – silent non-convergence is never returned.

The cap is a runaway guard: the search terminates by construction for a non-degenerate objective (monotone wR2 descent + the radius floor), so the cap only guards pathological ridge-walking on (near-)degenerate landscapes. Its default of 2000 is calibrated on the quartz anchor under the integrated recipe (slowest legitimate search: 1288 passes across 99 rotations, so 2000 has headroom); raise it via config if a dataset with shallower minima trips it.

optimize_thickness: per-rotation specimen-thickness calibration by grid search.

A Plan -> Plan step that replaces each rotation’s thickness with the value that best matches its observed pattern. The specimen’s 3D shape is irregular, so each orientation presents a different beam path length; this optimizes that length per rotation rather than assuming one shared thickness.

For each orientation it evaluates n_steps candidate thicknesses spaced evenly from min_thickness to max_thickness and keeps the candidate with the lowest scaling-optimised weighted R-factor (wR2). All candidates are simulated in a single forward Bloch pass: the expensive eigendecomposition depends only on the orientation and the structure factors, while thickness enters only the cheap propagation tail, so scoring 100 thicknesses costs barely more than scoring one (score_orientation_per_thickness()).

The captured refinement is read-only context the step never mutates; the simulation inside is deterministic and depends only on its inputs, so it is ordinary computation, not a side effect.

The search is an evenly-spaced (np.linspace) grid of candidate thicknesses, per-candidate wR2 via the scaling factor, then the per-rotation minimum.

Plan-agnostic: replace(op, thickness=...) swaps the thickness on either an OrientationPlan or a CoupledOrientationPlan (whose _solve_segmented reads the top-level thickness, ignoring the stale sub-plan copies), so a coupled plan is optimized unchanged.

diffBloch.preprocess.steps.optimize_thickness.optimize_thickness(refinement: RefinementSetup, grid: ThicknessGrid, *, method: SolverMethod = 'matrix_exp', device: Device | None = None, max_batch: int | None = None, logger: Logger = NULL_LOGGER, absorption: Absorption = NO_ABSORPTION, scores: ScoresFn = wr2_scores, residual: str = 'wr2') → PlanStep

Return a Plan -> Plan step optimizing each rotation’s thickness by grid search.

scores (default wr2_scores()) is the per-thickness metric the grid search argmins over – pass cfg.loss_metrics.to_scores() to search the same residual the gradient refinement stage minimises (to_scores()). Execution-only like method: the resolved ExperimentConfig.loss_metrics already rides in dataset_config_digest(). residual (default "wr2") is the display name for scores – pass cfg.loss_metrics.residual alongside it so ThicknessOptimized reports the score under its real name.

refinement (constraint spec, ASU expansion, atomic numbers, seeded params) is captured read-only and rejoined to the geometry Plan via build_engine(); the orientation-invariant F_gb is computed once and reused across every orientation. Each rotation is then assigned the lowest-wR2 of grid.n_steps candidate thicknesses spaced evenly from grid.min_thickness to grid.max_thickness (inclusive, Angstroms). grid is a pre-validated ThicknessGrid (invalid bounds are unrepresentable, so this function never re-validates); method configures the engine’s solver.

device (default None = CPU) places the grid search’s forward solve on the given accelerator by moving the seed params there; the engine co-locates every invariant onto the param device at the use site. Execution-only (kept out of the recipe identity), exactly as in optimize_orientation().

max_batch (default None) caps the matrix_exp propagator block. None lets each solve derive a memory-safe block from its beam count – it matters most here because the grid search evaluates grid.n_steps thicknesses at once, so a wide coupled segment’s (C, T, N, N) propagator can be tens of GiB if left unbounded. Raise it to fill a larger GPU. The bound matches the unbounded solve to machine precision (memory only) and is execution-only, like device.

logger (default the null sink) receives a ThicknessOptimized per rotation as its grid search completes – the progress stream for this phase (mirroring optimize_orientation); the memory-heavy thickness search is otherwise silent under a console logger.

integrate_rocking_curve: bake each rotation’s rocking-curve tilt set into the geometry.

A Plan -> Plan step that replaces each rotation’s single static geometry with N tilted sub-orientations spanning the integration semi-angle – the forward model then sums |psi|^2 over the tilts (an incoherent rotation-frame integration; see diffBloch.core.products.BlochSolution.integrate()). The rocking curve is the scientific enabling structure of the rotation-electron-diffraction forward model, so it is a composable, toggleable step rather than baked into from_experiment: composing it in with rocking.sampling == 1 (a single angle-0 tilt) is the identity, so leaving it out leaves the Plan unchanged.

It is pure geometry – no engine, no structure factors, no refinement: the tilts depend only on the fixed RockingCurve and each settled nominal orientation, so they are prebuilt into the Plan exactly like the per-orientation beam plans. Ordered last in the pipeline (after select_beams / optimize_orientation / optimize_thickness, which score on the fast single-solve): the fits settle the nominal orientation and the one shared beam set, then this bakes the integration geometry those results are held fixed at, reusing that beam set across every tilt.

Building the tilt matrices and integrating their intensities is exposed here as one composable unit, rather than being wired in unconditionally.

diffBloch.preprocess.steps.rocking_curve.integrate_rocking_curve(rocking: RockingCurve) → PlanStep

Return a Plan -> Plan step baking each rotation’s rocking-curve tilt geometry.

rocking is a pre-validated RockingCurve (invalid bounds are unrepresentable, so this never re-validates): rocking.sampling tilts about the goniometer axis spanning +/- rocking.integration.semiangle degrees (rocking.integration.geometry selects the sweep). The tilt matrices are orientation-independent, so they are generated once and left-multiplied onto every rotation’s nominal orientation (R_tilt @ orientation); each rotation is rebuilt with its N sub-orientations sharing its one existing beam set.

Convergence testing: grow a simulation-accuracy knob until the diffraction pattern stops moving.

A convergence sweep is self-referential – unlike optimize_orientation / optimize_thickness (which match the simulation to observed data), it asks whether two consecutive simulations still differ, so it is a numerical resolution study (has the calculation stopped depending on the knob?), run before and orthogonally to the accuracy fit. This module has three layers:

• simulation_rfactor() – the measurement: (previous, current) -> float, the mean per-orientation R-factor between two Plans’ simulations (0 when they are identical).

• converge_scalar() – the parameter-agnostic driver: given a build(value) -> object and a measure it clicks a scalar knob upward until two consecutive builds stop changing (the R-factor drops below threshold), or a hard cap raises. It knows nothing about beams (or even Plans); adapters instantiate it.

• converge_beams() – the beam-window adapter: a Plan -> Plan step that widens integration_semiangle until the pattern stabilises, re-running select_beams from the seed.

• converge_sampling() – the forward-model lever: refines the rocking-curve tilt count (rocking_curve_sampling) until the integrated pattern stabilises. Independent of the beam levers (the tilt count does not touch the Fgb support), so it composes as an ordinary lever.

simulation_converged() wraps simulation_rfactor() with a threshold to give the boolean ConvergenceCheck that iterate_until() drives to a fixpoint.

The convergence sweep grows each numeric knob (beam-pool radius, excitation window, tilt count) and stops the first time the per-orientation Bragg R-factor between consecutive simulations drops below r_factor_threshold – a deliberately simple stopping rule (no patience, no skip-null).

diffBloch.preprocess.steps.convergence.converge_beams(selection: BeamSelection, refinement: RefinementSetup, tolerance: ConvergenceTolerance, *, step: float, method: SolverMethod = 'matrix_exp') → PlanStep

Return a Plan -> Plan step: widen integration_semiangle until the pattern stabilises.

The window lever of beam-set convergence (the physically primary “how many near-Ewald beams”). Each candidate re-runs select_beams() from the incoming seed Plan at a wider integration_semiangle – selecting from the fixed seed each time, not from the previous (already-pruned) candidate, so widening can admit beams a narrower window dropped. The sweep starts at selection.integration.semiangle and clicks up by step (degrees) until converge_scalar() settles the pattern (first sub-threshold step wins); rsg / dsg are held fixed. step must be positive.

diffBloch.preprocess.steps.convergence.converge_sampling(rocking: RockingCurve, refinement: RefinementSetup, tolerance: ConvergenceTolerance, *, step: float, method: SolverMethod = 'matrix_exp') → PlanStep

Return a Plan -> Plan step: refine the rocking-curve tilt count until the pattern stops.

The forward-model convergence lever (independent of the two beam levers): each candidate bakes the rocking-curve integration geometry at a finer rocking_curve_sampling (the tilt count) via integrate_rocking_curve(), so the summed |psi|^2 over tilts approaches the continuous rotation-frame integral. The sweep starts at rocking.sampling and clicks up by step (rounded to a whole tilt count) until converge_scalar() settles the pattern (first sub-threshold step wins): it settles when a finer tilt grid stops moving the integrated intensities. Only sampling is swept – the tilt span (rocking.integration.semiangle) and rocking.integration.geometry are held fixed. Re-integrating from the incoming seed each step (integrate_rocking_curve rebuilds tilts from each nominal orientation, discarding any prior tilts) makes the sweep independent of the seed’s tilt state.

step must be positive. Unlike the beam levers this needs no grid guard (the tilt count does not touch the Fgb support) and does not couple to them, so it composes as an ordinary extra lever.

diffBloch.preprocess.steps.convergence.converge_scalar(build: Callable[[float], T], measure: Callable[[T, T], float], tolerance: ConvergenceTolerance, *, start: float, step: float, accept_converged_candidate: bool = True) → T

Grow a scalar knob until two consecutive builds stop changing; return the converged object.

The parameter-agnostic convergence driver – it knows nothing about beams or Plans. build(value) rebuilds the object at a knob value; measure(previous, candidate) is the consecutive-output R-factor (0 when identical). Starting from start and clicking by step each iteration, it stops at the first candidate whose R-factor against the previous build is below tolerance.r_factor_threshold. By default it returns that candidate; accept_converged_candidate=False retains the previous value instead. The stopping rule is deliberately simple – the first dip stops the sweep, and an unchanged build (R = 0) counts as converged – with no patience and no null-step handling. Raises RuntimeError if tolerance.max_iterations steps pass without a dip below threshold (silent non-convergence is never returned, matching iterate_until()).

diffBloch.preprocess.steps.convergence.simulation_converged(refinement: RefinementSetup, tolerance: ConvergenceTolerance, *, method: SolverMethod = 'matrix_exp') → ConvergenceCheck

Return a (previous, current) -> bool check: have consecutive simulations stabilised?

Thin threshold wrapper over simulation_rfactor(): the mean per-orientation R-factor is compared against tolerance.r_factor_threshold. This is the boolean ConvergenceCheck that iterate_until() drives to a fixpoint (the cross-lever composition); converge_scalar() uses the underlying float measure directly, applying the same threshold inline.

diffBloch.preprocess.steps.convergence.simulation_rfactor(refinement: RefinementSetup, *, method: SolverMethod = 'matrix_exp', comparison_hkl: tuple[Tensor, ...] | None = None) → SimulationRfactor

Return (previous, current) -> float: the mean consecutive-simulation R-factor.

refinement (the read-only structure context) is captured and rejoined to each Plan via build_engine(); method configures the solver. The returned measure simulates both Plans, computes the scale-optimised rbragg R-factor between them on each orientation’s shared reflections, and averages over orientations. The comparison is a control-flow decision, not a gradient path, so the simulated intensities are detached. It is 0 exactly when the two Plans produce identical simulations (the null-step signal the sweep skips).

The two Plans must describe the same orientations in the same order (a convergence step rebuilds each orientation, changing only its beam set), and each pair must share at least one reflection (the retained 000 guarantees this in practice).

Coverage sweep: grow a beam knob to the minimum that recovers the most matched reflections.

The second convergence operation (the first is self-stability in convergence.py). Where a converge_* sweep asks whether two consecutive simulations agree, a coverage sweep asks a cheaper, purely geometric question: how many of each orientation’s observed reflections does the current beam set actually include? Growing a beam knob admits more beams; a candidate is accepted only while it increases that matched count, and the sweep stops at the first knob step that buys no new match – the minimal beam set that still covers the data.

• plan_coverage() – the objective: sum over orientations |beam_hkl & observed hkl|. It is a pure function of the Plan (no engine, no structure factors): a “match” is set membership (an experimental hkl present in the filtered simulated beam set, with no intensity gate).

• maximize_scalar() – the parameter-agnostic driver: click a scalar knob upward, keep the build while the objective strictly increases, return the last build at the first non-increase, or raise at a hard cap. The match-count analogue of converge_scalar().

• cover_beams() – the Plan -> Plan adapter for the Klar window (integration_semiangle) lever.

The two convergence operations are two kinds of objective: coverage maximises observed matches (sequential per-knob sweeps, accepting a candidate only while the match count increases, capped), while self-stability (convergence.py) settles simulations.

diffBloch.preprocess.steps.coverage.cover_beams(selection: BeamSelection, *, step: float, max_iterations: int = 100) → PlanStep

Return a Plan -> Plan step: widen the Klar window to the minimum that maximises coverage.

The window (integration_semiangle) lever of the coverage sweep: each candidate re-runs select_beams() from the incoming seed at a wider window, and maximize_scalar() keeps widening while plan_coverage() strictly increases, stopping at the first window that admits no new matched reflection. step must be positive.

diffBloch.preprocess.steps.coverage.maximize_scalar(build: Callable[[float], T], objective: Callable[[T], float], *, start: float, step: float, max_iterations: int = 100) → T

Grow a scalar knob while objective strictly increases; return the last accepted build.

The parameter-agnostic coverage driver – it knows nothing about beams or Plans. build(value) rebuilds the object at a knob value; objective(obj) is the score to maximise (for coverage, plan_coverage()). Starting from start and clicking by step, it keeps a candidate while its score is strictly greater than the best so far and returns the best at the first step that does not improve it (accept while candidate > best, else stop). Raises RuntimeError if max_iterations steps pass while the score is still increasing (the score never plateaus). max_iterations must be >= 1.

diffBloch.preprocess.steps.coverage.plan_coverage(plan: Plan) → int

Count matched reflections: sum over orientations |beam_hkl intersect observed hkl|.

A pure, engine-free measure of how much of the observed data the current beam set covers. A “match” is an observed reflection whose hkl is present in that orientation’s active beam set (set membership, no intensity threshold).

couple_beams: choose the rocking curve’s beam-coupling policy across its tilts.

A Plan -> Plan step selecting how each rotation couples beams over its rocking-curve tilts, a discriminated union rather than a boolean toggle:

• TiltIndependent – the default: one beam set (the select_beams active set) shared across every tilt. A no-op here (the shared set is already on the plan).

• UnionCoupling – a policy that partitions the tilts into contiguous chunks and gives each its own boundary-union beam set (build_coupling_segments()), replacing each OrientationPlan with a CoupledOrientationPlan the engine reassembles + reduces.

It is the tilt-dependent generalization of select_beams. The default app recipe does not use this step – it couples per trial inside optimize_orientation (coupling=...); couple_beams is the explicit composable step when a caller wants to settle a coupled Plan directly. It is only meaningful once select_beams has established the Klar-selected scored set (before it, an orientation’s alignment is the seed-pool intersection, too wide). It accepts either a plain OrientationPlan (the first application) or an already-coupled CoupledOrientationPlan (the re-couple), re-deriving each rotation’s segments from its current source fields (orientation / tilts / energy / u0). The candidate pool is the shared grid (plan.structure_factor_grid.structure_factor_hkl, which spans the coupling cap), and the tilt set is the one integrate_rocking_curve already baked, so this raises if a rotation has fewer than two tilts. The upstream tilt_reduction (e.g. a mosaicity broadening) is carried through unchanged.

Crucially it decouples the two reflection sets the plan otherwise conflates: the solve set expands to the excitation coupling union, while the scored set stays pinned to the pre-couple select_beams selection (op.alignment.hkl), intersected with the union.

diffBloch.preprocess.steps.coupling.couple_beams(policy: TiltIndependent | UnionCoupling | PerTiltCoupling) → PlanStep

Return a Plan -> Plan step applying the policy beam-coupling to every rotation.

policy is a CouplingPolicy selected by construction: TiltIndependent yields the identity (the shared beam set is kept), and UnionCoupling replaces each rotation with its per-chunk CoupledOrientationPlan. Pre-validated, so this never re-validates its bounds.

Orchestration and terminals

Simulation-convergence testing over g_max, sg_max, and tilt steps.

class diffBloch.preprocess.driver.ConvergenceState(g_max: float, sg_max: float, tilt_steps: int)

Bases: object

The three numerical controls varied by a convergence test.

g_max: float

sg_max: float

tilt_steps: int

diffBloch.preprocess.driver.converge_numerics(test: ConvergenceTest, rocking: RockingCurve, simulation: UnionCoupling | PerTiltCoupling, refinement: RefinementSetup, tolerance: ConvergenceTolerance, *, method: SolverMethod = 'matrix_exp', logger: Logger = NULL_LOGGER) → PlanStep

Return a step that converges g_max, sg_max, and rocking-curve tilt steps.

diffBloch.preprocess.driver.run_convergence(plan: Plan, state: ConvergenceState, test: ConvergenceTest, rocking: RockingCurve, simulation: UnionCoupling | PerTiltCoupling, refinement: RefinementSetup, tolerance: ConvergenceTolerance, *, method: SolverMethod = 'matrix_exp', logger: Logger = NULL_LOGGER) → tuple[Plan, ConvergenceState]

Run coordinate sweeps until all three simulation controls are self-stable.

The inference terminal: run the forward model over every rotation and score it, no refinement.

run_inference is the eval-only member of the preprocess pipeline’s terminal family (the other is engine.refine, which optimizes structure): it runs one forward Bloch pass per orientation under no_grad and reports a per-rotation RotationInference (the Bragg R-factor R_obs and two diagnostics).

Built entirely from the public forward spine – engine.simulate + core.products.align() + core.losses.rbragg()/core.losses.optimal_scale() – so callers never reach into engine internals. Preprocess is composed in optionally via the preprocess PlanStep; the solver is swappable via method.

class diffBloch.preprocess.inference.InferenceResult(per_rotation: tuple[RotationInference, ...])

Bases: object

Per-rotation forward-inference metrics for a whole Plan.

per_rotation: tuple[RotationInference, ...]

property n_evaluated: int

Rotations with a finite r_obs (i.e. at least one I > 3*sigma reflection).

property mean_wr2: float

Mean weighted-R2 over rotations with a finite value; nan when none has one.

The companion to mean_r_obs, filtered independently: a rotation can produce a finite score under one metric and not the other, so the two means need not share a denominator.

property mean_r_obs: float

Mean R_obs over the finite rotations.

nan when no rotation has a finite r_obs. The per-rotation R_obs values are averaged, skipping rotations with no reflections.

class diffBloch.preprocess.inference.RotationInference(r_obs: float, wr2: float, n_observed: int, n_beams: int)

Bases: object

One rotation’s forward-inference metrics.

r_obs is the scaling-optimised Bragg R-factor of calculated vs observed intensities over the reflections with I > 3*sigma (core.losses.rbragg); it is nan when no reflection passes that cut. n_observed counts those reflections and n_beams the active beam set – both diagnostics for why an r_obs is what it is.

r_obs: float

wr2: float

n_observed: int

n_beams: int

diffBloch.preprocess.inference.run_inference(plan: Plan, refinement: RefinementSetup, *, prepare: PlanStep = identity, method: SolverMethod = 'matrix_exp', device: Device | None = None, max_batch: int | None = None, absorption: Absorption = NO_ABSORPTION, logger: Logger = NULL_LOGGER) → InferenceResult

Run the forward model once per orientation and score each against its observed pattern.

First applies prepare to plan – one composed Plan -> Plan pipeline (compose the run’s steps with pipeline(), e.g. build_orientation_plans -> optimize_orientation -> optimize_thickness); it defaults to the identity (evaluate the plan as given). Then builds a RefinementEngine, simulates every orientation under no_grad with the swappable method solver, and returns per-rotation RotationInference.

Emits a RotationScored per rotation and one InferenceCompleted aggregate to logger (the NULL_LOGGER default discards them, so the returned value is unchanged whether or not a sink is attached). Attach a console/wandb logger at the boundary to watch per-rotation R_obs live – e.g. while chasing a residual.

device (default None = CPU, unchanged) runs the forward solve on the given accelerator: the seed params are moved there, and the engine co-locates every invariant onto the param device at the use site (RefinableParams.to()), so the whole eigensolve runs on-device. The scoring tail (align / optimal_scale) is device-safe (observed data is co-located there), so the returned R_obs is identical (to solver tolerance) across devices.

max_batch (default None) caps the matrix_exp propagator block on the terminal solve; None lets the engine pick a memory-safe block per beam count. Execution-only (memory), like device. See build_engine().

Checkpoint / resume

Serialize a settled Plan to a portable .npz and read it back (source persisted, compiled geometry rebuilt on load). The infer CLI checkpoints/resumes against this plus the plan.lock provenance in diffBloch.config.manifest.

Serialize a preprocessed Plan to a portable .npz checkpoint and read it back.

The persistence primitive behind the run program’s checkpoint/resume: serialize the whole ``Plan``, never per-facet state. It exploits the plan types’ own design – both OrientationPlan and CoupledOrientationPlan separate their source / rebuild inputs (orientation / tilts / thickness / beam set(s) / observed pattern / energy / u0 / tilt_reduction, plus the segmented plan’s per-chunk (union_hkl, covered_tilt_indices) and pinned scored set) from their built geometry (beam_plans incl. the heavy StructureFactorGather, alignment, the union + union_beam_index), and StructureFactorGrid.from_cell + .build rebuild the built parts from the source. So we persist only the source and rebuild the derived geometry on read – a stored gather is a pure function of the beam set + grid and could only desync.

Format: one .npz (numpy is already a core dependency; portable, non-pickle hence auditable, and it holds the ragged per-rotation/per-segment arrays natively as index-keyed entries). Scalars, structure, the per-orientation kind discriminant + tilt_reduction discriminant, and the plan’s provenance (the recipe that produced it) ride in a JSON string stored as the reserved __meta__ array, so the file is self-describing. Read uses allow_pickle=False: a checkpoint is data, never code.

This lives in preprocess (where Plan lives), not io – io is the input-record layer below core / engine / preprocess, so a Plan serializer there would invert the dependency. The plan.lock provenance that binds a checkpoint to the inputs + config + software version that produced it is a separate, data-free concern in config.manifest.

diffBloch.preprocess.serialize.read_plan(path: str | Path) → Plan

Read a .npz checkpoint from write_plan(), rebuilding the built geometry.

diffBloch.preprocess.serialize.write_plan(plan: Plan, path: str | Path) → None

Write plan to path as a portable .npz checkpoint (source arrays + JSON meta).