Engine

The refinement engine: compiled per-orientation geometry plans, the forward simulation spine, the objective, the hard molecular constraints and soft penalties composed into it, and the (deliberately quarantined) imperative optimization loop.

Scientific composition is done with typed Python values, not config: build a problem with build_refinement_problem and add hard constraints (e.g. with_hydrogen_riding) or soft penalties.

Refinement-invariant geometry plans: the shared scattering grid and per-orientation bundles.

These are the static (refinement-invariant) inputs the RefinementEngine composes. The grid is owned once by StructureFactorGrid and reused by both structure_factors and every BeamPlan, so the two sides cannot silently disagree on the Fgb support (the difference-support constraint is validated when the beam plans are built).

class diffBloch.engine.plan.OrientationPlan(orientation: Tensor, tilts: Tensor, energy: float, u0: float, thickness: Tensor, beam_hkl: Tensor, beam_plans: tuple[BeamPlan, ...], pattern: PatternBatch, alignment: AlignmentPlan, tilt_reduction: PlainSum | MosaicSmoothed = PlainSum())

Bases: object

The refinement-invariant plans for a single rotation/orientation.

Self-describing: it carries both its source / rebuild inputs (orientation, energy, u0, thickness – what preprocess steps like select_beams / optimize_orientation / optimize_thickness consume to rebuild) and the built geometry (beam_plan, alignment – what engine.simulate consumes). Source and built geometry are only ever set together by build(), so they cannot desync. orientation is the source of truth; the lab-frame basis is derived from it, never stored. tilts (N, 3, 3) is the rocking-curve integration tilt set (source): N goniometer sub-orientations, each built into the matching entry of beam_plans (N = len(beam_plans)). The default is a single identity tilt (1, 3, 3) – one static solve; a longer set is baked by integrate_rocking_curve and summed as |psi|^2 over the tilts by the engine. thickness (T,) is the specimen’s thickness for this rotation (its beam path length at this tilt), held fixed during refinement. It is seeded from the sample thickness and later replaced by the best-fitting value optimize_thickness finds. The forward model uses it for this orientation unless the caller is refining thickness directly (see _thickness_for()).

orientation: Tensor

tilts: Tensor

energy: float

u0: float

thickness: Tensor

beam_hkl: Tensor

beam_plans: tuple[BeamPlan, ...]

pattern: PatternBatch

alignment: AlignmentPlan

tilt_reduction: PlainSum | MosaicSmoothed = PlainSum()

classmethod build(grid: StructureFactorGrid, beam_hkl: NDArray[int64], pattern: PatternBatch, *, energy: float, thickness: Tensor | NDArray[float64] | Sequence[float], u0: float = 0.0, orientation: Tensor | NDArray[float64] | None = None, tilts: NDArray[float64] | None = None, tilt_reduction: PlainSum | MosaicSmoothed = PLAIN_SUM, gather: StructureFactorGather | None = None, validate: bool = True, build_alignment: bool = True) → OrientationPlan

Assemble an orientation’s plans against the shared grid (enforces grid coupling).

orientation (3, 3) is the crystal orientation matrix for this rotation; the lab-frame reciprocal cell is derived from it and the grid’s real-space cell via orientation_basis(grid.cell, orientation) = reciprocal_cell(cell @ orientation.T) and drives g -> Sg / Mii only. When None the orientation is the identity and the shared grid.reciprocal_basis is used directly (the untilted / single-orientation case), making that path identical to the unoriented build. The rotation convention is derived upstream in preprocess (see preprocess.orientation); the Fgb gather is keyed on grid.structure_factor_hkl and is unaffected.

orientation accepts either a NumPy array or a Tensor (e.g. a prior plan’s stored orientation), so a later Plan -> Plan rebuild can pass old_plan.orientation directly without ad-hoc conversion.

thickness (T,) is required: this rotation’s frozen per-rotation conditioning, coerced to a 1-D float64 tensor. A rebuild threads old_plan.thickness through unchanged; optimize_thickness bakes the single gridsearch winner (1,).

tilts (N, 3, 3) is the optional rocking-curve integration set: N goniometer rotations, each left-multiplying orientation (R_tilt @ orientation) into its own built beam_plan, sharing this orientation’s one beam set. None (the default) is a single identity tilt, so beam_plans has length 1 and the untilted path is the plain static solve; integrate_rocking_curve passes the tilt matrices from rocking_curve_tilts().

tilt_reduction selects how the engine reduces the tilt sub-solutions over the rocking curve: PlainSum (the default) sums them; MosaicSmoothed applies mosaicity broadening first. It is a rebuild-preserved attribute (geometry-independent), resolved during dataset setup.

gather may be a precomputed StructureFactorGather for this beam set against the shared grid. The F-gather is basis- and orientation-free, so all N tilts here share one, and a caller rebuilding this plan over a fixed beam set (rocking integration, orientation-search trials) passes the seed plan’s gather (op.beam_plans[0].gather) to skip re-deriving it on every rebuild – the dominant preprocess cost. When None it is built once here and shared across the tilts.

validate (default True) is forwarded to build_structure_factor_gather() when it builds the gather here – pass False only on a hot rebuild loop whose grid coverage an upstream g_max guard already guarantees. Ignored when a precomputed gather is supplied.

with_orientation(grid: StructureFactorGrid, orientation: Tensor | NDArray[float64]) → OrientationPlan

Rebuild this plan at a new orientation, reusing everything else (F-gather included).

The pure rebuild verb an orientation search needs: same beam set, tilts, thickness, reduction, pattern / alignment – only orientation changes, so only the orientation-dependent beam bases are recomputed while the orientation-free StructureFactorGather (shared across the tilts) is reused via gather=. grid is required because the plan does not own the shared support the bases derive from; the caller threads its Plan.structure_factor_grid. This makes a hexagonal-search trial one call and keeps the plan (not the fit) the source of truth for the beam set.

class diffBloch.engine.plan.StructureFactorGrid(structure_factor_hkl: Tensor, cell: Tensor, reciprocal_basis: Tensor, gpts: tuple[int, int, int], cell_volume: float, g_max: float)

Bases: object

The shared Fgb support grid, owned once and reused by structure factors and beam plans.

Three reciprocal-space radii are easy to conflate; they are distinct concerns:

• solve cutoff – the beams that couple in one Bloch solve (|g| <= solve_g_max; the coupled fit’s UnionCoupling.g_max).

• structure-factor support – the Fgb grid this class holds. It must cover every beam difference g_j - g_i, which reaches 2 * solve_g_max, so the support radius is ~2x the solve cutoff. from_cell_for_beam_cutoff() derives it from the solve cutoff.

• scored cutoff – ScoredHklSelection.g_max, the reflections compared in the objective. It selects what enters the loss, not what solves. The default app recipe reuses the solve cutoff for it (see _trial_coupling); a custom composition may set a separate radius.

structure_factor_hkl (G, 3) are the Miller indices Fgb is tabulated on (|g| <= g_max); cell (3, 3) is the real-space basis and reciprocal_basis (3, 3) / cell_volume the metric it derives (kept together; only from_cell() constructs them, so they cannot desync). gpts is the ravel box. g_max is the structure-factor support radius and must span the beam difference support or beam-plan construction raises.

structure_factor_hkl: Tensor

cell: Tensor

reciprocal_basis: Tensor

gpts: tuple[int, int, int]

cell_volume: float

g_max: float

classmethod from_cell(cell: NDArray[float64], g_max: float) → StructureFactorGrid

Build the grid from a real-space cell (3, 3) and a structure-factor g_max.

g_max here is the structure-factor support radius – the grid must already span the beam difference support. Callers who know their solve cutoff rather than the support radius should use from_cell_for_beam_cutoff(), which derives the support for them.

classmethod from_cell_for_beam_cutoff(cell: NDArray[float64], solve_g_max: float) → StructureFactorGrid

Build the grid from the solve cutoff – the radius of the beams in one Bloch solve.

A beam set bounded by |g| <= solve_g_max produces dynamical-matrix terms F(g_j - g_i) whose differences reach |g_j - g_i| <= 2 * solve_g_max (triangle inequality), so the structure-factor support must cover 2 * solve_g_max.

The 2x is fundamental, not an implementation artifact: any coupled Bloch solve that gathers F(g - h) from a tabulated grid needs the table out to twice the beam cutoff, so the factor cannot be designed away. The API is arranged so the caller declares the physical solve cutoff and the 2x support is derived here, rather than declaring the doubled radius and halving it internally (self.g_max = sf_g_max / 2) – an error-prone double-entry. solve_g_max is the beam/coupling cutoff, distinct from the scoring-resolution cutoff (ScoredHklSelection.g_max), which selects reflections for the objective, not the solve.

The support radius is 2 * solve_g_max + _SUPPORT_MARGIN. The half-Angstrom headroom reconciles the coupling filter’s orientation metric (u_matrix is not guaranteed exactly orthonormal – see preprocess.orientation) with this reciprocal_cell metric of the same authoritative cell, so a coupled beam difference near 2 * solve_g_max in the former still lands inside the grid in the latter. The shell only enlarges the (unused-at-the-margin) SF table; it changes neither the coupled beam set nor the scored set.

class diffBloch.engine.plan.SegmentPlan(plan: OrientationPlan, cover: Tensor, union_beam_index: Tensor)

Bases: object

One coupled tilt-chunk of a CoupledOrientationPlan: a sub-plan + reassembly map.

plan is an ordinary OrientationPlan over the segment’s own (smaller) beam set, solved at just the tilts this chunk covers (plan.tilts are the covered tilt matrices, so len(plan.beam_plans) == len(cover)). cover (C,) are the segment’s global rocking-curve tilt indices (contiguous; disjoint across a rotation’s segments; tiling every tilt once). union_beam_index (n_seg,) maps each of the segment’s beams to its column in the parent plan’s union beam set, so the segment’s per-tilt intensities scatter onto the shared rocking curve before the tilt reduction runs on the whole curve.

plan: OrientationPlan

cover: Tensor

union_beam_index: Tensor

class diffBloch.engine.plan.CoupledOrientationPlan(orientation: Tensor, tilts: Tensor, energy: float, u0: float, thickness: Tensor, beam_hkl: Tensor, segments: tuple[SegmentPlan, ...], pattern: PatternBatch, alignment: AlignmentPlan, tilt_reduction: PlainSum | MosaicSmoothed = PlainSum())

Bases: object

A rotation whose rocking curve couples a different beam set per tilt chunk.

The tilt-dependent generalization of OrientationPlan: instead of one beam set shared across all tilts, the curve is partitioned into SegmentPlan chunks, each solving its own boundary-union beam set over its covered tilts (see diffBloch.preprocess.coupling.build_coupling_segments()). The engine solves each segment and reassembles every reflection’s per-tilt intensity onto the shared union beam axis before reducing over tilts (diffBloch.engine.forward.RefinementEngine._solve()), returning an ordinary BlochSolution over that union – so align / scoring stay identical to the tilt-independent path. Reassembling before the reduction is required: the mosaicity sample span can cover more tilts than any single chunk holds.

beam_hkl (N_union, 3) is the union of every segment’s beams (deduplicated, sorted, and always including 000); pattern / alignment bridge that union to the observed reflections; tilts (N, 3, 3) is the full rocking-curve set (N the total tilt count); tilt_reduction is carried over unchanged from the orientation this was coupled from (so a mosaicity broadening set upstream still applies). orientation / energy / u0 / thickness mirror OrientationPlan as the rotation’s frozen conditioning.

orientation: Tensor

tilts: Tensor

energy: float

u0: float

thickness: Tensor

beam_hkl: Tensor

segments: tuple[SegmentPlan, ...]

pattern: PatternBatch

alignment: AlignmentPlan

tilt_reduction: PlainSum | MosaicSmoothed = PlainSum()

classmethod build(grid: StructureFactorGrid, segments: Sequence[tuple[NDArray[int64], Sequence[int]]], pattern: PatternBatch, *, energy: float, thickness: Tensor | NDArray[float64] | Sequence[float], u0: float, orientation: Tensor | NDArray[float64], tilts: NDArray[float64], tilt_reduction: PlainSum | MosaicSmoothed = PLAIN_SUM, scored_hkl: NDArray[int64] | None = None, gathers: Sequence[StructureFactorGather] | None = None, validate: bool = True) → CoupledOrientationPlan

Assemble a segmented plan from (beam_hkl, cover) chunks against the shared grid.

Each segments entry is one chunk’s beam set (n_seg, 3) and the global tilt indices it covers; tilts (N, 3, 3) is the full rocking-curve set the covers index into. The union beam set is the sorted, deduplicated concatenation of every chunk’s beams (000 is present because each chunk’s coupling always includes it); each chunk is built into an OrientationPlan over its beam set and covered tilts (sharing the rotation’s orientation / energy / u0 / thickness), and its union_index records where its beams sit in the union.

scored_hkl (S, 3) pins the scored reflection set (via build_alignment_plan’s restrict_to): the union is the enlarged solve set, but scoring stays on this set intersected with the union – the select_beams selection couple_beams hands in, so expanding the solve does not drag scoring onto the union’s weak beams. None scores the whole pattern ∩ union (the tilt-independent behaviour).

gathers optionally supplies one precomputed StructureFactorGather per segment (same order as segments), threaded into each chunk’s OrientationPlan.build() to skip re-deriving the orientation-free F-gather – the dominant cost. A rebuild at a new orientation over the same segments (with_orientation()) passes the seed plan’s per-segment gathers; None builds each fresh here.

validate (default True) is forwarded to each segment’s OrientationPlan.build() when it builds a gather (i.e. where gathers is None) – pass False only on a hot per-trial re-couple whose grid coverage an upstream g_max guard already guarantees.

with_orientation(grid: StructureFactorGrid, orientation: Tensor | NDArray[float64]) → CoupledOrientationPlan

Rebuild at a new orientation, reusing the segments’ beams, covers, and F-gathers.

The segmented counterpart of OrientationPlan.with_orientation(): the segment partition (each chunk’s beam set + covered tilts), the union, the pinned scored set (alignment.hkl, idempotent under the intersection since it is already a subset of the union), and every chunk’s StructureFactorGather are carried over; only the orientation-dependent bases recompute. This lets a fit tilt an already-coupled plan trial-by-trial at ~eigensolve cost (no re-gather, no re-coupling): the frozen-union fit. grid is threaded from the caller’s Plan.structure_factor_grid.

diffBloch.engine.plan.mean_plan_thickness(plan: Sequence[OrientationPlan | CoupledOrientationPlan]) → Tensor

Return the mean physical thickness across a settled orientation plan.

Forward composition spine: raw structure parameters -> diffraction -> scalar objective.

A RefinementEngine holds the refinement-invariant plans (constraint spec, ASU-expansion plan, the shared scattering grid, and one OrientationPlan per rotation) and maps RefinableParams to a differentiable objective:

constrain -> expand ASU -> structure_factors (Fgb on the shared grid)

-> per orientation: build_bloch_system -> propagate -> intensities -> align -> loss

objective_value / simulate are pure and differentiable; run_refinement_model() delegates to the quarantined imperative loop in diffBloch.engine.refine. Engines are assembled from explicit per-orientation beam sets that preprocessing has already selected.

class diffBloch.engine.forward.ForwardContext(thickness: Tensor | None = None)

Bases: object

Forward-model values supplied by refinement model components.

Deliberately narrow: apparent thickness is the only value admitted, since it is the only value a component currently supplies. Scale/background/damage fields should be added only when consumed.

thickness: Tensor | None = None

type diffBloch.engine.forward.LossFn = Callable[[AlignedIntensities], Tensor]

class diffBloch.engine.forward.ModelComponent(*args, **kwargs)

Bases: Protocol

A trainable/differentiable model component that can feed the forward simulation.

property key: str

Stable component parameter-tree key for validation and optimizer grouping.

initial_params(*, dtype: dtype, device: device) → Mapping[str, Tensor]

Return initial parameter tensors for this component.

forward_context(params: Mapping[str, Tensor], *, rotation_index: int, orientation: OrientationPlan | CoupledOrientationPlan) → ForwardContext

Return this component’s values for one orientation.

class diffBloch.engine.forward.RotationMetrics(rotation_index: int, wr2: float, r_obs: float, n_matched: int)

Bases: object

One rotation’s scaling-optimised wR2/R_obs for a settled model snapshot.

Report/plot use (RefinementEngine.per_rotation_metrics()), not the objective: each metric independently re-optimises its own intensity scale (diffBloch.core.losses.optimal_scale()), exactly as refinement_metrics/the training objective do, so wr2/r_obs here match what those report elsewhere. rotation_index is the original zero-based PETS rotation index.

rotation_index: int

wr2: float

r_obs: float

n_matched: int

type diffBloch.engine.forward.ScoresFn = Callable[[AlignedIntensities], Tensor]

class diffBloch.engine.forward.StructureComponent(initial: RefinableParams, constraints: tuple[ConstraintTransform, ...] = ())

Bases: object

The physical-structure component of a refinement model.

A thin wrapper around the structure refinement inputs: initial is the RefinableParams, and constraints is the tuple of hard molecular transforms applied after the crystallographic constrain. Structure-local hard parameterizations belong here.

initial: RefinableParams

constraints: tuple[ConstraintTransform, ...] = ()

class diffBloch.engine.forward.RefinementEngine(spec: ~diffBloch.params.ConstraintSpec, asu_plan: ~diffBloch.core.symmetry.AsuExpansionPlan, numbers: ~torch.Tensor, grid: ~diffBloch.engine.plan.StructureFactorGrid, orientations: tuple[~diffBloch.engine.plan.OrientationPlan | ~diffBloch.engine.plan.CoupledOrientationPlan, ...], loss: ~diffBloch.engine.forward.LossFn, scores: ~diffBloch.engine.forward.ScoresFn = <function wr2_scores>, method: ~diffBloch.core.solver.SolverMethod = 'matrix_exp', max_batch: int | None = None, absorption: ~diffBloch.specs.Absorption = Absorption(enabled=False), active_structure_factor_indices: ~torch.Tensor | None = None, profile: bool = False, checkpoint_activations: bool = True)

Bases: object

Forward from raw structure parameters to a differentiable scalar objective.

Holds the refinement-invariant context: the constraint spec, the ASU-expansion asu_plan, the ASU atomic numbers, the shared grid, the per-rotation orientations (each carrying its own frozen thickness), the per-orientation loss, and the propagation method.

spec: ConstraintSpec

asu_plan: AsuExpansionPlan

numbers: Tensor

grid: StructureFactorGrid

orientations: tuple[OrientationPlan | CoupledOrientationPlan, ...]

loss: LossFn

scores() → Tensor

Per-thickness scaling-optimised weighted-R2 (shape (T,)).

The calculated intensities come off the dynamical solve on an arbitrary structure-factor scale, while the observed are PETS intensities on their own scale. Compared raw (w_rbragg_loss()), wR2 is denominator-dominated and parks near ~1 with a vanishing gradient, so a gradient refinement cannot descend it. Every call therefore re-fits the multiplicative intensity scale independently for every thickness through optimal_scale(). The selected grid branch remains differentiable in its calculated intensities (torch.min routes the gradient through the winning candidate); only a boundary where the winning grid point changes is piecewise-smooth.

method: SolverMethod = 'matrix_exp'

max_batch: int | None = None

absorption: Absorption = Absorption(enabled=False)

active_structure_factor_indices: Tensor | None = None

profile: bool = False

checkpoint_activations: bool = True

simulate(params: RefinableParams) → tuple[BlochSolution, ...]

Return the calculated BlochSolution for every orientation (no loss).

fgb(params: RefinableParams) → Tensor

The calculated structure factors F_gb on the shared grid.

The orientation-invariant part of the forward model: compute once and reuse across orientations (e.g. when scoring many trial orientations of one structure).

structure_factor_values(params: RefinableParams, indices: Tensor) → Tensor

Calculate Fgb only for selected rows of the shared support grid.

score_orientation(orientation: OrientationPlan | CoupledOrientationPlan, fgb: Tensor) → Tensor

This engine’s configured scores for one orientation against its observed pattern.

Runs the forward Bloch simulation for orientation from a precomputed fgb (fgb()), aligns calculated vs observed intensities, and reduces by self.scores (default wR2, grid-searching the intensity scale via diffBloch.core.losses.optimal_scale()). With multiple thicknesses the best-fitting thickness’s score is returned – thickness is a nuisance when scoring orientation. This is the objective optimize_orientation minimises.

score_orientation_per_thickness(orientation: OrientationPlan | CoupledOrientationPlan, fgb: Tensor) → Tensor

This engine’s configured scores for each of the orientation’s thicknesses ((T,)).

One forward Bloch simulation from a precomputed fgb (fgb()) covers all T thicknesses at once: the expensive eigendecomposition depends only on the orientation and fgb, while thickness enters only the cheap propagation tail. The calculated and observed intensities are aligned once (alignment is thickness-independent), then reduced by self.scores (default wr2_scores()) – the same per-metric function self.loss sums for the gradient objective, so this search and the refinement stage share one ExperimentConfig.objective.

optimize_thickness grid-searches this vector and bakes the lowest-scoring thickness; score_orientation() collapses it with .min() (thickness is a nuisance there).

Runs under torch.no_grad(): this is search scoring (optimize_orientation / optimize_thickness grid search + argmin), never backpropagated – every caller consumes a detached scalar. Without it the T-thickness solve builds an autograd graph whose retained matrix_exp intermediates accumulate across every propagator block, defeating the max_batch memory bound (propagate()) and OOMing a wide coupled segment on the GPU. Grad-off does not change the returned scores.

objective_value(params: RefinableParams, penalties: tuple[PenaltyTerm, ...] = (), constraints: tuple[ConstraintTransform, ...] = ()) → ObjectiveValue

Return the objective as a scalar total plus named scalar components.

Differentiable in params. The objective is composed in a fixed order that separates hard constraints (enforced transforms) from soft penalties (additive terms):

raw RefinableParams
  -> crystallographic constraints (constrain / ConstraintSpec):
       site-symmetry position projector, ADP equalities, positivity/bounded transforms
  -> PhysicalState
  -> molecular hard constraints  (the `constraints` ConstraintTransform layer)
  -> diffraction term + soft penalties
  -> scalar objective

self.physical_state(params) applies the crystallographic constraints; each ConstraintTransform in constraints then reparameterizes that state in tuple order (duplicate names rejected), so both the diffraction term and the penalties see the transformed state. The "diffraction" component is always present; penalties add weighted soft-penalty components (bond-length, etc.) without making the optimizer know their details. Hydrogen riding (HydrogenRiding) is one such ConstraintTransform; a soft penalty and a hard constraint are distinct – a constraint reparameterizes, it is not a cost term.

objective_value_model(model: RefinementModel, *, penalties: tuple[PenaltyTerm, ...] = ()) → ObjectiveValue

Return the objective for a RefinementModel.

Structure-only models delegate to objective_value() for exact behavior parity. Models with components use the same objective order, but let components supply forward-context values such as per-orientation thickness before falling back to the built-in thickness path.

refinement_metrics(model: RefinementModel) → tuple[float, int, int, int, int]

Return mean R_obs and reflection counts for one refinement-model snapshot.

Counts are over PETS rows in the selected rotations: matched rows enter the diffraction alignment, unmatched rows do not; strong/weak split matched rows at I > 3 sigma.

per_rotation_metrics(model: RefinementModel) → tuple[RotationMetrics, ...]

Per-rotation wR2/R_obs for one refinement-model snapshot (report/plot use).

Same loop shape as refinement_metrics() (component-aware thickness, one thickness per rotation picked by the wR2 that is actually minimised during training – not by R_obs, which is reported but never optimised), returning every rotation’s pair instead of an aggregate.

physical_state(params: RefinableParams) → PhysicalState

Return bounded physical ASU quantities for params.

class diffBloch.engine.forward.RefinementModel(structure: StructureComponent, components: tuple[ModelComponent, ...] = (), component_params: Mapping[str, Mapping[str, Tensor]] = mappingproxy({}))

Bases: object

Trainable refinement model value, currently structure-only.

The model is the value optimized against a static RefinementEngine. Non-structure components provide forward-context values such as apparent thickness.

structure: StructureComponent

components: tuple[ModelComponent, ...] = ()

component_params: Mapping[str, Mapping[str, Tensor]] = mappingproxy({})

class diffBloch.engine.forward.ModelRefinementResult(model: RefinementModel, losses: Tensor, best_model: RefinementModel, best_step: int, selection_losses: Tensor | None = None, history: tuple[~diffBloch.observability.RefinementStep, ...]=(), reflection_counts: Mapping[str, int]=<factory>, artifacts: Mapping[str, str]=<factory>, objective_manifest: ObjectiveManifest | None = None)

Bases: object

The outcome of optimizing a refinement model.

Returns the optimized model value. params/best_params properties preserve the structure-only convenience used by the default app path and existing tests.

model: RefinementModel

losses: Tensor

best_model: RefinementModel

best_step: int

selection_losses: Tensor | None = None

history: tuple[RefinementStep, ...] = ()

reflection_counts: Mapping[str, int]

artifacts: Mapping[str, str]

objective_manifest: ObjectiveManifest | None = None

property params: RefinableParams

Final refined structure parameters.

property best_params: RefinableParams

Best recorded structure parameters.

property best_loss: float

The loss used to select best_step.

By default this is the training objective. When run_refinement_model is given a selection engine (the app’s held-out validation split), it is the corresponding selection objective instead.

class diffBloch.engine.forward.RefinementProblem(penalties: tuple[PenaltyTerm, ...] = ())

Bases: object

Objective-side scientific composition for one refinement run.

The trainable model value is passed separately to run_refinement_model. The problem records only additive objective terms such as bond, angle, and planarity penalties.

penalties: tuple[PenaltyTerm, ...] = ()

diffBloch.engine.forward.build_refinement_model(*, initial: RefinableParams, constraints: tuple[ConstraintTransform, ...] = (), components: tuple[ModelComponent, ...] = (), component_params: Mapping[str, Mapping[str, Tensor]] = MappingProxyType({})) → RefinementModel

Construct a RefinementModel.

With no components supplied it is behavior-equivalent to structure-only refinement; supplied components contribute forward-context values (e.g. apparent thickness) to that structure core.

diffBloch.engine.forward.build_refinement_problem(*, penalties: tuple[PenaltyTerm, ...] = ()) → RefinementProblem

Construct objective-side refinement composition data.

diffBloch.engine.forward.run_refinement_model(engine: RefinementEngine, model: RefinementModel, problem: RefinementProblem, *, trainable: TrainableSpec, steps: int, optimizer: OptimizerName = 'adam', lr: float = 1e-3, logger: Logger = NULL_LOGGER, verbose: bool = False, profile: bool = False, selection_engine: RefinementEngine | None = None) → ModelRefinementResult

Optimize a refinement model against the supplied engine/static context and problem terms.

Before the first step this emits an ObjectiveManifest naming the penalties (with weights), constraints, and components the run actually composed, and returns the same value on the result. This is the one place that holds the problem and the model together, so it is the only place that can state the objective’s composition; the structure-only run_refinement() receives a bare objective callable and therefore declares nothing.

verbose (“verbose refinement”) additionally reports one RefinementOrientationStep per rotation per step (wr2/r_obs/ diff_loss), for diagnosing which orientations drive the epoch mean reported by the ordinary RefinementStep. Off by default: the per-rotation stream is n_orientations``x louder and is a diagnosis tool, not the default reporting shape. It is execution-only, like ``logger itself – it changes what gets reported, never the objective or the optimizer trajectory, so it is a function argument / CLI flag, not an experiment.yaml field (a config field would enter the preprocess/refinement identity for no scientific reason).

profile logs per-phase wall time (structure factors, each rotation’s solve, backward, optimizer step) via stdlib diagnostics logging (logging.getLogger(__name__), level INFO, the "profile: "-prefixed lines) – see diffBloch.engine.forward._Timer. Execution-only and off by default: it forces a CUDA sync around every measured block, which is itself real overhead, so it is a diagnosis tool for one run, not something to leave on. engine.profile must also be set (build_engine() threads it through) for the structure-factor/solve breakdown; this flag alone only times backward/optimizer-step.

selection_engine is an optional held-out objective used only to choose best_model / best_step. It does not contribute gradients or alter the optimizer trajectory. The app uses it for refinement.split.train_test validation selection; without it, best selection remains the training objective. It costs one extra no-grad forward pass per step over the held-out rotations – roughly val_frac of a training forward, paid every epoch with no opt-out – and it changes which objective RefinementCompleted reports (see that event’s selection field).

Hard molecular constraints on the physical state.

A ConstraintTransform is a hard reparameterization of the bounded PhysicalState, distinct from a soft BondLengthPenalty: a constraint rewrites the state so an invariant holds exactly at every optimizer step (e.g. deriving hydrogen positions from their parents), rather than adding a cost the optimizer may trade against. Constraints are applied in the refinement objective after constrain and before the diffraction term, so the diffraction and the penalties see the transformed state.

class diffBloch.engine.constraints.ConstraintTransform(*args, **kwargs)

Bases: Protocol

A hard constraint applied to the physical state during refinement.

name identifies the transform for deterministic ordering and duplicate rejection; apply returns a new PhysicalState with the constraint enforced (no in-place mutation), so gradients flow through the reparameterization to the free parameters it depends on.

name is a read-only member: constraints are immutable value-types (frozen dataclasses), so the protocol must not demand a settable attribute (a settable one excludes frozen implementations).

property name: str

A stable identifier for deterministic ordering and duplicate rejection.

apply(state: PhysicalState) → PhysicalState

Return a new physical state with this constraint enforced.

class diffBloch.engine.constraints.HydrogenRiding(name: str, h_index: Tensor, parent_index: Tensor, offset: Tensor, u_iso_scale: Tensor)

Bases: object

Ride each hydrogen on its parent heavy atom – a hard constraint on the physical state.

Constant-offset riding: a hydrogen’s position is its parent’s position plus a fixed fractional offset taken from the input structure, and its displacement is u_iso_scale times the parent’s uij_star. Both are reparameterizations – the H rows are overwritten so they carry no independent degree of freedom, and gradients flow through to the parent – so hydrogens stay in the forward scattering while tracking the refined heavy-atom frame. Pair with TrainableSpec(positions=exclude_elements("H"), adp=exclude_elements("H")) so the H rows are never optimizer leaves.

Limitations (the model deliberately kept simple; upgrades slot into this same transform): the fixed offset does not re-point H when the parent’s local frame rotates (only translation is followed); the uij_star scale is exact for an isotropic parent (all-Uiso) and approximate for an anisotropic one.

name: str

h_index: Tensor

parent_index: Tensor

offset: Tensor

u_iso_scale: Tensor

apply(state: PhysicalState) → PhysicalState

diffBloch.engine.constraints.perceive_hydrogen_riding(structure: StructureRecord, *, cutoff_scale: float = 1.2, cutoff_margin_angstrom: float = 0.1, u_iso_scale: float = 1.2) → HydrogenRiding | None

Build a HydrogenRiding from a structure: each H rides its nearest bonded heavy atom.

For every hydrogen (atomic number 1) the parent is the nearest heavy atom within the covalent cutoff cutoff_scale * (r_H + r_heavy) + cutoff_margin_angstrom (Cartesian distance via the cell). The stored offset is the fractional parent->H vector from the input structure; the molecule is assumed ASU-contiguous, so no minimum-image wrapping is applied (matching the penalties layer).

Riding is for general-position hydrogens only: it overwrites the H coordinate after the crystallographic projector, so a special-position H would be pushed off its site-symmetry manifold. A hydrogen on a special position is therefore rejected. Returns None when the structure has no hydrogens; raises when a hydrogen has no heavy neighbour within the cutoff.

diffBloch.engine.constraints.with_hydrogen_riding(structure: StructureRecord, trainable: TrainableSpec) → tuple[TrainableSpec, tuple[ConstraintTransform, ...]]

Compose hydrogen riding onto a base trainable selection (Python/API scientific composition).

Riding derives every hydrogen from its parent heavy atom each step – position (constant parent->H offset) and Uiso (scaled from the parent) – so the hydrogens must not also be optimizer leaves. This freezes them (excludes H from both positions and adp) and perceives the HydrogenRiding constraint from the structure geometry, returning the (trainable, constraints) pair to place on StructureComponent via build_refinement_model(). This is expressed here in Python, not in config: it is scientific composition, not a stable default-path knob.

When the structure has no hydrogens the constraint tuple is empty and the freeze is harmless (no H rows to exclude), so the call is safe to apply unconditionally.

Soft refinement penalties evaluated on the bounded physical ASU state.

class diffBloch.engine.penalties.BondLengthPenalty(pairs: Tensor, target_angstrom: Tensor, sigma_angstrom: Tensor, frac_to_cart: Tensor, weight: float = 1.0, name: str = 'bond', criterion: Literal['mse', 'flat_bottom_l1'] = 'mse')

Bases: object

Bond-length soft penalties in Cartesian Angstrom units.

The current ASU positions are fractional coordinates, so the penalty owns the invariant fractional-to-Cartesian cell matrix. pairs indexes ASU atom rows; target_angstrom and sigma_angstrom carry the penalty target and tolerance for each pair. The raw loss is the mean squared normalized bond-distance residual by default. flat_bottom_l1 is a robust criterion: zero inside the sigma tolerance and linear outside.

pairs: Tensor

target_angstrom: Tensor

sigma_angstrom: Tensor

frac_to_cart: Tensor

weight: float = 1.0

name: str = 'bond'

criterion: Literal['mse', 'flat_bottom_l1'] = 'mse'

value(state: PhysicalState) → Tensor

Return the raw mean squared normalized bond-distance residual.

diffBloch.engine.penalties.perceive_bond_length_penalty(structure: StructureRecord, *, include_hydrogen: bool = False, sigma_angstrom: float = 0.02, cutoff_scale: float = 1.2, cutoff_margin_angstrom: float = 0.1, weight: float = 1.0, criterion: Literal['mse', 'flat_bottom_l1'] = 'mse') → BondLengthPenalty

Perceive ASU-contiguous bonds and tether them to the starting distances.

This is an explicit source/builder layer, separate from the pure penalty value. It assumes the bonded molecule is contiguous in the ASU and deliberately does not do minimum-image wrapping. The perceived target for each bond is the current Cartesian distance in the input structure; an external source of literature bond targets and sigmas could supply them instead.

Named LossFn/ScoresFn builders: the objective terms ExperimentConfig.objective picks.

Each adapts a pure diffBloch.core.losses intensity comparison (which reduces over the reflection axis, yielding a (T,) per-thickness loss) into two shapes: a ScoresFn (AlignedIntensities -> (T,) Tensor, one score per thickness) and the corresponding LossFn (AlignedIntensities -> scalar, the ScoresFn summed over thickness). RefinementEngine uses both off the same underlying metric – scores for the orientation/thickness preprocessing search (score_orientation/score_orientation_per_thickness, which need a per-thickness vector to argmin over) and loss for the differentiable gradient-refinement objective – so picking one ExperimentConfig.objective genuinely drives the whole pipeline, not just the gradient stage. Saves every caller from rewriting lambda a: mse(a.calculated, a.observed).sum() and gives each objective a named, importable home – e.g. RefinementEngine(loss=rbragg_loss, scores=robs_scores, ...).

diffBloch.engine.losses.l1_loss(aligned: AlignedIntensities) → Tensor

Per-orientation L1 loss term, summed over thicknesses to a scalar.

diffBloch.engine.losses.mse_loss(aligned: AlignedIntensities) → Tensor

Per-orientation MSE loss term, summed over thicknesses to a scalar.

diffBloch.engine.losses.rbragg_loss(aligned: AlignedIntensities) → Tensor

Scaling-optimised Bragg R(obs) objective. Sums robs_scores() over thickness.

diffBloch.engine.losses.robs_scores(aligned: AlignedIntensities) → Tensor

Per-thickness scaling-optimised Bragg R(obs) (shape (T,)).

Refits the intensity scale per thickness exactly like wr2_scores(), but against the rbragg() metric instead of the default w_rbragg – the same R_obs the app reports elsewhere (refinement_metrics()).

diffBloch.engine.losses.wr2_loss(aligned: AlignedIntensities) → Tensor

Scaling-optimised weighted-R2 – the default refinement and orientation-search objective.

Sums wr2_scores() over the thickness axis to a scalar; see there for the scale-fit.

diffBloch.engine.losses.wr2_scores(aligned: AlignedIntensities) → Tensor

Per-thickness scaling-optimised weighted-R2 (shape (T,)).

The calculated intensities come off the dynamical solve on an arbitrary structure-factor scale, while the observed are PETS intensities on their own scale. Compared raw (w_rbragg_loss()), wR2 is denominator-dominated and parks near ~1 with a vanishing gradient, so a gradient refinement cannot descend it. Every call therefore re-fits the multiplicative intensity scale independently for every thickness through optimal_scale(). The selected grid branch remains differentiable in its calculated intensities (torch.min routes the gradient through the winning candidate); only a boundary where the winning grid point changes is piecewise-smooth.

diffBloch.engine.losses.w_rbragg_loss(aligned: AlignedIntensities) → Tensor

Per-orientation weighted-R2 term (default mu), summed over thicknesses to a scalar.

Raw: no calc<->obs scaling. Correct only where the caller has already put calculated on the observed scale; for the refinement objective use wr2_loss(), which is the build_engine() default.

The imperative refinement loop – the one deliberately stateful corner of an otherwise pure core.

torch.optim optimizers mutate .grad and leaf tensors in place and carry internal state, so a training loop cannot be a pure function. This module quarantines that imperativeness behind a functional contract: run_refinement() takes the engine’s pure objective_value callable and the caller’s parameters, clones the target fields into fresh requires_grad leaves (the rest become detached constants), steps a chosen backend, and returns a new detached RefinementResult. The caller’s parameters are never touched. core/ stays free of torch.optim entirely.

class diffBloch.engine.refine.AtomSelection(mode: Literal['all', 'none'], element_include: tuple[str, ...] = (), element_exclude: tuple[str, ...] = ())

Bases: object

A coarse atom/parameter selection for one trainable group.

The current modes distinguish whole-group all vs none. The value can be extended with element or index filters without reintroducing stringly-typed refinement targets.

mode: Literal['all', 'none']

element_include: tuple[str, ...] = ()

element_exclude: tuple[str, ...] = ()

classmethod all() → AtomSelection

Select every present parameter in the group.

classmethod include_elements(*symbols: str) → AtomSelection

Select only atoms whose element symbol is listed.

classmethod exclude_elements(*symbols: str) → AtomSelection

Select all atoms except those whose element symbol is listed.

classmethod none() → AtomSelection

Select no parameters in the group.

property selects_any: bool

Whether this selection unlocks the corresponding trainable group.

property has_element_filter: bool

Whether this selection needs ASU atomic numbers to resolve per-row leaves.

class diffBloch.engine.refine.ObjectiveComponent(raw: Tensor, weight: float = 1.0)

Bases: object

One named refinement objective term.

raw is the scientifically meaningful scalar diagnostic (for example, a bond penalty before weighting). weight scales that diagnostic into the optimizer-facing contribution.

raw: Tensor

weight: float = 1.0

property contribution: Tensor

The weighted scalar contribution this component adds to the objective total.

class diffBloch.engine.refine.ObjectiveValue(components: Mapping[str, ObjectiveComponent], diagnostics: Mapping[str, float] = MappingProxyType({}), per_rotation: Sequence[Mapping[str, float]] = ())

Bases: object

A scalar refinement objective plus named scalar components.

total is computed from component contributions so reporting and optimization cannot silently drift. components is a read-only mapping whose values retain both raw diagnostics and weights for future penalty reporting. per_rotation is optional per-rotation diagnostics (rotation_index/wr2/r_obs/diff_loss) for verbose per-orientation reporting; empty for an objective with no rotation structure. It is plain data, computed unconditionally alongside the existing epoch-mean diagnostics (no extra forward cost) – whether it is actually reported is an execution-only choice made at the refinement loop, not here.

total: Tensor

components: Mapping[str, ObjectiveComponent]

diagnostics: Mapping[str, float]

per_rotation: tuple[Mapping[str, float], ...]

type diffBloch.engine.refine.OptimizerName = Literal['adam', 'adamw', 'lbfgs']

class diffBloch.engine.refine.RefinementResult(params: RefinableParams, losses: Tensor, best_params: RefinableParams, best_step: int)

Bases: object

The outcome of a refinement run (all tensors detached).

params are the final parameters after the last step; losses (steps,) is the per-step training curve (each entry is the objective before that step’s update); best_params / best_step snapshot the lowest recorded loss, for early-stopping callers.

params: RefinableParams

losses: Tensor

best_params: RefinableParams

best_step: int

property best_loss: float

The lowest recorded (pre-update) loss.

class diffBloch.engine.refine.PenaltyTerm(*args, **kwargs)

Bases: Protocol

A soft refinement penalty evaluated on the physical ASU state.

Penalties are objective components, not hard constraints: loss returns the raw scientific diagnostic and weight scales it into the optimizer-facing contribution. Concrete terms own their invariant context (for example metric/cell, connectivity, targets, and sigmas) instead of bloating PhysicalState with every possible penalty input.

name: str

weight: float

value(state: PhysicalState) → Tensor

Return this penalty’s raw scalar loss for the current physical state.

class diffBloch.engine.refine.TrainableSpec(positions: AtomSelection = <factory>, adp: AtomSelection = <factory>, occupancy: AtomSelection = <factory>)

Bases: object

Explicit selection of which parameter groups are trainable in a refinement problem.

positions: AtomSelection

adp: AtomSelection

occupancy: AtomSelection

classmethod positions_and_adp() → TrainableSpec

The default: refine positions and ADPs when present.

diffBloch.engine.refine.run_refinement(objective_value: Callable[[RefinableParams], ObjectiveValue], params: RefinableParams, *, steps: int, trainable: TrainableSpec, optimizer: OptimizerName, lr: float, atomic_numbers: Tensor | None = None, logger: Logger = NULL_LOGGER) → RefinementResult

Optimize the trainable parameter groups to minimise objective_value(params).total.

Functional contract over an unavoidably imperative core: the caller’s params are never mutated. Selected fields become fresh requires_grad leaves (non-selected fields detached constants); a backend steps them for steps iterations via a closure (which unifies LBFGS’ re-evaluation with Adam/AdamW). Trainable selections map through _TRAINABLE_FIELDS; a selected group with no present parameter, or zero steps, raises.

logger receives a RefinementStep per iteration and one RefinementCompleted at the end; the default NULL_LOGGER makes emission a no-op, so the returned result is unchanged. Step events include the structured ObjectiveValue components as numeric diagnostics, making diffraction/penalty tradeoffs inspectable.