"""``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
:meth:`~diffBloch.engine.RefinementEngine.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 :func:`_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 (:meth:`~diffBloch.engine.plan.OrientationPlan.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
(:func:`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`` (:class:`~diffBloch.specs.NelderMeadSearch`, off by default)
guards against exactly this by dividing the comparison score by the matched count
(:func:`_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 :class:`OrientationPlan` and the
:class:`~diffBloch.engine.plan.CoupledOrientationPlan`, so an already-segmented plan is optimized
under its frozen union.
"""
from __future__ import annotations
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import replace
import numpy as np
import torch
from numpy.typing import NDArray
from scipy.optimize import minimize
from torch import Tensor
from diffBloch.core.crystal import orientation_basis
from diffBloch.core.dynamical import (
StructureFactorGather,
build_structure_factor_gather,
grid_source_indices,
)
from diffBloch.core.reciprocal import g_vectors
from diffBloch.core.solver import SolverMethod
from diffBloch.engine import RefinementEngine, ScoresFn, wr2_scores
from diffBloch.engine.plan import CoupledOrientationPlan, OrientationPlanLike, StructureFactorGrid
from diffBloch.observability import (
NULL_LOGGER,
Logger,
OrientationOptimizationStarted,
OrientationOptimizationSummary,
OrientationOptimized,
)
from diffBloch.params import Device
from diffBloch.preprocess.coupling import build_coupling_segments
from diffBloch.preprocess.experiment import RefinementSetup
from diffBloch.preprocess.orientation import goniometer_rotation
from diffBloch.preprocess.pipeline import PlanStep, as_step
from diffBloch.preprocess.plan import Plan, require_built_plans, unique_hkl_count
from diffBloch.preprocess.scoring import active_structure_factor_indices, build_engine
from diffBloch.preprocess.steps.beams import klar_beam_mask
from diffBloch.specs import (
NO_ABSORPTION,
Absorption,
NelderMeadSearch,
TrialCoupling,
assert_grid_covers_coupling,
)
__all__ = ["optimize_orientation"]
[docs]
def 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
:class:`~diffBloch.observability.OrientationOptimized` reports the score under its real name.
``scores`` (default :func:`~diffBloch.engine.wr2_scores`) is the per-thickness metric
:meth:`~diffBloch.engine.forward.RefinementEngine.score_orientation` searches -- pass
``cfg.loss_metrics.to_scores()`` to search the same residual the gradient refinement stage
minimises (:func:`~diffBloch.config.schema.LossMetricsConfig.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 :func:`~diffBloch.config.manifest.dataset_config_digest`).
``refinement`` (constraint spec, ASU expansion, atomic numbers, seeded params) is captured
read-only and rejoined to the geometry ``Plan`` via :func:`build_engine`; the
orientation-invariant ``F_gb`` is computed once and reused across every orientation and trial.
``search`` is a pre-validated :class:`~diffBloch.specs.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
:class:`~diffBloch.specs.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
(:func:`~diffBloch.core.dynamical.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
:func:`build_engine`.
``logger`` receives an :class:`~diffBloch.observability.OrientationOptimized` per rotation as its
search completes (the optimization is the run's long phase, so this is the progress stream); the
default :data:`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 :func:`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.
"""
if workers < 1:
raise ValueError("workers must be >= 1")
def run(plan: Plan) -> Plan:
# Coverage guard for the coupled path: the grid sphere must span the beam-difference
# support so a per-trial gather cannot address a reflection outside it -- the precondition
# for running those gathers validate=False (their silent-zero gap has no runtime backstop).
# O(1), orientation-independent, always on -- fails at setup, not deep in the search.
if coupling is not None:
assert_grid_covers_coupling(coupling.policy, plan.structure_factor_grid.g_max)
engine = build_engine(
plan,
refinement,
method=method,
max_batch=max_batch,
absorption=absorption,
scores=scores,
)
params = refinement.params if device is None else refinement.params.to(device)
fgb = engine.fgb(params)
def refine(op: OrientationPlanLike) -> tuple[OrientationPlanLike, float, int, int]:
trial_fgb: Tensor | Callable[[OrientationPlanLike], Tensor] = fgb
if coupling is not None:
cached = fgb.clone()
known = torch.zeros(cached.shape[0], dtype=torch.bool, device=cached.device)
assert engine.active_structure_factor_indices is not None
known[engine.active_structure_factor_indices.to(device=known.device)] = True
def lazy_fgb(candidate: OrientationPlanLike) -> Tensor:
indices = active_structure_factor_indices(
(candidate,), plan.structure_factor_grid.gpts
).to(device=cached.device)
missing = indices[~known.index_select(0, indices)]
if missing.numel():
with torch.no_grad():
values = engine.structure_factor_values(params, missing)
cached.index_copy_(0, missing, values)
known[missing] = True
return cached
trial_fgb = lazy_fgb
return _refine_one(
engine,
trial_fgb,
plan,
op,
search=search,
coupling=coupling,
validate=validate,
)
built = require_built_plans(plan)
logger.report(OrientationOptimizationStarted(total_rotations=len(built)))
results_by_index: dict[int, tuple[OrientationPlanLike, float, int, int]] = {}
cap = search.max_iterations
def report(
index: int,
result: tuple[OrientationPlanLike, float, int, int],
) -> None:
fitted, score, n_trials, n_passes = result
results_by_index[index] = result
pattern_index = fitted.alignment.pattern_index
n_matched = int(pattern_index.shape[0])
logger.report(
OrientationOptimized(
rotation_index=fitted.pattern.rotation_index,
score=score,
residual=residual,
n_matched_hkl=n_matched,
n_trials=n_trials,
n_passes=n_passes,
pass_cap=cap,
)
)
if workers > 1:
pool = ThreadPoolExecutor(max_workers=workers)
try:
futures = {pool.submit(refine, op): index for index, op in enumerate(built)}
for future in as_completed(futures): # emit progress as searches finish
index = futures[future]
report(index, future.result())
finally:
# An early abort surfaces as a logger.report() raising mid-loop. Cancel the
# not-yet-started searches and don't block on the in-flight ones, so the abort
# actually saves the remaining budget instead of draining every queued rotation
# first (the whole point of stopping early). The `with`-block's default
# shutdown(wait=True) would run them all before the exception surfaced. Rotations
# already executing cannot be interrupted, so up to `workers` still finish; the
# queued remainder is dropped. On normal completion nothing is pending -- a no-op.
pool.shutdown(wait=False, cancel_futures=True)
else:
for index, op in enumerate(built):
report(index, refine(op))
ordered_results = tuple(results_by_index[i] for i in range(len(built)))
ordered = tuple(result[0] for result in ordered_results)
def strong_matched_hkl(op: OrientationPlanLike) -> Tensor:
pattern_index = op.alignment.pattern_index
strong = op.pattern.intensities[pattern_index] > 3.0 * op.pattern.sigmas[pattern_index]
return op.alignment.hkl[strong]
logger.report(
OrientationOptimizationSummary(
n_orientations=len(ordered_results),
mean_score=sum(result[1] for result in ordered_results) / len(ordered_results),
residual=residual,
unique_matched_hkl=unique_hkl_count(op.alignment.hkl for op in ordered),
unique_strong_hkl=unique_hkl_count(strong_matched_hkl(op) for op in ordered),
unique_observed_hkl=unique_hkl_count(op.pattern.hkl for op in ordered),
total_trials=sum(result[2] for result in ordered_results),
max_passes=max(result[3] for result in ordered_results),
)
)
return replace(plan, orientations=ordered)
# search rides in the config digest too, but coupling is a composition-site kwarg (not config),
# so it MUST be in the recipe identity; workers/logger/device are execution-only (device shifts
# output only to solver tolerance -- see the docstring -- so it stays out of the identity).
return as_step(
"optimize_orientation",
{"search": search, "coupling": coupling, "absorption": absorption},
run,
)
def _comparable_score(score: float, plan: OrientationPlanLike, search: NelderMeadSearch) -> float:
"""``score``, or ``score`` normalised by matched-reflection count when the search opts in.
A trial's SCORED set can vary in size (:func:`_coupled_trial` re-selects it per trial, and the
residual metrics have no term correcting for that), so comparing the raw ``score`` lets a trial
win by drifting to geometry that matches a smaller, easier subset rather than by optimizing
better. ``search.penalize_fewer_reflections`` (default ``False``, matching the reference
implementation's own disabled-by-default normalisation) opts into dividing by the matched
count: a trial with fewer matched reflections then needs a proportionally lower ``score`` to be
accepted. ``0`` matched reflections scores ``inf`` (never accepted; ``score`` itself is
typically ``nan`` there already). With the flag off this is the identity on ``score``.
"""
if not search.penalize_fewer_reflections:
return score
n_matched = int(plan.alignment.pattern_index.shape[0])
return score / n_matched if n_matched > 0 else float("inf")
def _refine_one(
engine: RefinementEngine,
fgb: Tensor | Callable[[OrientationPlanLike], Tensor],
plan: Plan,
op: OrientationPlanLike,
*,
search: NelderMeadSearch,
coupling: TrialCoupling | None,
validate: bool = True,
) -> tuple[OrientationPlanLike, float, int, int]:
"""Local Nelder-Mead search over the goniometer correction ``(alpha, beta, omega)``.
Every trial composes directly off the fixed seed orientation: ``seed_orientation @
goniometer_rotation(alpha, beta, omega)`` -- the search explores the whole ``step_size``
neighbourhood at once rather than annealing a step down. ``scipy.optimize.minimize`` runs
``method="Nelder-Mead"`` from a fixed initial simplex of edge length ``search.step_size``
around ``(alpha, beta, omega) = (0, 0, 0)``, exactly mirroring the reference implementation
this port is checked against. ``n_passes`` is scipy's reported iteration count (``result.nit``).
"""
grid = plan.structure_factor_grid
n_trials = 0
# One rotation's search revisits the same excitation unions across many nearby trials, so the
# beam set across this rotation's trials.
gather_cache: dict[bytes, StructureFactorGather] = {}
seed_orientation = np.asarray(op.orientation, dtype=np.float64)
def build_trial(orientation: NDArray[np.float64]) -> OrientationPlanLike:
if coupling is None:
return op.with_orientation(grid, orientation)
return _coupled_trial(grid, op, orientation, coupling, gather_cache, validate=validate)
def objective(params: NDArray[np.float64]) -> float:
nonlocal n_trials
alpha, beta, omega = params
orientation = seed_orientation @ goniometer_rotation(alpha, beta, omega)
trial = build_trial(orientation)
n_trials += 1
trial_fgb = fgb(trial) if callable(fgb) else fgb
raw = float(engine.score_orientation(trial, trial_fgb))
# scipy minimises this directly, so the (opt-in) fewer-reflections guard lives here: a
# trial cannot win merely by drifting to geometry that matches a smaller, easier subset.
return _comparable_score(raw, trial, search)
step = search.step_size
initial_simplex = np.array(
[
[0.0, 0.0, 0.0],
[step, 0.0, 0.0],
[0.0, step, 0.0],
[0.0, 0.0, step],
]
)
result = minimize(
objective,
x0=np.zeros(3),
method="Nelder-Mead",
options={
"initial_simplex": initial_simplex,
"maxiter": search.max_iterations,
"xatol": search.x_tolerance,
"fatol": search.f_tolerance,
},
)
alpha, beta, omega = result.x
best_orientation = seed_orientation @ goniometer_rotation(alpha, beta, omega)
current = build_trial(best_orientation)
n_trials += 1
current_fgb = fgb(current) if callable(fgb) else fgb
# result.fun is the comparable (penalized) score minimised above; report the plain score
# instead (self.scores, under whichever residual ExperimentConfig.loss_metrics configures).
score = float(engine.score_orientation(current, current_fgb))
return current, score, n_trials, int(result.nit)
def _coupled_trial(
grid: StructureFactorGrid,
op: OrientationPlanLike,
orientation: NDArray[np.float64],
coupling: TrialCoupling,
gather_cache: dict[bytes, StructureFactorGather] | None = None,
*,
validate: bool = True,
) -> CoupledOrientationPlan:
"""Re-couple the solve union and re-select the scored set, both at ``orientation`` (one trial).
One objective evaluation: (1) ``build_coupling_segments`` re-derives the per-tilt-segment
excitation union at ``orientation`` (the SOLVE set); (2) ``coupling.scored`` (the Klar
rsg/dsg window + resolution cap) is re-applied to that fresh union's own lab-frame geometry at
this orientation, giving the SCORED set its own fresh selection every trial -- mirroring the
reference implementation's per-trial ``filter_hkls`` -- rather than pinning it to the seed's.
The observed ``pattern``, ``thickness``, and ``tilt_reduction`` are carried from ``op``
unchanged. The atomic ``F_gb`` is untouched either way; ``gather_cache`` (keyed by a segment's
beam-set bytes) reuses the orientation-free per-segment F-gathers across a search's trials --
identical beam set, identical gather -- collapsing the per-trial rebuild cost.
``validate`` (default ``True``) is forwarded to each cache-miss
:func:`~diffBloch.core.dynamical.build_structure_factor_gather`; ``False`` skips its O(N^2)
integrity checks on the hot path (safe under the caller's coverage guard). It reaches only that
build -- the ``CoupledOrientationPlan.build`` below always receives the precomputed
``gathers``, so its own ``validate`` never triggers a rebuild here.
"""
cell = np.asarray(grid.cell, dtype=np.float64)
tilts = np.asarray(op.tilts, dtype=np.float64)
segments = build_coupling_segments(
coupling.policy,
np.asarray(grid.structure_factor_hkl, dtype=np.int64),
cell=cell,
orientation=orientation,
tilts=tilts,
energy=op.energy,
u0=op.u0,
)
union_hkl = np.unique(np.concatenate([segment.union_hkl for segment in segments]), axis=0)
basis = orientation_basis(cell, orientation)
g = g_vectors(union_hkl, basis)
keep = klar_beam_mask(
g,
energy=op.energy,
u0=op.u0,
rsg=coupling.scored.klar.rsg,
dsg=coupling.scored.klar.dsg,
semiangle=coupling.scored.klar.integration.semiangle,
geometry=coupling.scored.klar.integration.geometry,
)
keep &= np.linalg.norm(g, axis=1) <= coupling.scored.g_max
keep |= (union_hkl == 0).all(axis=1) # 000 anchors psi0; retained when present
scored_hkl = union_hkl[keep]
gathers = None
if gather_cache is not None:
structure_factor_hkl = np.asarray(grid.structure_factor_hkl)
# The gather's grid-side ravel is identical for every segment (same
# structure_factor_hkl/gpts), so build
# it once here and reuse across this trial's segment builds -- re-raveling the support grid
# per segment was the residual per-trial cost after the |g|<cap pre-filter. Lazy: only paid
# when a segment actually misses the cache (an unchanged union rebuilds nothing).
source: NDArray[np.int64] | None = None
gathers = []
for segment in segments:
key = np.ascontiguousarray(segment.union_hkl).tobytes()
gather = gather_cache.get(key)
if gather is None:
if source is None:
source = grid_source_indices(structure_factor_hkl, grid.gpts)
gather = build_structure_factor_gather(
structure_factor_hkl,
segment.union_hkl,
grid.gpts,
validate=validate,
structure_factor_indices=source,
)
gather_cache[key] = gather
gathers.append(gather)
return CoupledOrientationPlan.build(
grid,
[(segment.union_hkl, segment.covered_tilt_indices) for segment in segments],
op.pattern,
energy=op.energy,
thickness=op.thickness,
u0=op.u0,
orientation=orientation,
tilts=tilts,
tilt_reduction=op.tilt_reduction,
scored_hkl=scored_hkl,
gathers=gathers,
)