diff --git a/src/quantem/core/fitting/__init__.py b/src/quantem/core/fitting/__init__.py new file mode 100644 index 00000000..5b687b7b --- /dev/null +++ b/src/quantem/core/fitting/__init__.py @@ -0,0 +1,21 @@ +from quantem.core.fitting.background import DCBackground as DCBackground +from quantem.core.fitting.background import GaussianBackground as GaussianBackground +from quantem.core.fitting.base import Component as Component +from quantem.core.fitting.base import Model as Model +from quantem.core.fitting.base import ModelContext as ModelContext +from quantem.core.fitting.base import OriginND as OriginND +from quantem.core.fitting.base import Parameter as Parameter +from quantem.core.fitting.diffraction import DiskTemplate as DiskTemplate +from quantem.core.fitting.diffraction import SyntheticDiskLattice as SyntheticDiskLattice + +__all__ = [ + "Component", + "DCBackground", + "DiskTemplate", + "GaussianBackground", + "Model", + "ModelContext", + "OriginND", + "Parameter", + "SyntheticDiskLattice", +] \ No newline at end of file diff --git a/src/quantem/core/fitting/background.py b/src/quantem/core/fitting/background.py new file mode 100644 index 00000000..84f86cb3 --- /dev/null +++ b/src/quantem/core/fitting/background.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +from typing import Any, Sequence + +import torch +from torch import nn + +from quantem.core.fitting.base import OriginND, RenderComponent, RenderContext + + +class DCBackground(RenderComponent): + def __init__( + self, + *, + intensity: float | int | Sequence[float | int | None] = 0.0, + name: str = "dc_background", + constraint_params: dict[str, Any] | None = None, + ): + """ + Build a constant background component. + + Notes + ----- + Validity is enforced via hard constraints/parameter bounds. Forward + intentionally avoids hard clamps for gradient flow. + """ + super().__init__() + self.name = str(name) + intensity_init, intensity_lo, intensity_hi = self.parse_bounded_init( + intensity, name="intensity" + ) + self.intensity_raw = nn.Parameter(torch.tensor(intensity_init, dtype=torch.float32)) + bounded_lo = 0.0 if intensity_lo is None else max(float(intensity_lo), 0.0) + self.register_parameter_bounds("intensity_raw", bounded_lo, intensity_hi) + if constraint_params is not None: + self.apply_constraint_params(constraint_params, strict=True) + self._enforce_parameter_bounds() + + def forward(self, ctx: RenderContext) -> torch.Tensor: + """ + Render constant background from raw trainable intensity. + + Notes + ----- + Validity is enforced via hard constraints/parameter bounds, not via + forward-time hard clamps. + """ + inten = self.intensity_raw.to(device=ctx.device, dtype=ctx.dtype) + return torch.ones(ctx.shape, device=ctx.device, dtype=ctx.dtype) * inten + + def forward_batched( + self, + ctx: RenderContext, + *, + intensity_raw_b: torch.Tensor, + ) -> torch.Tensor: + B = intensity_raw_b.shape[0] + return intensity_raw_b.view(B, 1, 1).expand(B, ctx.shape[0], ctx.shape[1]).to( + device=ctx.device, dtype=ctx.dtype + ) + + +class GaussianBackground(RenderComponent): + def __init__( + self, + *, + sigma: float | int | Sequence[float | int | None] = (40.0, 5.0, None), + sigma_col: float | int | Sequence[float | int | None] | None = None, + intensity: float | int | Sequence[float | int | None] = 0.0, + origin: OriginND | None = None, + origin_key: str = "origin", + name: str = "gaussian_background", + constraint_params: dict[str, Any] | None = None, + ): + """ + Build a Gaussian background component centered at origin. + + Parameters + ---------- + sigma : float or sequence + Row-direction (and column-direction when ``sigma_col`` is None) width. + Accepts ``x``, ``(x0, delta)``, or ``(x0, lo, hi)``. + sigma_col : float or sequence or None + Column-direction width for an anisotropic (elliptical) Gaussian. + When ``None`` (default) the background is isotropic and ``sigma`` + governs both axes. Providing a value enables independent per-axis + widths, which is useful for tilted datasets where the transmitted-disk + background is elongated. + + Notes + ----- + ``sigma_raw``, ``sigma_col_raw`` (when present), and ``intensity_raw`` + validity is enforced via hard constraints/parameter bounds. + """ + super().__init__() + self.name = str(name) + self.origin = origin + self.origin_key = str(origin_key) + sigma_init, sigma_lo, sigma_hi = self.parse_bounded_init(sigma, name="sigma") + intensity_init, intensity_lo, intensity_hi = self.parse_bounded_init( + intensity, name="intensity" + ) + self.sigma_raw = nn.Parameter(torch.tensor(sigma_init, dtype=torch.float32)) + sigma_bounded_lo = 1e-6 if sigma_lo is None else max(float(sigma_lo), 1e-6) + self.register_parameter_bounds("sigma_raw", sigma_bounded_lo, sigma_hi) + + if sigma_col is not None: + sc_init, sc_lo, sc_hi = self.parse_bounded_init(sigma_col, name="sigma_col") + self.sigma_col_raw = nn.Parameter(torch.tensor(sc_init, dtype=torch.float32)) + sc_bounded_lo = 1e-6 if sc_lo is None else max(float(sc_lo), 1e-6) + self.register_parameter_bounds("sigma_col_raw", sc_bounded_lo, sc_hi) + else: + self.sigma_col_raw = None + + self.intensity_raw = nn.Parameter(torch.tensor(intensity_init, dtype=torch.float32)) + intensity_bounded_lo = 0.0 if intensity_lo is None else max(float(intensity_lo), 0.0) + self.register_parameter_bounds("intensity_raw", intensity_bounded_lo, intensity_hi) + if constraint_params is not None: + self.apply_constraint_params(constraint_params, strict=True) + self._enforce_parameter_bounds() + + def set_origin(self, origin: OriginND) -> None: + self.origin = origin + + def forward(self, ctx: RenderContext) -> torch.Tensor: + if self.origin is None: + raise RuntimeError("GaussianBackground requires an OriginND instance.") + + rr = torch.arange(ctx.shape[0], device=ctx.device, dtype=ctx.dtype)[:, None] + cc = torch.arange(ctx.shape[1], device=ctx.device, dtype=ctx.dtype)[None, :] + r0, c0 = self.origin.coords[0], self.origin.coords[1] + + sigma_r = self.sigma_raw.to(device=ctx.device, dtype=ctx.dtype) + sigma_c = ( + self.sigma_col_raw.to(device=ctx.device, dtype=ctx.dtype) + if self.sigma_col_raw is not None + else sigma_r + ) + inten = self.intensity_raw.to(device=ctx.device, dtype=ctx.dtype) + r2 = (rr - r0) ** 2 / (sigma_r * sigma_r) + (cc - c0) ** 2 / (sigma_c * sigma_c) + return inten * torch.exp(-0.5 * r2) + + def forward_batched( + self, + ctx: RenderContext, + *, + sigma_raw_b: torch.Tensor, + sigma_col_raw_b: torch.Tensor | None = None, + intensity_raw_b: torch.Tensor, + origin_coords_b: torch.Tensor, + ) -> torch.Tensor: + B = sigma_raw_b.shape[0] + rr = torch.arange(ctx.shape[0], device=ctx.device, dtype=ctx.dtype).view(1, ctx.shape[0], 1) + cc = torch.arange(ctx.shape[1], device=ctx.device, dtype=ctx.dtype).view(1, 1, ctx.shape[1]) + r0 = origin_coords_b[:, 0].view(B, 1, 1) + c0 = origin_coords_b[:, 1].view(B, 1, 1) + sigma_r = sigma_raw_b.view(B, 1, 1) + sigma_c = sigma_col_raw_b.view(B, 1, 1) if sigma_col_raw_b is not None else sigma_r + inten = intensity_raw_b.view(B, 1, 1) + r2 = (rr - r0) ** 2 / (sigma_r * sigma_r) + (cc - c0) ** 2 / (sigma_c * sigma_c) + return inten * torch.exp(-0.5 * r2) \ No newline at end of file diff --git a/src/quantem/core/fitting/base.py b/src/quantem/core/fitting/base.py new file mode 100644 index 00000000..ce2688fd --- /dev/null +++ b/src/quantem/core/fitting/base.py @@ -0,0 +1,1164 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal, Self, Sequence, cast + +import numpy as np +import torch +from torch import nn +from tqdm import tqdm + +from quantem.core.ml.optimizer_mixin import ( + OptimizerMixin, +) + + +def parse_bounded_init( + value: float | int | Sequence[float | int | None], *, name: str +) -> tuple[float, float | None, float | None]: + """ + Parse a scalar or bounded initializer specification. + + Parameters + ---------- + value : float | int | Sequence[float | int | None] + Accepted forms: + - ``x`` -> init ``x`` with no bounds. + - ``(x0, delta)`` -> init ``x0`` with bounds ``[x0-|delta|, x0+|delta|]``. + - ``(x0, lo, hi)`` -> init ``x0`` with explicit bounds. + name : str + Parameter name used in error messages. + + Returns + ------- + tuple[float, float | None, float | None] + Parsed ``(init, lo, hi)``. + + Raises + ------ + ValueError + If the sequence form is invalid, contains required ``None`` entries, + has invalid ordering, or ``init`` lies outside explicit bounds. + """ + if not isinstance(value, (list, tuple, np.ndarray)): + x = float(cast(float | int, value)) + return x, None, None + + seq = list(value) + if len(seq) == 0: + raise ValueError(f"{name} cannot be empty.") + if seq[0] is None: + raise ValueError(f"{name} initial value cannot be None.") + x0 = float(cast(float | int, seq[0])) + + if len(seq) == 1: + return x0, None, None + if len(seq) == 2: + if seq[1] is None: + raise ValueError(f"{name} delta cannot be None.") + delta = abs(float(cast(float | int, seq[1]))) + return x0, x0 - delta, x0 + delta + if len(seq) == 3: + if seq[1] is None or seq[2] is None: + raise ValueError(f"{name} bounds cannot contain None.") + lo = float(cast(float | int, seq[1])) + hi = float(cast(float | int, seq[2])) + if lo > hi: + raise ValueError(f"{name} has invalid bounds: lo ({lo}) > hi ({hi}).") + if x0 < lo or x0 > hi: + raise ValueError(f"{name} initial value {x0} is outside bounds [{lo}, {hi}].") + return x0, lo, hi + + raise ValueError(f"{name} must be scalar, (x0, delta), or (x0, lo, hi).") + + +@dataclass +class RenderContext: + shape: tuple[int, ...] + device: torch.device + dtype: torch.dtype + mask: torch.Tensor | None = None + fields: dict[str, Any] = field(default_factory=dict) + + +class OriginND(nn.Module): + def __init__(self, *, ndim: int, init: Sequence[float]): + super().__init__() + if int(ndim) <= 0: + raise ValueError("ndim must be >= 1.") + if len(init) != int(ndim): + raise ValueError("init length must match ndim.") + self.ndim = int(ndim) + self.coords = nn.Parameter(torch.as_tensor(init, dtype=torch.float32).reshape(self.ndim)) + + +class RenderComponent(OptimizerMixin,nn.Module): + DEFAULT_HARD_CONSTRAINTS: dict[str, Any] = {} + DEFAULT_SOFT_CONSTRAINTS: dict[str, Any] = {} + DEFAULT_CONSTRAINT_CONFIG: dict[str, Any] = {} + + DEFAULT_OPTIMIZER: str = "adam" + DEFAULT_LR: float = 1e-2 + DEFAULT_SCHEDULER_TYPE: str = "none" + + def __init__(self) -> None: + nn.Module.__init__(self) + + self._optimizer = None + self._scheduler = None + self._optimizer_params = {} + self._scheduler_params = {} + self.hard_constraints: dict[str, Any] = dict(self.DEFAULT_HARD_CONSTRAINTS) + self.soft_constraints: dict[str, Any] = dict(self.DEFAULT_SOFT_CONSTRAINTS) + self.constraint_config: dict[str, Any] = dict(self.DEFAULT_CONSTRAINT_CONFIG) + self.parameter_bounds: dict[str, tuple[float | None, float | None]] = {} + + optimizer_params = { + "type": self.DEFAULT_OPTIMIZER, + "lr": self.DEFAULT_LR + } + self.optimizer_params = optimizer_params + + scheduler_params = { + "type": self.DEFAULT_SCHEDULER_TYPE + } + self.scheduler_params = scheduler_params + + @staticmethod + def parse_bounded_init( + value: float | int | Sequence[float | int | None], *, name: str + ) -> tuple[float, float | None, float | None]: + """ + Parse bounded initializer forms into ``(init, lo, hi)``. + + Parameters + ---------- + value : float | int | Sequence[float | int | None] + Scalar, ``(x0, delta)``, or ``(x0, lo, hi)``. + name : str + Parameter name used in error messages. + + Returns + ------- + tuple[float, float | None, float | None] + Parsed ``(init, lo, hi)``. + """ + return parse_bounded_init(value, name=name) + + def register_parameter_bounds( + self, parameter_name: str, lo: float | None, hi: float | None + ) -> None: + """ + Register hard bounds for a trainable parameter. + + Parameters + ---------- + parameter_name : str + Name of an ``nn.Parameter`` attribute on this component. + lo : float | None + Lower bound, or ``None`` for unbounded lower side. + hi : float | None + Upper bound, or ``None`` for unbounded upper side. + + Returns + ------- + None + + Raises + ------ + ValueError + If ``lo > hi``. + """ + if lo is not None and hi is not None and float(lo) > float(hi): + raise ValueError(f"Invalid bounds for {parameter_name}: lo ({lo}) > hi ({hi}).") + self.parameter_bounds[str(parameter_name)] = ( + None if lo is None else float(lo), + None if hi is None else float(hi), + ) + + def _enforce_parameter_bounds(self) -> None: + """ + Clamp registered parameters in-place to configured bounds. + + Returns + ------- + None + + Raises + ------ + AttributeError + If a registered parameter attribute is missing. + TypeError + If a registered attribute is not an ``nn.Parameter``. + """ + if not self.parameter_bounds: + return + with torch.no_grad(): + for param_name, (lo, hi) in self.parameter_bounds.items(): + if not hasattr(self, param_name): + raise AttributeError( + f"Parameter '{param_name}' is not an attribute of {self.__class__.__name__}." + ) + param = getattr(self, param_name) + if not isinstance(param, nn.Parameter): + raise TypeError( + f"Attribute '{param_name}' on {self.__class__.__name__} is not an nn.Parameter." + ) + if lo is None and hi is None: + continue + if lo is None: + assert hi is not None + param.clamp_(max=float(hi)) + elif hi is None: + assert lo is not None + param.clamp_(min=float(lo)) + else: + param.clamp_(min=float(lo), max=float(hi)) + + def _set_constraints( + self, + current: dict[str, Any], + defaults: dict[str, Any], + constraints: dict[str, Any], + *, + strict: bool, + ) -> None: + if strict: + unknown = [k for k in constraints if k not in defaults] + if unknown: + keys = ", ".join(str(k) for k in unknown) + raise KeyError(f"Unknown constraint keys: {keys}") + current.update(constraints) + + def set_hard_constraints(self, constraints: dict[str, Any], strict: bool = True) -> None: + self._set_constraints( + self.hard_constraints, self.DEFAULT_HARD_CONSTRAINTS, constraints, strict=strict + ) + + def set_soft_constraints(self, constraints: dict[str, Any], strict: bool = True) -> None: + self._set_constraints( + self.soft_constraints, self.DEFAULT_SOFT_CONSTRAINTS, constraints, strict=strict + ) + + def apply_constraint_params(self, params: dict[str, Any], strict: bool = True) -> None: + if not isinstance(params, dict): + raise TypeError("constraint params must be a dict.") + if "hard" in params or "soft" in params: + hard = params.get("hard") + soft = params.get("soft") + if hard is not None: + if not isinstance(hard, dict): + raise TypeError("constraint params 'hard' value must be a dict.") + self.set_hard_constraints(hard, strict=strict) + if soft is not None: + if not isinstance(soft, dict): + raise TypeError("constraint params 'soft' value must be a dict.") + self.set_soft_constraints(soft, strict=strict) + return + + hard_updates: dict[str, Any] = {} + soft_updates: dict[str, Any] = {} + unknown: dict[str, Any] = {} + for k, v in params.items(): + if k in self.DEFAULT_HARD_CONSTRAINTS: + hard_updates[k] = v + elif k in self.DEFAULT_SOFT_CONSTRAINTS: + soft_updates[k] = v + else: + unknown[k] = v + + if unknown and strict: + keys = ", ".join(str(k) for k in unknown.keys()) + raise KeyError(f"Unknown constraint keys for {self.__class__.__name__}: {keys}") + if unknown: + soft_updates.update(unknown) + if hard_updates: + self.set_hard_constraints(hard_updates, strict=strict) + if soft_updates: + self.set_soft_constraints(soft_updates, strict=strict) + + def effective_soft_constraints(self, params: dict[str, Any] | None = None) -> dict[str, Any]: + effective = dict(self.soft_constraints) + if isinstance(params, dict): + effective.update(params) + return effective + + def enforce_hard_constraints(self, ctx: RenderContext) -> None: + self._enforce_parameter_bounds() + + def forward(self, ctx: RenderContext) -> torch.Tensor: + raise NotImplementedError + + def constraint_loss( + self, ctx: RenderContext, params: dict[str, Any] | None = None + ) -> torch.Tensor: + return torch.zeros((), device=ctx.device, dtype=ctx.dtype) + + def get_optimization_parameters(self) -> dict[str, list]: + params = [p for p in self.parameters() if p.requires_grad] + if not params: + return {} + return {'default': params} + + def initialize_optimizer(self, + optimizer_params: dict[str, Any] | None = None, + scheduler_params: dict[str, Any] | None = None, + num_iter: int | None = None, + ) -> None: + trainable_params = (self.get_optimization_parameters()) + if not trainable_params: + self._optimizer = None + self._scheduler = None + return + if optimizer_params is not None: + self.optimizer_params = optimizer_params + else: + self.optimizer_params = { + "type": self.DEFAULT_OPTIMIZER, + "lr": self.DEFAULT_LR + } + self.set_optimizer(self.optimizer_params) + + if scheduler_params is not None: + self.scheduler_params = scheduler_params + else: + self.scheduler_params = { + "type": self.DEFAULT_SCHEDULER_TYPE + } + self.set_scheduler(self.scheduler_params, num_iter = num_iter) + return + + + def _infer_optimizer_rebuild_params(self) -> Any: + if self.optimizer_params: + return self.optimizer_params + if self.optimizer is not None: + opt_type: str | type[torch.optim.Optimizer] + if isinstance(self.optimizer, torch.optim.AdamW): + opt_type = "adamw" + elif isinstance(self.optimizer, torch.optim.Adam): + opt_type = "adam" + elif isinstance(self.optimizer, torch.optim.SGD): + opt_type = "sgd" + else: + opt_type = type(self.optimizer) + lr = float( + self.optimizer.param_groups[0].get( + "lr", getattr(self, "DEFAULT_LR", self.DEFAULT_LR) + ) + ) + return {"type": opt_type, "lr": lr} + return { + "type": getattr(self, "DEFAULT_OPTIMIZER_TYPE", self.DEFAULT_OPTIMIZER), + "lr": float(getattr(self, "DEFAULT_LR", self.DEFAULT_LR)), + } + + def _infer_scheduler_rebuild_params(self) -> Any: + if self.scheduler_params: + return self.scheduler_params + return { + "type": self.DEFAULT_SCHEDULER_TYPE, + } + + def _rebuild_optimizer_after_trainability_change(self) -> None: + trainable_params = (self.get_optimization_parameters()) + if not trainable_params: + self._optimizer = None + self._scheduler = None + return + # rebuild_params = self._infer_optimizer_rebuild_params() + # rebuild_params_scheduler = self._infer_scheduler_rebuild_params() + self.set_optimizer(self.optimizer_params) + self.set_scheduler(self.scheduler_params) + + def initialize_constraint_config(self, config: dict[str, Any], strict: bool = True) -> None: + if not hasattr(self, 'constraint_config'): + if strict: + raise AttributeError( + f"{self.__class__.__name__} does not have constraint_config attribute" + ) + return + if not isinstance(config, dict): + raise TypeError("constraint config must be a dict.") + + unknown: dict[str, Any] = {} + for k, v in config.items(): + if k in self.DEFAULT_CONSTRAINT_CONFIG: + self.constraint_config[k] = v + else: + unknown[k] = v + + if unknown and strict: + keys = ", ".join(str(k) for k in unknown.keys()) + raise KeyError(f"Unknown constraint keys for {self.__class__.__name__}: {keys}") + return + + +class AdditiveRenderModel(nn.Module): # step all otpimzers + def __init__(self, *, origin: nn.Module, components: list[RenderComponent]): + super().__init__() + self.origin = origin + self.components = nn.ModuleList(components) + + def forward(self, ctx: RenderContext) -> torch.Tensor: + if len(self.components) == 0: + return torch.zeros(ctx.shape, device=ctx.device, dtype=ctx.dtype) + out = self.components[0](ctx) + for component in self.components[1:]: + out = out + component(ctx) + return out + + def _component_constraint_name(self, component: RenderComponent, idx: int) -> str: + name = getattr(component, "name", None) + if isinstance(name, str) and name: + return name + class_name = component.__class__.__name__ + if class_name: + return class_name + return f"component_{idx}" + + def apply_constraint_params( + self, constraint_params: dict[str, Any], strict: bool = True + ) -> None: + if not isinstance(constraint_params, dict): + raise TypeError("constraint_params must be a dict.") + source = constraint_params.get("components") + component_map = source if isinstance(source, dict) else constraint_params + for target, params in component_map.items(): + if not isinstance(params, dict): + if strict: + raise TypeError(f"Constraint params for '{target}' must be a dict.") + continue + target_str = str(target) + name_matches: list[RenderComponent] = [] + class_matches: list[RenderComponent] = [] + for idx, module in enumerate(self.components): + component = cast(RenderComponent, module) + if self._component_constraint_name(component, idx) == target_str: + name_matches.append(component) + if component.__class__.__name__ == target_str: + class_matches.append(component) + targets = name_matches if name_matches else class_matches + if not targets: + if strict: + raise KeyError(f"No matching component for constraint target '{target_str}'.") + continue + for component in targets: + component.apply_constraint_params(params, strict=strict) + + def apply_hard_constraints(self, ctx: RenderContext) -> None: + for module in self.components: + component = cast(RenderComponent, module) + component.enforce_hard_constraints(ctx) + + def total_constraint_loss(self, ctx: RenderContext, **kwargs: Any) -> torch.Tensor: + from quantem.core.fitting.diffraction import SyntheticDiskLattice + loss = torch.zeros((), device=ctx.device, dtype=ctx.dtype) + for module in self.components: + component = cast(RenderComponent, module) + if isinstance(component, SyntheticDiskLattice): + loss = loss + component.constraint_loss(ctx, neighbor_target=kwargs.get("neighbor_target")) + else: + loss = loss + component.constraint_loss(ctx) + return loss + + def initilize_independant_optimizers(self, + individual_optimizers: dict[str, dict[str, Any]] | None = None, + individual_schedulers: dict[str, dict[str, Any]] | None = None, + num_iter: int | None = None, + ) -> None: + for idx, module in enumerate(self.components): + component = cast(RenderComponent, module) + component_name = self._component_constraint_name(component, idx) + + component_optimizer_params = None + component_scheduler_params = None + if individual_optimizers is not None: + if component_name in individual_optimizers: + component_optimizer_params = individual_optimizers[component_name] + elif component.__class__.__name__ in individual_optimizers : + component_optimizer_params = individual_optimizers[component.__class__.__name__] + + if individual_schedulers is not None: + if component_name in individual_schedulers: + component_scheduler_params = individual_schedulers[component_name] + elif component.__class__.__name__ in individual_schedulers: + component_scheduler_params = individual_schedulers[component.__class__.__name__] + + component.initialize_optimizer(component_optimizer_params, component_scheduler_params, num_iter=num_iter) + + def set_independant_optimizer_params(self, + individual_optimizer_params: dict[str, dict[str, Any]], + individual_scheduler_params: dict[str, dict[str, Any]] | None = None, + num_iter: int | None = None, + ) -> None: + for component_name, param in individual_optimizer_params.items(): + scheduler_params = None + if individual_scheduler_params is not None and component_name in individual_scheduler_params: + scheduler_params = individual_scheduler_params[component_name] + + component = self._resolve_component_by_name(str(component_name)) + component.initialize_optimizer(param,scheduler_params,num_iter) + + def rebuild_independant_optimizers(self) -> None: + for module in self.components: + component = cast(RenderComponent, module) + component.initialize_optimizer() + + def step_optimizers(self) -> None: + for module in self.components: + component = cast(RenderComponent, module) + component.step_optimizer() + + def step_schedulers(self, loss: float | None = None) -> None: + for module in self.components: + component = cast(RenderComponent, module) + if hasattr(component, 'step_scheduler'): + try: + component.step_scheduler(loss) + except (AttributeError, TypeError): + pass + + def zero_grad_optimizers(self) -> None: + for module in self.components: + component = cast(RenderComponent, module) + component.zero_optimizer_grad() + + def _iter_named_components(self) -> list[tuple[str, RenderComponent]]: + """ + Return canonical component names paired with components. + + Returns + ------- + list[tuple[str, RenderComponent]] + ``(name, component)`` entries using the model's canonical naming + rule. Names fall back to class-name/index behavior when ``.name`` is + missing. + + Raises + ------ + RuntimeError + If the model is not defined. + """ + entries: list[tuple[str, RenderComponent]] = [] + for idx, module in enumerate(self.components): + component = cast(RenderComponent, module) + name = self._component_constraint_name(component, idx) + entries.append((name, component)) + return entries + + def _resolve_component_by_name(self, component_name: str) -> RenderComponent: + target = str(component_name) + for resolved_name, component in self._iter_named_components(): + if resolved_name == target: + return component + + for resolved_name, component in self._iter_named_components(): + if component.__class__.__name__ == target: + return component + + known = ", ".join(self.get_component_names()) + raise KeyError(f"Component not found: {target}. Known components: {known}") + + def get_component_names(self) -> list[str]: + """ + Return canonical component names. + + Returns + ------- + list[str] + Canonical component names. + """ + return [name for name, _ in self._iter_named_components()] + + def get_component_by_name(self, component_name: str) -> RenderComponent: + return self._resolve_component_by_name(component_name) + + + def set_component_trainable( + self, + component_name: str, + enabled: bool, + rebuild_optimizer: bool = True, + num_iter: int | None = None, + ) -> None: + """ + Enable or disable optimization for all parameters in one component. + + Parameters + ---------- + component_name : str + Resolved component name. + enabled : bool + If ``True``, mark component parameters trainable. + rebuild_optimizer : bool, optional + If ``True``, rebuild optimizer param groups after toggling. + + Returns + ------- + None + + Raises + ------ + RuntimeError + If the model is not defined. + KeyError + If ``component_name`` is unknown. + + Notes + ----- + When rebuilding, the optimizer is reconstructed from stored optimizer + parameters if available, otherwise inferred from the current optimizer + type and learning rate, else defaults. Scheduler state is cleared + predictably by setting scheduler type to ``"none"``. + """ + component = self._resolve_component_by_name(component_name) + for _, param in component.named_parameters(recurse=True): + param.requires_grad_(bool(enabled)) + if rebuild_optimizer: + component._rebuild_optimizer_after_trainability_change() + + def set_parameter_trainable( + self, + component_name: str, + parameter_name: str, + enabled: bool, + rebuild_optimizer: bool = True, + num_iter: int | None = None, + ) -> None: + """ + Enable or disable optimization for one component parameter. + + Parameters + ---------- + component_name : str + Resolved component name. + parameter_name : str + Parameter name from ``component.named_parameters()``. + enabled : bool + If ``True``, mark parameter trainable. + rebuild_optimizer : bool, optional + If ``True``, rebuild optimizer param groups after toggling. + + Returns + ------- + None + + Raises + ------ + RuntimeError + If the model is not defined. + KeyError + If ``component_name`` or ``parameter_name`` is unknown. + + Notes + ----- + When rebuilding, scheduler state is cleared by setting scheduler type + to ``"none"``. + """ + component = self._resolve_component_by_name(component_name) + params = dict(component.named_parameters(recurse=True)) + if parameter_name not in params: + known = ", ".join(sorted(params.keys())) + raise KeyError( + f"Parameter '{parameter_name}' not found in component '{component_name}'. " + f"Known parameters: {known}" + ) + params[parameter_name].requires_grad_(bool(enabled)) + if rebuild_optimizer: + component._rebuild_optimizer_after_trainability_change() + + def set_parameters_trainable( + self, + component_name: str, + parameter_names: list[str], + enabled: bool, + rebuild_optimizer: bool = True, + num_iter: int | None = None, + ) -> None: + """ + Enable or disable optimization for multiple component parameters. + + Parameters + ---------- + component_name : str + Resolved component name. + parameter_names : list[str] + Parameter names from ``component.named_parameters()``. + enabled : bool + If ``True``, mark parameters trainable. + rebuild_optimizer : bool, optional + If ``True``, rebuild optimizer param groups after toggling. + + Returns + ------- + None + + Raises + ------ + RuntimeError + If the model is not defined. + KeyError + If any parameter name is unknown. + """ + component = self._resolve_component_by_name(component_name) + params = dict(component.named_parameters(recurse=True)) + missing = [name for name in parameter_names if name not in params] + if missing: + known = ", ".join(sorted(params.keys())) + raise KeyError( + f"Unknown parameters for component '{component_name}': {', '.join(missing)}. " + f"Known parameters: {known}" + ) + for name in parameter_names: + params[name].requires_grad_(bool(enabled)) + if rebuild_optimizer: + component._rebuild_optimizer_after_trainability_change() + + def get_component_trainable(self, component_name: str) -> dict[str, bool]: + """ + Return trainability flags for one component's parameters. + + Parameters + ---------- + component_name : str + Resolved component name. + + Returns + ------- + dict[str, bool] + Mapping of parameter name to ``requires_grad``. + + Raises + ------ + RuntimeError + If the model is not defined. + KeyError + If ``component_name`` is unknown. + """ + component = self._resolve_component_by_name(component_name) + return {name: bool(param.requires_grad) for name, param in component.named_parameters()} + + def apply_constraint_configs( + self, constraint_configs: dict[str, Any], strict: bool = True + ) -> None: + for component_name, param in constraint_configs.items(): + component = self._resolve_component_by_name(str(component_name)) + component.initialize_constraint_config(param, strict=strict) + + + + + + + + +@dataclass +class FitResult: + losses: list[float] + lrs: list[float] + final_loss: float + num_steps: int + metrics: dict[str, list[float]] = field(default_factory=dict) + +class BCEMSELoss(nn.Module): + + def __init__( + self, + gamma: float = 1, + percentile: float = 0.95, + ): + super().__init__() + self.gamma = gamma + self.percentile = percentile + + def forward(self, pred, target): + mse_loss = torch.mean((pred - target) ** 2) + + threshold_pred = torch.quantile(pred, self.percentile) + threshold_target = torch.quantile(target, self.percentile) + + # Create binary masks + target_binary = (target > threshold_target).float() + pred_binary = (pred > threshold_pred).float() + + epsilon = 1e-7 + pred_binary_clamped = torch.clamp(pred_binary, epsilon, 1.0 - epsilon) + bce_loss = -torch.mean( + target_binary * torch.log(pred_binary_clamped) + + (1 - target_binary) * torch.log(1 - pred_binary_clamped) + ) + + # Combined loss + total_loss = mse_loss + self.gamma ** bce_loss + + return total_loss + +class SqrtMSELoss(nn.Module): + def __init__( + self, + gamma: float = 0.25, + ): + super().__init__() + self.gamma = gamma + self.mse_fn = torch.nn.MSELoss(reduction="mean") + + def forward(self, pred, target): + eps = 1 + pred_modified = (pred - pred.min() + eps) ** self.gamma + target_modified = (target - target.min() + eps) ** self.gamma + loss = self.mse_fn(pred_modified, target_modified) + return loss + +class LogMSELoss(nn.Module): + def __init__( + self, + ): + super().__init__() + self.mse_fn = torch.nn.MSELoss(reduction="mean") + + def forward(self, pred, target): + return self.mse_fn(torch.log(1+pred), torch.log(1+target)) + +class WeightedMSELoss(nn.Module): + """Signal-weighted MSE: residuals weighted by ``(target / mean target)**power + floor``. + + Power-compressed data (``**0.25`` / sqrt preprocessing) inflates the noisy near-zero + pixels; weighting by the target intensity puts the emphasis back on the diffraction + disks, where the lattice-position information is. Weights are normalized to mean 1. + """ + + def __init__(self, power: float = 1.0, floor: float = 0.05): + super().__init__() + self.power = float(power) + self.floor = float(floor) + + def weights(self, target: torch.Tensor) -> torch.Tensor: + dims = tuple(range(target.ndim - 2, target.ndim)) + t = target.clamp(min=0) + w = (t / t.mean(dims, keepdim=True).clamp(min=1e-12)) ** self.power + self.floor + return w / w.mean(dims, keepdim=True) + + def forward(self, pred, target): + return (self.weights(target) * (pred - target) ** 2).mean() + + +class FitBase(OptimizerMixin): + + def __init__(self): + super().__init__() + # Core wiring + # self.loss_fn = torch.nn.L1Loss(reduction="mean") + # self.loss_fn = torch.nn.MSELoss(reduction="mean") + self.loss_fn = SqrtMSELoss() + self.model: AdditiveRenderModel | None = None + self.ctx: RenderContext | None = None + + # State/checkpoints + self.state_initialized: dict[str, torch.Tensor] | None = None + + # Histories/results + self.fit_history: dict[str, FitResult] = {} + + # self.multi_optimizers: dict[str, torch.optim.Optimizer] | None = None + # self.use_mutiple_optimizers: bool = False + + def get_optimization_parameters(self) -> Any: + if self.model is None: + return [] + + return [p for p in self.model.parameters() if p.requires_grad] + + @property + def state_current(self) -> dict[str, torch.Tensor] | None: + if self.model is None: + return None + return self._get_model_state_dict_copy() + + @property + def render_initialized(self) -> np.ndarray: + if self.state_initialized is None: + raise RuntimeError("initialized state is unavailable. Call .define_model(...) first.") + return self._render_state_array(self.state_initialized) + + @property + def render_current(self) -> np.ndarray: + if self.model is None or self.ctx is None: + raise RuntimeError("Call .define_model(...) first.") + return self.model(self.ctx).detach().cpu().numpy() + + def set_loss(self, loss_fn: str = "sqrtmse", gamma: float = 0.25, power: float = 1.0, floor: float = 0.05): + mode_in = loss_fn.strip().lower() + if mode_in in {"mse"}: + self.loss_fn = torch.nn.MSELoss(reduction="mean") + elif mode_in in {"sqrtmse"}: + self.loss_fn = SqrtMSELoss(gamma = gamma) + elif mode_in in {"logmse"}: + self.loss_fn = LogMSELoss() + elif mode_in in {"l1"}: + self.loss_fn = torch.nn.L1Loss(reduction="mean") + elif mode_in in {"wmse"}: + self.loss_fn = WeightedMSELoss(power=power, floor=floor) + else: + raise ValueError( + "loss function must be mse, sqrtmse, logmse, l1 or wmse" + ) + + def reset( + self, + reset_to: Literal["initialized"] = "initialized", + ) -> Self: + if reset_to != "initialized": + raise ValueError("FitBase.reset only supports reset_to='initialized'.") + if self.state_initialized is None: + raise RuntimeError("initialized state is unavailable. Call .define_model(...) first.") + self._load_model_state_dict_copy(self.state_initialized) + self._clear_fit_history_all() + return self + + def fit_render( + self, + *, + target: torch.Tensor, + n_steps: int, + constraint_weight: float = 1.0, + constraint_params: dict[str, Any] | None = None, + constraint_config_params: dict[str, Any] | None = None, + optimizer_params: dict[str, dict[str, Any]] | None = None, + scheduler_params: dict[str, dict[str, Any]] | None = None, + progress: bool = False, + run_key: str = "default", + **kwargs: Any, + ) -> FitResult: + """ + Fit model parameters to a target render. + + Parameters + ---------- + target : torch.Tensor + Target tensor to fit. + n_steps : int + Number of optimization steps. + constraint_weight : float, optional + Multiplier applied to the summed soft-constraint loss. + constraint_params : dict[str, Any] | None, optional + Optional constraint updates applied once to matching components before + optimization starts. If ``None``, existing component constraints are reused. + optimizer_params : dict | None, optional + Optimizer configuration override for this call. + scheduler_params : dict | None, optional + Scheduler configuration override for this call. + progress : bool, optional + If ``True``, display a progress bar. + run_key : str, optional + History key used to store/append fit metrics. + **kwargs : Any + Forwarded to internal forward/loss hooks. + + Returns + ------- + FitResult + Fit history and final loss metadata for this run key. + + Raises + ------ + RuntimeError + If model/context are undefined. + + Notes + ----- + Hard constraints are applied after each optimizer step. + """ + if self.model is None or self.ctx is None: + raise RuntimeError("Model and context are not defined for fitting.") + if constraint_params is not None: + self.model.apply_constraint_params(constraint_params, strict=True) + if constraint_config_params is not None: + self.model.apply_constraint_configs(constraint_config_params, strict=True) + + n_steps = int(n_steps) + self.model.initilize_independant_optimizers( + optimizer_params, + scheduler_params, + num_iter=n_steps + ) + + pbar = tqdm(range(n_steps), desc="Fit render", disable=not progress) + loss_vals = torch.empty(n_steps, device=self.ctx.device) + + losses: list[float] = [] + lrs: list[float] = [] + for step in pbar: + self.model.zero_grad_optimizers() + pred = self._forward_for_fit(target=target, **kwargs) + data_loss = self._fidelity_loss(pred, target, **kwargs) + constraint_loss = self._constraint_loss(pred, target, **kwargs) + total_loss = data_loss + constraint_weight * constraint_loss + total_loss.backward() + self.model.step_optimizers() + if self.model is None or self.ctx is None: + raise RuntimeError("Model and context are not defined for fitting.") + self.model.apply_hard_constraints(self.ctx) + total_loss_value = (total_loss.detach()) + self.model.step_schedulers(total_loss_value) + # losses.append(total_loss_value) + loss_vals[step] = total_loss_value + + first_lr = 0.0 + if len(self.model.components) > 0: + first_comp = cast(RenderComponent, self.model.components[0]) + if hasattr(first_comp, 'get_current_lr'): + try: + first_lr = float(first_comp.get_current_lr()) + except (AttributeError, TypeError): + first_lr = 0.0 + lrs.append(first_lr) + losses = loss_vals.cpu().tolist() + + key = str(run_key) + if key in self.fit_history: + prev = self.fit_history[key] + prev.losses.extend(losses) + prev.lrs.extend(lrs) + prev.final_loss = prev.losses[-1] if prev.losses else float("nan") + prev.num_steps = len(prev.losses) + result = prev + else: + result = FitResult( + losses=losses, + lrs=lrs, + final_loss=(losses[-1] if losses else float("nan")), + num_steps=n_steps, + ) + self.fit_history[key] = result + return result + + + def _clone_state_dict(self, state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return {k: v.detach().clone() for k, v in state.items()} + + def _get_model_state_dict_copy(self) -> dict[str, torch.Tensor]: + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + return self._clone_state_dict(self.model.state_dict()) + + def _load_model_state_dict_copy(self, state: dict[str, torch.Tensor]) -> None: + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + self.model.load_state_dict(self._clone_state_dict(state), strict=True) + + def _clear_fit_history_all(self) -> None: + self.fit_history.clear() + + def _clear_fit_history_run(self, run_key: str) -> None: + self.fit_history.pop(str(run_key), None) + + def _render_state_array(self, state: dict[str, torch.Tensor]) -> np.ndarray: + if self.model is None or self.ctx is None: + raise RuntimeError("Call .define_model(...) first.") + live = self._get_model_state_dict_copy() + try: + self._load_model_state_dict_copy(state) + arr = self.model(self.ctx).detach().cpu().numpy() + finally: + self._load_model_state_dict_copy(live) + return arr + + def _forward_for_fit(self, *, target: torch.Tensor, **kwargs: Any) -> torch.Tensor: + if self.model is None or self.ctx is None: + raise RuntimeError("Model and context are not defined for fitting.") + return self.model(self.ctx) + + def _fidelity_loss( + self, pred: torch.Tensor, target: torch.Tensor, **kwargs: Any + ) -> torch.Tensor: + if self.ctx is not None and self.ctx.mask is not None: + valid = self.ctx.mask.bool() + return self.loss_fn(pred[valid], target[valid]) + return self.loss_fn(pred, target) + + def _constraint_loss( + self, + pred: torch.Tensor, + target: torch.Tensor, + **kwargs: Any, + ) -> torch.Tensor: + if self.model is None or self.ctx is None: + raise RuntimeError("Model and context are not defined for fitting.") + return self.model.total_constraint_loss(self.ctx, **kwargs) + + def set_component_trainable( + self, + component_name: str, + enabled: bool, + rebuild_optimizer: bool = True, + num_iter: int | None = None + ) -> None: + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + self.model.set_component_trainable(component_name, enabled, rebuild_optimizer, num_iter) + + def set_parameter_trainable( + self, + component_name: str, + parameter_name: str, + enabled: bool, + rebuild_optimizer: bool = True, + num_iter: int | None = None + ) -> None: + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + self.model.set_parameter_trainable(component_name, parameter_name, enabled, rebuild_optimizer, num_iter) + + def set_parameters_trainable( + self, + component_name: str, + parameter_names: list[str], + enabled: bool, + rebuild_optimizer: bool = True, + num_iter: int | None = None + ) -> None: + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + self.model.set_parameters_trainable(component_name, parameter_names, enabled, rebuild_optimizer, num_iter) + + def get_component_trainable(self, component_name: str) -> dict[str, bool]: + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + return self.model.get_component_trainable(component_name) + + def get_component_names(self) -> list[str]: + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + return self.model.get_component_names() + + def _iter_named_components(self) -> list[tuple[str, RenderComponent]]: + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + return self.model._iter_named_components() + + def _resolve_component_by_name(self, component_name: str) -> RenderComponent: + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + return self.model._resolve_component_by_name(component_name) + + def get_component_by_name(self, component_name: str) -> RenderComponent: + return self._resolve_component_by_name(component_name) + + def _rebuild_optimizer_after_trainability_change(self) -> None: + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + self.model.rebuild_independant_optimizers() + + def apply_constraint_config_params( + self, constraint_configs: dict[str, Any], strict: bool = True + ) -> None: + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + self.model.apply_constraint_configs(constraint_configs, strict=strict) + + +Component = RenderComponent +ModelContext = RenderContext +Model = AdditiveRenderModel +Parameter = nn.Parameter \ No newline at end of file diff --git a/src/quantem/core/fitting/diffraction.py b/src/quantem/core/fitting/diffraction.py new file mode 100644 index 00000000..18a232d5 --- /dev/null +++ b/src/quantem/core/fitting/diffraction.py @@ -0,0 +1,1302 @@ +from __future__ import annotations + +from typing import Any, Iterable, Sequence, cast + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +from quantem.core.fitting.base import OriginND, RenderComponent, RenderContext + + +def _splat_patch( + out: torch.Tensor, + *, + r0: torch.Tensor, + c0: torch.Tensor, + patch_vals: torch.Tensor, + dr: torch.Tensor, + dc: torch.Tensor, + scale: torch.Tensor, +) -> None: + h, w = out.shape + r = r0 + dr + c = c0 + dc + + r_base = torch.floor(r) + c_base = torch.floor(c) + fr = r - r_base + fc = c - c_base + r0i = r_base.to(torch.long) + c0i = c_base.to(torch.long) + + w00 = (1.0 - fr) * (1.0 - fc) + w01 = (1.0 - fr) * fc + w10 = fr * (1.0 - fc) + w11 = fr * fc + v = patch_vals * scale + + def put(rr: torch.Tensor, cc: torch.Tensor, ww: torch.Tensor) -> None: + keep = (rr >= 0) & (rr < h) & (cc >= 0) & (cc < w) + if torch.any(keep): + out.index_put_((rr[keep], cc[keep]), v[keep] * ww[keep], accumulate=True) + + put(r0i, c0i, w00) + put(r0i, c0i + 1, w01) + put(r0i + 1, c0i, w10) + put(r0i + 1, c0i + 1, w11) + + +def _gaussian_tap_weights( + frac: torch.Tensor, taps: torch.Tensor, sigma: float +) -> torch.Tensor: + """ + Normalized 1D Gaussian weights from fractional positions to integer taps. + + Parameters + ---------- + frac : torch.Tensor + Fractional parts in ``[0, 1)``, any shape ``(...,)``. + taps : torch.Tensor + Integer tap offsets, shape ``(T,)``. + sigma : float + Gaussian width in pixels. + + Returns + ------- + torch.Tensor + Weights of shape ``(..., T)`` summing to 1 along the last axis, so the + splatted flux is conserved and (unlike bilinear) the effective blur is + independent of the subpixel phase — no integer-pixel "locking" minima. + """ + d = taps.reshape((1,) * frac.ndim + (-1,)) - frac.unsqueeze(-1) + w = torch.exp(-(d * d) / (2.0 * sigma * sigma)) + return w / w.sum(dim=-1, keepdim=True) + + +def _smooth_tap_weights( + frac: torch.Tensor, taps: torch.Tensor, width: float +) -> torch.Tensor: + """ + Normalized finite-support smoothstep splat weights. + + A compact alternative to the Gaussian tap weights with NO tails: the weight + is ``smoothstep(1 - |d|/width)`` (``smoothstep(x) = 3x^2 - 2x^3``), which is + exactly zero for ``|d| >= width`` and, for ``width > 1``, still nonzero at + the neighbouring integer taps when centred -- so the effective blur stays + phase-independent (no integer-pixel locking) while the kernel has finite + bandwidth. The disk edge itself is carried by the template, so this kernel + only interpolates sub-pixel position and provides anti-lock. + """ + d = (taps.reshape((1,) * frac.ndim + (-1,)) - frac.unsqueeze(-1)).abs() + x = torch.clamp(1.0 - d / float(width), min=0.0) + w = x * x * (3.0 - 2.0 * x) + return w / w.sum(dim=-1, keepdim=True) + + +def _kernel_taps_weights(frac, width, kernel, device, dtype): + """Return (taps, weights) for the requested finite render kernel.""" + if kernel == "smoothstep": + K = max(2, int(np.ceil(float(width)))) + taps = torch.arange(-K + 1, K + 1, device=device, dtype=dtype) + return taps, _smooth_tap_weights(frac, taps, width) + K = max(2, int(np.ceil(3.0 * float(width)))) + taps = torch.arange(-K + 1, K + 1, device=device, dtype=dtype) + return taps, _gaussian_tap_weights(frac, taps, width) + + +def _splat_patch_gaussian( + out: torch.Tensor, + *, + r0: torch.Tensor, + c0: torch.Tensor, + patch_vals: torch.Tensor, + dr: torch.Tensor, + dc: torch.Tensor, + scale: torch.Tensor, + sigma: float, + kernel: str = "gaussian", +) -> None: + """Finite-kernel analogue of ``_splat_patch`` (phase-independent blur).""" + h, w = out.shape + r = r0 + dr + c = c0 + dc + r_base = torch.floor(r) + c_base = torch.floor(c) + fr = r - r_base + fc = c - c_base + r0i = r_base.to(torch.long) + c0i = c_base.to(torch.long) + v = patch_vals * scale + + taps, wr = _kernel_taps_weights(fr, sigma, kernel, out.device, out.dtype) + _, wc = _kernel_taps_weights(fc, sigma, kernel, out.device, out.dtype) + taps_i = taps.to(torch.long) + + for i in range(taps_i.numel()): + rr = r0i + taps_i[i] + r_ok = (rr >= 0) & (rr < h) + for j in range(taps_i.numel()): + cc = c0i + taps_i[j] + keep = r_ok & (cc >= 0) & (cc < w) + if torch.any(keep): + ww = wr[:, i] * wc[:, j] + out.index_put_( + (rr[keep], cc[keep]), v[keep] * ww[keep], accumulate=True + ) + + +def _splat_patch_batched_gaussian( + shape: tuple[int, int], + *, + r0: torch.Tensor, + c0: torch.Tensor, + vals: torch.Tensor, + device: torch.device, + dtype: torch.dtype, + sigma: float, + kernel: str = "gaussian", +) -> torch.Tensor: + """Finite-kernel analogue of ``_splat_patch_batched``.""" + h, w = int(shape[0]), int(shape[1]) + B, N = r0.shape + r_base = torch.floor(r0) + c_base = torch.floor(c0) + fr = r0 - r_base + fc = c0 - c_base + r0i = r_base.to(torch.long) + c0i = c_base.to(torch.long) + + taps, wr = _kernel_taps_weights(fr, sigma, kernel, device, dtype) + _, wc = _kernel_taps_weights(fc, sigma, kernel, device, dtype) + taps_i = taps.to(torch.long) + + out_flat = torch.zeros(B, h * w, device=device, dtype=dtype) + for i in range(taps_i.numel()): + rr = r0i + taps_i[i] + r_ok = (rr >= 0) & (rr < h) + rr_c = rr.clamp(0, h - 1) + for j in range(taps_i.numel()): + cc = c0i + taps_i[j] + keep = r_ok & (cc >= 0) & (cc < w) + weighted = wr[:, :, i] * wc[:, :, j] * vals * keep.to(dtype) + out_flat.scatter_add_(1, rr_c * w + cc.clamp(0, w - 1), weighted) + return out_flat.reshape(B, h, w) + + +def _splat_patch_batched( + shape: tuple[int, int], + *, + r0: torch.Tensor, + c0: torch.Tensor, + vals: torch.Tensor, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + h, w = int(shape[0]), int(shape[1]) + B, N = r0.shape + + r_base = torch.floor(r0) + c_base = torch.floor(c0) + fr = r0 - r_base + fc = c0 - c_base + r0i = r_base.to(torch.long) + c0i = c_base.to(torch.long) + + w00 = (1.0 - fr) * (1.0 - fc) + w01 = (1.0 - fr) * fc + w10 = fr * (1.0 - fc) + w11 = fr * fc + + rr = torch.stack([r0i, r0i, r0i + 1, r0i + 1], dim=0) + cc = torch.stack([c0i, c0i + 1, c0i, c0i + 1], dim=0) + ww = torch.stack([w00, w01, w10, w11], dim=0) + + keep = (rr >= 0) & (rr < h) & (cc >= 0) & (cc < w) + weighted = ww * vals.unsqueeze(0) * keep.to(dtype) + + rr_c = rr.clamp(0, h - 1) + cc_c = cc.clamp(0, w - 1) + flat_idx = rr_c * w + cc_c + + flat_idx_b = flat_idx.permute(1, 0, 2).reshape(B, -1) + weighted_b = weighted.permute(1, 0, 2).reshape(B, -1) + + out_flat = torch.zeros(B, h * w, device=device, dtype=dtype) + out_flat.scatter_add_(1, flat_idx_b, weighted_b) + return out_flat.reshape(B, h, w) + + +class DiskTemplate(RenderComponent): + DEFAULT_HARD_CONSTRAINTS: dict[str, bool] = { + "force_center": False, + "force_positive": True, + "force_norm": True, # force range [0,1] + "force_shrinkage": False, + "force_cutoff": False, + "force_circular_mask": True, + } + DEFAULT_SOFT_CONSTRAINTS: dict[str, float] = { + "tv_weight": 0.0, + "cutoff_weight": 0.0, + "circular_weight": 0.0, + } + DEFAULT_CONSTRAINT_CONFIG: dict[str, float] = { + "soft_cutoff_threshold": 0.0, + "soft_cutoff_target_ratio": 0.1, + "hard_cutoff_threshold": 0.35, + "shrinkage_amount": 0.25, + "circular_mask_radius_fraction": 0.95, + "circular_mask_sharpness": 0, + "soft_circular_mask": False, + } + + def __init__( + self, + *, + name: str, + array: np.ndarray, + refine_all_pixels: bool = False, + normalize: str = "none", + origin: OriginND | None = None, + origin_key: str = "origin", + intensity: float | Sequence[float] = 1.0, + render_sigma: float | None = None, + render_kernel: str = "gaussian", + constraint_params: dict[str, Any] | None = None, + constraint_config: dict[str, Any] | None = None, + ): + """ + Build a disk template renderer centered at the shared origin. + + Parameters + ---------- + intensity : float | Sequence[float], optional + Trainable scalar amplitude applied to the rendered template. + Accepts ``x``, ``(x0, delta)``, or ``(x0, lo, hi)``. + render_sigma : float | None, optional + If set, splat template pixels with a normalized Gaussian kernel of + this width (px) instead of bilinear interpolation. The Gaussian + blur is identical at every subpixel phase, which removes the + integer-pixel "locking" minima of bilinear splatting. Applies to + this template and to any lattice that renders it. + + Returns + ------- + None + + Raises + ------ + ValueError + If ``array`` is not 2D or if ``normalize`` is unsupported. + + Notes + ----- + ``template_raw`` controls template shape and ``intensity_raw`` controls + center-disk amplitude. + """ + super().__init__() + self.name = str(name) + self.refine_all_pixels = bool(refine_all_pixels) + self.render_sigma = None if render_sigma is None else float(render_sigma) + # "gaussian" (tails) or "smoothstep" (finite support, no tails); with + # smoothstep, render_sigma is the half-support in px (needs > 1 to + # anti-lock). The disk edge is carried by the template either way. + self.render_kernel = str(render_kernel) + self.origin = origin + self.origin_key = str(origin_key) + intensity_init, intensity_lo, intensity_hi = self.parse_bounded_init( + intensity, name="intensity" + ) + self.intensity_raw = nn.Parameter(torch.tensor(intensity_init, dtype=torch.float32)) + if intensity_lo is not None or intensity_hi is not None: + self.register_parameter_bounds("intensity_raw", intensity_lo, intensity_hi) + + a = np.asarray(array, dtype=np.float32) + if a.ndim != 2: + raise ValueError("DiskTemplate.array must be 2D.") + if normalize == "max": + s = float(np.max(a)) + if s > 0.0: + a = a / s + elif normalize == "mean": + s = float(np.mean(a)) + if s != 0.0: + a = a / s + elif normalize != "none": + raise ValueError("normalize must be one of: 'none', 'max', 'mean'.") + + template = torch.as_tensor(a, dtype=torch.float32) + self.template_raw = nn.Parameter(template.clone(), requires_grad=self.refine_all_pixels) + + ht, wt = int(template.shape[0]), int(template.shape[1]) + rr, cc = np.mgrid[0:ht, 0:wt] + rr = rr.astype(np.float32) - (ht - 1) * 0.5 + cc = cc.astype(np.float32) - (wt - 1) * 0.5 + self.register_buffer("dr", torch.as_tensor(rr.ravel(), dtype=torch.float32)) + self.register_buffer("dc", torch.as_tensor(cc.ravel(), dtype=torch.float32)) + + self.constraint_config = self.DEFAULT_CONSTRAINT_CONFIG.copy() + if constraint_config is not None: + self.constraint_config.update(constraint_config) + + if constraint_params is not None: + self.apply_constraint_params(constraint_params, strict=True) + if bool(self.hard_constraints.get("force_positive", False)): + self._enforce_positivity() + if bool(self.hard_constraints.get("force_shrinkage", False)) and bool(self.hard_constraints.get("force_positive", True)): + raise RuntimeWarning("Setting shrinkage true and positivity false might cause negative values in disk template") + + @classmethod + def from_array( + cls, + *, + name: str, + array: np.ndarray, + refine_all_pixels: bool = False, + normalize: str = "none", + origin: OriginND | None = None, + origin_key: str = "origin", + intensity: float | Sequence[float] = 1.0, + render_sigma: float | None = None, + render_kernel: str = "gaussian", + constraint_params: dict[str, Any] | None = None, + constraint_config: dict[str, Any] | None = None, + + ) -> "DiskTemplate": + return cls( + name=name, + array=array, + refine_all_pixels=refine_all_pixels, + normalize=normalize, + origin=origin, + origin_key=origin_key, + intensity=intensity, + render_sigma=render_sigma, + render_kernel=render_kernel, + constraint_params=constraint_params, + constraint_config=constraint_config, + + ) + + def set_origin(self, origin: OriginND) -> None: + self.origin = origin + + def set_intensity(self, value: float | int) -> None: + """Assign ``intensity_raw`` in-place.""" + with torch.no_grad(): + self.intensity_raw.copy_(torch.as_tensor(float(value), dtype=self.intensity_raw.dtype)) + + def patch_values(self) -> torch.Tensor: + return self.template_raw.reshape(-1) + + def patch_offsets(self) -> tuple[torch.Tensor, torch.Tensor]: + return cast(torch.Tensor, self.dr), cast(torch.Tensor, self.dc) + + def add_patch( + self, out: torch.Tensor, *, r0: torch.Tensor, c0: torch.Tensor, scale: torch.Tensor + ) -> None: + vals = self.patch_values().to(device=out.device, dtype=out.dtype) + dr = cast(torch.Tensor, self.dr).to(device=out.device, dtype=out.dtype) + dc = cast(torch.Tensor, self.dc).to(device=out.device, dtype=out.dtype) + if self.render_sigma is not None: + _splat_patch_gaussian( + out, r0=r0, c0=c0, patch_vals=vals, dr=dr, dc=dc, scale=scale, + sigma=self.render_sigma, + kernel=self.render_kernel, + ) + else: + _splat_patch(out, r0=r0, c0=c0, patch_vals=vals, dr=dr, dc=dc, scale=scale) + + def forward(self, ctx: RenderContext) -> torch.Tensor: + """ + Render template at origin with scalar amplitude ``intensity_raw``. + + Parameters + ---------- + ctx : RenderContext + Rendering context. + + Returns + ------- + torch.Tensor + Rendered center disk image. + """ + out = torch.zeros(ctx.shape, device=ctx.device, dtype=ctx.dtype) + if self.origin is None: + raise RuntimeError("DiskTemplate.forward() requires an OriginND instance.") + r0, c0 = self.origin.coords[0], self.origin.coords[1] + scale = self.intensity_raw.to(device=ctx.device, dtype=ctx.dtype) + self.add_patch(out, r0=r0, c0=c0, scale=scale) + return out + + def forward_batched( + self, + ctx: RenderContext, + *, + template_raw_b: torch.Tensor, + intensity_raw_b: torch.Tensor, + origin_coords_b: torch.Tensor, + ) -> torch.Tensor: + B = template_raw_b.shape[0] + N = int(cast(torch.Tensor, self.dr).numel()) + dr = cast(torch.Tensor, self.dr).to(device=ctx.device, dtype=ctx.dtype) + dc = cast(torch.Tensor, self.dc).to(device=ctx.device, dtype=ctx.dtype) + r0 = origin_coords_b[:, 0:1] + dr.unsqueeze(0) + c0 = origin_coords_b[:, 1:2] + dc.unsqueeze(0) + vals = template_raw_b.reshape(B, N) * intensity_raw_b.view(B, 1) + if self.render_sigma is not None: + return _splat_patch_batched_gaussian( + ctx.shape, r0=r0, c0=c0, vals=vals, device=ctx.device, + dtype=ctx.dtype, sigma=self.render_sigma, kernel=self.render_kernel, + ) + return _splat_patch_batched( + ctx.shape, r0=r0, c0=c0, vals=vals, device=ctx.device, dtype=ctx.dtype + ) + + def _center_disk(self) -> None: + with torch.no_grad(): + template = self.template_raw + h, w = int(template.shape[0]), int(template.shape[1]) + weights = torch.clamp(template, min=0.0) + mass = torch.sum(weights) + if float(mass.detach().cpu()) <= 1e-12: + return + rr = torch.arange(h, device=template.device, dtype=template.dtype)[:, None] + cc = torch.arange(w, device=template.device, dtype=template.dtype)[None, :] + com_r = torch.sum(weights * rr) / mass + com_c = torch.sum(weights * cc) / mass + target_r = torch.as_tensor((h - 1) * 0.5, device=template.device, dtype=template.dtype) + target_c = torch.as_tensor((w - 1) * 0.5, device=template.device, dtype=template.dtype) + shift_r = target_r - com_r + shift_c = target_c - com_c + denom_h = max(h - 1, 1) + denom_w = max(w - 1, 1) + ty = -2.0 * shift_r / float(denom_h) + tx = -2.0 * shift_c / float(denom_w) + theta = torch.as_tensor( + [[1.0, 0.0, tx], [0.0, 1.0, ty]], + device=template.device, + dtype=template.dtype, + )[None, ...] + src = template[None, None, :, :] + grid = F.affine_grid(theta, [1, 1, h, w], align_corners=True) + shifted = F.grid_sample( + src, + grid, + mode="bilinear", + padding_mode="zeros", + align_corners=True, + )[0, 0] + self.template_raw.copy_(shifted) + + def _enforce_positivity(self) -> None: + with torch.no_grad(): + self.template_raw.clamp_(min=0.0) + self.intensity_raw.clamp_(min=0.0) + + def _enforce_norm(self) -> None: # pick value and cut off 5 percent of mean, or every iteration shrinkage, every iteration just subtract a valye of 0.01 + with torch.no_grad(): + self.template_raw -= self.template_raw.min() + self.template_raw /= self.template_raw.max() + + def _enforce_shrinkage(self) -> None: + with torch.no_grad(): + self.template_raw -= self.constraint_config["shrinkage_amount"] + + def _enforce_cutoff(self) -> None: + with torch.no_grad(): + mean_val = torch.max(self.template_raw) * self.constraint_config["hard_cutoff_threshold"] + self.template_raw[self.template_raw <= mean_val] = 0 + + def _enforce_circular_mask(self) -> None: + with torch.no_grad(): + h, w = self.template_raw.shape + radius = (min(h, w) / 2.0) * self.constraint_config["circular_mask_radius_fraction"] + + r = torch.arange(-h/2, h/2, device=self.template_raw.device, dtype=self.template_raw.dtype) + c = torch.arange(-w/2, w/2, device=self.template_raw.device, dtype=self.template_raw.dtype) + rr, cc = torch.meshgrid(r, c, indexing='ij') + circle_matrix = torch.sqrt(rr**2 + cc**2) + + if self.constraint_config["soft_circular_mask"]: + mask = torch.sigmoid(self.constraint_config["circular_mask_sharpness"]*(radius-circle_matrix)) + else: + mask = circle_matrix <= radius + self.template_raw *= mask + + + def enforce_hard_constraints(self, ctx: RenderContext) -> None: + if bool(self.hard_constraints.get("force_center", False)): + self._center_disk() + if bool(self.hard_constraints.get("force_cutoff", False)): + self._enforce_cutoff() + if bool(self.hard_constraints.get("force_circular_mask", False)): + self._enforce_circular_mask() + if bool(self.hard_constraints.get("force_shrinkage", False)): + self._enforce_shrinkage() + if bool(self.hard_constraints.get("force_positive", False)): + self._enforce_positivity() + if bool(self.hard_constraints.get("force_norm", False)): # could be put in positivity + self._enforce_norm() + + super().enforce_hard_constraints(ctx) + + def constraint_loss( + self, ctx: RenderContext, params: dict[str, object] | None = None + ) -> torch.Tensor: + cfg = self.effective_soft_constraints(cast(dict[str, object] | None, params)) + tv_weight = float(cfg.get("tv_weight", 0.0)) + cutoff_weight = float(cfg.get("cutoff_weight", 0.0)) + circular_weight = float(cfg.get("circular_weight", 0.0)) + + if tv_weight <= 0.0: + tv_weight = 0.0 + if cutoff_weight <= 0.0: + cutoff_weight = 0.0 + if circular_weight <= 0.0: + circular_weight = 0.0 + + # tv loss calculation + template = self.template_raw.to(device=ctx.device, dtype=ctx.dtype) + tv_r = ( + torch.mean(torch.abs(template[1:, :] - template[:-1, :])) + if template.shape[0] > 1 + else torch.zeros((), device=ctx.device, dtype=ctx.dtype) + ) + tv_c = ( + torch.mean(torch.abs(template[:, 1:] - template[:, :-1])) + if template.shape[1] > 1 + else torch.zeros((), device=ctx.device, dtype=ctx.dtype) + ) + tv_loss = torch.as_tensor(tv_weight, device=ctx.device, dtype=ctx.dtype) * (tv_r + tv_c) + + # Cutoff loss calculation + num_px = torch.prod(torch.tensor(template.shape)) + px_under_threshold = torch.sum(template <= torch.mean(template)*self.constraint_config["soft_cutoff_threshold"])/num_px + cutoff_loss = torch.as_tensor(cutoff_weight, device=ctx.device, dtype=ctx.dtype) * torch.maximum( + px_under_threshold-self.constraint_config["soft_cutoff_target_ratio"], torch.as_tensor(0.0, device=ctx.device, dtype=ctx.dtype)) + + # circular loss calculation + h, w = template.shape + radius = (min(h, w) / 2.0) * self.constraint_config["circular_mask_radius_fraction"] + + r = torch.arange(h, device=ctx.device, dtype=ctx.dtype) - h / 2.0 + c = torch.arange(w, device=ctx.device, dtype=ctx.dtype) - w / 2.0 + rr, cc = torch.meshgrid(r, c, indexing='ij') + + circle_mask = torch.sqrt(rr**2 + cc**2) + + dist_from_radius = torch.abs(circle_mask - radius) + dist_from_radius = torch.relu(dist_from_radius) + circular_err = torch.mean(dist_from_radius * template) + + circular_loss = torch.as_tensor(circular_weight, device=ctx.device, dtype=ctx.dtype) * circular_err + + return cutoff_loss + tv_loss + circular_loss + + def constraint_loss_batched( + self, + ctx: RenderContext, + *, + template_raw_b: torch.Tensor, + params: dict[str, object] | None = None, + ) -> torch.Tensor: + """ + Per-sample analogue of ``constraint_loss`` for stacked templates. + + Parameters + ---------- + template_raw_b : torch.Tensor + Stacked templates with shape ``(B, H_t, W_t)``. + + Returns + ------- + torch.Tensor + Per-sample soft-constraint losses with shape ``(B,)``. Identical + semantics to ``constraint_loss`` on each slice, but reductions are + taken over ``dim=(1, 2)`` only. + """ + cfg = self.effective_soft_constraints(cast(dict[str, object] | None, params)) + tv_weight = max(float(cfg.get("tv_weight", 0.0)), 0.0) + cutoff_weight = max(float(cfg.get("cutoff_weight", 0.0)), 0.0) + circular_weight = max(float(cfg.get("circular_weight", 0.0)), 0.0) + + template = template_raw_b.to(device=ctx.device, dtype=ctx.dtype) + B, h, w = template.shape + + if h > 1: + tv_r = torch.mean(torch.abs(template[:, 1:, :] - template[:, :-1, :]), dim=(1, 2)) + else: + tv_r = torch.zeros(B, device=ctx.device, dtype=ctx.dtype) + if w > 1: + tv_c = torch.mean(torch.abs(template[:, :, 1:] - template[:, :, :-1]), dim=(1, 2)) + else: + tv_c = torch.zeros(B, device=ctx.device, dtype=ctx.dtype) + tv_loss = tv_weight * (tv_r + tv_c) + + # Per-sample cutoff (note: hard `<=` is non-differentiable; matches serial). + per_sample_mean = template.mean(dim=(1, 2), keepdim=True) + thresh = per_sample_mean * float(self.constraint_config["soft_cutoff_threshold"]) + frac_under = (template <= thresh).to(dtype=ctx.dtype).mean(dim=(1, 2)) + target_ratio = float(self.constraint_config["soft_cutoff_target_ratio"]) + cutoff_loss = cutoff_weight * torch.relu(frac_under - target_ratio) + + # Per-sample circular: mask depends only on (H, W), so build once. + radius = (min(h, w) / 2.0) * float(self.constraint_config["circular_mask_radius_fraction"]) + r = torch.arange(h, device=ctx.device, dtype=ctx.dtype) - h / 2.0 + c = torch.arange(w, device=ctx.device, dtype=ctx.dtype) - w / 2.0 + rr, cc = torch.meshgrid(r, c, indexing="ij") + circle_mask = torch.sqrt(rr * rr + cc * cc) + dist_from_radius = torch.relu(torch.abs(circle_mask - radius)) # (H, W) + circular_err = (dist_from_radius.unsqueeze(0) * template).mean(dim=(1, 2)) + circular_loss = circular_weight * circular_err + + return tv_loss + cutoff_loss + circular_loss + + def get_optimization_parameters(self) -> dict[str, list[torch.nn.Parameter]]: + params = [] + for name, param in self.named_parameters(recurse=True): + if not name.startswith('origin.') and param.requires_grad: + params.append(param) + if not params: + return {} + return {'default': params} + + + + +class SyntheticDiskLattice(RenderComponent): + DEFAULT_HARD_CONSTRAINTS: dict[str, bool] = { + "force_positive_intensity": True, + } + DEFAULT_SOFT_CONSTRAINTS: dict[str, float] = { + "consistency_weight": 0.0, + "consistency_window": 1, + } + + def __init__( + self, + *, + name: str, + disk: DiskTemplate, + u_row: float | Sequence[float], + u_col: float | Sequence[float], + v_row: float | Sequence[float], + v_col: float | Sequence[float], + u_max: int = 0, + v_max: int = 0, + intensity_0: float | Sequence[float] = 0.0, + intensity_row: float | Sequence[float] = 0.0, + intensity_col: float | Sequence[float] = 0.0, + intensity_row_row: float | Sequence[float] = 0.0, + intensity_col_col: float | Sequence[float] = 0.0, + intensity_row_col: float | Sequence[float] = 0.0, + per_disk_intensity: bool = False, + per_disk_slopes: bool = True, + max_intensity_order: int | None = None, + default_pattern_intensity_order: int | None = None, + center_intensity_0: float | Sequence[float] | None = None, + exclude_indices: Iterable[tuple[int, int]] | None = None, + boundary_px: float = 0.0, + min_frac_inside_mask: float | None = None, + max_slope_ratio: float | None = None, + own_template: bool = False, + intensity_softplus_beta: float | None = None, + slope_l2_weight: float = 0.0, + origin: OriginND | None = None, + origin_key: str = "origin", + constraint_params: dict[str, Any] | None = None, + ): + """ + Build a synthetic disk lattice renderer. + + Parameters + ---------- + u_row, u_col, v_row, v_col : float | Sequence[float | int | None] + Lattice basis parameters. Accept ``x``, ``(x0, delta)``, or + ``(x0, lo, hi)``. + intensity_0 : float | Sequence[float], optional + Baseline intensity for included lattice disks. Accepts ``x``, + ``(x0, delta)``, or ``(x0, lo, hi)``. + intensity_row, intensity_col, intensity_row_row, intensity_col_col, intensity_row_col : + float | Sequence[float | int | None] + Intensity polynomial parameters. Accept ``x``, ``(x0, delta)``, or + ``(x0, lo, hi)``. + center_intensity_0 : float | Sequence[float] | None, optional + Optional center-disk baseline. Accepts ``x``, ``(x0, delta)``, or + ``(x0, lo, hi)`` and routes by center ownership rules. + max_slope_ratio : float | None, optional + If set, cap the per-disk linear slope magnitude at + ``ratio * i0 / support_radius`` (projected after each step). + ``1.0`` keeps the intensity ramp nonnegative everywhere (no + saturation); larger values allow the rendered disk to saturate at + zero over part of its support while bounding how far the lit + region's centroid can shift. ``None`` leaves slopes unconstrained. + exclude_indices : Iterable[tuple[int, int]] | None, optional + Lattice indices excluded from rendering. By default, ``(0, 0)`` is + excluded. To include center explicitly, pass ``exclude_indices`` that + does not contain ``(0, 0)``. + + Returns + ------- + None + + Raises + ------ + ValueError + If ``center_intensity_0`` is provided with center included while + ``per_disk_intensity=False``. + + Notes + ----- + Center-intensity routing is explicit: + - Center excluded (default): ``center_intensity_0`` maps to + ``disk.intensity_raw``. + - Center included with ``per_disk_intensity=True``: + ``center_intensity_0`` maps to lattice center ``i0_raw`` entry. + In this case ``disk.intensity_raw`` is set to ``0`` to avoid + duplicate center ownership when disk is rendered standalone. + """ + super().__init__() + self.name = str(name) + self.disk = disk + self.origin = origin + self.origin_key = str(origin_key) + self.per_disk_intensity = bool(per_disk_intensity) + self.u_max = int(u_max) + self.v_max = int(v_max) + self.boundary_px = float(boundary_px) + # Drop a lattice disk from the render/fit when less than this fraction + # of its (circular) template patch falls inside ctx.mask -- i.e. the disk + # has too little illuminated data to constrain it. None disables it. + self.min_frac_inside_mask = ( + None if min_frac_inside_mask is None else float(min_frac_inside_mask) + ) + self.max_slope_ratio = None if max_slope_ratio is None else float(max_slope_ratio) + # When True, ``disk`` is the lattice's OWN template (a distinct + # DiskTemplate from the center-beam component), so its ``template_raw`` + # is trained via this lattice and shaped by the disk's constraints. When + # False (default) the lattice shares the center disk's template. + self.own_template = bool(own_template) + # None -> hard clamp max(f, 0); float -> softplus(beta*f)/beta (smooth, + # gradient-alive positivity that approaches zero without crossing it). + self.intensity_softplus_beta = ( + None if intensity_softplus_beta is None else float(intensity_softplus_beta) + ) + # Soft L2 prior on the per-disk linear slopes (ir, ic). Unlike the hard + # max_slope_ratio cap this shrinks slopes toward zero proportionally to + # the data misfit, so it auto-adapts: slopes relax to ~0 where there is + # no real intensity asymmetry (e.g. precession-averaged data) and grow + # only where the data supports them. + self.slope_l2_weight = float(slope_l2_weight) + + if max_intensity_order is None: + max_intensity_order = 1 if bool(per_disk_slopes) else 0 + self.max_intensity_order = int(max_intensity_order) + if self.max_intensity_order < 0 or self.max_intensity_order > 2: + raise ValueError("max_intensity_order must be 0, 1, or 2.") + + if default_pattern_intensity_order is None: + default_pattern_intensity_order = self.max_intensity_order + self.default_pattern_intensity_order = int(default_pattern_intensity_order) + + u_row_init, u_row_lo, u_row_hi = self.parse_bounded_init(u_row, name="u_row") + u_col_init, u_col_lo, u_col_hi = self.parse_bounded_init(u_col, name="u_col") + v_row_init, v_row_lo, v_row_hi = self.parse_bounded_init(v_row, name="v_row") + v_col_init, v_col_lo, v_col_hi = self.parse_bounded_init(v_col, name="v_col") + self.u_row = nn.Parameter(torch.tensor(u_row_init, dtype=torch.float32)) + if u_row_lo is not None or u_row_hi is not None: + self.register_parameter_bounds("u_row", u_row_lo, u_row_hi) + self.u_col = nn.Parameter(torch.tensor(u_col_init, dtype=torch.float32)) + if u_col_lo is not None or u_col_hi is not None: + self.register_parameter_bounds("u_col", u_col_lo, u_col_hi) + self.v_row = nn.Parameter(torch.tensor(v_row_init, dtype=torch.float32)) + if v_row_lo is not None or v_row_hi is not None: + self.register_parameter_bounds("v_row", v_row_lo, v_row_hi) + self.v_col = nn.Parameter(torch.tensor(v_col_init, dtype=torch.float32)) + if v_col_lo is not None or v_col_hi is not None: + self.register_parameter_bounds("v_col", v_col_lo, v_col_hi) + + exclude = {(0, 0)} if exclude_indices is None else set(exclude_indices) + center_included = (0, 0) not in exclude + if center_intensity_0 is not None and center_included and not self.per_disk_intensity: + raise ValueError( + "center_intensity_0 with center included requires per_disk_intensity=True, " + "or exclude (0,0) and use DiskTemplate intensity ownership." + ) + uv: list[tuple[int, int]] = [] + for u in range(-self.u_max, self.u_max + 1): + for v in range(-self.v_max, self.v_max + 1): + if (u, v) not in exclude: + uv.append((u, v)) + uv_t = ( + torch.as_tensor(uv, dtype=torch.long) if uv else torch.zeros((0, 2), dtype=torch.long) + ) + self.register_buffer("uv_indices", uv_t) + + n_uv = int(uv_t.shape[0]) + i0_init, i0_lo, i0_hi = self.parse_bounded_init(intensity_0, name="intensity_0") + if center_intensity_0 is None: + i0_center, i0_center_lo, i0_center_hi = i0_init, None, None + else: + i0_center, i0_center_lo, i0_center_hi = self.parse_bounded_init( + center_intensity_0, name="center_intensity_0" + ) + if center_intensity_0 is not None and not center_included: + self.disk.set_intensity(float(i0_center)) + if i0_center_lo is not None or i0_center_hi is not None: + self.disk.register_parameter_bounds("intensity_raw", i0_center_lo, i0_center_hi) + ir_init, ir_lo, ir_hi = self.parse_bounded_init(intensity_row, name="intensity_row") + ic_init, ic_lo, ic_hi = self.parse_bounded_init(intensity_col, name="intensity_col") + irr_init, irr_lo, irr_hi = self.parse_bounded_init( + intensity_row_row, name="intensity_row_row" + ) + icc_init, icc_lo, icc_hi = self.parse_bounded_init( + intensity_col_col, name="intensity_col_col" + ) + irc_init, irc_lo, irc_hi = self.parse_bounded_init( + intensity_row_col, name="intensity_row_col" + ) + self._center_i0_bounds: tuple[float | None, float | None] | None = None + self._center_i0_index: int | None = None + + if self.per_disk_intensity: + i0_values = torch.full((n_uv,), float(i0_init), dtype=torch.float32) + if n_uv > 0: + center_mask = (uv_t[:, 0] == 0) & (uv_t[:, 1] == 0) + i0_values[center_mask] = float(i0_center) + if center_intensity_0 is not None and bool(torch.any(center_mask)): + self._center_i0_index = int(torch.nonzero(center_mask, as_tuple=False)[0, 0]) + self._center_i0_bounds = (i0_center_lo, i0_center_hi) + self.i0_raw = nn.Parameter(i0_values) + if i0_lo is not None or i0_hi is not None: + self.register_parameter_bounds("i0_raw", i0_lo, i0_hi) + if center_intensity_0 is not None and center_included: + self.disk.set_intensity(0.0) + if self.max_intensity_order >= 1: + self.ir = nn.Parameter(torch.full((n_uv,), float(ir_init), dtype=torch.float32)) + self.ic = nn.Parameter(torch.full((n_uv,), float(ic_init), dtype=torch.float32)) + if ir_lo is not None or ir_hi is not None: + self.register_parameter_bounds("ir", ir_lo, ir_hi) + if ic_lo is not None or ic_hi is not None: + self.register_parameter_bounds("ic", ic_lo, ic_hi) + else: + self.ir = None + self.ic = None + if self.max_intensity_order >= 2: + self.irr = nn.Parameter(torch.full((n_uv,), float(irr_init), dtype=torch.float32)) + self.icc = nn.Parameter(torch.full((n_uv,), float(icc_init), dtype=torch.float32)) + self.irc = nn.Parameter(torch.full((n_uv,), float(irc_init), dtype=torch.float32)) + if irr_lo is not None or irr_hi is not None: + self.register_parameter_bounds("irr", irr_lo, irr_hi) + if icc_lo is not None or icc_hi is not None: + self.register_parameter_bounds("icc", icc_lo, icc_hi) + if irc_lo is not None or irc_hi is not None: + self.register_parameter_bounds("irc", irc_lo, irc_hi) + else: + self.irr = None + self.icc = None + self.irc = None + else: + self.i0_raw = nn.Parameter(torch.tensor(i0_init, dtype=torch.float32)) + if i0_lo is not None or i0_hi is not None: + self.register_parameter_bounds("i0_raw", i0_lo, i0_hi) + self.ir = nn.Parameter(torch.tensor(ir_init, dtype=torch.float32)) + if ir_lo is not None or ir_hi is not None: + self.register_parameter_bounds("ir", ir_lo, ir_hi) + self.ic = nn.Parameter(torch.tensor(ic_init, dtype=torch.float32)) + if ic_lo is not None or ic_hi is not None: + self.register_parameter_bounds("ic", ic_lo, ic_hi) + self.irr = nn.Parameter(torch.tensor(irr_init, dtype=torch.float32)) + if irr_lo is not None or irr_hi is not None: + self.register_parameter_bounds("irr", irr_lo, irr_hi) + self.icc = nn.Parameter(torch.tensor(icc_init, dtype=torch.float32)) + if icc_lo is not None or icc_hi is not None: + self.register_parameter_bounds("icc", icc_lo, icc_hi) + self.irc = nn.Parameter(torch.tensor(irc_init, dtype=torch.float32)) + if irc_lo is not None or irc_hi is not None: + self.register_parameter_bounds("irc", irc_lo, irc_hi) + if constraint_params is not None: + self.apply_constraint_params(constraint_params, strict=True) + if bool(self.hard_constraints.get("force_positive_intensity", False)): + self._enforce_positive_intensity_params() + + def set_origin(self, origin: OriginND) -> None: + self.origin = origin + + def _slope_support_radius(self) -> float: + """Maximum pixel radius of the disk-template support (for slope caps).""" + template = self.disk.template_raw.detach() + dr = cast(torch.Tensor, self.disk.dr) + dc = cast(torch.Tensor, self.disk.dc) + radius = torch.sqrt(dr * dr + dc * dc) + support = template.reshape(-1) > 1e-3 * template.max().clamp(min=1e-12) + if bool(support.any()): + return float(radius[support].max()) + return float(radius.max()) + + def _enforce_positive_intensity_params(self) -> None: + """ + Project base intensity parameter(s) to nonnegative values. + + Notes + ----- + ``i0_raw`` is projected to ``>= 0`` after optimizer steps. The forward + path saturates the per-pixel intensity polynomial at zero, so a steep + tilt clips part of the disk dark instead of rendering negative counts. + If ``max_slope_ratio`` is set, the linear slope pair ``(ir, ic)`` is + additionally capped at ``ratio * i0 / support_radius`` per disk, which + bounds how far the clipped disk's lit centroid can wander. + """ + with torch.no_grad(): + self.i0_raw.clamp_(min=0.0) + if ( + self.max_slope_ratio is not None + and self.ir is not None + and self.ic is not None + ): + radius = self._slope_support_radius() + if radius > 0.0: + slope = torch.sqrt(self.ir * self.ir + self.ic * self.ic) + limit = self.max_slope_ratio * self.i0_raw / radius + scale = torch.clamp(limit / slope.clamp(min=1e-12), max=1.0) + self.ir.mul_(scale) + self.ic.mul_(scale) + + def enforce_hard_constraints(self, ctx: RenderContext) -> None: + if bool(self.hard_constraints.get("force_positive_intensity", False)): + self._enforce_positive_intensity_params() + if self._center_i0_bounds is not None and self._center_i0_index is not None: + with torch.no_grad(): + lo, hi = self._center_i0_bounds + idx = self._center_i0_index + if lo is not None: + self.i0_raw[idx].clamp_(min=float(lo)) + if hi is not None: + self.i0_raw[idx].clamp_(max=float(hi)) + # Shape the lattice's own template with the disk's hard constraints + # (center/norm/circular mask), the same way the center disk is shaped. + if self.own_template: + self.disk.enforce_hard_constraints(ctx) + super().enforce_hard_constraints(ctx) + + + def _mask_keep( + self, ctx: RenderContext, centers_r: torch.Tensor, centers_c: torch.Tensor + ) -> torch.Tensor | None: + """ + Per-disk keep mask from ``min_frac_inside_mask``. + + Returns a boolean tensor shaped like ``centers_r`` that is True where at + least ``min_frac_inside_mask`` of the disk's circular template patch lands + on a True pixel of ``ctx.mask``; returns None when the filter is disabled + or no mask is set. Off-frame patch pixels count as outside the mask. + """ + if self.min_frac_inside_mask is None or ctx.mask is None: + return None + mask = ctx.mask + h, w = int(mask.shape[0]), int(mask.shape[1]) + dr = cast(torch.Tensor, self.disk.dr).to(device=ctx.device) + dc = cast(torch.Tensor, self.disk.dc).to(device=ctx.device) + tv = self.disk.patch_values().detach() + supp = tv > 0.5 * tv.max().clamp(min=1e-12) + drs = dr[supp] + dcs = dc[supp] + if drs.numel() == 0: + return None + rr = (centers_r.reshape(-1)[:, None] + drs[None, :]).round().to(torch.long) + cc = (centers_c.reshape(-1)[:, None] + dcs[None, :]).round().to(torch.long) + in_frame = (rr >= 0) & (rr < h) & (cc >= 0) & (cc < w) + mval = mask[rr.clamp(0, h - 1), cc.clamp(0, w - 1)].to(torch.bool) & in_frame + frac = mval.to(torch.float32).mean(dim=1) + return (frac >= self.min_frac_inside_mask).reshape(centers_r.shape) + + def _apply_positivity(self, inten: torch.Tensor) -> torch.Tensor: + """Map the per-pixel intensity polynomial to nonnegative values. + + Hard ``max(f, 0)`` clips part of a tilted disk dark but leaves a dead + gradient in the clipped region; softplus (``intensity_softplus_beta`` + set) is smooth and keeps the slope gradient alive everywhere while still + approaching zero without crossing it. + """ + if self.intensity_softplus_beta is not None: + beta = self.intensity_softplus_beta + return F.softplus(beta * inten) / beta + return (inten) + + def forward(self, ctx: RenderContext) -> torch.Tensor: + if self.origin is None: + raise RuntimeError("SyntheticDiskLattice requires an OriginND instance.") + + out = torch.zeros(ctx.shape, device=ctx.device, dtype=ctx.dtype) + uv_indices = cast(torch.Tensor, self.uv_indices) + if torch.numel(uv_indices) == 0: + return out + + uv = torch.as_tensor(uv_indices, device=ctx.device) + u = uv[:, 0].to(dtype=ctx.dtype) + v = uv[:, 1].to(dtype=ctx.dtype) + r0, c0 = self.origin.coords[0], self.origin.coords[1] + centers_r = r0 + u * self.u_row + v * self.v_row + centers_c = c0 + u * self.u_col + v * self.v_col + + b = torch.as_tensor(self.boundary_px, device=ctx.device, dtype=ctx.dtype) + keep = (centers_r >= b) & (centers_r <= (ctx.shape[0] - 1) - b) + keep = keep & (centers_c >= b) & (centers_c <= (ctx.shape[1] - 1) - b) + mk = self._mask_keep(ctx, centers_r, centers_c) + if mk is not None: + keep = keep & mk + if not torch.any(keep): + return out + + centers_r = centers_r[keep] + centers_c = centers_c[keep] + keep_idx = torch.nonzero(keep, as_tuple=False).reshape(-1) + num_disks = centers_r.shape[0] + + active_order = int( + ctx.fields.get( + "lattice_intensity_order_override", self.default_pattern_intensity_order + ) + ) + active_order = max(0, min(active_order, self.max_intensity_order)) + + dr = cast(torch.Tensor, self.disk.dr).to(device=ctx.device, dtype=ctx.dtype) + dc = cast(torch.Tensor, self.disk.dc).to(device=ctx.device, dtype=ctx.dtype) + patch_vals = self.disk.patch_values().to(device=ctx.device, dtype=ctx.dtype) + num_pixels = patch_vals.shape[0] + + if self.per_disk_intensity: + i0 = self.i0_raw[keep_idx][:, None] + inten = i0.expand(-1, num_pixels) + + if active_order >= 1 and self.ir is not None: + ir = self.ir[keep_idx][:, None] + ic = self.ic[keep_idx][:, None] + inten = inten + ir * dr[None, :] + ic * dc[None, :] + + if active_order >= 2 and self.irr is not None: + irr = self.irr[keep_idx][:, None] + icc = self.icc[keep_idx][:, None] + irc = self.irc[keep_idx][:, None] + inten = inten + irr * (dr*dr)[None, :] + icc * (dc*dc)[None, :] + irc * (dr*dc)[None, :] + else: + inten = self.i0_raw if isinstance(self.i0_raw, torch.Tensor) else self.i0_raw + if active_order >= 1: + inten = inten + self.ir * centers_r + self.ic * centers_c + if active_order >= 2: + inten = inten + self.irr * centers_r**2 + self.icc * centers_c**2 + self.irc * centers_r * centers_c + + # Positivity: hard clip, or smooth softplus if a beta is configured. + inten = self._apply_positivity(inten) + inten = inten[:, None].expand(-1, num_pixels) if inten.ndim == 1 else inten.expand(num_disks, num_pixels) + total_pixels = num_disks * num_pixels + r0_all = centers_r[:, None].expand(-1, num_pixels).reshape(total_pixels) + c0_all = centers_c[:, None].expand(-1, num_pixels).reshape(total_pixels) + dr_all = dr[None, :].expand(num_disks, -1).reshape(total_pixels) + dc_all = dc[None, :].expand(num_disks, -1).reshape(total_pixels) + vals_all = (patch_vals[None, :] * inten).reshape(total_pixels) + if self.disk.render_sigma is not None: + _splat_patch_gaussian( + out, r0=r0_all, c0=c0_all, patch_vals=vals_all, + dr=dr_all, dc=dc_all, scale=torch.ones_like(vals_all), + sigma=self.disk.render_sigma, + kernel=self.disk.render_kernel, + ) + else: + _splat_patch( + out, + r0=r0_all, + c0=c0_all, + patch_vals=vals_all, + dr=dr_all, + dc=dc_all, + scale=torch.ones_like(vals_all) + ) + return out + + def forward_batched( + self, + ctx: RenderContext, + *, + u_row_b: torch.Tensor, + u_col_b: torch.Tensor, + v_row_b: torch.Tensor, + v_col_b: torch.Tensor, + i0_raw_b: torch.Tensor, + ir_b: torch.Tensor | None, + ic_b: torch.Tensor | None, + irr_b: torch.Tensor | None, + icc_b: torch.Tensor | None, + irc_b: torch.Tensor | None, + template_raw_b: torch.Tensor, + origin_coords_b: torch.Tensor, + ) -> torch.Tensor: + if self.origin is None: + raise RuntimeError("SyntheticDiskLattice requires an OriginND instance.") + + B = u_row_b.shape[0] + uv = cast(torch.Tensor, self.uv_indices).to(device=ctx.device) + K = int(uv.shape[0]) + if K == 0: + return torch.zeros(B, ctx.shape[0], ctx.shape[1], device=ctx.device, dtype=ctx.dtype) + + u = uv[:, 0].to(dtype=ctx.dtype) + v = uv[:, 1].to(dtype=ctx.dtype) + + r0_kb = ( + origin_coords_b[:, 0:1] + + u.unsqueeze(0) * u_row_b.unsqueeze(1) + + v.unsqueeze(0) * v_row_b.unsqueeze(1) + ) + c0_kb = ( + origin_coords_b[:, 1:2] + + u.unsqueeze(0) * u_col_b.unsqueeze(1) + + v.unsqueeze(0) * v_col_b.unsqueeze(1) + ) + + bb = torch.as_tensor(self.boundary_px, device=ctx.device, dtype=ctx.dtype) + keep = (r0_kb >= bb) & (r0_kb <= (ctx.shape[0] - 1) - bb) + keep = keep & (c0_kb >= bb) & (c0_kb <= (ctx.shape[1] - 1) - bb) + mk = self._mask_keep(ctx, r0_kb, c0_kb) + if mk is not None: + keep = keep & mk + keep_f = keep.to(dtype=ctx.dtype) + + active_order = int( + ctx.fields.get( + "lattice_intensity_order_override", self.default_pattern_intensity_order + ) + ) + active_order = max(0, min(active_order, self.max_intensity_order)) + + dr = cast(torch.Tensor, self.disk.dr).to(device=ctx.device, dtype=ctx.dtype) + dc = cast(torch.Tensor, self.disk.dc).to(device=ctx.device, dtype=ctx.dtype) + N_pix = int(dr.numel()) + patch_vals = template_raw_b.reshape(B, N_pix) + + if self.per_disk_intensity: + inten = i0_raw_b.unsqueeze(2).expand(B, K, N_pix) + if active_order >= 1 and ir_b is not None and ic_b is not None: + inten = inten + ir_b.unsqueeze(2) * dr.view(1, 1, N_pix) + ic_b.unsqueeze(2) * dc.view(1, 1, N_pix) + if active_order >= 2 and irr_b is not None and icc_b is not None and irc_b is not None: + inten = ( + inten + + irr_b.unsqueeze(2) * (dr * dr).view(1, 1, N_pix) + + icc_b.unsqueeze(2) * (dc * dc).view(1, 1, N_pix) + + irc_b.unsqueeze(2) * (dr * dc).view(1, 1, N_pix) + ) + else: + inten = i0_raw_b.view(B, 1, 1).expand(B, K, N_pix).clone() + if active_order >= 1 and ir_b is not None and ic_b is not None: + inten = inten + ir_b.view(B, 1, 1) * r0_kb.unsqueeze(2) + ic_b.view(B, 1, 1) * c0_kb.unsqueeze(2) + if active_order >= 2 and irr_b is not None and icc_b is not None and irc_b is not None: + inten = ( + inten + + irr_b.view(B, 1, 1) * (r0_kb * r0_kb).unsqueeze(2) + + icc_b.view(B, 1, 1) * (c0_kb * c0_kb).unsqueeze(2) + + irc_b.view(B, 1, 1) * (r0_kb * c0_kb).unsqueeze(2) + ) + + # Positivity: hard clip, or smooth softplus if a beta is configured. + inten = self._apply_positivity(inten) + inten = inten * keep_f.unsqueeze(2) + vals = patch_vals.unsqueeze(1) * inten + + r0_full = (r0_kb.unsqueeze(2) + dr.view(1, 1, N_pix)).reshape(B, K * N_pix) + c0_full = (c0_kb.unsqueeze(2) + dc.view(1, 1, N_pix)).reshape(B, K * N_pix) + vals_full = vals.reshape(B, K * N_pix) + + if self.disk.render_sigma is not None: + return _splat_patch_batched_gaussian( + ctx.shape, r0=r0_full, c0=c0_full, vals=vals_full, + device=ctx.device, dtype=ctx.dtype, sigma=self.disk.render_sigma, kernel=self.disk.render_kernel, + ) + return _splat_patch_batched( + ctx.shape, r0=r0_full, c0=c0_full, vals=vals_full, + device=ctx.device, dtype=ctx.dtype, + ) + + def get_optimization_parameters(self) -> dict[str, list[torch.nn.Parameter]]: + params = [] + for name, param in self.named_parameters(recurse=True): + if name.startswith('disk.'): + # The lattice's own template is trained here; other disk params + # (intensity, its origin) are not owned by the lattice. + if self.own_template and name == 'disk.template_raw' and param.requires_grad: + params.append(param) + continue + if param.requires_grad: + params.append(param) + if not params: + return {} + return {'default': params} + + def constraint_loss( + self, + ctx: RenderContext, + params: dict[str, object] | None = None, + neighbor_target: dict[str, torch.Tensor] | None = None, + ) -> torch.Tensor: + if self.slope_l2_weight <= 0.0 or self.ir is None or self.ic is None: + l2_loss = torch.zeros((), device=ctx.device, dtype=ctx.dtype) + else: + ir = self.ir.to(device=ctx.device, dtype=ctx.dtype) + ic = self.ic.to(device=ctx.device, dtype=ctx.dtype) + l2_loss = self.slope_l2_weight * torch.mean(ir * ir + ic * ic) + + cfg = self.effective_soft_constraints(cast(dict[str, object] | None, params)) + consistency_weight = max(float(cfg.get("consistency_weight", 0.0)), 0.0) + consistency_loss = torch.zeros((), device=ctx.device, dtype=ctx.dtype) + if consistency_weight > 0.0 and neighbor_target is not None: + for attr, key in (("u_row", "u_row"), ("u_col", "u_col"), ("v_row", "v_row"), ("v_col", "v_col")): + tgt = neighbor_target.get(key) + if tgt is None or tgt != tgt: # NaN check without importing math + continue + val = getattr(self, attr).to(device=ctx.device, dtype=ctx.dtype) + consistency_loss = consistency_loss + consistency_weight * (val - tgt) ** 2 + + return consistency_loss + l2_loss + + def constraint_loss_batched( + self, + ctx: RenderContext, + *, + u_row_b: torch.Tensor | None = None, + u_col_b: torch.Tensor | None = None, + v_row_b: torch.Tensor | None = None, + v_col_b: torch.Tensor | None = None, + ir_b: torch.Tensor | None = None, + ic_b: torch.Tensor | None = None, + neighbor_target: dict[str, torch.Tensor] | None = None, + params: dict[str, object] | None = None, + ) -> torch.Tensor: + """Per-sample analogue of ``constraint_loss`` for stacked (batched) lattices.""" + ref = next((t for t in (u_row_b, ir_b) if t is not None), None) + B = ref.shape[0] if ref is not None else 1 + out = torch.zeros(B, device=ctx.device, dtype=ctx.dtype) + + if self.slope_l2_weight > 0.0 and ir_b is not None and ic_b is not None: + out = out + self.slope_l2_weight * (ir_b * ir_b + ic_b * ic_b) + + cfg = self.effective_soft_constraints(cast(dict[str, object] | None, params)) + consistency_weight = max(float(cfg.get("consistency_weight", 0.0)), 0.0) + + if consistency_weight > 0.0 and neighbor_target is not None: + def penalty(x, tgt): + if x is None or tgt is None: + return torch.zeros(B, device=ctx.device, dtype=ctx.dtype) + valid = ~torch.isnan(tgt) + safe_tgt = torch.where(valid, tgt, x.detach()) + return torch.where(valid, (x - safe_tgt) ** 2, torch.zeros_like(x)) + + out = out + consistency_weight * ( + penalty(u_row_b, neighbor_target.get("u_row")) + + penalty(u_col_b, neighbor_target.get("u_col")) + + penalty(v_row_b, neighbor_target.get("v_row")) + + penalty(v_col_b, neighbor_target.get("v_col")) + ) + + return out \ No newline at end of file diff --git a/src/quantem/core/utils/imaging_utils.py b/src/quantem/core/utils/imaging_utils.py index 9507bc53..29e8ea5b 100644 --- a/src/quantem/core/utils/imaging_utils.py +++ b/src/quantem/core/utils/imaging_utils.py @@ -6,7 +6,7 @@ import numpy as np import torch from numpy.typing import NDArray -from scipy.ndimage import gaussian_filter +from scipy.ndimage import gaussian_filter, map_coordinates from quantem.core.utils.utils import generate_batches @@ -1180,3 +1180,207 @@ def radially_project_fourier_tensor( array_1d = array_1d[0] return q_bins_out, array_1d + +def _parabolic_peak(v) -> float: + denom = 4.0 * v[1] - 2.0 * v[2] - 2.0 * v[0] + if denom == 0: + return 0.0 + return float((v[2] - v[0]) / denom) + + +def _upsampled_correlation_numpy( + imageCorr: NDArray, + upsampleFactor: int, + xyShift: NDArray, +) -> NDArray: + xyShift = np.round(xyShift * float(upsampleFactor)) / float(upsampleFactor) + globalShift = math.floor(math.ceil(upsampleFactor * 1.5) / 2.0) + upsampleCenter = float(globalShift) - (float(upsampleFactor) * xyShift) + + im_up = dft_upsample( + np.conj(imageCorr), upsampleFactor, (float(upsampleCenter[0]), float(upsampleCenter[1])) + ) + imageCorrUpsample = np.conj(im_up) + + flat_idx = int(np.argmax(imageCorrUpsample.real)) + r = flat_idx // imageCorrUpsample.shape[1] + c = flat_idx % imageCorrUpsample.shape[1] + + dx = 0.0 + dy = 0.0 + patch = imageCorrUpsample.real[r - 1 : r + 2, c - 1 : c + 2] + if patch.shape == (3, 3): + dx = _parabolic_peak(patch[:, 1]) + dy = _parabolic_peak(patch[1, :]) + + xySubShift = np.array([float(r), float(c)], dtype=float) - float(globalShift) + xyShift = xyShift + (xySubShift + np.array([dx, dy], dtype=float)) / float(upsampleFactor) + + return xyShift + + +def weighted_cross_correlation_shift( + im_ref=None, + im=None, + *, + cc=None, + weight_real=None, + upsample_factor: int = 1, + max_shift=None, + fft_input: bool = False, + fft_output: bool = False, + return_shifted_image: bool = False, +): + """ + Weighted peak selection + DFT subpixel refinement for Fourier cross-correlation. + + Provide either: + - im_ref and im (real-space images, or Fourier-domain if fft_input=True), OR + - cc (the Fourier-domain cross-spectrum), where cc = F_ref * conj(F_im) + + The weight is applied ONLY in real-space correlation to choose the peak location, + but the subpixel refinement uses the true (unweighted) cross-spectrum `cc`. + + Returns + ------- + shift_rc : tuple[float, float] + (d_row, d_col) shift to apply to `im` to align it to `im_ref`. + shifted : ndarray (optional) + If return_shifted=True: shifted image. If fft_output=True returns FFT (corner-centered), + else returns real-space image. + """ + if cc is None: + if im_ref is None or im is None: + raise ValueError("Provide either `cc` or both `im_ref` and `im`.") + F_ref = np.asarray(im_ref) if fft_input else np.fft.fft2(np.asarray(im_ref)) + F_im = np.asarray(im) if fft_input else np.fft.fft2(np.asarray(im)) + cc = F_ref * np.conj(F_im) + else: + cc = np.asarray(cc) + F_im = None + + cc_real = np.fft.ifft2(cc).real + M, N = cc_real.shape + + if weight_real is not None: + w = np.asarray(weight_real) + if w.shape != cc_real.shape: + raise ValueError( + f"weight_real.shape={w.shape} must match correlation shape {cc_real.shape}." + ) + cc_pick = cc_real * w + else: + cc_pick = cc_real + + if max_shift is not None: + fr = np.fft.fftfreq(M) * M + fc = np.fft.fftfreq(N) * N + mask = fr[:, None] ** 2 + fc[None, :] ** 2 > float(max_shift) ** 2 + cc_pick = cc_pick.copy() + cc_pick[mask] = -np.inf + + flat_idx = int(np.argmax(cc_pick)) + x0 = flat_idx // N + y0 = flat_idx % N + + x_inds = [((x0 + dx) % M) for dx in (-1, 0, 1)] + y_inds = [((y0 + dy) % N) for dy in (-1, 0, 1)] + vx = cc_pick[x_inds, y0] + vy = cc_pick[x0, y_inds] + + dx = _parabolic_peak(vx) + dy = _parabolic_peak(vy) + + x0 = np.round((float(x0) + float(dx)) * 2.0) / 2.0 + y0 = np.round((float(y0) + float(dy)) * 2.0) / 2.0 + xy_shift = np.array([x0, y0], dtype=float) + + if upsample_factor > 2: + xy_shift = _upsampled_correlation_numpy(cc, int(upsample_factor), xy_shift) + + dr = ((xy_shift[0] + M / 2) % M) - M / 2 + dc = ((xy_shift[1] + N / 2) % N) - N / 2 + shift_rc = (float(dr), float(dc)) + + if not return_shifted_image: + return shift_rc + + if im is None: + raise ValueError( + "return_shifted_image=True requires `im` (or its FFT via fft_input=True)." + ) + + if F_im is None: + F_im = np.asarray(im) if fft_input else np.fft.fft2(np.asarray(im)) + + kr = np.fft.fftfreq(M)[:, None] + kc = np.fft.fftfreq(N)[None, :] + phase_ramp = np.exp(-2j * np.pi * (kr * shift_rc[0] + kc * shift_rc[1])) + F_im_shifted = F_im * phase_ramp + + if fft_output: + return shift_rc, F_im_shifted + return shift_rc, np.fft.ifft2(F_im_shifted).real + +def rotate_image( + im, + rotation_deg: float, + origin: tuple[float, float] | None = None, + clockwise: bool = True, + interpolation: str = "bilinear", + mode: str = "constant", + cval: float = 0.0, +): + """Rotate an array about a pixel origin using bilinear/bicubic interpolation.""" + im = np.asarray(im) + if im.ndim < 2: + raise ValueError("im must have at least 2 dimensions") + + H, W = im.shape[-2], im.shape[-1] + if origin is None: + r0 = float(H // 2) + c0 = float(W // 2) + else: + r0 = float(origin[0]) + c0 = float(origin[1]) + + interp = str(interpolation).lower() + if interp in {"bilinear", "linear"}: + order = 1 + elif interp in {"bicubic", "cubic"}: + order = 3 + else: + raise ValueError("interpolation must be 'bilinear' or 'bicubic'") + + theta = float(np.deg2rad(rotation_deg)) + if not clockwise: + theta = -theta + + ct = float(np.cos(theta)) + st = float(np.sin(theta)) + + r_out, c_out = np.meshgrid( + np.arange(H, dtype=np.float64), + np.arange(W, dtype=np.float64), + indexing="ij", + ) + + c_rel = c_out - c0 + r_rel = r_out - r0 + + c_in = ct * c_rel + st * r_rel + c0 + r_in = -st * c_rel + ct * r_rel + r0 + + coords = np.vstack((r_in.ravel(), c_in.ravel())) + + if im.ndim == 2: + out = map_coordinates(im, coords, order=order, mode=mode, cval=cval) + return out.reshape(H, W) + + prefix = im.shape[:-2] + n = int(np.prod(prefix)) if prefix else 1 + im_flat = im.reshape(n, H, W) + out_flat = np.empty((n, H * W), dtype=np.result_type(im_flat.dtype, np.float64)) + for i in range(n): + out_flat[i] = map_coordinates(im_flat[i], coords, order=order, mode=mode, cval=cval) + return out_flat.reshape(*prefix, H, W) \ No newline at end of file diff --git a/src/quantem/diffraction/__init__.py b/src/quantem/diffraction/__init__.py index 166c4672..32e344b5 100644 --- a/src/quantem/diffraction/__init__.py +++ b/src/quantem/diffraction/__init__.py @@ -6,8 +6,11 @@ from quantem.diffraction.phase import PhaseMap as PhaseMap from quantem.diffraction.reverse_monte_carlo import ReverseMonteCarlo as ReverseMonteCarlo from quantem.diffraction.strain import StrainMap as StrainMap +from quantem.diffraction.strain_autocorrelation import StrainMapAutocorrelation as StrainMapAutocorrelation +from quantem.diffraction.model_fitting import ModelDiffraction as ModelDiffraction from quantem.diffraction import bloch as bloch from quantem.diffraction import calibration as calibration from quantem.diffraction import crystal_rotations as crystal_rotations from quantem.diffraction import digital_dark_field as digital_dark_field +from quantem.diffraction import ellipticity as ellipticity from quantem.diffraction import illumination as illumination diff --git a/src/quantem/diffraction/bragg_vectors.py b/src/quantem/diffraction/bragg_vectors.py index f9ccc18d..8e872928 100644 --- a/src/quantem/diffraction/bragg_vectors.py +++ b/src/quantem/diffraction/bragg_vectors.py @@ -9,6 +9,7 @@ import torch from numpy.typing import NDArray +from quantem.core import config from quantem.core.datastructures.dataset2d import Dataset2d from quantem.core.datastructures.dataset4dstem import Dataset4dstem from quantem.core.datastructures.vector import Vector @@ -85,14 +86,16 @@ class BraggVectors(AutoSerialize): def __init__( self, dataset: Dataset4dstem, - device: str = "cpu", + device: str | None = None, _token: object | None = None, ): if _token is not self._token: raise RuntimeError("Use BraggVectors.from_dataset() to instantiate this class.") super(BraggVectors, self).__init__() self.dataset = dataset - self.device = device + self.device = config.validate_device( + config.get("device") if device is None else device + )[0] self.peaks: Vector | None = None self.calibration = None @@ -109,6 +112,7 @@ def __init__( self.reference_ab: np.ndarray | None = None self.reference_qpos: np.ndarray | None = None self.reference_intensity: np.ndarray | None = None + self.max_peak_shift: float | None = None self.g1_array: np.ndarray | None = None self.g2_array: np.ndarray | None = None @@ -124,7 +128,7 @@ def __init__( @classmethod def from_dataset( - cls, dataset: Dataset4dstem, *, device: str = "cpu", name: str | None = None + cls, dataset: Dataset4dstem, *, device: str | None = None, name: str | None = None ) -> "BraggVectors": """Create a BraggVectors workflow bound to a 4D-STEM dataset. @@ -132,8 +136,9 @@ def from_dataset( ---------- dataset : Dataset4dstem The 4D-STEM dataset to analyze. - device : str, default="cpu" + device : str, optional Torch device used for detection (e.g. ``"cpu"`` or ``"cuda"``). + Defaults to ``config.get("device")``. name : str, optional If given, sets ``dataset.name``. @@ -327,7 +332,7 @@ def make_template_from_data( def make_template_from_probe( self, - probe: NDArray | torch.Tensor, + probe: NDArray | torch.Tensor | Dataset2d, center: tuple[float, float] | None = None, subtract_mean: bool = False, ) -> "BraggVectors": @@ -335,7 +340,7 @@ def make_template_from_probe( Parameters ---------- - probe : np.ndarray or torch.Tensor + probe : np.ndarray or torch.Tensor or Dataset2d Probe image; must match the diffraction-pattern shape. center : tuple of float, optional ``(row, col)`` probe center rolled to the origin; defaults to the @@ -350,6 +355,8 @@ def make_template_from_probe( BraggVectors ``self``, for method chaining. """ + if isinstance(probe, Dataset2d): + probe = probe.array probe_t = torch.as_tensor(probe, dtype=torch.float, device=self.device) if center is None and tuple(probe_t.shape) == tuple(self.dataset.shape[-2:]): center = probe_centroid(probe_t) @@ -559,10 +566,10 @@ def detect_disks( def correct_peak_origins( self, - origins: NDArray, + origins: NDArray | None = None, origin_ref: NDArray | tuple[float, float] | None = None, *, - inplace: bool = False, + modify_in_place: bool = False, ) -> "BraggVectors": """Shift the detected peak coordinates so all positions share one origin. @@ -580,24 +587,30 @@ def correct_peak_origins( Parameters ---------- - origins : np.ndarray + origins : np.ndarray, optional ``(scan_row, scan_col, 2)`` per-position diffraction origins in - detector pixels (row, col). + detector pixels (row, col). Defaults to :meth:`measure_origins` with + its own defaults. Measure them yourself and pass the result when you + need a non-default ``search_radius`` or ``center`` -- e.g. a beam that + is not centred on the detector -- or want to check the plane fit with + :meth:`plot_origin_fit` first. origin_ref : array-like of float, optional ``(row, col)`` common origin the corrected peaks are referred to. Defaults to the scan-mean of ``origins``. - inplace : bool, default=False + modify_in_place : bool, default=False If ``True``, correct ``self`` instead of returning a corrected copy. Returns ------- BraggVectors The workflow holding the corrected peaks (a new instance unless - ``inplace=True``); its :attr:`bvm` is recomputed. + ``modify_in_place=True``); its :attr:`bvm` is recomputed. """ if self.peaks is None: raise ValueError("Run detect_disks() before correct_peak_origins().") scan_shape = tuple(int(v) for v in self.dataset.shape[:2]) + if origins is None: + origins = self.measure_origins() origins = np.asarray(origins, dtype=float) if origins.shape != scan_shape + (2,): raise ValueError(f"origins must have shape {scan_shape + (2,)}, got {origins.shape}.") @@ -605,7 +618,7 @@ def correct_peak_origins( origin_ref = origins.mean(axis=(0, 1)) origin_ref = np.asarray(origin_ref, dtype=float).reshape(2) - if inplace: + if modify_in_place: bv = self else: bv = type(self)(dataset=self.dataset, device=self.device, _token=type(self)._token) @@ -827,6 +840,10 @@ def _candidate(idx: int) -> NDArray: self.origin = np.asarray(origin_rc, dtype=float).reshape(2) self.g1 = np.asarray(g1_rc, dtype=float).reshape(2) self.g2 = np.asarray(g2_rc, dtype=float).reshape(2) + self.reference_ab = None + self.reference_qpos = None + self.reference_intensity = None + self.max_peak_shift = None if plot or returnfig: fig, ax = plot_basis_vectors( @@ -845,6 +862,7 @@ def _candidate(idx: int) -> NDArray: def index_peaks( self, *, + max_peak_shift: float | None = None, plot: bool = True, returnfig: bool = False, **show_kwargs, @@ -864,11 +882,18 @@ def index_peaks( With ``plot=True`` (default) the reference lattice is drawn over the Bragg vector map, each site ringed and labelled with its ``(a, b)`` index. The ring color encodes how far the picked candidate sits from its ideal lattice site - ``origin + a*g1 + b*g2``, so a mis-picked or duplicate candidate (which rings - far from zero offset) is easy to spot. + ``origin + a*g1 + b*g2``, scaled to ``max_peak_shift`` -- the radius + :meth:`fit_lattice` matches within -- so a mis-picked or duplicate candidate + rings far from zero offset, and one that would later be rejected reads at + full scale here. Parameters ---------- + max_peak_shift : float, optional + Inclusion radius in pixels that :meth:`fit_lattice` matches within, + stored on :attr:`max_peak_shift`. The plot's ring colors are scaled to + it, so the figure shows the tolerance that will actually be applied. + Defaults to ``0.5 * min(|g1|, |g2|)`` -- half the shorter lattice spacing. plot : bool, default=True If ``True``, show the reference lattice via :func:`~quantem.diffraction.bragg_vectors_visualization.plot_reference_lattice`. @@ -905,6 +930,11 @@ def index_peaks( self.reference_ab = ab[keep] self.reference_qpos = cand_rc[keep] self.reference_intensity = cand_int[keep] + if max_peak_shift is None: + max_peak_shift = 0.5 * float( + min(np.linalg.norm(self.g1), np.linalg.norm(self.g2)) + ) + self.max_peak_shift = float(max_peak_shift) if not (plot or returnfig): return self @@ -916,6 +946,7 @@ def index_peaks( self.origin, self.g1, self.g2, + max_peak_shift=self.max_peak_shift, **show_kwargs, ) if returnfig: @@ -966,8 +997,11 @@ def fit_lattice( Inclusion radius in pixels: a detected peak is kept only if it lands within this distance of its nearest *ideal* lattice site (``origin + a*g1 + b*g2``), excluding peaks that stray too far from where - the best-fit lattice predicts. Defaults to ``0.5 * min(|g1|, |g2|)`` — - half the shorter lattice spacing. + the best-fit lattice predicts. Defaults to the value :meth:`index_peaks` + stored (itself ``0.5 * min(|g1|, |g2|)`` -- half the shorter lattice + spacing -- unless given there). Passing a different value here warns, + because the reference-lattice plot's ring colors are scaled to the + stored one. progressbar : bool, default=True If ``True``, show a tqdm progress bar over the scan positions. plot : bool, default=True @@ -997,6 +1031,8 @@ def fit_lattice( g2 = np.asarray(self.g2, dtype=float).reshape(2) ideal_qpos = o[None, :] + ref_ab[:, 0:1] * g1[None, :] + ref_ab[:, 1:2] * g2[None, :] + if max_peak_shift is None: + max_peak_shift = self.max_peak_shift if max_peak_shift is None: max_peak_shift = 0.5 * float(min(np.linalg.norm(g1), np.linalg.norm(g2))) @@ -1261,6 +1297,9 @@ def show_template( if radius is None: radius = radius_est crop = (float(center[0]), float(center[1]), float(crop_factor) * float(radius)) + dr = float(center[0]) - dp_mean.shape[0] / 2.0 + dc = float(center[1]) - dp_mean.shape[1] / 2.0 + print(f"mean diffraction (beam {dr:+.1f}, {dc:+.1f} px from detector center)") fig, ax = plot_template( dp_mean, @@ -1494,6 +1533,9 @@ def peak_histogram(self, *, returnfig: bool = False, **kwargs): **kwargs Extra keyword arguments forwarded to :func:`~quantem.diffraction.bragg_vectors_visualization.plot_bvm`. + **counts_kwargs : dict, optional + Keyword arguments for the peak-count panel (e.g. ``cmap``, ``norm``, + ``cbar``), forwarded to ``plot_bvm``. Returns ------- diff --git a/src/quantem/diffraction/bragg_vectors_visualization.py b/src/quantem/diffraction/bragg_vectors_visualization.py index a105909a..953a9bd8 100644 --- a/src/quantem/diffraction/bragg_vectors_visualization.py +++ b/src/quantem/diffraction/bragg_vectors_visualization.py @@ -564,15 +564,13 @@ def plot_bvm( ``(fig, ax)`` with ``ax`` a length-2 array of axes. """ from quantem.core.visualization import show_2d - fig, ax = plt.subplots(1, 2, figsize=figsize) - + plotting_kwargs.setdefault("title", "Bragg vector map") show_2d( np.asarray(bvm), figax=(fig, ax[0]), cmap=cmap, norm=norm, - title="Bragg vector map", **plotting_kwargs, ) counts_plot_kwargs = { @@ -593,6 +591,7 @@ def plot_reference_lattice( origin: np.ndarray, g1: np.ndarray, g2: np.ndarray, + max_peak_shift: float | None = None, *, cmap: str = "gray", norm: str | dict = "log_auto", @@ -605,9 +604,9 @@ def plot_reference_lattice( Each indexed reference site is a ring labelled with its ``(a, b)`` Miller index; the chosen origin is a green marker and ``g1`` (red) / ``g2`` (blue) are arrows from it. The ring color encodes how far the picked candidate sits from its *ideal* - lattice site ``origin + a*g1 + b*g2`` (the colorbar reads pixels), scaled to half - the shorter lattice spacing — the default :meth:`fit_lattice` match radius. A - mis-picked or duplicate candidate stands out as a ring far from zero offset (bright + lattice site ``origin + a*g1 + b*g2`` (the colorbar reads pixels), scaled to + ``max_peak_shift`` -- the radius :meth:`BraggVectors.fit_lattice` matches + within. A mis-picked or duplicate candidate stands out as a ring far from zero offset (bright color), sitting off the regular grid, or carrying an index that breaks the pattern. Parameters @@ -626,6 +625,10 @@ def plot_reference_lattice( First lattice vector as a ``(row, col)`` offset from ``origin``. g2 : np.ndarray Second lattice vector as a ``(row, col)`` offset from ``origin``. + max_peak_shift : float, optional + Radius in pixels the ring colors are scaled to -- the match radius + :meth:`BraggVectors.fit_lattice` will use. Defaults to + ``0.5 * min(|g1|, |g2|)``. cmap : str, default="gray" Colormap for the Bragg vector map; gray keeps the colored overlays legible. norm : str or dict, default="log_auto" @@ -674,7 +677,11 @@ def plot_reference_lattice( # color the rings by that offset, scaled to half the shorter lattice spacing (the # default fit_lattice match radius) so the colorbar previews which candidates sit # near the inclusion tolerance. - radius = 0.5 * float(min(np.hypot(*g1), np.hypot(*g2))) + radius = ( + 0.5 * float(min(np.hypot(*g1), np.hypot(*g2))) + if max_peak_shift is None + else float(max_peak_shift) + ) cmap_offset = plt.get_cmap("plasma") norm_offset = Normalize(vmin=0.0, vmax=radius if radius > 0 else 1.0) diff --git a/src/quantem/diffraction/model_fitting.py b/src/quantem/diffraction/model_fitting.py new file mode 100644 index 00000000..37fee590 --- /dev/null +++ b/src/quantem/diffraction/model_fitting.py @@ -0,0 +1,2382 @@ +from __future__ import annotations + +from typing import Any, Literal, Sequence, cast + +import numpy as np +import torch +import torch.nn.functional as F +from scipy.ndimage import shift as ndi_shift +from scipy.signal.windows import tukey +from tqdm import tqdm + +from quantem.core.datastructures import Dataset2d, Dataset3d, Dataset4d, Dataset4dstem +from quantem.core.fitting.background import DCBackground, GaussianBackground +from quantem.core.fitting.base import ( + AdditiveRenderModel, + FitBase, + LogMSELoss, + OriginND, + RenderComponent, + RenderContext, + SqrtMSELoss, + WeightedMSELoss, +) +from quantem.core.fitting.diffraction import DiskTemplate, SyntheticDiskLattice +from quantem.core.io.serialize import AutoSerialize +from quantem.core.ml.optimizer_mixin import OptimizerParams +from quantem.core.utils.imaging_utils import cross_correlation_shift +from quantem.diffraction.model_fitting_visualizations import ModelDiffractionVisualizations +from quantem.diffraction.strain import StrainMap + + +def _parse_init(value: float | int | Sequence[float | int | None], *, name: str) -> float: + if isinstance(value, (list, tuple, np.ndarray)): + if len(value) == 0: + raise ValueError(f"{name} cannot be empty.") + if value[0] is None: + raise ValueError(f"{name} initial value cannot be None.") + return float(value[0]) + return float(cast(float | int, value)) + + +class ModelDiffraction(ModelDiffractionVisualizations, FitBase, AutoSerialize): + _token = object() + DEFAULT_LR = 5e-2 + DEFAULT_OPTIMIZER_TYPE = "adam" + + def __init__(self, dataset: Any, _token: object | None = None): + if _token is not self._token: + raise RuntimeError("Use ModelDiffraction.from_dataset() or .from_file().") + AutoSerialize.__init__(self) + FitBase.__init__(self) + + # Dataset/input references + self.dataset = dataset + self.image_ref: np.ndarray | None = None + self.preprocess_shifts: np.ndarray | None = None + self.index_shape: tuple[int, ...] | None = None + self.target_mean: torch.Tensor | None = None + + # Diffraction-specific state/checkpoints + self.state_mean_refined: dict[str, torch.Tensor] | None = None + self.mean_refined: bool = False + + self.state_individual_refined: np.ndarray | None = None + self.individual_refined: bool = False + + # Misc metadata + self.metadata: dict[str, Any] = {} + + self.g1_array: np.ndarray | None = None + self.g2_array: np.ndarray | None = None + + self.real_space = False + + @classmethod + def from_dataset( + cls, dataset: Dataset2d | Dataset3d | Dataset4d | Dataset4dstem | Any + ) -> "ModelDiffraction": + if isinstance(dataset, (Dataset2d, Dataset3d, Dataset4d, Dataset4dstem)): + return cls(dataset=dataset, _token=cls._token) + raise TypeError( + "from_dataset expects a Dataset2d, Dataset3d, Dataset4d, or Dataset4dstem instance." + ) + + @property + def components(self) -> torch.nn.ModuleList: + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + return self.model.components + + def get_component(self, name: str) -> RenderComponent: + """ + Return a live model component by resolved name. + + Parameters + ---------- + name : str + Resolved component name. + + Returns + ------- + RenderComponent + The live component object. + + Raises + ------ + RuntimeError + If the model is not defined. + KeyError + If no component matches ``name``. + """ + return self._resolve_component_by_name(name) + + def get_rendered_component(self, name: str) -> np.ndarray: + """ + Render a component and return a NumPy array. + + Parameters + ---------- + name : str + Resolved component name. + + Returns + ------- + np.ndarray + Rendered component image. + + Raises + ------ + RuntimeError + If model/context are not defined. + KeyError + If no component matches ``name``. + """ + if self.ctx is None: + raise RuntimeError("Call .define_model(...) first.") + ctx = self.ctx + component = self._resolve_component_by_name(name) + rendered = component(ctx) + return rendered.detach().cpu().numpy() + + def get_rendered_disk_template(self, name: str | None = None) -> np.ndarray: + """ + Return a DiskTemplate patch as a numpy array--not rendered onto the full frame. + + Parameters + ---------- + name : str | None, optional + DiskTemplate component name. If omitted, requires exactly one DiskTemplate. + + Returns + ------- + np.ndarray + Template-sized array from ``template_raw``. + + Raises + ------ + RuntimeError + If model/context are not defined, no DiskTemplate exists, or multiple + DiskTemplates exist when ``name`` is omitted. + TypeError + If a named component exists but is not a DiskTemplate. + """ + if self.ctx is None or self.model is None: + raise RuntimeError("Call .define_model(...) first.") + + if name is not None: + component = self._resolve_component_by_name(name) + if not isinstance(component, DiskTemplate): + raise TypeError(f"Component '{name}' is not a DiskTemplate.") + return component.template_raw.detach().cpu().numpy() + matches = [m for m in self.model.components if isinstance(m, DiskTemplate)] + if len(matches) == 0: + raise RuntimeError("No DiskTemplate components found.") + if len(matches) > 1: + raise RuntimeError("Multiple DiskTemplate components found; pass name explicitly.") + disk = cast(DiskTemplate, matches[0]) + return disk.template_raw.detach().cpu().numpy() + + def set_disk_template_trainable( + self, enabled: bool, name: str | None = None, rebuild_optimizer: bool = True + ) -> None: + """ + Toggle DiskTemplate ``template_raw`` trainability. + + Parameters + ---------- + enabled : bool + If ``True``, enable optimization of ``template_raw``. + name : str | None, optional + DiskTemplate component name. If ``None``, applies to all DiskTemplate + components in the current model. + rebuild_optimizer : bool, optional + If ``True``, rebuild optimizer param groups after toggling. + + Returns + ------- + None + + Raises + ------ + KeyError + If ``name`` does not match any component. + RuntimeError + If model is not defined or no DiskTemplate components are found. + TypeError + If ``name`` resolves to a non-DiskTemplate component. + + Notes + ----- + This toggles only ``template_raw.requires_grad``. Other DiskTemplate + parameters (for example ``intensity_raw``) are unchanged. When + ``rebuild_optimizer=True``, optimizer param groups are rebuilt to match + current ``requires_grad`` flags. + """ + if self.model is None: + raise RuntimeError("Call .define_model(...) first.") + + if name is not None: + component = self._resolve_component_by_name(name) + if not isinstance(component, DiskTemplate): + raise TypeError(f"Component '{name}' is not a DiskTemplate.") + self.set_parameter_trainable( + name, + "template_raw", + enabled=enabled, + rebuild_optimizer=rebuild_optimizer, + ) + return + + disk_names = [ + component_name + for component_name, component in self._iter_named_components() + if isinstance(component, DiskTemplate) + ] + if len(disk_names) == 0: + raise RuntimeError("No DiskTemplate components found.") + + for disk_name in disk_names: + self.set_parameter_trainable( + disk_name, + "template_raw", + enabled=enabled, + rebuild_optimizer=False, + ) + if rebuild_optimizer: + self._rebuild_optimizer_after_trainability_change() + + def get_component_constraints(self, name: str) -> dict[str, dict[str, Any]]: + component = self._resolve_component_by_name(name) + return { + "hard": dict(component.hard_constraints), + "soft": dict(component.soft_constraints), + } + + def get_overlay_coordinates(self) -> tuple[np.ndarray, np.ndarray]: + """ + Return origin and lattice disk-center coordinates for overlay plotting. + + Parameters + ---------- + None + + Returns + ------- + origin_rc : np.ndarray + Origin coordinate array with shape ``(2,)`` as ``(row, col)``. + disk_centers_rc : np.ndarray + Disk-center array with shape ``(N, 2)`` as ``(row, col)``. + + Raises + ------ + RuntimeError + If model/context are not defined. + + Notes + ----- + Coordinates are computed from current model parameters without mutating state. + Boundary filtering matches ``SyntheticDiskLattice.forward`` behavior. + """ + if self.model is None or self.ctx is None: + raise RuntimeError("Call .define_model(...) first.") + + with torch.no_grad(): + origin = cast(OriginND, self.model.origin) + origin_rc = origin.coords[:2].detach().cpu().numpy().astype(np.float32, copy=False) + + centers: list[np.ndarray] = [] + for module in self.model.components: + component = cast(RenderComponent, module) + if not isinstance(component, SyntheticDiskLattice): + continue + if component.origin is None: + continue + uv_indices = cast(torch.Tensor, component.uv_indices) + if torch.numel(uv_indices) == 0: + continue + + uv = torch.as_tensor(uv_indices, device=self.ctx.device) + u = uv[:, 0].to(dtype=self.ctx.dtype) + v = uv[:, 1].to(dtype=self.ctx.dtype) + r0, c0 = component.origin.coords[0], component.origin.coords[1] + centers_r = r0 + u * component.u_row + v * component.v_row + centers_c = c0 + u * component.u_col + v * component.v_col + + b = torch.as_tensor( + component.boundary_px, device=self.ctx.device, dtype=self.ctx.dtype + ) + keep = (centers_r >= b) & (centers_r <= (self.ctx.shape[0] - 1) - b) + keep = keep & (centers_c >= b) & (centers_c <= (self.ctx.shape[1] - 1) - b) + mk = component._mask_keep(self.ctx, centers_r, centers_c) + if mk is not None: + keep = keep & mk + if torch.any(keep): + rc = torch.stack((centers_r[keep], centers_c[keep]), dim=1) + centers.append(rc.detach().cpu().numpy().astype(np.float32, copy=False)) + + if centers: + disk_centers_rc = np.concatenate(centers, axis=0) + else: + disk_centers_rc = np.zeros((0, 2), dtype=np.float32) + + return origin_rc, disk_centers_rc + + def preprocess( + self, + *, + align: bool = False, + origins: Any = None, + edge_blend: float = 8.0, + upsample_factor: int = 32, + max_shift: float | None = None, + shift_order: int = 1, + gamma: float = 0.5, + mode: str = "linear", + rows=None, + cols = None, + ) -> "ModelDiffraction": + arr = np.asarray(self.dataset.array) + self.mask = np.ones(arr.shape[:2]) + + if arr.ndim < 2: + raise ValueError("dataset.array must have at least 2 dimensions.") + mode_in = mode.strip().lower() + if mode_in in {"linear", "patterson", "paterson", "acf", "autocorrelation"}: + mode_norm = "linear" + elif mode_in in {"log", "cepstrum", "cepstral"}: + mode_norm = "log" + elif mode_in in {"gamma", "power", "sqrt"}: + mode_norm = "gamma" + else: + raise ValueError( + "mode must be 'linear', 'log', or 'gamma' (aliases: 'patterson'->'linear', 'cepstrum'/'cepstral'->'log')." + ) + + self.metadata["mode"] = mode_norm + if mode_norm == "gamma": + self.metadata["gamma"] = gamma + + if mode_norm == "linear": + arr = arr + elif mode_norm == "log": + arr = np.log1p(arr) + elif mode_norm == "gamma": + arr = np.power(np.clip(arr, 0.0, None), self.metadata["gamma"]) + else: + raise RuntimeError("Unreachable: normalized mode mapping failed.") + self.dataset.array = np.asarray(arr) + + if rows is None and cols is None: + rows = range(self.dataset.shape[0]) + cols = range(self.dataset.shape[1]) + elif rows is not None and cols is None: + cols = range(self.dataset.shape[1]) + elif rows is None and cols is not None: + rows = range(self.dataset.shape[0]) + else: + rows = rows + cols = cols + + if isinstance(rows, int): + rows = np.array([rows]).astype(int) + else: + rows = np.asarray(rows).astype(int) + + if isinstance(cols, int): + cols = np.array([cols]).astype(int) + else: + cols = np.asarray(cols).astype(int) + + arr = arr[np.ix_(rows, cols)] + + h, w = arr.shape[-2], arr.shape[-1] + self.index_shape = tuple(arr.shape[:-2]) + + stack = arr.reshape((-1, h, w)).astype(np.float32, copy=False) + n = stack.shape[0] + + if origins is not None: + origins = np.asarray(origins, dtype=float) + scan_shape = tuple(int(v) for v in self.dataset.shape[:2]) + if origins.shape != scan_shape + (2,): + raise ValueError( + f"origins must have shape {scan_shape + (2,)}, got {origins.shape}." + ) + origins_sub = origins[np.ix_(rows, cols)].reshape(-1, 2) + shifts = (origins_sub.mean(axis=0) - origins_sub).astype(np.float32) + aligned = np.empty_like(stack, dtype=np.float32) + for i in range(n): + aligned[i] = ndi_shift( + stack[i], + shift=(float(shifts[i, 0]), float(shifts[i, 1])), + order=int(shift_order), + mode="nearest", + prefilter=False, + ) + self.image_ref = np.mean(aligned, axis=0) + self.preprocess_shifts = shifts.reshape(self.index_shape + (2,)) + return self + + if not align or n <= 1: + self.image_ref = np.mean(stack, axis=0) + self.preprocess_shifts = None + return self + + alpha_r = 0.0 if edge_blend <= 0 else min(1.0, 2.0 * float(edge_blend) / float(h)) + alpha_c = 0.0 if edge_blend <= 0 else min(1.0, 2.0 * float(edge_blend) / float(w)) + window = tukey(h, alpha=alpha_r)[:, None] * tukey(w, alpha=alpha_c)[None, :] + window = window.astype(np.float32, copy=False) + + shifts = np.zeros((n, 2), dtype=np.float32) + fft_ref = np.fft.fft2(window * stack[0]) + for i in range(1, n): + fft_i = np.fft.fft2(window * stack[i]) + drc, fft_shift = cross_correlation_shift( + fft_ref, + fft_i, + upsample_factor=int(upsample_factor), + max_shift=max_shift, + fft_input=True, + fft_output=True, + return_shifted_image=True, + ) + if not isinstance(drc, (list, tuple, np.ndarray)) or len(drc) < 2: + raise RuntimeError("cross_correlation_shift returned an invalid shift vector.") + shifts[i, 0] = float(drc[0]) + shifts[i, 1] = float(drc[1]) + fft_ref = fft_ref * (i / (i + 1)) + fft_shift / (i + 1) + + shifts -= np.mean(shifts, axis=0, keepdims=True) + aligned = np.empty_like(stack, dtype=np.float32) + for i in range(n): + aligned[i] = ndi_shift( + stack[i], + shift=(float(shifts[i, 0]), float(shifts[i, 1])), + order=int(shift_order), + mode="nearest", + prefilter=False, + ) + + self.image_ref = np.mean(aligned, axis=0) + self.preprocess_shifts = shifts.reshape(self.index_shape + (2,)) + return self + + def define_model( + self, + *, + origin_row: float | Sequence[float], + origin_col: float | Sequence[float], + components: list[RenderComponent], + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + mask: np.ndarray | torch.Tensor | None = None, + origin_key: str = "origin", + ) -> "ModelDiffraction": + if self.image_ref is None: + self.preprocess() + if self.image_ref is None: + raise RuntimeError("image_ref not available.") + + h, w = int(self.image_ref.shape[0]), int(self.image_ref.shape[1]) + dev = torch.device(device) if device is not None else torch.device("cpu") + dt = dtype if dtype is not None else torch.float32 + + mask_t = None + if mask is not None: + mask_t = ( + mask.to(device=dev, dtype=dt) + if torch.is_tensor(mask) + else torch.as_tensor(mask, device=dev, dtype=dt) + ) + if tuple(mask_t.shape) != (h, w): + raise ValueError("mask must have shape (H, W).") + + origin = OriginND( + ndim=2, + init=[ + _parse_init(origin_row, name="origin_row"), + _parse_init(origin_col, name="origin_col"), + ], + ) + origin._quantem_origin_key = str(origin_key) # type: ignore[attr-defined] + + for component in components: + if hasattr(component, "set_origin"): + component.set_origin(origin) # type: ignore[misc] + elif hasattr(component, "origin") and getattr(component, "origin") is None: + component.origin = origin # type: ignore[attr-defined] + + self.model = AdditiveRenderModel(origin=origin, components=list(components)).to( + device=dev, dtype=dt + ) + self.ctx = RenderContext(shape=(h, w), device=dev, dtype=dt, mask=mask_t, fields={}) + self.target_mean = torch.as_tensor(self.image_ref, device=dev, dtype=dt) + + s0 = self._get_model_state_dict_copy() + self.state_initialized = s0 + self.state_mean_refined = None + self.mean_refined = False + self._clear_fit_history_all() + self.remove_optimizer() + return self + + def fit_mean_diffraction_pattern( + self, + *, + n_steps: int = 200, + reset: bool | Literal["initialized", "mean_refined"] = False, + optimizer_params: dict | None = None, + scheduler_params: dict | None = None, + constraint_weight: float = 1.0, + constraint_params: dict[str, Any] | None = None, + constraint_config_params: dict[str, Any] | None = None, + progress: bool = True, + ) -> "ModelDiffraction": + """ + Fit the mean diffraction pattern. + + Parameters + ---------- + n_steps : int, optional + Number of optimization steps. + reset : bool | Literal["initialized", "mean_refined"], optional + Reset behavior before fitting. + optimizer_params : dict | None, optional + Optimizer override for this fit call. + scheduler_params : dict | None, optional + Scheduler override for this fit call. + constraint_weight : float, optional + Global multiplier for soft-constraint loss. + constraint_params : dict[str, Any] | None, optional + Optional constraint updates applied once to components before fitting. + If ``None``, previously assigned constraints are reused. + progress : bool, optional + If ``True``, show progress bar. + + Returns + ------- + ModelDiffraction + Self, with updated fit state and history. + + Raises + ------ + RuntimeError + If model/context/target are not defined. + ValueError + If ``reset`` has an unsupported value. + + Notes + ----- + Constraint assignments persist on components across fit calls. + """ + if self.model is None or self.ctx is None or self.target_mean is None: + raise RuntimeError("Call .define_model(...) first.") + if reset is True: + self.reset("initialized") + elif isinstance(reset, str): + if reset not in ("initialized", "mean_refined"): + raise ValueError("reset must be False, True, 'initialized', or 'mean_refined'.") + self.reset(reset_to=cast(Literal["initialized", "mean_refined"], reset)) + elif reset not in (False,): + raise ValueError("reset must be False, True, 'initialized', or 'mean_refined'.") + + self.fit_render( + # target=torch.tensor(self.dataset[30,30].array.astype("float32")), + target=self.target_mean, + n_steps=int(n_steps), + constraint_weight=float(constraint_weight), + constraint_params=constraint_params, + constraint_config_params=constraint_config_params, + optimizer_params=optimizer_params, + scheduler_params=scheduler_params, + progress=bool(progress), + run_key="mean", + ) + + s_fit = self._get_model_state_dict_copy() + self.state_mean_refined = self._clone_state_dict(s_fit) + self.mean_refined = True + return self + + def reset( + self, + reset_to: Literal["initialized", "mean_refined", "individual"] = "mean_refined", + reset_history: bool = True, + individual_row: int = 0, + individual_col: int = 0, + ) -> "ModelDiffraction": + if reset_to == "initialized": + state = self.state_initialized + if state is None: + raise RuntimeError( + "initialized state is unavailable. Call .define_model(...) first." + ) + if reset_history: + self._clear_fit_history_all() + elif reset_to == "mean_refined": + state = self.state_mean_refined + if state is None: + raise RuntimeError( + "mean_refined state is unavailable. Run .fit_mean_diffraction_pattern(...) first." + ) + if reset_history: + mean_hist = self.fit_history.get("mean") + self._clear_fit_history_all() + if mean_hist is not None: + self.fit_history["mean"] = mean_hist + elif reset_to == "individual": + if self.state_individual_refined is None: + raise ValueError("individual states is unavalible. Run fit_individual_diffraction_pattern(....) first") + if (individual_row >= self.state_individual_refined.shape[0]) or (individual_col >= self.state_individual_refined.shape[1]): + raise ValueError("row and column values not in range") + state = self.state_individual_refined[individual_row, individual_col] + if state is None: + raise RuntimeError( + f"No refined state for position ({individual_row}, {individual_col}). " + "Run fit_individual_diffraction_pattern(...) for that row and column first." + ) + if reset_history: + self._clear_fit_history_all() + else: + raise ValueError("reset_to must be 'initialized' or 'mean_refined' or 'individual'.") + + self._load_model_state_dict_copy(state) + return self + + def fit_individual_diffraction_pattern( + self, + *, + rows=None, + cols = None, + n_steps: int = 200, + reset: bool | Literal["initialized", "mean_refined"], + optimizer_params: dict | None = None, + scheduler_params: dict | None = None, + constraint_weight: float = 1.0, + constraint_params: dict[str, Any] | None = None, + constraint_config_params: dict[str, Any] | None = None, + edge_weight: float = 0.0, + progress: bool = True, + batch_size: int | None = None, + frozen_components: list[str] | str | None = None, + sample_trainability: dict[str, Any] | None = None, + save_to_gpu: bool = True, + **_compat_kwargs: Any, + ) -> "ModelDiffraction": + if batch_size is not None: + return self.fit_individual_diffraction_pattern_batched( + rows=rows, + cols=cols, + batch_size=int(batch_size), + n_steps=int(n_steps), + reset=cast(Literal["initialized", "mean_refined"], reset), + optimizer_params=optimizer_params, + scheduler_params=scheduler_params, + constraint_weight=float(constraint_weight), + constraint_params=constraint_params, + constraint_config_params=constraint_config_params, + edge_weight=float(edge_weight), + frozen_components=frozen_components, + sample_trainability=sample_trainability, + progress=progress, + save_to_gpu = save_to_gpu + ) + + if self.model is None or self.ctx is None or self.target_mean is None: + raise RuntimeError("Call .define_model(...) first.") + if reset not in ("initialized", "mean_refined"): + raise ValueError("reset must be initialized', or 'mean_refined'.") + self.reset(reset_to=cast(Literal["initialized", "mean_refined"], reset)) + if not isinstance(self.dataset, Dataset4d): + raise ValueError("Dataset must be Dataset4d or Dataset4dstem for fit_individual_diffraction_pattern") + + scan_r = self.dataset.shape[0] + scan_c = self.dataset.shape[1] + + if rows is None and cols is None: + rows = range(self.dataset.shape[0]) + cols = range(self.dataset.shape[1]) + elif rows is not None and cols is None: + cols = range(self.dataset.shape[1]) + elif rows is None and cols is not None: + rows = range(self.dataset.shape[0]) + else: + rows = rows + cols = cols + + if isinstance(rows, int): + rows = np.array([rows]).astype(int) + else: + rows = np.asarray(rows).astype(int) + + if isinstance(cols, int): + cols = np.array([cols]).astype(int) + else: + cols = np.asarray(cols).astype(int) + + self.state_individual_refined = np.full(shape=(scan_r, scan_c), fill_value=None, dtype=object) + + if progress: + pbar = tqdm(total=rows.shape[0] * cols.shape[0], desc="Fit individual") + + for r in rows: + for c in cols: + # print(self.dataset.array[r,c].shape) + self.reset(reset_to=cast(Literal["initialized", "mean_refined"], reset), reset_history=False) + nt_single = _neighbor_uv_target_single(self.state_individual_refined, r, c) + self.fit_render( + target=torch.as_tensor(self.dataset.array[r,c],device=self.ctx.device,dtype=self.ctx.dtype), + n_steps=int(n_steps), + constraint_weight=float(constraint_weight), + constraint_params=constraint_params, + optimizer_params=optimizer_params, + scheduler_params=scheduler_params, + progress=False, + run_key=f"individual_{r}_{c}", + neighbor_target=nt_single, + ) + + s_fit = self._get_model_state_dict_copy() + self.state_individual_refined[r,c] = self._clone_state_dict(s_fit) + # self.reset(reset_to=cast(Literal["initialized", "mean_refined"], reset), reset_history=False) + if progress: + pbar.update(1) + if progress: + pbar.close() + + self.individual_refined=True + self.g1_array = None + self.g2_array = None + return self + + def fit_individual_diffraction_pattern_batched( + self, + *, + rows: Any = None, + cols: Any = None, + batch_size: int = 16, + n_steps: int = 200, + reset: Literal["initialized", "mean_refined"] = "mean_refined", + optimizer_params: dict | None = None, + scheduler_params: dict | None = None, + constraint_weight: float = 1.0, + constraint_params: dict[str, Any] | None = None, + constraint_config_params: dict[str, Any] | None = None, + edge_weight: float = 0.0, + frozen_components: list[str] | str | None = None, + sample_trainability: dict[str, Any] | None = None, + progress: bool = True, + save_to_gpu: bool = True, + ) -> "ModelDiffraction": + """ + Per-pattern fit, vectorized across a batch dimension on a single GPU. + + See ``fit_individual_diffraction_pattern`` for argument semantics. The + batched version runs ``batch_size`` patterns in parallel per optimizer + step, with per-sample stacked parameters and per-sample Adam moments. + + Soft constraint losses are added per-sample via + ``model.total_constraint_loss``-style accumulation (currently only + ``DiskTemplate`` contributes). Hard constraints and parameter bounds + are still enforced between steps. + + Parameters + ---------- + scheduler_params : dict | None, optional + Either a single spec dict (``{"type": "cosine", "t_max": N}``) + applied to every parameter group, or a dict keyed by component + name/class. Supported types: ``none``, ``cosine`` (a.k.a. + ``cosine_annealing``), ``linear``, ``exponential``. ``plateau`` + and ``cyclic`` fall back to constant LR with a warning. + constraint_weight : float + Global multiplier on the per-sample soft-constraint loss. + constraint_config_params : dict | None + Applied to the model once before the fit (``apply_constraint_configs``). + frozen_components : list[str] | str | None + Component name(s) whose parameters are frozen across all samples. + Resolved via ``AdditiveRenderModel._component_constraint_name``; + both instance names (``"disk0"``) and class names (``"DiskTemplate"``) + work. + sample_trainability : dict[str, ArrayLike] | None + Per-canonical-key (e.g. ``"disk.template_raw"``) boolean array of + length ``len(positions)``; ``False`` entries freeze that parameter + for the matching scan position. + """ + if self.model is None or self.ctx is None or self.target_mean is None: + raise RuntimeError("Call .define_model(...) first.") + if not isinstance(self.dataset, Dataset4d): + raise ValueError("Dataset must be Dataset4d or Dataset4dstem.") + if reset not in ("initialized", "mean_refined"): + raise ValueError("reset must be 'initialized' or 'mean_refined'.") + + # Choose initialization state and load into the live model so we can + # read current parameter values + constraint settings off the modules. + if reset == "mean_refined": + if self.state_mean_refined is None: + raise RuntimeError("mean_refined state is unavailable. Run fit_mean_diffraction_pattern first.") + init_state = self._clone_state_dict(self.state_mean_refined) + else: + if self.state_initialized is None: + raise RuntimeError("initialized state is unavailable. Call define_model first.") + init_state = self._clone_state_dict(self.state_initialized) + self._load_model_state_dict_copy(init_state) + + if constraint_params is not None: + self.model.apply_constraint_params(constraint_params, strict=True) + if constraint_config_params is not None: + self.model.apply_constraint_configs(constraint_config_params, strict=True) + + scan_r = int(self.dataset.shape[0]) + scan_c = int(self.dataset.shape[1]) + rows_arr, cols_arr = _resolve_rows_cols_for_batched(rows, cols, scan_r, scan_c) + positions: list[tuple[int, int]] = [(int(r), int(c)) for r in rows_arr for c in cols_arr] + if len(positions) == 0: + return self + + if self.state_individual_refined is None or self.state_individual_refined.shape != (scan_r, scan_c): + self.state_individual_refined = np.full(shape=(scan_r, scan_c), fill_value=None, dtype=object) + + ctx = self.ctx + components_list = list(self.model.components) + plan = _BatchedPlan.from_model(self.model, components_list, optimizer_params or {}) + plan.set_scheduler_params(scheduler_params) + + # Build per-position trainability arrays (default True everywhere). + n_positions = len(positions) + is_trainable_pos: dict[str, np.ndarray] = { + key: np.ones(n_positions, dtype=bool) for key in plan.lrs + } + frozen_keys = plan.resolve_component_keys(frozen_components) + for key in frozen_keys: + is_trainable_pos[key][:] = False + hard_skip_keys: set[str] = set(frozen_keys) + if sample_trainability is not None: + for key, arr in sample_trainability.items(): + if key not in is_trainable_pos: + raise KeyError( + f"sample_trainability key '{key}' is not a known stacked-param. " + f"Known: {sorted(is_trainable_pos)}" + ) + mask = np.asarray(arr, dtype=bool).reshape(-1) + if mask.shape != (n_positions,): + raise ValueError( + f"sample_trainability['{key}'] has shape {mask.shape}, " + f"expected ({n_positions},)." + ) + is_trainable_pos[key] = mask + + loss_fn = self.loss_fn + + fidelity_mask: torch.Tensor | None = None + fidelity_valid_count: torch.Tensor | None = None + if ctx.mask is not None: + fidelity_mask = ctx.mask.bool().to(ctx.device) + fidelity_valid_count = fidelity_mask.sum().clamp(min=1).to(ctx.dtype) + + total_steps = len(positions) * n_steps + pbar = tqdm(total=total_steps, desc="Fit individual (batched)", disable=not progress) + + + dataset_gpu: torch.Tensor | None = None + dev = self.ctx.device + + if save_to_gpu and str(dev) != "cpu": + try: + print(f"Loading data to {dev}...", flush=True) + dataset_gpu = torch.as_tensor( + self.dataset.array, + device=dev, + dtype=ctx.dtype + ) + except (RuntimeError, torch.cuda.OutOfMemoryError): + dataset_gpu = None + + + for start in range(0, len(positions), int(batch_size)): + chunk = positions[start:start + int(batch_size)] + B = len(chunk) + + if dataset_gpu is not None: + rows_batch = [r for r, c in chunk] + cols_batch = [c for r, c in chunk] + targets = dataset_gpu[rows_batch, cols_batch] + else: + targets = torch.stack( + [ + torch.as_tensor(self.dataset.array[r, c], device=ctx.device, dtype=ctx.dtype) + for (r, c) in chunk + ], + dim=0, + ) + + stacked = plan.build_stacked_params(B) + + neighbor_target = _neighbor_uv_target( + self.state_individual_refined, chunk, plan.lat_names[0], ctx.device, ctx.dtype, + ) + + opt_state: dict[str, dict[str, torch.Tensor]] = { + name: ( + {"m": torch.zeros_like(p.detach()), "v": torch.zeros_like(p.detach())} + if plan.opt_type in ("adam", "adamw") + else {"buf": torch.zeros_like(p.detach())} + ) + for name, p in stacked.items() + } + + # Per-chunk trainability slice as (B,) bool tensors on device. + chunk_slice = slice(start, start + B) + chunk_trainable: dict[str, torch.Tensor] = {} + mixed_trainable_keys: list[str] = [] + init_snapshots: dict[str, torch.Tensor] = {} + for key in stacked: + if key in is_trainable_pos: + mask_np = is_trainable_pos[key][chunk_slice] + mask = torch.as_tensor(mask_np, device=ctx.device, dtype=torch.bool) + chunk_trainable[key] = mask + # A "mixed" key has some frozen samples and some trainable. + # We snapshot the init value so frozen samples can be + # restored after hard constraints / Adam. + if not bool(mask.all()) and not bool((~mask).all()): + mixed_trainable_keys.append(key) + init_snapshots[key] = stacked[key].detach().clone() + + is_sqrt = isinstance(loss_fn, SqrtMSELoss) + is_log = isinstance(loss_fn, LogMSELoss) + is_l1 = isinstance(loss_fn, torch.nn.L1Loss) + # signal-weighted MSE: per-pixel weights fixed from the targets for this chunk + w_pix = loss_fn.weights(targets) if isinstance(loss_fn, WeightedMSELoss) else None + ew = float(edge_weight) + if is_sqrt: + sqrt_gamma = float(loss_fn.gamma) + sqrt_eps = 1.0 + if fidelity_mask is not None: + tgt_min = targets.masked_fill(~fidelity_mask, float("inf")).amin(dim=(1, 2), keepdim=True) + else: + tgt_min = targets.amin(dim=(1, 2), keepdim=True) + tgt_mod = (targets - tgt_min + sqrt_eps) ** sqrt_gamma + elif is_log: + tgt_mod = torch.log1p(targets) + else: + tgt_mod = targets + if ew > 0.0: + # Edge term compares spatial gradients of the (compressed) images. + # Disk-position information is concentrated at the aperture edges, + # while per-disk intensity tilts act on disk interiors, so this + # term anchors the lattice geometry against tilt-shift mimicry. + tgt_dr = tgt_mod[:, 1:, :] - tgt_mod[:, :-1, :] + tgt_dc = tgt_mod[:, :, 1:] - tgt_mod[:, :, :-1] + for step in range(int(n_steps)): + pred = plan.batched_forward(ctx, stacked) + + if is_sqrt: + if fidelity_mask is not None: + pred_min = pred.masked_fill(~fidelity_mask, float("inf")).amin(dim=(1, 2), keepdim=True) + else: + pred_min = pred.amin(dim=(1, 2), keepdim=True) + pred_mod = (pred - pred_min + sqrt_eps) ** sqrt_gamma + elif is_log: + pred_mod = torch.log1p(pred) + else: + pred_mod = pred + + diff_mod = pred_mod - tgt_mod + val = diff_mod.abs() if is_l1 else diff_mod * diff_mod + if w_pix is not None: + val = val * w_pix + if fidelity_mask is not None: + per_sample_loss = (val * fidelity_mask).sum(dim=(1, 2)) / fidelity_valid_count + else: + per_sample_loss = val.mean(dim=(1, 2)) + + if ew > 0.0: + pr_dr = pred_mod[:, 1:, :] - pred_mod[:, :-1, :] + pr_dc = pred_mod[:, :, 1:] - pred_mod[:, :, :-1] + per_sample_loss = per_sample_loss + ew * ( + ((pr_dr - tgt_dr) ** 2).mean(dim=(1, 2)) + + ((pr_dc - tgt_dc) ** 2).mean(dim=(1, 2)) + ) + + # Add per-sample soft constraint loss. + if float(constraint_weight) != 0.0: + constraint_per_sample = plan.batched_constraint_loss(ctx, stacked, neighbor_target=neighbor_target) + per_sample_loss = per_sample_loss + float(constraint_weight) * constraint_per_sample + + total_loss = per_sample_loss.sum() + + grads_list = list(torch.autograd.grad(total_loss, list(stacked.values()), allow_unused=True)) + + # Apply trainability masks: zero out per-sample gradient entries + # for frozen samples/params before the Adam step. + for i, (name, g) in enumerate(zip(stacked.keys(), grads_list)): + if g is None: + continue + mask_b = chunk_trainable.get(name) + if mask_b is None: + continue + if bool(mask_b.all()): + continue + # Broadcast (B,) over remaining dims. + view_shape = (g.shape[0],) + (1,) * (g.ndim - 1) + grads_list[i] = g * mask_b.view(view_shape).to(dtype=g.dtype) + + t = step + 1 + current_lrs = plan.lr_at_step(t, int(n_steps)) + _optimizer_step_inplace( + plan.opt_type, stacked, tuple(grads_list), opt_state, + current_lrs, plan.opt_hparams, chunk_trainable, t, + ) + + with torch.no_grad(): + plan.apply_hard_constraints(stacked, skip_keys=hard_skip_keys) + for key in stacked: + if key not in chunk_trainable: + continue + mask = chunk_trainable[key] + if bool(mask.all()): + continue + if key not in init_snapshots: + init_snapshots[key] = stacked[key].detach().clone() + view_shape = (mask.shape[0],) + (1,) * (stacked[key].ndim - 1) + keep = mask.view(view_shape) + stacked[key].data.copy_( + torch.where(keep, stacked[key].data, init_snapshots[key]) + ) + + pbar.update(B) + + # Unstack each sample back into a state_dict and store + for b, (r, c) in enumerate(chunk): + sample_state = plan.build_sample_state_dict(init_state, stacked, b) + self.state_individual_refined[r, c] = sample_state + + pbar.close() + self.individual_refined = True + self.g1_array = None + self.g2_array = None + return self + + def get_individual_uv_vectors(self) -> "ModelDiffraction": + scan_r = self.dataset.shape[0] + scan_c = self.dataset.shape[1] + + self.g1_array = np.empty(shape=(scan_r, scan_c, 2)) + self.g2_array = np.empty(shape=(scan_r, scan_c, 2)) + if self.state_individual_refined is None: + raise RuntimeError("Call .fit_individual_diffraction_pattern(...) on all patterns first.") + for r in range(scan_r): + for c in range(scan_c): + pos_state = self.state_individual_refined[r,c] + if pos_state is None: + self.g1_array[r,c,:] = None + self.g2_array[r,c,:] = None + continue + for key in pos_state.keys(): + key_parts = key.split('.') + if(key_parts[-1] == 'u_row'): + self.g1_array[r,c,0] = pos_state[key] + if(key_parts[-1] == 'u_col'): + self.g1_array[r,c,1] = pos_state[key] + if(key_parts[-1] == 'v_row'): + self.g2_array[r,c,0] = pos_state[key] + if(key_parts[-1] == 'v_col'): + self.g2_array[r,c,1] = pos_state[key] + + return self + + def render_individual_pattern(self, row, col): + if self.state_individual_refined is None: + raise RuntimeError( + "individual_refined_state is unavalible. Run fit_individual_diffraction_pattern(...) first." + ) + if self.dataset.shape[0] <= row or self.dataset.shape[1] <= col: + raise ValueError("individual row or column outside bounds of dataset") + if row < 0 or col < 0: + raise ValueError("individual row or column outside bounds of dataset") + if self.state_individual_refined[row, col] is None: + raise RuntimeError( + "individual_refined_state is not avalible for given row and column. Run fit_individual_diffraction_pattern(...) for that row and column." + ) + return self._render_state_array(self.state_individual_refined[row, col]) + + + def create_mask( + self, + use_radial_method: bool = False, + exclusion_radius_fraction: float = 0.1, + plot: bool = True, + figsize: tuple[float, float] = (5, 4), + ): + """Compute the per-position weight :attr:`mask` from the fitted lattice signal. + + Builds a ``(scan_row, scan_col)`` weight in ``[0, 1]`` measuring the lattice + signal at each position -- the model-fitting analogue of + :attr:`BraggVectors.mask_weight` and the cepstral ``mask_weight``. It is passed to + :meth:`calculate_strain_map`, where higher-weight positions contribute more to the + reference lattice and weak/vacuum positions are down-weighted. + + The raw signal (summed intensity of the non-central fitted disks, or the + diffracted intensity outside a central disk for the radial method) is min-max + normalized to ``[0, 1]``; no contrast windowing or smoothing is applied. Display + contrast is applied via the ``mask_range`` argument of + :meth:`StrainMap.plot_strain`, :meth:`StrainMap.update_reference`, and + :meth:`StrainMap.estimate_strain_precision` (e.g. ``mask_range=(0.37, 0.5)``): + weights at/below ``low`` render black, at/above ``high`` render full color. The + weight-map plot shows the distribution so an appropriate ``mask_range`` can be + chosen. + + Parameters + ---------- + use_radial_method : bool, default=False + If ``True``, weight by the total diffracted intensity *outside* a central + disk (radius ``exclusion_radius_fraction`` of the detector width). If + ``False`` (default, recommended), weight by the summed intensity of the + fitted non-central lattice disks -- requires + :meth:`fit_individual_diffraction_pattern` to have been run. + exclusion_radius_fraction : float, default=0.1 + Central-disk radius (fraction of detector width) excluded by the radial + method. + plot : bool, default=True + If ``True``, show the resulting per-position weight map (so a separate + ``plt.imshow(mask)`` cell is unnecessary). + figsize : tuple of float, default=(5, 4) + Figure size in inches for the weight-map plot. + + Returns + ------- + ModelDiffraction + ``self``, with :attr:`mask` set. + """ + if not isinstance(self.dataset, (Dataset4d, Dataset4dstem)): + raise ValueError("Dataset must be Dataset4d or Dataset4dstem.") + + scan_r = self.dataset.shape[0] + scan_c = self.dataset.shape[1] + self.i0_sum_array = np.zeros(shape=(scan_r, scan_c)) + + if use_radial_method: + center_y, center_x = np.array(self.dataset.shape[:-2]) / 2 + y, x = np.ogrid[: self.dataset.shape[-2], : self.dataset.shape[-1]] + radius_map = np.sqrt((x - center_x) ** 2 + (y - center_y) ** 2) + exclusion_radius = exclusion_radius_fraction * self.dataset.shape[-1] + outside_mask = radius_map > exclusion_radius + for r in range(scan_r): + for c in range(scan_c): + self.i0_sum_array[r, c] = np.sum(self.dataset.array[r, c][outside_mask]) + else: + if self.state_individual_refined is None: + raise RuntimeError("Call .fit_individual_diffraction_pattern(...) first.") + for r in range(scan_r): + for c in range(scan_c): + pos_state = self.state_individual_refined[r, c] + if pos_state is None: + continue + i0_raw = None + uv_indices = None + for key in pos_state.keys(): + if key.endswith("i0_raw"): + i0_raw = pos_state[key].cpu().numpy() + if key.endswith("uv_indices"): + uv_indices = pos_state[key].cpu().numpy() + if i0_raw is None or uv_indices is None: + continue + is_not_center = ~((uv_indices[:, 0] == 0) & (uv_indices[:, 1] == 0)) + self.i0_sum_array[r, c] = np.sum(i0_raw[is_not_center]) + + # Min-max normalization to [0, 1], no contrast windowing (set mask_range at + # display). Subtracting the floor -- rather than only dividing by the max -- keeps + # the weight from saturating near 1.0 when every position carries a baseline + # lattice intensity, and matches the cepstral _amplitude_mask_weight. Degenerate + # (constant / non-finite) input falls back to uniform full weight. + lo = np.nanmin(self.i0_sum_array) + hi = np.nanmax(self.i0_sum_array) + if np.isfinite(lo) and np.isfinite(hi) and hi > lo: + self.mask = (self.i0_sum_array - lo) / (hi - lo) + else: + self.mask = np.ones_like(self.i0_sum_array) + + if plot: + import matplotlib.pyplot as plt + + fig, ax = plt.subplots(1, 1, figsize=figsize) + handle = ax.imshow(self.mask, cmap="gray", vmin=0.0, vmax=1.0) + ax.set_title("Per-position lattice weight (mask)") + ax.set_xlabel("scan column") + ax.set_ylabel("scan row") + fig.colorbar(handle, ax=ax, fraction=0.046, pad=0.04) + fig.tight_layout() + + return self + + def calculate_strain_map( + self, + g1_ref: np.ndarray | None = None, + g2_ref: np.ndarray | None = None, + mask: np.ndarray | None = None, + q_to_r_rotation_ccw_deg: float | None = None, + q_transpose: bool | None = None, + calculation_metric: str = "median", + ) -> StrainMap: + """Build a :class:`StrainMap` from the fitted per-position lattice vectors. + + Mirrors :meth:`BraggVectors.calculate_strain_map` and + :meth:`StrainMapAutocorrelation.calculate_strain_map` so the downstream strain + cells (``plot_strain``, ``update_reference``, ``estimate_strain_precision``) are + identical across the correlation, cepstral, and model-fitting workflows. + + Parameters + ---------- + g1_ref : np.ndarray, optional + ``(2,)`` reference for the first lattice vector. Defaults to the median over + the scan inside :class:`StrainMap`. + g2_ref : np.ndarray, optional + ``(2,)`` reference for the second lattice vector. Defaults to the median over + the scan inside :class:`StrainMap`. + mask : np.ndarray, optional + ``(scan_row, scan_col)`` per-position weighting used when computing the + reference lattice. Defaults to :attr:`mask` from :meth:`create_mask` (the + lattice signal strength), so strong, well-fit positions dominate the + reference. + calculation_metric : {"median", "mean"}, default="median" + Statistic used for the reference lattice (weighted by ``mask``). Stored on + the returned :class:`StrainMap` and reused by later + :meth:`~StrainMap.update_reference` calls. + + Returns + ------- + StrainMap + A strain map initialized from the fitted lattice vectors. + """ + if self.g1_array is None or self.g2_array is None: + self.get_individual_uv_vectors() + if not isinstance(self.dataset, (Dataset4d, Dataset4dstem)): + raise ValueError("Dataset must be Dataset4d or Dataset4dstem.") + + if mask is None: + mask = self.mask + + default_units = None + default_sampling = None + if hasattr(self.dataset, 'units'): + if isinstance(self.dataset.units, (tuple, list)): + default_units = str(self.dataset.units[0]) + else: + default_units = str(self.dataset.units) + if hasattr(self.dataset, 'sampling'): + if isinstance(self.dataset.sampling, (tuple, list, np.ndarray)): + default_sampling = float(self.dataset.sampling[0]) + else: + default_sampling = float(self.dataset.sampling) + + parent_rot = self.dataset.metadata.get("q_to_r_rotation_ccw_deg", None) + parent_tr = self.dataset.metadata.get("q_transpose", None) + + used_parent = False + if q_to_r_rotation_ccw_deg is None and parent_rot is not None: + q_to_r_rotation_ccw_deg = parent_rot + used_parent = True + if q_transpose is None and parent_tr is not None: + q_transpose = parent_tr + used_parent = True + + if used_parent: + import warnings + + warnings.warn( + "StrainMapAutocorrelation.preprocess: using parent Dataset4dstem metadata " + f"(q_to_r_rotation_ccw_deg={q_to_r_rotation_ccw_deg or 0.0}, " + f"q_transpose={q_transpose or False}).", + UserWarning, + ) + + if q_to_r_rotation_ccw_deg is None or q_transpose is None: + import warnings + + q_to_r_rotation_ccw_deg = 0.0 if q_to_r_rotation_ccw_deg is None else q_to_r_rotation_ccw_deg + q_transpose = False if q_transpose is None else q_transpose + warnings.warn( + "StrainMapAutocorrelation.preprocess: setting q_to_r_rotation_ccw_deg=0.0 and q_transpose=False.", + UserWarning, + ) + + self.metadata["q_to_r_rotation_ccw_deg"] = q_to_r_rotation_ccw_deg + self.metadata["q_transpose"] = q_transpose + + return StrainMap( + g1_array = self.g1_array, + g2_array = self.g2_array, + ds_shape = self.dataset.shape, + real_space = self.real_space, + g1_ref = g1_ref, + g2_ref = g2_ref, + mask = mask, + ds_sampling=default_sampling, + ds_units = default_units, + q_to_r_rotation_ccw_deg = q_to_r_rotation_ccw_deg, + q_transpose = q_transpose, + calculation_metric = calculation_metric, + ) + + @property + def render_mean_refined(self) -> np.ndarray: + if self.state_mean_refined is None: + raise RuntimeError( + "mean_refined state is unavailable. Run .fit_mean_diffraction_pattern(...) first." + ) + return self._render_state_array(self.state_mean_refined) + + +def _resolve_rows_cols_for_batched( + rows: Any, cols: Any, scan_r: int, scan_c: int +) -> tuple[np.ndarray, np.ndarray]: + if rows is None: + rows = range(scan_r) + if cols is None: + cols = range(scan_c) + if isinstance(rows, int): + rows_arr = np.array([rows], dtype=int) + else: + rows_arr = np.asarray(list(rows), dtype=int) + if isinstance(cols, int): + cols_arr = np.array([cols], dtype=int) + else: + cols_arr = np.asarray(list(cols), dtype=int) + return rows_arr, cols_arr + +def _neighbor_uv_target(state_individual_refined, positions, lat_prefix, device, dtype, radius=1): + scan_r, scan_c = state_individual_refined.shape + keys = [f"{lat_prefix}.u_row", f"{lat_prefix}.u_col", f"{lat_prefix}.v_row", f"{lat_prefix}.v_col"] + out = {k: [] for k in keys} + for (r, c) in positions: + vals = {k: [] for k in keys} + for dr in range(-radius, radius + 1): + for dc in range(-radius, radius + 1): + if dr == 0 and dc == 0: + continue + rr, cc = r + dr, c + dc + if not (0 <= rr < scan_r and 0 <= cc < scan_c): + continue + neighbor_state = state_individual_refined[rr, cc] + if neighbor_state is None: + continue + for k in keys: + if k in neighbor_state: + vals[k].append(float(neighbor_state[k])) + for k in keys: + out[k].append(float(np.mean(vals[k])) if vals[k] else float("nan")) + return {k: torch.tensor(v, device=device, dtype=dtype) for k, v in out.items()} + +def _neighbor_uv_target_single(state_individual_refined, r, c, radius=1): + """Non-batched analogue of _neighbor_uv_target: looks up neighbors by key + *suffix* rather than exact prefix, since state_individual_refined here + holds raw nn.Module.state_dict() entries (full module-path keys), not the + _BatchedPlan naming convention.""" + scan_r, scan_c = state_individual_refined.shape + suffixes = ("u_row", "u_col", "v_row", "v_col") + vals = {s: [] for s in suffixes} + for dr in range(-radius, radius + 1): + for dc in range(-radius, radius + 1): + if dr == 0 and dc == 0: + continue + rr, cc = r + dr, c + dc + if not (0 <= rr < scan_r and 0 <= cc < scan_c): + continue + if not (0 <= rr < scan_r and 0 <= cc < scan_c): + continue + neighbor_state = state_individual_refined[rr, cc] + if neighbor_state is None: + continue + for s in suffixes: + key = next((k for k in neighbor_state if k.endswith(f".{s}")), None) + if key is not None: + vals[s].append(float(neighbor_state[key])) + return {s: (float(np.mean(vals[s])) if vals[s] else float("nan")) for s in suffixes} + +def is_outlier(state_individual_refined, r, c, deviation_threshold, radius=1): + state = state_individual_refined[r, c] + if state is None: + return False, None + nt = _neighbor_uv_target_single(state_individual_refined, r, c, radius=radius) + cur = _extract_uv_single(state) + devs = [abs(cur[k] - nt[k]) for k in cur if nt[k] == nt[k]] + if not devs: + return False, None + return max(devs) > deviation_threshold, nt + +def redo_position(model_diff, r, c, neighbor_target, n_steps=500, + optimizer_params=None, scheduler_params=None): + model_diff._load_model_state_dict_copy(model_diff.state_individual_refined[r, c]) + model_diff.fit_render( + target=torch.as_tensor(model_diff.dataset.array[r, c], device=model_diff.ctx.device, dtype=model_diff.ctx.dtype), + n_steps=n_steps, optimizer_params=optimizer_params, scheduler_params=scheduler_params, + progress=False, run_key=f"refit_{r}_{c}", + neighbor_target=neighbor_target, + ) + model_diff.state_individual_refined[r, c] = model_diff._get_model_state_dict_copy() + +def refit_outlier_positions(model_diff, deviation_threshold, radius=1, n_steps=500, + optimizer_params=None, scheduler_params=None, consistency_weight=1.0): + scan_r, scan_c = model_diff.state_individual_refined.shape + lat = next(c for c in model_diff.model.components if type(c).__name__ == "SyntheticDiskLattice") + old_weight = lat.soft_constraints.get("consistency_weight", 0.0) + lat.apply_constraint_params({"consistency_weight": consistency_weight}) + + n_refit = 0 + for r in range(scan_r): + for c in range(scan_c): + outlier, nt = is_outlier(model_diff.state_individual_refined, r, c, deviation_threshold, radius=radius) + if outlier: + redo_position(model_diff, r, c, nt, n_steps=n_steps, + optimizer_params=optimizer_params, scheduler_params=scheduler_params) + n_refit += 1 + + lat.apply_constraint_params({"consistency_weight": old_weight}) + model_diff.get_individual_uv_vectors() + print(f"refit {n_refit} / {scan_r*scan_c} outlier positions") + return model_diff + +def _extract_uv_single(state_dict): + """Pull a position's own (u_row, u_col, v_row, v_col) out of one state dict, + matching by key suffix so it works regardless of key-naming convention + (batched 'lat.u_row' or raw state_dict 'components.N.u_row').""" + suffixes = ("u_row", "u_col", "v_row", "v_col") + out = {} + for s in suffixes: + key = next((k for k in state_dict if k.endswith(f".{s}")), None) + out[s] = float(state_dict[key]) if key is not None else float("nan") + return out + +def _lr_for_component( + component: RenderComponent, + component_idx: int, + optimizer_params: dict[str, Any], + model: AdditiveRenderModel, +) -> float: + name = model._component_constraint_name(component, component_idx) + if name in optimizer_params: + return float(optimizer_params[name].get("lr", component.DEFAULT_LR)) + class_name = component.__class__.__name__ + if class_name in optimizer_params: + return float(optimizer_params[class_name].get("lr", component.DEFAULT_LR)) + return float(component.DEFAULT_LR) + + +def _opt_spec_for_component( + component: RenderComponent, + component_idx: int, + optimizer_params: dict[str, Any], + model: AdditiveRenderModel, +): + """Resolve the full optimizer spec (type + hyperparameters) for one component. + + Mirrors ``_lr_for_component``'s name-then-class lookup, but returns the parsed + ``OptimizerParams`` spec instead of just the learning rate, so batched fitting can + honor a requested optimizer ``type`` (e.g. ``"sgd"``) instead of silently always + running Adam. + """ + name = model._component_constraint_name(component, component_idx) + class_name = component.__class__.__name__ + if name in optimizer_params: + d = dict(optimizer_params[name]) + elif class_name in optimizer_params: + d = dict(optimizer_params[class_name]) + else: + return OptimizerParams.Adam(lr=component.DEFAULT_LR) + d.pop("param_lr", None) # batched-fit per-parameter LR overrides, read by _BatchedPlan + d.setdefault("type", "adam") + d.setdefault("lr", component.DEFAULT_LR) + return OptimizerParams.parse_dict(d) + + +def _param_lr_for_component( + component: RenderComponent, + component_idx: int, + optimizer_params: dict[str, Any], + model: AdditiveRenderModel, +) -> dict[str, float]: + """Per-parameter LR overrides, e.g. ``{"lat0": {"lr": 2e-2, "param_lr": {"i0_raw": 0.1}}}``.""" + name = model._component_constraint_name(component, component_idx) + spec = optimizer_params.get(name, optimizer_params.get(component.__class__.__name__, {})) + return {str(k): float(v) for k, v in dict(spec.get("param_lr", {}) or {}).items()} + + +def _masked(t: torch.Tensor, mask_expanded: torch.Tensor | None) -> torch.Tensor: + return t if mask_expanded is None else t * mask_expanded + + +def _adam_step_inplace( + stacked: dict[str, torch.Tensor], + grads: tuple[torch.Tensor, ...], + adam_state: dict[str, dict[str, torch.Tensor]], + lrs: dict[str, float], + hparams: dict[str, dict[str, Any]], + chunk_trainable: dict[str, torch.Tensor], + t: int, + decoupled_wd: bool = False, +) -> None: + """ + In-place Adam/AdamW step with per-sample trainability support. + + Frozen samples (mask=False) will: + 1. Not accumulate gradient statistics in moments + 2. Not receive parameter updates -- including from weight decay, which is + explicitly masked too (unlike a plain zeroed gradient, weight decay would + otherwise still pull frozen samples toward zero every step). + """ + with torch.no_grad(): + for (name, p), g in zip(stacked.items(), grads): + if g is None: + continue + + hp = hparams.get(name, {}) + beta1, beta2 = hp.get("betas", (0.9, 0.999)) + eps = hp.get("eps", 1e-8) + wd = hp.get("weight_decay", 0.0) + + st = adam_state[name] + mask_b = chunk_trainable.get(name) + mask_expanded = None + if mask_b is not None and not bool(mask_b.all()): + view_shape = (g.shape[0],) + (1,) * (g.ndim - 1) + mask_expanded = mask_b.view(view_shape).to(dtype=g.dtype) + + g_eff = _masked(g, mask_expanded) + if wd != 0: + if decoupled_wd: + lr_ = float(lrs.get(name, 1e-2)) + p.data.sub_(_masked(lr_ * wd * p.data, mask_expanded)) + else: + g_eff = g_eff + _masked(wd * p.data, mask_expanded) + + st["m"].mul_(beta1).add_(g_eff, alpha=1.0 - beta1) + st["v"].mul_(beta2).addcmul_(g_eff, g_eff, value=1.0 - beta2) + + m_hat = st["m"] / (1.0 - beta1 ** t) + v_hat = st["v"] / (1.0 - beta2 ** t) + + lr = float(lrs.get(name, 1e-2)) + p.data.addcdiv_(m_hat, v_hat.sqrt().add_(eps), value=-lr) + + +def _sgd_step_inplace( + stacked: dict[str, torch.Tensor], + grads: tuple[torch.Tensor, ...], + sgd_state: dict[str, dict[str, torch.Tensor]], + lrs: dict[str, float], + hparams: dict[str, dict[str, Any]], + chunk_trainable: dict[str, torch.Tensor], + t: int, +) -> None: + """In-place SGD step (with optional momentum/dampening/nesterov/weight_decay), + matching ``torch.optim.SGD`` exactly, with the same per-sample trainability + masking as :func:`_adam_step_inplace` (including for the weight-decay term).""" + with torch.no_grad(): + for (name, p), g in zip(stacked.items(), grads): + if g is None: + continue + + hp = hparams.get(name, {}) + momentum = hp.get("momentum", 0.0) + dampening = hp.get("dampening", 0.0) + wd = hp.get("weight_decay", 0.0) + nesterov = hp.get("nesterov", False) + + mask_b = chunk_trainable.get(name) + mask_expanded = None + if mask_b is not None and not bool(mask_b.all()): + view_shape = (g.shape[0],) + (1,) * (g.ndim - 1) + mask_expanded = mask_b.view(view_shape).to(dtype=g.dtype) + + d_p = _masked(g, mask_expanded) + if wd != 0: + d_p = d_p + _masked(wd * p.data, mask_expanded) + if momentum != 0: + buf = sgd_state[name]["buf"] + if t == 1: + buf.copy_(d_p) + else: + buf.mul_(momentum).add_(d_p, alpha=1.0 - dampening) + d_p = d_p.add(buf, alpha=momentum) if nesterov else buf + + lr = float(lrs.get(name, 1e-2)) + p.data.add_(d_p, alpha=-lr) + + +def _optimizer_step_inplace( + opt_type: str, + stacked: dict[str, torch.Tensor], + grads: tuple[torch.Tensor, ...], + opt_state: dict[str, dict[str, torch.Tensor]], + lrs: dict[str, float], + hparams: dict[str, dict[str, Any]], + chunk_trainable: dict[str, torch.Tensor], + t: int, +) -> None: + """Dispatch a single in-place optimizer step for the batched fit loop.""" + if opt_type in ("adam", "adamw"): + _adam_step_inplace( + stacked, grads, opt_state, lrs, hparams, chunk_trainable, t, + decoupled_wd=(opt_type == "adamw"), + ) + elif opt_type == "sgd": + _sgd_step_inplace(stacked, grads, opt_state, lrs, hparams, chunk_trainable, t) + else: + raise NotImplementedError(f"Batched fit does not support optimizer type '{opt_type}'.") + +class _BatchedPlan: + """Resolved layout for the batched per-pattern fit: component refs, lrs, and helpers.""" + + def __init__(self) -> None: + self.origin: OriginND | None = None + self.disk: DiskTemplate | None = None + self.dcbg: DCBackground | None = None + self.disk_idx: int | None = None + self.dcbg_idx: int | None = None + self.lrs: dict[str, float] = {} + self.opt_type: str = "adam" + self.opt_hparams: dict[str, dict[str, Any]] = {} + self.scheduler_specs: dict[str, dict[str, Any]] = {} + self.component_keys: dict[str, list[str]] = {} + self.lats: list[SyntheticDiskLattice] = [] + self.lat_idxs: list[int] = [] + self.lat_names: list[str] = [] + self.lat_disk_key: dict[str, str] = {} + self.gaussbgs: list[GaussianBackground] = [] + self.gaussbg_idxs: list[int] = [] + self.gaussbg_names: list[str] = [] + + self._trainable_flags: dict[str, bool] = {} + + @classmethod + def from_model( + cls, + model: AdditiveRenderModel, + components_list: list[Any], + optimizer_params: dict[str, Any], + ) -> "_BatchedPlan": + self = cls() + self.origin = cast(OriginND, model.origin) + + for idx, comp in enumerate(components_list): + if isinstance(comp, DiskTemplate) and self.disk is None: + self.disk = comp + self.disk_idx = idx + elif isinstance(comp, DCBackground) and self.dcbg is None: + self.dcbg = comp + self.dcbg_idx = idx + elif isinstance(comp, GaussianBackground) : + self.gaussbgs.append(comp) + self.gaussbg_idxs.append(idx) + elif isinstance(comp, SyntheticDiskLattice): + self.lats.append(comp) + self.lat_idxs.append(idx) + else: + raise TypeError( + f"Batched fit does not yet support component type {type(comp).__name__} " + f"at index {idx} (or duplicate of an already-handled type)." + ) + + # Track trainability for each parameter + if self.disk is not None and self.disk_idx is not None: + lr = _lr_for_component(self.disk, self.disk_idx, optimizer_params, model) + name = model._component_constraint_name(self.disk, self.disk_idx) + + # Check if parameters are trainable + self._trainable_flags["disk.template_raw"] = self.disk.template_raw.requires_grad + self._trainable_flags["disk.intensity_raw"] = self.disk.intensity_raw.requires_grad + + # Only add to lrs if trainable + if self._trainable_flags["disk.template_raw"]: + self.lrs["disk.template_raw"] = lr + if self._trainable_flags["disk.intensity_raw"]: + self.lrs["disk.intensity_raw"] = lr + + self.component_keys[name] = ["disk.template_raw", "disk.intensity_raw"] + self.component_keys[self.disk.__class__.__name__] = list(self.component_keys[name]) + + if self.dcbg is not None and self.dcbg_idx is not None: + lr = _lr_for_component(self.dcbg, self.dcbg_idx, optimizer_params, model) + name = model._component_constraint_name(self.dcbg, self.dcbg_idx) + + self._trainable_flags["dcbg.intensity_raw"] = self.dcbg.intensity_raw.requires_grad + if self._trainable_flags["dcbg.intensity_raw"]: + self.lrs["dcbg.intensity_raw"] = lr + + self.component_keys[name] = ["dcbg.intensity_raw"] + self.component_keys[self.dcbg.__class__.__name__] = list(self.component_keys[name]) + + if self.gaussbgs: + seen_prefixes: set[str] = set() + for gaussbg, gaussbg_idx in zip(self.gaussbgs, self.gaussbg_idxs): + lr = _lr_for_component(gaussbg, gaussbg_idx, optimizer_params, model) + canonical_name = model._component_constraint_name(gaussbg, gaussbg_idx) + key_prefix = "gaussbg" if len(self.gaussbgs) == 1 else canonical_name + if key_prefix in seen_prefixes: + raise ValueError( + f"Multiple GaussianBackground components resolve to the same name " + f"'{key_prefix}'. Pass a distinct name=... to each GaussianBackground." + ) + seen_prefixes.add(key_prefix) + self.gaussbg_names.append(key_prefix) + + gaussbg_keys: list[str] = [] + + sigma_key = f"{key_prefix}.sigma_raw" + self._trainable_flags[sigma_key] = gaussbg.sigma_raw.requires_grad + if self._trainable_flags[sigma_key]: + self.lrs[sigma_key] = lr + gaussbg_keys.append(sigma_key) + + if gaussbg.sigma_col_raw is not None: + sigma_col_key = f"{key_prefix}.sigma_col_raw" + self._trainable_flags[sigma_col_key] = gaussbg.sigma_col_raw.requires_grad + if self._trainable_flags[sigma_col_key]: + self.lrs[sigma_col_key] = lr + gaussbg_keys.append(sigma_col_key) + + intensity_key = f"{key_prefix}.intensity_raw" + self._trainable_flags[intensity_key] = gaussbg.intensity_raw.requires_grad + if self._trainable_flags[intensity_key]: + self.lrs[intensity_key] = lr + gaussbg_keys.append(intensity_key) + + self.component_keys[canonical_name] = gaussbg_keys + cls_key = gaussbg.__class__.__name__ + bucket = self.component_keys.setdefault(cls_key, []) + for k in gaussbg_keys: + if k not in bucket: + bucket.append(k) + + if self.lats: + first_lr = None + for lat, lat_idx in zip(self.lats, self.lat_idxs): + lr = _lr_for_component(lat, lat_idx, optimizer_params, model) + param_lr = _param_lr_for_component(lat, lat_idx, optimizer_params, model) + if first_lr is None: + first_lr = lr + canonical_name = model._component_constraint_name(lat, lat_idx) + key_prefix = "lat" if len(self.lats) == 1 else canonical_name + self.lat_names.append(key_prefix) + + # Resolve which tracked disk this lattice shares (by identity). + if self.disk is not None and lat.disk is self.disk: + self.lat_disk_key[key_prefix] = "disk" + else: + raise NotImplementedError( + f"Batched fit requires lattice '{canonical_name}' to share the " + f"single tracked DiskTemplate component; found a different/untracked disk." + ) + + lat_keys = [] + for k in ("u_row", "u_col", "v_row", "v_col", "i0_raw", "ir", "ic", "irr", "icc", "irc"): + full_key = f"{key_prefix}.{k}" + param = getattr(lat, k, None) + if param is not None: + is_trainable = param.requires_grad if isinstance(param, torch.nn.Parameter) else False + self._trainable_flags[full_key] = is_trainable + if is_trainable: + self.lrs[full_key] = param_lr.get(k, lr) + lat_keys.append(full_key) + + self.component_keys[canonical_name] = lat_keys + cls_key = lat.__class__.__name__ + bucket = self.component_keys.setdefault(cls_key, []) + for k in lat_keys: + if k not in bucket: + bucket.append(k) + + # Origin uses the first lattice's LR (origin is shared/global). + self._trainable_flags["origin.coords"] = self.origin.coords.requires_grad + if self._trainable_flags["origin.coords"]: + self.lrs["origin.coords"] = first_lr + elif self.gaussbgs: + intensity_key = f"{self.gaussbg_names[0]}.intensity_raw" + if self._trainable_flags.get(intensity_key, False): + self.lrs["origin.coords"] = self.lrs[intensity_key] + else: + if self.origin.coords.requires_grad: + self.lrs["origin.coords"] = 1e-2 + + # Default schedulers: constant LR for every key + self.scheduler_specs = {k: {"type": "none"} for k in self.lrs} + + # Resolve optimizer type/hyperparameters. All components must agree on the + # optimizer class -- mirroring OptimizerMixin.set_optimizer's own constraint -- + # since a single manual step function runs once per training step over every + # stacked parameter. Keys with no explicit spec fall back to that type's + # library defaults inside _adam_step_inplace/_sgd_step_inplace (e.g. + # "origin.coords", which never has its own optimizer_params entry). + spec_types: set[type] = set() + for idx, comp in enumerate(components_list): + spec = _opt_spec_for_component(comp, idx, optimizer_params, model) + spec_types.add(type(spec)) + canonical_name = model._component_constraint_name(comp, idx) + for key in self.component_keys.get(canonical_name, []): + self.opt_hparams[key] = spec.params() + if len(spec_types) > 1: + raise ValueError( + "All components must use the same optimizer type for batched fitting; " + f"got {sorted(t.__name__ for t in spec_types)}." + ) + if spec_types: + self.opt_type = next(iter(spec_types))._name + + return self + + # ... (keep set_scheduler_params, lr_at_step, batched_constraint_loss, resolve_component_keys as-is) + + def build_stacked_params(self, B: int) -> dict[str, torch.Tensor]: + """Build stacked params only for trainable parameters, expand frozen ones as needed.""" + out: dict[str, torch.Tensor] = {} + assert self.origin is not None + + # Only stack trainable params + if self._trainable_flags.get("origin.coords", False): + out["origin.coords"] = self._stack(self.origin.coords, B) + + if self.disk is not None: + if self._trainable_flags.get("disk.template_raw", False): + out["disk.template_raw"] = self._stack(self.disk.template_raw, B) + if self._trainable_flags.get("disk.intensity_raw", False): + out["disk.intensity_raw"] = self._stack(self.disk.intensity_raw, B) + + if self.dcbg is not None: + if self._trainable_flags.get("dcbg.intensity_raw", False): + out["dcbg.intensity_raw"] = self._stack(self.dcbg.intensity_raw, B) + + if self.gaussbgs is not None: + for gaussbg_name, gaussbg in zip(self.gaussbg_names, self.gaussbgs): + sigma_key = f"{gaussbg_name}.sigma_raw" + if self._trainable_flags.get(sigma_key, False): + out[sigma_key] = self._stack(gaussbg.sigma_raw, B) + sigma_col_key = f"{gaussbg_name}.sigma_col_raw" + if self._trainable_flags.get(sigma_col_key, False) and gaussbg.sigma_col_raw is not None: + out[sigma_col_key] = self._stack(gaussbg.sigma_col_raw, B) + intensity_key = f"{gaussbg_name}.intensity_raw" + if self._trainable_flags.get(intensity_key, False): + out[intensity_key] = self._stack(gaussbg.intensity_raw, B) + + if self.lats is not None: + for lat_name, lat in zip(self.lat_names, self.lats): + for attr in ("u_row", "u_col", "v_row", "v_col", "i0_raw"): + t = getattr(lat, attr) + key = f"{lat_name}.{attr}" + if t is not None and self._trainable_flags.get(key, False): + out[key] = self._stack(t, B) + for attr in ("ir", "ic", "irr", "icc", "irc"): + t = getattr(lat, attr, None) + key = f"{lat_name}.{attr}" + if t is not None and self._trainable_flags.get(key, False): + out[key] = self._stack(t, B) + + return out + + @staticmethod + def _stack(p: torch.Tensor, B: int) -> torch.Tensor: + x = p.detach().clone().unsqueeze(0).expand(B, *p.shape).contiguous() + x.requires_grad_(True) + return x + + def batched_forward( + self, ctx: RenderContext, stacked: dict[str, torch.Tensor] + ) -> torch.Tensor: + """Batched forward with frozen parameter support.""" + B = next(iter(stacked.values())).shape[0] if stacked else 1 + + # Get origin (might be frozen) + origin_b = stacked.get("origin.coords") + if origin_b is None: + # Origin is frozen - expand from component + origin_b = self.origin.coords.unsqueeze(0).expand(B, *self.origin.coords.shape).clone() + + pred = torch.zeros(B, ctx.shape[0], ctx.shape[1], device=ctx.device, dtype=ctx.dtype) + + if self.disk is not None: + # Handle potentially frozen disk params + template_b = stacked.get("disk.template_raw") + if template_b is None: + template_b = self.disk.template_raw.unsqueeze(0).expand(B, *self.disk.template_raw.shape).clone() + + intensity_b = stacked.get("disk.intensity_raw") + if intensity_b is None: + intensity_b = self.disk.intensity_raw.unsqueeze(0).expand(B).clone() + + pred = pred + self.disk.forward_batched( + ctx, + template_raw_b=template_b, + intensity_raw_b=intensity_b, + origin_coords_b=origin_b, + ) + + if self.dcbg is not None: + intensity_b = stacked.get("dcbg.intensity_raw") + if intensity_b is None: + intensity_b = self.dcbg.intensity_raw.unsqueeze(0).expand(B).clone() + pred = pred + self.dcbg.forward_batched(ctx, intensity_raw_b=intensity_b) + + for gaussbg_name, gaussbg in zip(self.gaussbg_names, self.gaussbgs): + sigma_b = stacked.get(f"{gaussbg_name}.sigma_raw") + if sigma_b is None: + sigma_b = gaussbg.sigma_raw.unsqueeze(0).expand(B).clone() + + sigma_col_b = stacked.get(f"{gaussbg_name}.sigma_col_raw") + if sigma_col_b is None and gaussbg.sigma_col_raw is not None: + sigma_col_b = gaussbg.sigma_col_raw.unsqueeze(0).expand(B).clone() + + intensity_b = stacked.get(f"{gaussbg_name}.intensity_raw") + if intensity_b is None: + intensity_b = gaussbg.intensity_raw.unsqueeze(0).expand(B).clone() + + pred = pred + gaussbg.forward_batched( + ctx, + sigma_raw_b=sigma_b, + sigma_col_raw_b=sigma_col_b, + intensity_raw_b=intensity_b, + origin_coords_b=origin_b, + ) + + if self.lats: + # Template (from disk) — resolved once, shared by every lattice. + template_b = stacked.get("disk.template_raw") + if template_b is None and self.disk is not None: + template_b = self.disk.template_raw.unsqueeze(0).expand(B, *self.disk.template_raw.shape).clone() + + for lat_name, lat in zip(self.lat_names, self.lats): + u_row_b = stacked.get(f"{lat_name}.u_row") + if u_row_b is None: + u_row_b = lat.u_row.unsqueeze(0).expand(B).clone() + + u_col_b = stacked.get(f"{lat_name}.u_col") + if u_col_b is None: + u_col_b = lat.u_col.unsqueeze(0).expand(B).clone() + + v_row_b = stacked.get(f"{lat_name}.v_row") + if v_row_b is None: + v_row_b = lat.v_row.unsqueeze(0).expand(B).clone() + + v_col_b = stacked.get(f"{lat_name}.v_col") + if v_col_b is None: + v_col_b = lat.v_col.unsqueeze(0).expand(B).clone() + + i0_raw_b = stacked.get(f"{lat_name}.i0_raw") + if i0_raw_b is None: + i0_raw_b = lat.i0_raw.unsqueeze(0).expand(B, *lat.i0_raw.shape).clone() + + ir_b = stacked.get(f"{lat_name}.ir") + if ir_b is None and lat.ir is not None: + ir_b = lat.ir.unsqueeze(0).expand(B, *lat.ir.shape).clone() + + ic_b = stacked.get(f"{lat_name}.ic") + if ic_b is None and lat.ic is not None: + ic_b = lat.ic.unsqueeze(0).expand(B, *lat.ic.shape).clone() + + irr_b = stacked.get(f"{lat_name}.irr") + if irr_b is None and lat.irr is not None: + irr_b = lat.irr.unsqueeze(0).expand(B, *lat.irr.shape).clone() + + icc_b = stacked.get(f"{lat_name}.icc") + if icc_b is None and lat.icc is not None: + icc_b = lat.icc.unsqueeze(0).expand(B, *lat.icc.shape).clone() + + irc_b = stacked.get(f"{lat_name}.irc") + if irc_b is None and lat.irc is not None: + irc_b = lat.irc.unsqueeze(0).expand(B, *lat.irc.shape).clone() + + pred = pred + lat.forward_batched( + ctx, + u_row_b=u_row_b, + u_col_b=u_col_b, + v_row_b=v_row_b, + v_col_b=v_col_b, + i0_raw_b=i0_raw_b, + ir_b=ir_b, + ic_b=ic_b, + irr_b=irr_b, + icc_b=icc_b, + irc_b=irc_b, + template_raw_b=template_b, + origin_coords_b=origin_b, + ) + return pred + + def apply_hard_constraints( + self, + stacked: dict[str, torch.Tensor], + skip_keys: set[str] | None = None, + ) -> None: + """Apply batched hard constraints, respecting frozen parameters.""" + skip_keys = skip_keys or set() + + # Add frozen params to skip_keys + frozen_keys = {k for k, trainable in self._trainable_flags.items() if not trainable} + skip_keys = skip_keys | frozen_keys + + # Parameter bounds (always elementwise; safe on any shape) + if self.disk is not None: + for pname, (lo, hi) in self.disk.parameter_bounds.items(): + key = f"disk.{pname}" + if key in stacked and key not in skip_keys: + self._clamp_bounds_inplace(stacked[key], lo, hi) + if self.dcbg is not None: + for pname, (lo, hi) in self.dcbg.parameter_bounds.items(): + key = f"dcbg.{pname}" + if key in stacked and key not in skip_keys: + self._clamp_bounds_inplace(stacked[key], lo, hi) + + for gaussbg_name, gaussbg in zip(self.gaussbg_names, self.gaussbgs): + for pname, (lo, hi) in gaussbg.parameter_bounds.items(): + key = f"{gaussbg_name}.{pname}" + if key in stacked and key not in skip_keys: + self._clamp_bounds_inplace(stacked[key], lo, hi) + + if self.lats is not None: + for lat_name, lat in zip(self.lat_names, self.lats): + i0_key = f"{lat_name}.i0_raw" + if i0_key not in skip_keys and bool(lat.hard_constraints.get("force_positive_intensity", False)): + i0 = stacked.get(i0_key) + if i0 is not None: + i0.clamp_(min=0.0) + + ir_key, ic_key = f"{lat_name}.ir", f"{lat_name}.ic" + ir = stacked.get(ir_key) + ic = stacked.get(ic_key) + if ( + lat.max_slope_ratio is not None + and ir is not None + and ic is not None + and ir_key not in skip_keys + and ic_key not in skip_keys + ): + i0_for_slope = stacked.get(i0_key) + if i0_for_slope is None: + i0_for_slope = lat.i0_raw.detach().to(ir.device, ir.dtype).expand_as(ir) + radius = lat._slope_support_radius() + if radius > 0.0: + slope = torch.sqrt(ir * ir + ic * ic) + limit = float(lat.max_slope_ratio) * i0_for_slope / radius + scale = torch.clamp(limit / slope.clamp(min=1e-12), max=1.0) + ir.mul_(scale) + ic.mul_(scale) + + disk_template_frozen = "disk.template_raw" in skip_keys + disk_intensity_frozen = "disk.intensity_raw" in skip_keys + + # DiskTemplate composite hard constraints + if self.disk is not None: + template = stacked.get("disk.template_raw") + intensity = stacked.get("disk.intensity_raw") + cfg = self.disk.constraint_config + force_positive = bool(self.disk.hard_constraints.get("force_positive", False)) + if not disk_template_frozen and template is not None and intensity is not None: + if bool(self.disk.hard_constraints.get("force_center", False)): + self._batched_center_disk(template) + if bool(self.disk.hard_constraints.get("force_cutoff", False)): + self._batched_enforce_cutoff(template, cfg) + if bool(self.disk.hard_constraints.get("force_circular_mask", False)): + self._batched_enforce_circular_mask(template, cfg) + if bool(self.disk.hard_constraints.get("force_shrinkage", False)): + template.sub_(float(cfg.get("shrinkage_amount", 0.25))) + if force_positive: + template.clamp_(min=0.0) + if not disk_intensity_frozen and intensity is not None: + intensity.clamp_(min=0.0) + if bool(self.disk.hard_constraints.get("force_norm", False)): + self._batched_enforce_norm(template) + if force_positive and not disk_intensity_frozen and intensity is not None: + intensity.clamp_(min=0.0) + + for lat_name, lat in zip(self.lat_names, self.lats): + key = f"{lat_name}.i0_raw" + if key not in skip_keys and bool(lat.hard_constraints.get("force_positive_intensity", False)): + i0 = stacked.get(key) + if i0 is not None: + i0.clamp_(min=0.0) + + @staticmethod + def _clamp_bounds_inplace(t: torch.Tensor, lo: float | None, hi: float | None) -> None: + if lo is None and hi is None: + return + if lo is None: + t.clamp_(max=float(hi)) # type: ignore[arg-type] + elif hi is None: + t.clamp_(min=float(lo)) + else: + t.clamp_(min=float(lo), max=float(hi)) + + @staticmethod + def _batched_center_disk(template_b: torch.Tensor) -> None: + # template_b: (B, H_t, W_t) + B, h, w = template_b.shape + weights = template_b.clamp(min=0.0) + mass = weights.sum(dim=(1, 2)) + if not torch.any(mass > 1e-12): + return + rr = torch.arange(h, device=template_b.device, dtype=template_b.dtype).view(1, h, 1) + cc = torch.arange(w, device=template_b.device, dtype=template_b.dtype).view(1, 1, w) + safe_mass = mass.clamp(min=1e-12) + com_r = (weights * rr).sum(dim=(1, 2)) / safe_mass + com_c = (weights * cc).sum(dim=(1, 2)) / safe_mass + target_r = (h - 1) * 0.5 + target_c = (w - 1) * 0.5 + shift_r = target_r - com_r # (B,) + shift_c = target_c - com_c + denom_h = max(h - 1, 1) + denom_w = max(w - 1, 1) + ty = -2.0 * shift_r / float(denom_h) + tx = -2.0 * shift_c / float(denom_w) + zeros = torch.zeros_like(tx) + ones = torch.ones_like(tx) + theta = torch.stack( + [ + torch.stack([ones, zeros, tx], dim=1), + torch.stack([zeros, ones, ty], dim=1), + ], + dim=1, + ) # (B, 2, 3) + src = template_b.unsqueeze(1) # (B, 1, H, W) + grid = F.affine_grid(theta, [B, 1, h, w], align_corners=True) + shifted = F.grid_sample(src, grid, mode="bilinear", padding_mode="zeros", align_corners=True)[:, 0] + # Only shift samples with nonzero mass; others left as-is + do_shift = (mass > 1e-12).view(B, 1, 1) + template_b.copy_(torch.where(do_shift, shifted, template_b)) + + @staticmethod + def _batched_enforce_norm(template_b: torch.Tensor) -> None: + mins = template_b.amin(dim=(1, 2), keepdim=True) + template_b.sub_(mins) + maxs = template_b.amax(dim=(1, 2), keepdim=True).clamp(min=1e-12) + template_b.div_(maxs) + + @staticmethod + def _batched_enforce_cutoff(template_b: torch.Tensor, cfg: dict[str, Any]) -> None: + thresh_ratio = float(cfg.get("hard_cutoff_threshold", 0.35)) + maxs = template_b.amax(dim=(1, 2), keepdim=True) + thresh = maxs * thresh_ratio + mask = template_b <= thresh + template_b.masked_fill_(mask, 0.0) + + @staticmethod + def _batched_enforce_circular_mask(template_b: torch.Tensor, cfg: dict[str, Any]) -> None: + B, h, w = template_b.shape + radius = (min(h, w) / 2.0) * float(cfg.get("circular_mask_radius_fraction", 0.95)) + r = torch.arange(-h / 2, h / 2, device=template_b.device, dtype=template_b.dtype) + c = torch.arange(-w / 2, w / 2, device=template_b.device, dtype=template_b.dtype) + rr, cc = torch.meshgrid(r, c, indexing="ij") + circle = torch.sqrt(rr * rr + cc * cc) + if bool(cfg.get("soft_circular_mask", False)): + sharpness = float(cfg.get("circular_mask_sharpness", 0)) + mask2d = torch.sigmoid(sharpness * (radius - circle)) + else: + mask2d = (circle <= radius).to(dtype=template_b.dtype) + template_b.mul_(mask2d.view(1, h, w)) + + def build_sample_state_dict( + self, + init_state: dict[str, torch.Tensor], + stacked: dict[str, torch.Tensor], + b: int, + ) -> dict[str, torch.Tensor]: + """Extract sample b from stacked params, handling frozen params.""" + out = {k: v.detach().clone() for k, v in init_state.items()} + + # Origin: top-level key plus any 'components.X.origin.coords' or + # 'components.X.disk.origin.coords' that PyTorch state_dict registers + if "origin.coords" in stacked: + origin_val = stacked["origin.coords"][b].detach().clone() + else: + # Origin was frozen - use the init value + origin_val = init_state.get("origin.coords") + if origin_val is not None: + origin_val = origin_val.detach().clone() + + if origin_val is not None: + for key in list(out.keys()): + if key == "origin.coords" or key.endswith(".origin.coords"): + out[key] = origin_val.clone() + + # DiskTemplate (shared with lat.disk) + if self.disk is not None and self.disk_idx is not None: + if "disk.template_raw" in stacked: + t_val = stacked["disk.template_raw"][b].detach().clone() + else: + # Template was frozen + t_val = init_state.get(f"components.{self.disk_idx}.template_raw") + if t_val is not None: + t_val = t_val.detach().clone() + + if "disk.intensity_raw" in stacked: + i_val = stacked["disk.intensity_raw"][b].detach().clone() + else: + # Intensity was frozen + i_val = init_state.get(f"components.{self.disk_idx}.intensity_raw") + if i_val is not None: + i_val = i_val.detach().clone() + + if t_val is not None: + for key in list(out.keys()): + if key.endswith(".template_raw"): + out[key] = t_val.clone() + + if i_val is not None: + keys = [f"components.{self.disk_idx}.intensity_raw"] + for lat_name, lat_idx in zip(self.lat_names, self.lat_idxs): + if self.lat_disk_key.get(lat_name) == "disk": + keys.append(f"components.{lat_idx}.disk.intensity_raw") + for key in keys: + if key in out: + out[key] = i_val.clone() + + # DCBackground + if self.dcbg is not None and self.dcbg_idx is not None: + if "dcbg.intensity_raw" in stacked: + out[f"components.{self.dcbg_idx}.intensity_raw"] = stacked["dcbg.intensity_raw"][b].detach().clone() + else: + # Was frozen - keep init value (already in out) + pass + + # GaussianBackground + for gaussbg_name, gaussbg_idx in zip(self.gaussbg_names, self.gaussbg_idxs): + sigma_key = f"{gaussbg_name}.sigma_raw" + if sigma_key in stacked: + out[f"components.{gaussbg_idx}.sigma_raw"] = stacked[sigma_key][b].detach().clone() + sigma_col_key = f"{gaussbg_name}.sigma_col_raw" + if sigma_col_key in stacked: + out[f"components.{gaussbg_idx}.sigma_col_raw"] = stacked[sigma_col_key][b].detach().clone() + intensity_key = f"{gaussbg_name}.intensity_raw" + if intensity_key in stacked: + out[f"components.{gaussbg_idx}.intensity_raw"] = stacked[intensity_key][b].detach().clone() + + # SyntheticDiskLattice scalar + tensor params + if self.lats is not None and self.lat_idxs is not None: + for lat_name, lat_idx in zip(self.lat_names, self.lat_idxs): + for attr in ("u_row", "u_col", "v_row", "v_col", "i0_raw", "ir", "ic", "irr", "icc", "irc"): + key = f"components.{lat_idx}.{attr}" + stacked_key = f"{lat_name}.{attr}" + if key in out and stacked_key in stacked: + out[key] = stacked[stacked_key][b].detach().clone() + # else: was frozen, keep init value + + return out + + def set_scheduler_params(self, scheduler_params: Any) -> None: + """Configure per-key scheduler specs.""" + if scheduler_params is None: + return + if not isinstance(scheduler_params, dict): + return + + def _normalize(spec: dict[str, Any]) -> dict[str, Any]: + out = dict(spec) + t = str(out.pop("type", out.pop("name", "none"))).lower() + if t in ("cosine", "cosineannealing"): + t = "cosine_annealing" + if t == "cosine_annealing": + if "t_max" in out: + out["T_max"] = out.pop("t_max") + if "T_max" in out and out["T_max"] is not None: + out["T_max"] = int(float(out["T_max"])) + if "eta_min" in out: + out["eta_min"] = float(out["eta_min"]) + elif t == "exponential": + if "gamma" in out: + out["gamma"] = float(out["gamma"]) + elif t == "linear": + for k in ("start_factor", "end_factor", "total_iters"): + if k in out and out[k] is not None: + out[k] = float(out[k]) if k != "total_iters" else int(float(out[k])) + out["type"] = t + return out + + if "type" in scheduler_params or "name" in scheduler_params: + spec = _normalize(cast(dict[str, Any], scheduler_params)) + for k in self.scheduler_specs: + self.scheduler_specs[k] = dict(spec) + return + + for comp_name, spec in scheduler_params.items(): + if not isinstance(spec, dict): + continue + norm = _normalize(spec) + param_keys = self.component_keys.get(str(comp_name), []) + for k in param_keys: + if k in self.scheduler_specs: + self.scheduler_specs[k] = dict(norm) + + def lr_at_step(self, step: int, n_steps: int) -> dict[str, float]: + """Evaluate the analytic LR schedule for each stacked-param key at step.""" + out: dict[str, float] = {} + for key, base_lr in self.lrs.items(): + spec = self.scheduler_specs.get(key, {"type": "none"}) + t = str(spec.get("type", "none")).lower() + if t == "cosine_annealing": + T_max = int(spec.get("T_max") or n_steps) + if T_max < 1: + T_max = 1 + eta_min = float(spec.get("eta_min", 1e-7)) + s = max(step - 1, 0) + import math + lr = eta_min + 0.5 * (base_lr - eta_min) * (1.0 + math.cos(math.pi * s / T_max)) + out[key] = float(lr) + elif t == "exponential": + gamma = float(spec.get("gamma", 1.0)) + out[key] = float(base_lr * (gamma ** max(step - 1, 0))) + elif t == "linear": + start = float(spec.get("start_factor", 1.0 / 3.0)) + end = float(spec.get("end_factor", 1.0)) + total = int(spec.get("total_iters", n_steps) or n_steps) + if total < 1: + total = 1 + s = min(max(step - 1, 0), total) + factor = start + (end - start) * (s / total) + out[key] = float(base_lr * factor) + else: + out[key] = float(base_lr) + return out + + def batched_constraint_loss( + self, + ctx: RenderContext, + stacked: dict[str, torch.Tensor], + neighbor_target=None + ) -> torch.Tensor: + B = next(iter(stacked.values())).shape[0] if stacked else 1 + out = torch.zeros(B, device=ctx.device, dtype=ctx.dtype) + if self.disk is not None: + template_b = stacked.get("disk.template_raw") + if template_b is not None: + out = out + self.disk.constraint_loss_batched(ctx, template_raw_b=template_b) + for lat_name, lat in zip(self.lat_names, self.lats): + nt = None + if neighbor_target is not None: + nt = {s: neighbor_target.get(f"{lat_name}.{s}") for s in ("u_row", "u_col", "v_row", "v_col")} + out = out + lat.constraint_loss_batched( + ctx, + u_row_b=stacked.get(f"{lat_name}.u_row"), + u_col_b=stacked.get(f"{lat_name}.u_col"), + v_row_b=stacked.get(f"{lat_name}.v_row"), + v_col_b=stacked.get(f"{lat_name}.v_col"), + ir_b=stacked.get(f"{lat_name}.ir"), + ic_b=stacked.get(f"{lat_name}.ic"), + neighbor_target=nt, + ) + return out + + def resolve_component_keys(self, components: Any) -> list[str]: + """Resolve a name/list of component names to the stacked-param keys they own.""" + if components is None: + return [] + if isinstance(components, str): + components = [components] + keys: list[str] = [] + for name in components: + ks = self.component_keys.get(str(name)) + if ks is None: + raise KeyError( + f"Unknown component name '{name}' for batched fit. " + f"Known: {sorted(self.component_keys)}" + ) + for k in ks: + if k not in keys: + keys.append(k) + return keys \ No newline at end of file diff --git a/src/quantem/diffraction/model_fitting_visualizations.py b/src/quantem/diffraction/model_fitting_visualizations.py new file mode 100644 index 00000000..5fac6d9a --- /dev/null +++ b/src/quantem/diffraction/model_fitting_visualizations.py @@ -0,0 +1,497 @@ +from typing import TYPE_CHECKING, Any, Literal, cast + +import numpy as np +from matplotlib import gridspec +from matplotlib import pyplot as plt + +from quantem.core import config +from quantem.core.visualization import show_2d +from quantem.core.visualization.custom_normalizations import ( + CustomNormalization, + _resolve_normalization, +) + +if TYPE_CHECKING: + from quantem.diffraction.model_fitting import ModelDiffraction + + +class ModelDiffractionVisualizations: + def _plot_overlays( + self, + ax: Any, + origin_rc: np.ndarray, + disk_centers_rc: np.ndarray, + *, + overlay_origin: bool = True, + overlay_disks: bool = True, + origin_marker_kwargs: dict[str, Any] | None = None, + disk_marker_kwargs: dict[str, Any] | None = None, + ) -> None: + """ + Plot origin and disk-center overlays on an axis. + + Parameters + ---------- + ax : Any + Matplotlib axis receiving overlays. + origin_rc : np.ndarray + Origin coordinate as ``(row, col)``. + disk_centers_rc : np.ndarray + Disk centers as ``(N, 2)`` in ``(row, col)`` order. + overlay_origin : bool, optional + If ``True``, plot origin marker. + overlay_disks : bool, optional + If ``True``, plot disk-center markers. + origin_marker_kwargs : dict[str, Any] | None, optional + Matplotlib kwargs merged onto origin marker defaults. + disk_marker_kwargs : dict[str, Any] | None, optional + Matplotlib kwargs merged onto disk marker defaults. + + Returns + ------- + None + """ + colors = config.get("viz.colors.set") + if overlay_origin and origin_rc.shape == (2,): + kw_origin = { + "marker": "+", + "color": colors[0], + "markersize": 10, + "markeredgewidth": 3, + "linestyle": "None", + } + if origin_marker_kwargs is not None: + kw_origin.update(origin_marker_kwargs) + ax.plot(float(origin_rc[1]), float(origin_rc[0]), **kw_origin) + + if overlay_disks and disk_centers_rc.ndim == 2 and disk_centers_rc.shape[0] > 0: + kw_disks = { + "marker": "x", + "color": colors[1], + "markersize": 5, + "markeredgewidth": 2.0, + "linestyle": "None", + } + if disk_marker_kwargs is not None: + kw_disks.update(disk_marker_kwargs) + ax.plot(disk_centers_rc[:, 1], disk_centers_rc[:, 0], **kw_disks) + + def plot_losses( + self, + figax: tuple[Any, Any] | None = None, + plot_lrs: bool = True, + plot_individual: bool = False, + individual_row: int = 0, + individual_col: int = 0, + ) -> tuple[Any, Any]: + md = cast("ModelDiffraction", self) + colors = config.get("viz.colors.set") + loss_color = "k" + lr_color = colors[8] + + if figax is None: + fig, ax = plt.subplots() + else: + fig, ax = figax + if plot_individual: + mean_hist = md.fit_history.get(f"individual_{individual_row}_{individual_col}") + losses = np.asarray([] if mean_hist is None else mean_hist.losses, dtype=np.float64) + else: + mean_hist = md.fit_history.get("mean") + losses = np.asarray([] if mean_hist is None else mean_hist.losses, dtype=np.float64) + if losses.size == 0: + ax.text( + 0.5, + 0.5, + "No fit history available", + ha="center", + va="center", + transform=ax.transAxes, + ) + ax.set_xlabel("Iterations") + ax.set_ylabel("Loss") + if figax is None: + plt.tight_layout() + plt.show() + return fig, ax + + iters = np.arange(losses.size) + lines: list[Any] = [] + lines.extend(ax.semilogy(iters, losses, c=loss_color, lw=2, label="loss")) + ax.set_xlabel("Iterations") + ax.set_ylabel("Loss", color=loss_color) + ax.tick_params(axis="y", which="both", colors=loss_color) + ax.spines["left"].set_color(loss_color) + ax.set_xbound(-2, max(1, int(iters.max())) + 2) + + lrs = np.asarray([] if mean_hist is None else mean_hist.lrs, dtype=np.float64) + if plot_lrs and lrs.size > 0: + if lrs.size == losses.size and not np.allclose(lrs, lrs[0]): + ax_lr = ax.twinx() + ax.set_zorder(2) + ax_lr.set_zorder(1) + ax.patch.set_visible(False) + ax_lr.spines["left"].set_visible(False) + lines.extend( + ax_lr.semilogy(np.arange(lrs.size), lrs, c=lr_color, lw=2, ls="--", label="LR") + ) + ax_lr.set_ylabel("LR", color=lr_color) + ax_lr.tick_params(axis="y", which="both", colors=lr_color) + ax_lr.spines["right"].set_color(lr_color) + else: + ax.set_title(f"LR: {float(lrs[-1]):.2e}", fontsize=10) + + labels = [line.get_label() for line in lines] + if len(labels) > 1: + ax.legend(lines, labels, loc="upper right") + + if figax is None: + plt.tight_layout() + plt.show() + return fig, ax + + def visualize( + self, + *, + individual_loss: bool = False, + pattern_row: int = 0, + pattern_col: int = 0, + power: float = 1, + axsize: tuple[int, int] = (6, 6), + overlay: bool = True, + overlay_origin: bool = True, + overlay_disks: bool = True, + overlay_on: Literal["model", "both"] = "model", + origin_marker_kwargs: dict[str, Any] | None = None, + disk_marker_kwargs: dict[str, Any] | None = None, + figsize: tuple[int, int] = (12,7), + **kwargs + ) -> tuple[Any, Any]: + """ + Visualize fit losses with reference/model image panels. + + Parameters + ---------- + power : float, optional + Power-law display scaling. + cbar : bool, optional + If ``True``, draw colorbars. + axsize : tuple[int, int], optional + Axis size passed through to ``show_2d``. + overlay : bool, optional + If ``True``, draw coordinate overlays. + overlay_origin : bool, optional + If ``True``, include origin marker in overlays. + overlay_disks : bool, optional + If ``True``, include disk-center markers in overlays. + overlay_on : {"model", "both"}, optional + Which image panel(s) receive overlays. + origin_marker_kwargs : dict[str, Any] | None, optional + Marker kwargs override for origin marker. + disk_marker_kwargs : dict[str, Any] | None, optional + Marker kwargs override for disk-center markers. + + Returns + ------- + tuple[Any, Any] + ``(fig, axs)`` for further editing. + + Raises + ------ + RuntimeError + If model/context are not defined. + ValueError + If ``overlay_on`` is invalid. + """ + md = cast("ModelDiffraction", self) + + if md.image_ref is None: + md.preprocess() + if md.image_ref is None or md.model is None or md.ctx is None: + raise RuntimeError("Call .define_model(...) first.") + if md.dataset.shape[0] <= pattern_row or md.dataset.shape[1] <= pattern_col: + raise ValueError("individual row or column outside bounds of dataset") + + fig = plt.figure(figsize=figsize) + gs = gridspec.GridSpec(2, 1, height_ratios=[1, 2], hspace=0.3) + ax_top = fig.add_subplot(gs[0]) + md.plot_losses(figax=(fig, ax_top), plot_lrs=True, plot_individual=individual_loss,individual_row=pattern_row,individual_col=pattern_col) + + if individual_loss: + pred = md.render_individual_pattern(row=pattern_row, col=pattern_col) + ref = np.asarray(md.dataset.array[pattern_row, pattern_col], dtype=np.float32) + else: + ref = np.asarray(md.image_ref, dtype=np.float32) + pred = md.render_current + + refp = ref if power == 1.0 else np.maximum(ref, 0.0) ** float(power) + predp = pred if power == 1.0 else np.maximum(pred, 0.0) ** float(power) + kwargs.setdefault("vmin", float(min(refp.min(), predp.min()))) + kwargs.setdefault("vmax", float(max(refp.max(), predp.max()))) + + gs_bot = gridspec.GridSpecFromSubplotSpec(1, 2, subplot_spec=gs[1], wspace=0.15) + axs = np.array( + [fig.add_subplot(gs_bot[0, 0]), fig.add_subplot(gs_bot[0, 1])], dtype=object + ) + + kwargs.setdefault("cmap", config.get("viz.cmap")) + show_2d( + [refp, predp], + figax=(fig, axs), + title=["image_ref", "model"], + returnfig=False, + axsize=axsize, + **kwargs + ) + + if overlay: + if overlay_on not in ("model", "both"): + raise ValueError("overlay_on must be 'model' or 'both'.") + origin_rc, disk_centers_rc = md.get_overlay_coordinates() + axes = [axs[1]] if overlay_on == "model" else [axs[0], axs[1]] + for ax in axes: + self._plot_overlays( + ax, + origin_rc, + disk_centers_rc, + overlay_origin=overlay_origin, + overlay_disks=overlay_disks, + origin_marker_kwargs=origin_marker_kwargs, + disk_marker_kwargs=disk_marker_kwargs, + ) + + mean_hist = md.fit_history.get("mean") + if individual_loss: + mean_hist = md.fit_history.get(f"individual_{pattern_row}_{pattern_row}") + if mean_hist is not None and len(mean_hist.losses) > 0: + fig.suptitle( + f"Final loss: {mean_hist.losses[-1]:.3e} | Iters: {len(mean_hist.losses)}", + fontsize=13, + y=0.98, + ) + plt.show() + return fig, axs + + def plot_model( + self, + *, + plot_mean_model: bool = False, + plot_individual_model: bool = False, + pattern_row: int = 0, + pattern_col: int= 0, + power: float = 0.25, + returnfig: bool = False, + axsize: tuple[int, int] = (6, 6), + overlay: bool = True, + overlay_origin: bool = True, + overlay_disks: bool = True, + overlay_on: Literal["model", "both"] = "model", + origin_marker_kwargs: dict[str, Any] | None = None, + disk_marker_kwargs: dict[str, Any] | None = None, + **kwargs, + ) -> tuple[Any, Any] | None: + """ + Plot reference and indiviual diffraction images. + + Parameters + ---------- + power : float, optional + Power-law display scaling. + returnfig : bool, optional + If ``True``, return ``(fig, ax)``. + axsize : tuple[int, int], optional + Axis size passed through to ``show_2d``. + overlay : bool, optional + If ``True``, draw coordinate overlays. + overlay_origin : bool, optional + If ``True``, include origin marker in overlays. + overlay_disks : bool, optional + If ``True``, include disk-center markers in overlays. + overlay_on : {"model", "both"}, optional + Which image panel(s) receive overlays. + origin_marker_kwargs : dict[str, Any] | None, optional + Marker kwargs override for origin marker. + disk_marker_kwargs : dict[str, Any] | None, optional + Marker kwargs override for disk-center markers. + **_ : Any + Ignored extra kwargs for backward compatibility. + + Returns + ------- + tuple[Any, Any] | None + Figure/axes tuple when ``returnfig=True``; otherwise ``None``. + + Raises + ------ + RuntimeError + If model/context are not defined. + ValueError + If ``overlay_on`` is invalid. + """ + md = cast("ModelDiffraction", self) + if md.image_ref is None: + md.preprocess() + if md.image_ref is None or md.model is None or md.ctx is None: + raise RuntimeError("Call .define_model(...) first.") + if plot_mean_model and plot_individual_model: + raise RuntimeError("can only plot mean or plot individual, not both") + if plot_individual_model: + pred = md.render_individual_pattern(pattern_row, pattern_col) + ref = np.asarray(md.dataset.array[pattern_row, pattern_col], dtype=np.float32) + elif plot_mean_model: + ref = np.asarray(md.image_ref, dtype=np.float32) + pred = md.render_mean_refined + else: + ref = np.asarray(md.image_ref, dtype=np.float32) + pred = md.render_current + + refp = ref if power == 1.0 else np.maximum(ref, 0.0) ** float(power) + predp = pred if power == 1.0 else np.maximum(pred, 0.0) ** float(power) + # kwargs.setdefault("vmin", float(min(refp.min(), predp.min()))) + # kwargs.setdefault("vmax", float(max(refp.max(), predp.max()))) + + norm = kwargs.get("norm", None) + if norm is None: + kwargs.setdefault("vmin", float(min(refp.min(), predp.min()))) + kwargs.setdefault("vmax", float(max(refp.max(), predp.max()))) + else: + # Force both panels onto one shared interval, derived from the + # reference image (the model may be uninitialized -> bad scale). + cfg = _resolve_normalization(norm) + cnorm = CustomNormalization( + interval_type=cfg.interval_type, + stretch_type=cfg.stretch_type, + lower_quantile=cfg.lower_quantile, + upper_quantile=cfg.upper_quantile, + vmin=cfg.vmin, + vmax=cfg.vmax, + vcenter=cfg.vcenter, + half_range=cfg.half_range, + power=cfg.power, + logarithmic_index=cfg.logarithmic_index, + asinh_linear_range=cfg.asinh_linear_range, + ) + vmin_shared, vmax_shared = cnorm.interval.get_limits(np.asarray(refp)) + kwargs["norm"] = { + "interval_type": "manual", + "stretch_type": cfg.stretch_type, + "vmin": float(vmin_shared), + "vmax": float(vmax_shared), + "power": cfg.power, + "logarithmic_index": cfg.logarithmic_index, + "asinh_linear_range": cfg.asinh_linear_range, + } + + + t1 = kwargs.pop("title", "") + fig, ax = show_2d( + [refp, predp], + title=[t1 + " image_ref", t1 + " model"], + cmap=config.get("viz.cmap"), + cbar=True, + returnfig=True, + axsize=axsize, + **kwargs, + ) + if overlay: + if overlay_on not in ("model", "both"): + raise ValueError("overlay_on must be 'model' or 'both'.") + origin_rc, disk_centers_rc = md.get_overlay_coordinates() + axs_arr = np.asarray(ax, dtype=object).reshape(-1) + axes = [axs_arr[1]] if overlay_on == "model" else [axs_arr[0], axs_arr[1]] + for a in axes: + self._plot_overlays( + a, + origin_rc, + disk_centers_rc, + overlay_origin=overlay_origin, + overlay_disks=overlay_disks, + origin_marker_kwargs=origin_marker_kwargs, + disk_marker_kwargs=disk_marker_kwargs, + ) + if returnfig: + return fig, ax + return None + + def visualize_components( + self, + components: str | list[str], + *, + power: float = 0.25, + cbar: bool = False, + axsize: tuple[int, int] = (6, 6), + returnfig: bool = False, + overlay: bool = True, + overlay_origin: bool = True, + overlay_disks: bool = True, + origin_marker_kwargs: dict[str, Any] | None = None, + disk_marker_kwargs: dict[str, Any] | None = None, + ) -> tuple[Any, Any] | None: + """ + Render and display a summed component image. + + Parameters + ---------- + components : str | list[str] + Component name or list of component names. Multiple names are + composited by summing component renders. + power : float, optional + Power-law display scaling. + cbar : bool, optional + If ``True``, draw colorbars. + axsize : tuple[int, int], optional + Axis size passed through to ``show_2d``. + returnfig : bool, optional + If ``True``, return ``(fig, ax)``. + overlay : bool, optional + If ``True``, draw coordinate overlays on all component panels. + overlay_origin : bool, optional + If ``True``, include origin marker in overlays. + overlay_disks : bool, optional + If ``True``, include disk-center markers in overlays. + origin_marker_kwargs : dict[str, Any] | None, optional + Marker kwargs override for origin marker. + disk_marker_kwargs : dict[str, Any] | None, optional + Marker kwargs override for disk-center markers. + + Returns + ------- + tuple[Any, Any] | None + Figure/axes tuple when ``returnfig=True``; otherwise ``None``. + """ + md = cast("ModelDiffraction", self) + names = [components] if isinstance(components, str) else list(components) + if len(names) == 0: + raise ValueError("components must contain at least one component name.") + + rendered = [ + np.asarray(md.get_rendered_component(name), dtype=np.float32) for name in names + ] + summed = np.sum(np.stack(rendered, axis=0), axis=0) + summed_scaled = summed if power == 1.0 else np.maximum(summed, 0.0) ** float(power) + title = names[0] if len(names) == 1 else " + ".join(names) + + fig, ax = show_2d( + summed_scaled, + title=title, + cmap=config.get("viz.cmap"), + cbar=bool(cbar), + returnfig=True, + axsize=axsize, + ) + + if overlay: + origin_rc, disk_centers_rc = md.get_overlay_coordinates() + self._plot_overlays( + ax, + origin_rc, + disk_centers_rc, + overlay_origin=overlay_origin, + overlay_disks=overlay_disks, + origin_marker_kwargs=origin_marker_kwargs, + disk_marker_kwargs=disk_marker_kwargs, + ) + + if returnfig: + return fig, ax + return None \ No newline at end of file diff --git a/src/quantem/diffraction/strain.py b/src/quantem/diffraction/strain.py index cde29b50..24b601ba 100644 --- a/src/quantem/diffraction/strain.py +++ b/src/quantem/diffraction/strain.py @@ -487,7 +487,6 @@ def plot_strain( """ if arrow_style not in ("title", "legend"): raise ValueError("arrow_style must be 'title' or 'legend'") - e_rr = self.e_rr.array e_cc = self.e_cc.array e_rc = self.e_rc.array @@ -771,6 +770,11 @@ def _weighted_median(err_native: np.ndarray, factor: float) -> float: if returnfig: return fig, ax return result + + @property + def mask_weight(self) -> np.ndarray | None: + """Per-position fit weight (alias of :attr:`mask`, matching BraggVectors/StrainMapAutocorrelation).""" + return self.mask # ---- module-level fitting functions ---- @@ -1026,7 +1030,6 @@ def _rotate_strain_tensor( Row-column (shear) strain component. rotation_angle : float Frame rotation angle, in degrees. - Returns ------- tuple of np.ndarray diff --git a/src/quantem/diffraction/strain_autocorrelation.py b/src/quantem/diffraction/strain_autocorrelation.py new file mode 100644 index 00000000..b90535ea --- /dev/null +++ b/src/quantem/diffraction/strain_autocorrelation.py @@ -0,0 +1,2412 @@ +from __future__ import annotations + +import warnings +from typing import Any + +import matplotlib.pyplot as plt +import numpy as np +import torch +from numpy.typing import NDArray +from scipy.ndimage import distance_transform_edt + +from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.core.datastructures.dataset3d import Dataset3d +from quantem.core.datastructures.dataset4d import Dataset4d +from quantem.core.datastructures.dataset4dstem import Dataset4dstem +from quantem.core.io.serialize import AutoSerialize +from quantem.core.utils.imaging_utils import rotate_image +from quantem.core.utils.utils import electron_wavelength_angstrom +from quantem.core.utils.validators import ensure_valid_array +from quantem.core.visualization import ScalebarConfig, show_2d +from quantem.diffraction.strain import StrainMap + + +class StrainMapAutocorrelation(AutoSerialize): + """Cepstral / autocorrelation lattice fitting and strain mapping for 4D-STEM. + + An alternative to correlation-based :class:`~quantem.diffraction.bragg_vectors.BraggVectors` + that needs no disk template: every diffraction pattern is transformed into the + autocorrelation (real-space) domain, where the crystal periodicity shows up as a + lattice of sharp peaks, and those peaks are fit per scan position to recover the + local lattice vectors. Because the transform lives in real space, the peaks track + real-space lattice spacings (they *expand* under tension, opposite to Bragg disks), + so ``real_space=True`` -- :func:`~quantem.diffraction.strain._strain_tensor` then + reduces both this and the reciprocal-space Bragg path to the same deformation + gradient, and the two strain maps can be compared directly. + + Workflow (each step writes state consumed by the next): + + 1. :meth:`diffraction_mask` -- build a soft mask over the detector that masks pixels + below a threshold so the transform is dominated by the lattice. + 2. :meth:`preprocess` -- transform every pattern and average the magnitudes into a + mean transform image. Four ``mode`` choices set the intensity scaling applied + before the FFT: ``"linear"`` (Patterson / autocorrelation), ``"log"`` + (``log1p``), ``"log-min"`` (PC-STEM exit-wave power cepstrum, + ``log(I - min + log_offset)``), or ``"gamma"`` (power law). With + ``hann_window=True`` the scaled pattern is multiplied by a 2D Hann window before + the FFT, as in PC-STEM; this is required for ``method="cft"`` and ``refine_dft``.The detector->scan rotation + (``q_to_r_rotation_ccw_deg`` + ``q_transpose``) is read from the parent dataset + metadata, the single source of truth shared with the DPC/CoM and Bragg + workflows. + 3. :meth:`choose_lattice_vector` -- refine a hand-picked initial ``(g1, g2)`` basis + against the mean transform, optionally auto-detecting and re-fitting all peaks. + 4. :meth:`fit_lattice_vectors` -- the heavy step: transform and fit the lattice + vectors at every scan position into ``g1_array``/``g2_array`` of shape + ``(scan_row, scan_col, 2)``. The sub-pixel peak refinement is chosen with ``method``: ``"gaussian"`` + (isotropic Gaussian fit, optionally followed by DFT upsampling) or ``"cft"`` + (exact continuous-Fourier-transform maximum, PC-STEM style). + 5. :meth:`create_mask` -- compute the per-position weight :attr:`mask_weight` + (lattice signal strength) used to weight the reference lattice. + 6. :meth:`calculate_strain_map` -- hand the lattice vectors (and + :attr:`mask_weight`) to a :class:`~quantem.diffraction.strain.StrainMap`. + + Use :meth:`from_dataset` or :meth:`from_array` to construct an instance. + + Parameters + ---------- + dataset : Dataset4dstem + The 4D-STEM dataset to analyze. + input_data : Any, optional + The original object passed to the constructor (a dataset or array), retained + for provenance. + """ + + _token = object() + + # Cepstral / Patterson lattice vectors are measured in the autocorrelation + # (real-space) domain -- the peaks track real-space lattice spacings and expand + # under tension -- so real_space=True. _strain_tensor reduces both this and the + # reciprocal-space Bragg path to F.T, so the two strain maps agree. + real_space: bool = True + + def __init__( + self, + dataset: Dataset4dstem, + input_data: Any | None = None, + _token: object | None = None, + ): + """Private constructor; use :meth:`from_dataset` or :meth:`from_array`. + + Direct instantiation is blocked by the ``_token`` guard. The factory + classmethods are the supported entry points and document their arguments. + + Parameters + ---------- + dataset : Dataset4dstem + The 4D-STEM dataset to analyze. + input_data : Any, optional + The original object passed to a factory (kept for provenance). + _token : object, optional + Internal sentinel; must match the class token or a ``RuntimeError`` is + raised. + """ + if _token is not self._token: + raise RuntimeError( + "Use StrainMapAutocorrelation.from_dataset() or StrainMapAutocorrelation.from_array() to instantiate this class." + ) + # Explicit (two-arg) super() rather than the bare super(): the zero-arg form + # needs a compiler-created __class__ closure cell that is absent when this + # method's source is re-exec'd from a string (Jupyter autoreload), which would + # raise "super(): __class__ cell not found". + super(StrainMapAutocorrelation, self).__init__() + self.dataset = dataset + self.input_data = input_data + self.strain = None + self.metadata: dict[str, Any] = {} + self.transform: Dataset2d | None = None + self.transform_rotated: Dataset2d | None = None + + self.mean_img_peaks: NDArray | None = None + self.mean_img_weights: NDArray | None = None + self.mean_img_ab: NDArray | None = None + self.fft_window: NDArray | None = None + + self.mask_diffraction = np.ones(self.dataset.array.shape[2:]) + self.mask_diffraction_inv = np.zeros(self.dataset.array.shape[2:]) + + # initial basis from choose_lattice_vector(); per-position fits from + # fit_lattice_vectors(); per-position weight from create_mask(). + self.g1: np.ndarray | None = None + self.g2: np.ndarray | None = None + self.g1_peak_fit: Dataset3d | None = None + self.g2_peak_fit: Dataset3d | None = None + self.g1_array: np.ndarray | None = None + self.g2_array: np.ndarray | None = None + self.mask_weight: np.ndarray | None = None + self._gpu_cache: torch.Tensor | None = None + + @classmethod + def from_dataset(cls, dataset: Dataset4dstem, *, name: str | None = None) -> "StrainMapAutocorrelation": + """Create a cepstral strain workflow bound to a 4D-STEM dataset. + + Parameters + ---------- + dataset : Dataset4dstem + The 4D-STEM dataset to analyze. + name : str, optional + If given, sets ``dataset.name``. + + Returns + ------- + StrainMapAutocorrelation + A new workflow instance bound to ``dataset``. + """ + if not isinstance(dataset, Dataset4dstem): + raise TypeError("StrainMapAutocorrelation.from_dataset expects a Dataset4dstem instance.") + if name is not None: + dataset.name = name + return cls(dataset=dataset, input_data=dataset, _token=cls._token) + + @classmethod + def from_array(cls, array: NDArray, *, name: str = "strain_map_autocorrelation") -> "StrainMapAutocorrelation": + """Create a cepstral strain workflow from a raw 4D array. + + Parameters + ---------- + array : np.ndarray + 4D-STEM data with shape ``(scan_row, scan_col, dp_row, dp_col)``. + name : str, default="strain_map_autocorrelation" + Name for the wrapped :class:`Dataset4dstem`. + + Returns + ------- + StrainMapAutocorrelation + A new workflow instance wrapping the data in a :class:`Dataset4dstem`. + """ + arr = ensure_valid_array(array) + if arr.ndim != 4: + raise ValueError( + "StrainMapAutocorrelation.from_array expects a 4D array with shape (scan_r, scan_c, dp_r, dp_c)." + ) + ds4 = Dataset4dstem.from_array(arr, name=name) + return cls(dataset=ds4, input_data=array, _token=cls._token) + + def diffraction_mask( + self, + threshold=None, + threshold_percentile=50.0, + edge_blend=64.0, + plot_mask=True, + figsize=(8, 4), + ): + """Build a soft detector mask that removes dim detector regions. + + Pixels of the mean diffraction pattern below ``threshold`` (plus the detector + border) are treated as "outside" the useful signal. The kept region is feathered + with a raised-cosine taper of width ``edge_blend`` (via a Euclidean distance + transform) into :attr:`mask_diffraction` (multiplicative, in ``[0, 1]``), and a + complementary fill :attr:`mask_diffraction_inv` replaces the masked region with a + flat edge intensity. Both are applied to every pattern in :meth:`preprocess` and + :meth:`fit_lattice_vectors` so the transform is not dominated by the dim, noisy parts of the detector. The + central beam is bright and is therefore *kept*. Because the mask is thresholded on + the raw mean pattern, it can follow the disk lattice (disks kept, gaps masked); + check it with ``plot_mask=True``, since a lattice-shaped mask imprints the mean + lattice on every pattern's transform. + + Parameters + ---------- + threshold : float, optional + Absolute mean-intensity level below which detector pixels are masked out. If + ``None`` (default), the level is taken from ``threshold_percentile`` of the + mean diffraction pattern, so no manual intensity value is needed. + threshold_percentile : float, default=50.0 + Percentile (``0``-``100``) of the mean diffraction pattern used to set the + masking level when ``threshold`` is ``None``. The default keeps the brighter + half of the detector (disks and central beam) and masks the dimmer vacuum; + raise it to mask more aggressively. Ignored when ``threshold`` is given. + edge_blend : float, default=64.0 + Feather width in pixels of the raised-cosine taper between kept and masked + regions; larger values give a softer transition. + plot_mask : bool, default=True + If ``True``, show the raw mean pattern beside the masked/filled pattern (log + scaled) so the mask can be checked. + figsize : tuple of float, default=(8, 4) + Figure size in inches for the diagnostic plot. + + Returns + ------- + StrainMapAutocorrelation + ``self``, with :attr:`mask_diffraction` and :attr:`mask_diffraction_inv` set. + """ + dp_mean = np.mean(self.dataset.array, axis=(0, 1)) + if threshold is None: + threshold = np.percentile(dp_mean, threshold_percentile) + mask_init = dp_mean < threshold + mask_init[:, 0] = True + mask_init[0, :] = True + mask_init[:, -1] = True + mask_init[-1, :] = True + + # The feather cannot be wider than the kept region's inradius: otherwise no + # pixel reaches mask ~ 1 and the int_edge reduction below sees an empty + # selection (e.g. edge_blend=64 on a 128x128 detector). + edge_distance = distance_transform_edt(np.logical_not(mask_init)) + max_distance = float(np.max(edge_distance)) + if max_distance <= 0.0: + raise ValueError( + "No detector pixels survive the threshold; lower threshold or " + "threshold_percentile." + ) + if edge_blend > max_distance: + + warnings.warn( + f"edge_blend={edge_blend:g} exceeds the kept region's largest distance " + f"from the masked region ({max_distance:g} pixels); clamping to " + f"{max_distance:g}. Pass a smaller edge_blend to silence this warning.", + UserWarning, + stacklevel=2, + ) + edge_blend = max_distance + + self.mask_diffraction = np.sin( + np.clip( + edge_distance / edge_blend, + 0.0, + 1.0, + ) + * np.pi + / 2, + ) ** 2 + int_edge = np.min(dp_mean[self.mask_diffraction > 0.99]) + self.mask_diffraction_inv = (1 - self.mask_diffraction) * int_edge + + if plot_mask: + fig, ax = plt.subplots(1, 2, figsize=figsize) + ax[0].imshow( + np.log(np.maximum(dp_mean, np.min(dp_mean[dp_mean > 0]))), + cmap="gray", + ) + ax[1].imshow( + np.log( + dp_mean * self.mask_diffraction + self.mask_diffraction_inv, + ), + cmap="gray", + ) + + return self + + def preprocess( + self, + mode: str = "linear", + q_to_r_rotation_ccw_deg: float | None = None, + q_transpose: bool | None = None, + skip=None, + plot_transform: bool = True, + cropping_factor: float = 0.25, + gamma: float = 0.5, + hann_window: bool = False, + log_offset: float = 0.1, + log_offset_relative: bool = False, + **plot_kwargs: Any, + ) -> "StrainMapAutocorrelation": + """Transform every pattern into the autocorrelation domain and average it. + + For each diffraction pattern the masked/filled pattern (from + :meth:`diffraction_mask`) is intensity-scaled per ``mode``, Fourier transformed, + and its magnitude accumulated; the mean over all scan positions is stored + (fft-shifted, origin centered) as :attr:`transform`. A display copy rotated by + the detector->scan rotation is stored as :attr:`transform_rotated`. The crystal + periodicity appears as a lattice of peaks about the center, which later steps fit. + + The detector->scan rotation (``q_to_r_rotation_ccw_deg``) and transpose + (``q_transpose``) default to the parent dataset metadata -- the same source used + by the DPC/CoM and :class:`BraggVectors` workflows -- so the strain frame is + consistent across methods; pass them explicitly to override. + + Parameters + ---------- + mode : {"linear", "log", "log-min", "gamma"}, default="linear" + Intensity scaling applied before the FFT. ``"linear"`` is the Patterson / + autocorrelation (aliases ``"patterson"``, ``"acf"``, ``"autocorrelation"``); + ``"log"`` is ``log1p(I)`` (aliases ``"cepstrum"``, ``"cepstral"``), i.e. an + absolute offset of 1, so it also depends on the data's intensity units; + ``"log-min"`` is the PC-STEM exit-wave power cepstrum + ``log(I - min(I) + log_offset)`` (aliases ``"padgett"``, ``"log_min"``, + ``"pcepstrum"``); ``"gamma"`` raises intensity to the power ``gamma`` + (aliases ``"power"``, ``"sqrt"``). + q_to_r_rotation_ccw_deg : float, optional + Counter-clockwise detector->scan rotation in degrees for the rotated display + transform. Defaults to ``dataset.metadata["q_to_r_rotation_ccw_deg"]`` if + present, else ``0`` (with a warning). + q_transpose : bool, optional + Whether to transpose the detector axes before rotating the display transform. + Defaults to ``dataset.metadata["q_transpose"]`` if present, else ``False``. + skip : int, optional + If given, subsample the scan by this stride (``array[::skip, ::skip]``) when + building the mean transform -- a fast preview over fewer patterns. + plot_transform : bool, default=True + If ``True``, show the original and rotated mean transforms via + :meth:`plot_transform`. + cropping_factor : float, default=0.25 + Fraction of the transform width/height shown when ``plot_transform`` is + ``True`` (the lattice peaks sit near the center). + gamma : float, default=0.5 + Exponent for ``mode="gamma"`` (ignored otherwise). + hann_window : bool, default=False + If ``True``, multiply the intensity-scaled pattern by a 2D Hann window before + the FFT (after the scaling, as in PC-STEM). The window is stored in + :attr:`fft_window` and used by every later step. Required for + ``method="cft"`` and ``refine_dft=True``, which are otherwise pulled by + central-peak leakage. + log_offset : float, default=0.1 + Offset inside the log for ``mode="log-min"``: ``log(I - min(I) + offset)``. The + PC-STEM default 0.1 assumes intensities in counts; for normalized data it swamps + the signal, the log becomes nearly linear (a Patterson), and results change with + the data's intensity units. + log_offset_relative : bool, default=False + If ``True``, ``log_offset`` is in units of the dataset's mean min-shifted + intensity, so the transform does not depend on intensity scaling. For the SiGe + simulations, ``log_offset=4, log_offset_relative=True`` matches 0.1 on x1e3 data. + **plot_kwargs + Forwarded to :meth:`plot_transform`. + + Returns + ------- + StrainMapAutocorrelation + ``self``, with :attr:`transform` and :attr:`transform_rotated` set. + """ + mode_in = mode.strip().lower() + if mode_in in {"linear", "patterson", "paterson", "acf", "autocorrelation"}: + mode_norm = "linear" + elif mode_in in {"log", "cepstrum", "cepstral"}: + mode_norm = "log" + elif mode_in in {"log_min", "log-min", "padgett", "log_min_shift", "log-min-shift", "pcepstrum"}: + mode_norm = "log-min" + elif mode_in in {"gamma", "power", "sqrt"}: + mode_norm = "gamma" + else: + raise ValueError( + "mode must be 'linear', 'log', or 'gamma' (aliases: 'patterson'->'linear', 'cepstrum'/'cepstral'->'log')." + ) + + H, W = self.dataset.shape[2:] + if hann_window: + self.metadata["hann_window"] = True + self.fft_window = np.outer(np.hanning(H), np.hanning(W)) + else: + self.fft_window = np.ones((H,W)) + self.metadata["hann_window"] = False + + self.metadata["mode"] = mode_norm + if mode_norm == "gamma": + self.metadata["gamma"] = gamma + + qrow_unit = self.dataset.units[2] + qcol_unit = self.dataset.units[3] + + if qrow_unit in {"A^-1", "Å^-1", "1/A", "1/Å"}: + qrow_sampling_ang = self.dataset.sampling[2] + elif qrow_unit == "mrad": + wavelength = electron_wavelength_angstrom(self.dataset.metadata["energy"]) + qrow_sampling_ang = self.dataset.sampling[2] / 1000.0 / wavelength + else: + qrow_sampling_ang = 1.0 + qrow_unit = "pixels" + + if qcol_unit in {"A^-1", "Å^-1", "1/A", "1/Å"}: + qcol_sampling_ang = self.dataset.sampling[3] + elif qcol_unit == "mrad": + wavelength = electron_wavelength_angstrom(self.dataset.metadata["energy"]) + qcol_sampling_ang = self.dataset.sampling[3] / 1000.0 / wavelength + else: + qcol_sampling_ang = 1.0 + qcol_unit = "pixels" + + self.metadata["sampling_real"] = np.array( + ( + 1.0 / (qrow_sampling_ang * self.dataset.shape[2]), + 1.0 / (qcol_sampling_ang * self.dataset.shape[3]), + ), + dtype=float, + ) + + if qrow_unit == "pixels" and qcol_unit == "pixels": + self.metadata["real_units"] = "1/pixels" + else: + self.metadata["real_units"] = r"$\mathrm{\AA}$" + + parent_rot = self.dataset.metadata.get("q_to_r_rotation_ccw_deg", None) + parent_tr = self.dataset.metadata.get("q_transpose", None) + + used_parent = False + if q_to_r_rotation_ccw_deg is None and parent_rot is not None: + q_to_r_rotation_ccw_deg = parent_rot + used_parent = True + if q_transpose is None and parent_tr is not None: + q_transpose = parent_tr + used_parent = True + + if used_parent: + + warnings.warn( + "StrainMapAutocorrelation.preprocess: using parent Dataset4dstem metadata " + f"(q_to_r_rotation_ccw_deg={q_to_r_rotation_ccw_deg or 0.0}, " + f"q_transpose={q_transpose or False}).", + UserWarning, + ) + + if q_to_r_rotation_ccw_deg is None or q_transpose is None: + + q_to_r_rotation_ccw_deg = 0.0 if q_to_r_rotation_ccw_deg is None else q_to_r_rotation_ccw_deg + q_transpose = False if q_transpose is None else q_transpose + warnings.warn( + "StrainMapAutocorrelation.preprocess: setting q_to_r_rotation_ccw_deg=0.0 and q_transpose=False.", + UserWarning, + ) + + self.metadata["q_to_r_rotation_ccw_deg"] = q_to_r_rotation_ccw_deg + self.metadata["q_transpose"] = q_transpose + + arr = self.dataset.array if skip is None else self.dataset.array[::skip, ::skip] + dp = arr * self.mask_diffraction[None, None, :, :] + self.mask_diffraction_inv[None, None, :, :] + dp_min = float(np.min(self.dataset.array)) + if mode_norm == "log-min": + if log_offset_relative: + offset = float(log_offset * (np.mean(arr) - dp_min)) + else: + offset = float(log_offset) + if not np.isfinite(offset) or offset <= 0: + raise ValueError(f"log offset must be finite and > 0 (got {offset}).") + self.metadata["dp_min"] = dp_min + self.metadata["log_offset"] = offset + self.metadata["log_offset_relative"] = log_offset_relative + + if mode_norm == "linear": + dp_proc = dp + elif mode_norm == "log": + dp_proc = np.log1p(dp) + elif mode_norm == "gamma": + dp_proc = np.power(np.clip(dp, 0.0, None), self.metadata["gamma"]) + elif mode_norm == "log-min": + dp_proc = np.log(dp - dp_min + self.metadata["log_offset"]) + else: + raise RuntimeError("Unreachable: normalized mode mapping failed.") + + im = np.mean(np.abs(np.fft.fft2(dp_proc * self.fft_window)), axis=(0, 1)) + im = np.fft.fftshift(im) + + self.transform = Dataset2d.from_array( + im, + origin=(im.shape[0] // 2, im.shape[1] // 2), + sampling=(1.0, 1.0), + units=(qrow_unit, qcol_unit), + signal_units="intensity", + ) + + im_plot = self.transform.array + if self.metadata["q_transpose"]: + im_plot = im_plot.T + + self.transform_rotated = Dataset2d.from_array( + rotate_image( + im_plot, + self.metadata["q_to_r_rotation_ccw_deg"], + clockwise=False, + ), + origin=(im.shape[0] // 2, im.shape[1] // 2), + sampling=(1.0, 1.0), + units=(self.metadata["real_units"], self.metadata["real_units"]), + signal_units="intensity", + ) + + if plot_transform: + self.plot_transform(cropping_factor=cropping_factor, **plot_kwargs) + + return self + + def plot_transform( + self, + cropping_factor: float = 0.25, + scalebar_fraction: float = 0.25, + **plot_kwargs: Any, + ): + """Show the original and rotated mean transform images side by side. + + The color range is set from the brightest lattice peak (the global max of the + radially weighted transform) so the central DC peak does not wash out the panel, + and both panels are cropped to the central ``cropping_factor`` window where the + lattice peaks lie. A scale bar is drawn in real-space units. + + Parameters + ---------- + cropping_factor : float, default=0.25 + Fraction of the full transform width/height shown about the center. + scalebar_fraction : float, default=0.25 + Target scale-bar length as a fraction of the cropped view width (snapped to + a "nice" round value). + **plot_kwargs + Forwarded to :func:`~quantem.core.visualization.show_2d` (overriding the + defaults computed here). + + Returns + ------- + tuple + ``(fig, ax)`` from :func:`~quantem.core.visualization.show_2d`. + """ + if self.transform is None or self.transform_rotated is None: + raise ValueError("Run preprocess() first to compute transform images.") + + sampling = np.mean(self.metadata["sampling_real"]) + units = self.metadata.get("real_units", r"$\mathrm{\AA}$") + + W = self.transform.shape[1] + view_w_px = W * cropping_factor + target_units = scalebar_fraction * view_w_px * sampling + sb_len = _nice_length_units(target_units) + + kr = (np.arange(self.transform.shape[0], dtype=float) - self.transform.shape[0] // 2)[:, None] + kc = (np.arange(self.transform.shape[1], dtype=float) - self.transform.shape[1] // 2)[None, :] + qmag = np.sqrt(kr * kr + kc * kc) + im0 = self.transform.array + tmp = im0 * qmag + i0 = np.unravel_index(np.nanargmax(tmp), tmp.shape) + vmin = 0.0 + vmax = im0[i0] + + defaults = dict( + vmin=vmin, + vmax=vmax, + title=("Original Transform", "Rotated Transform"), + scalebar=ScalebarConfig( + sampling=sampling, + units=units, + length=sb_len if sb_len > 0 else None, + ), + ) + defaults.update(plot_kwargs) + + fig, ax = show_2d([self.transform, self.transform_rotated], **defaults) + + for a in _flatten_axes(ax): + _apply_center_crop_limits(a, self.transform.shape, cropping_factor) + + return fig, ax + + def plot_single_transform( + self, + row: int = 0, + col: int = 0, + cropping_factor: float = 0.25, + scalebar_fraction: float = 0.25, + **plot_kwargs: Any, + ): + """Show the transform of a single scan position with its fitted lattice vectors. + + Recomputes the transform for the pattern at ``(row, col)`` using the current + ``mode``, and overlays the per-position fitted ``g1``/``g2`` vectors (from + :meth:`fit_lattice_vectors`) plus any detected peaks. Useful for inspecting the + fit quality at a specific position. + + Parameters + ---------- + row : int, default=0 + Scan row index of the pattern to transform. + col : int, default=0 + Scan column index of the pattern to transform. + cropping_factor : float, default=0.25 + Fraction of the full transform width/height shown about the center. + scalebar_fraction : float, default=0.25 + Target scale-bar length as a fraction of the cropped view width. + **plot_kwargs + Forwarded to :func:`~quantem.core.visualization.show_2d`. + + Returns + ------- + None + Draws the figure; nothing is returned. + """ + if self.transform is None or self.transform_rotated is None: + raise ValueError("Run preprocess() first to compute transform images.") + if self.g1_peak_fit is None or self.g2_peak_fit is None: + raise ValueError("Run fit_lattice_vectors() first.") + + sampling = np.mean(self.metadata["sampling_real"]) + units = self.metadata.get("real_units", r"$\mathrm{\AA}$") + + W = self.transform.shape[1] + view_w_px = W * cropping_factor + target_units = scalebar_fraction * view_w_px * sampling + sb_len = _nice_length_units(target_units) + + kr = (np.arange(self.transform.shape[0], dtype=float) - self.transform.shape[0] // 2)[:, None] + kc = (np.arange(self.transform.shape[1], dtype=float) - self.transform.shape[1] // 2)[None, :] + qmag = np.sqrt(kr * kr + kc * kc) + im0 = self.transform.array + tmp = im0 * qmag + i0 = np.unravel_index(np.nanargmax(tmp), tmp.shape) + vmin = 0.0 + vmax = im0[i0] + + defaults = dict( + vmin=vmin, + vmax=vmax, + title=(f"Row={row} Col={col} Transform"), + scalebar=ScalebarConfig( + sampling=sampling, + units=units, + length=sb_len if sb_len > 0 else None, + ), + ) + defaults.update(plot_kwargs) + if row >= 0 and row < self.dataset.array.shape[0] and col >= 0 and col < self.dataset.array.shape[1]: + dp = self.dataset.array[row, col] * self.mask_diffraction + self.mask_diffraction_inv + else: + raise ValueError("row or column value out of bounds") + mode = self.metadata.get("mode", "linear").lower() + if mode == "gamma": + g = self.metadata["gamma"] + + + if mode == "linear": + im = np.fft.fftshift(np.abs(np.fft.fft2(dp * self.fft_window))) + elif mode == "log": + im = np.fft.fftshift(np.abs(np.fft.fft2(np.log1p(dp) * self.fft_window))) + elif mode == "log-min": + dp_min = self.metadata.get("dp_min", 0.0) + offset = self.metadata.get("log_offset", 0.1) + im = np.fft.fftshift(np.abs(np.fft.fft2(np.log(dp - dp_min + offset) * self.fft_window))) + elif mode == "gamma": + im = np.fft.fftshift(np.abs(np.fft.fft2(np.power(np.clip(dp, 0.0, None), g) * self.fft_window))) + else: + raise ValueError("metadata['mode'] must be 'linear', 'log', 'log-min', or 'gamma'") + + + fig, ax = show_2d(im, **defaults) + rot_ccw = self.metadata["q_to_r_rotation_ccw_deg"] + q_transpose = self.metadata["q_transpose"] + + _overlay_lattice_vectors( + ax=ax, + shape=self.transform.shape, + g1_rc= self.g1_peak_fit.array[row, col, :2], + g2_rc=self.g2_peak_fit.array[row, col, :2], + rot_ccw_deg=rot_ccw, + q_transpose=q_transpose, + peaks_plot=self.mean_img_peaks, + ) + + for a in _flatten_axes(ax): + _apply_center_crop_limits(a, self.transform.shape, cropping_factor) + + def choose_lattice_vector( + self, + *, + g1: tuple[float, float] | NDArray, + g2: tuple[float, float] | NDArray, + define_in_rotated: bool = False, + refine_gaussian: bool = True, + refine_all_peaks: bool = False, + refine_radius_px: float = 2.0, + gaussian_maxfev: int = 1000, + threshold_percentile: float = 0.9975, + min_peak_spacing: float = 0, + plot: bool = True, + cropping_factor: float = 0.25, + **plot_kwargs: Any, + ) -> "StrainMapAutocorrelation": + """Refine a hand-picked initial lattice basis against the mean transform. + + Takes an approximate basis ``(g1, g2)`` -- read off the transform plot by eye -- + and refines each vector to the nearest transform peak, storing the result in + :attr:`g1` and :attr:`g2`. These seed the per-position fit in + :meth:`fit_lattice_vectors`. Vectors are ``(row, col)`` offsets from the + transform center, in transform pixels. + + Parameters + ---------- + g1 : tuple of float or np.ndarray + Initial first lattice vector ``(d_row, d_col)`` relative to the center. + g2 : tuple of float or np.ndarray + Initial second lattice vector ``(d_row, d_col)`` relative to the center. + define_in_rotated : bool, default=False + If ``True``, ``g1``/``g2`` are given in the rotated (display) frame and are + converted back to the raw detector frame before fitting. + refine_gaussian : bool, default=True + If ``True``, refine each peak by a 2D isotropic Gaussian fit; otherwise use + the parabolic sub-pixel estimate only. + refine_all_peaks : bool, default=False + If ``True``, auto-detect all peaks above ``threshold_percentile`` and fit + the basis to the full set by weighted least squares (storing the detected + peaks/weights for reuse), rather than refining only ``g1`` and ``g2``. + refine_radius_px : float, default=2.0 + Half-width in pixels of the window used for sub-pixel/Gaussian refinement. + gaussian_maxfev : int, default=1000 + Maximum function evaluations for the Gaussian fit. + threshold_percentile : float, default=0.9975 + Intensity percentile (0--1) above which local maxima are kept as peaks when + ``refine_all_peaks=True``. + min_peak_spacing : float, default=0 + Minimum spacing in pixels between accepted peaks when + ``refine_all_peaks=True`` (0 disables the spacing filter). + plot : bool, default=True + If ``True``, show the transform with the refined ``g1``/``g2`` overlaid. + cropping_factor : float, default=0.25 + Fraction of the transform shown about the center when ``plot=True``. + **plot_kwargs + Forwarded to :meth:`plot_transform`. + + Returns + ------- + StrainMapAutocorrelation + ``self``, with :attr:`g1` and :attr:`g2` set (and the detected peaks/weights + when ``refine_all_peaks=True``). + """ + if self.transform is None or self.transform_rotated is None: + raise ValueError("Run preprocess() first to compute transform images.") + + g1_rc = np.asarray(g1, dtype=float).reshape(2) + g2_rc = np.asarray(g2, dtype=float).reshape(2) + + rot_ccw = self.metadata["q_to_r_rotation_ccw_deg"] + q_transpose = self.metadata["q_transpose"] + + if define_in_rotated: + g1_rc = _display_vec_to_raw(g1_rc, rotation_ccw_deg=rot_ccw, transpose=q_transpose) + g2_rc = _display_vec_to_raw(g2_rc, rotation_ccw_deg=rot_ccw, transpose=q_transpose) + + g1_fit_abs, g2_fit_abs, peaks, weights, peaks_ab = _refine_lattice_vectors( + self.transform.array, + g1_rc=g1_rc, + g2_rc=g2_rc, + radius_px=refine_radius_px, + refine_gaussian=refine_gaussian, + refine_dft=False, + refine_all_peaks=refine_all_peaks, + peaks=None, + weights=None, + maxfev=gaussian_maxfev, + threshold_percentile=threshold_percentile, + min_peak_spacing = min_peak_spacing, + warn_on_bound = True, + ) + + self.g1 = g1_fit_abs[:2] + self.g2 = g2_fit_abs[:2] + + if refine_all_peaks: + self.mean_img_peaks = peaks + self.mean_img_weights = weights + self.mean_img_ab = peaks_ab + + self.metadata["choose_define_in_rotated"] = define_in_rotated + self.metadata["choose_refine_gaussian"] = refine_gaussian + self.metadata["choose_refine_all_peaks"] = refine_all_peaks + self.metadata["choose_refine_radius_px"] = refine_radius_px + self.metadata["choose_gaussian_maxfev"] = gaussian_maxfev + self.metadata["choose_threshold_percentile"] = threshold_percentile + + if plot: + fig, ax = self.plot_transform(cropping_factor=cropping_factor, **plot_kwargs) + _overlay_lattice_vectors( + ax=ax, + shape=self.transform.shape, + g1_rc=self.g1, + g2_rc=self.g2, + rot_ccw_deg=rot_ccw, + q_transpose=q_transpose, + peaks_plot=self.mean_img_peaks, + ) + return self + + return self + + def fit_lattice_vectors( + self, + refine_gaussian: bool = True, + refine_dft: bool = False, + refine_all_peaks: bool = False, + refine_radius_px: float = 2.0, + upsample: int = 16, + gaussian_maxfev: int = 1000, + progressbar: bool = True, + device: str = "cpu", + batch_size: int | None = None, + batch_fit: bool = True, + save_to_gpu: bool = True, + method: str = "gaussian", + ) -> "StrainMapAutocorrelation": + """Fit the lattice vectors at every scan position (the heavy step). + + For each scan position the masked/filled pattern is transformed (same ``mode`` + as :meth:`preprocess`) and the lattice basis is refined from the + :meth:`choose_lattice_vector` seed ``(self.g1, self.g2)``. The fitted vectors are + written to :attr:`g1_array`/:attr:`g2_array` (shape ``(scan_row, scan_col, 2)``, + row/col components) for :meth:`calculate_strain_map`; the full fit records + (position, amplitude, width, background) are kept in :attr:`g1_peak_fit`/ + :attr:`g2_peak_fit` (shape ``(scan_row, scan_col, 5)``). + + With ``batch_fit=True`` the transforms and peak refinements are batched across + scan positions on ``device`` (the analogue of the correlation pipeline's + :func:`~quantem.diffraction.disk_detection.detect_disks_batch`); with + ``batch_fit=False`` each position is processed in turn with numpy/scipy. Every + ``method`` / ``refine_gaussian`` / ``refine_dft`` / ``refine_all_peaks`` + combination runs on both paths. ``cft``, the parabolic estimate and DFT + upsampling agree between the paths to rounding error. For ``method="gaussian"`` + the batched path uses a vectorized Levenberg-Marquardt solve (run on the CPU in + float64, where it is faster than on a GPU) in place of ``scipy.optimize.curve_fit``; + the two agree to ~1e-4 px on well-defined peaks, but can land in different local + minima on weak or non-Gaussian peaks, or with ``refine_radius_px`` > ~2.5. + + Parameters + ---------- + refine_dft : bool, default=False + If ``True``, after the Gaussian (or parabolic) estimate, refine the position + by DFT upsampling: |FT| of the windowed, scaled pattern is evaluated on a grid + of spacing ``1/upsample`` around the estimate. Ignored when ``method="cft"``. + Requires ``preprocess(hann_window=True)``. + refine_radius_px : float, default=2.0 + For ``method="gaussian"``: half-width of the Gaussian fit window (use ~1.5; + larger windows bias the fit). For ``method="cft"``: maximum distance the peak + may move from its starting pixel. + device : str, default="cpu" + Torch device for the batched path (e.g. ``"cuda:0"``). Ignored when + ``batch_fit=False``. + batch_size : int, optional + Number of scan positions per batch when ``batch_fit=True``. If ``None``, it is + chosen from the detector size. + batch_fit : bool, default=True + If ``True``, use the batched torch path on ``device``; if ``False``, process + positions one at a time with numpy/scipy (slow; mainly for checking). + save_to_gpu : bool, default=True + If ``True`` and ``device`` is a GPU, copy the whole dataset to the GPU once + (cached for later calls). Falls back to per-batch transfer if it does not fit. + method : {"gaussian", "cft"}, default="gaussian" + Sub-pixel peak refinement. ``"gaussian"`` fits an isotropic Gaussian to the + |FFT| pixels (optionally followed by ``refine_dft``); ``"cft"`` maximises the + exact continuous Fourier transform of the windowed pattern, as in PC-STEM, + and ignores ``refine_gaussian`` and ``refine_dft``. Use ``"cft"`` or + ``refine_dft`` with ``preprocess(hann_window=True)``. + + Returns + ------- + StrainMapAutocorrelation + ``self``, with :attr:`g1_array`, :attr:`g2_array`, :attr:`g1_peak_fit`, and + :attr:`g2_peak_fit` set. + """ + if self.g1 is None or self.g2 is None: + raise ValueError("Run choose_lattice_vector() first to set initial lattice vectors (self.g1, self.g2).") + if refine_all_peaks: + if self.mean_img_peaks is None or self.mean_img_weights is None: + raise ValueError("Run choose_lattice_vector() with refine_all_peaks=True to determine which peaks to fit") + if method not in ("gaussian", "cft"): + raise ValueError("method must be 'gaussian' or 'cft'.") + hann = bool(self.metadata.get("hann_window", False)) + if method == "gaussian" and refine_gaussian and refine_radius_px > 2.5: + warnings.warn("method='gaussian' with refine_radius_px > 2.5 is biased (0.8-1.5 px on the SiGe " + "sims) and, without a window, has multiple local minima; use ~1.5, or method='cft'.", + UserWarning, stacklevel=2) + if (method == "cft" or refine_dft) and not hann: + warnings.warn("method='cft' / refine_dft without preprocess(hann_window=True) is pulled by " + "central-peak leakage; enable the Hann window.", UserWarning, stacklevel=2) + if method == "gaussian" and refine_gaussian and not refine_dft and hann: + warnings.warn("hann_window=True with a plain Gaussian fit biases positions (up to ~0.5 px); " + "use method='cft' or refine_dft=True.", UserWarning, stacklevel=2) + + scan_r = self.dataset.shape[0] + scan_c = self.dataset.shape[1] + + self.g1_peak_fit = Dataset3d.from_shape( + (scan_r, scan_c, 5), + name="g1_peak_fit", + signal_units="mixed", + ) + self.g2_peak_fit = Dataset3d.from_shape( + (scan_r, scan_c, 5), + name="g2_peak_fit", + signal_units="mixed", + ) + + self.g1_array = np.zeros((scan_r, scan_c, 2)) + self.g2_array = np.zeros((scan_r, scan_c, 2)) + + g1_0 = np.asarray(self.g1, dtype=float).reshape(2) + g2_0 = np.asarray(self.g2, dtype=float).reshape(2) + + if batch_fit: + # Fast path: batch the transforms + self._fit_lattice_vectors_batched( + g1_0=g1_0, + g2_0=g2_0, + refine_gaussian=refine_gaussian, + refine_radius_px=refine_radius_px, + refine_dft=refine_dft, + upsample=upsample, + device=device, + progressbar=progressbar, + refine_all_peaks=refine_all_peaks, + peaks=self.mean_img_peaks if refine_all_peaks else None, + weights=self.mean_img_weights if refine_all_peaks else None, + batch_size = batch_size, + save_to_gpu = save_to_gpu, + method = method, + ) + else: + mode = self.metadata.get("mode", "linear").lower() + if mode == "gamma": + g = self.metadata["gamma"] + + it = np.ndindex(scan_r, scan_c) + if progressbar: + try: + from tqdm.auto import tqdm # type: ignore + + it = tqdm(it, total=scan_r * scan_c, desc="fit_lattice_vectors", leave=True) + except Exception: + pass + + for r, c in it: + dp = self.dataset.array[r, c] * self.mask_diffraction + self.mask_diffraction_inv + + if mode == "linear": + dp_proc = (dp) + elif mode == "log": + dp_proc = (np.log1p(dp)) + elif mode == "log-min": + dp_min = self.metadata.get("dp_min", 0.0) + dp_proc = (np.log(dp - dp_min + self.metadata.get("log_offset", 0.1))) + elif mode == "gamma": + dp_proc = (np.power(np.clip(dp, 0.0, None), g)) + else: + raise ValueError("metadata['mode'] must be 'linear', 'log', 'log-min', or 'gamma'") + im = np.fft.fftshift(np.abs(np.fft.fft2(dp_proc * self.fft_window))) + tr_np = dp_proc * self.fft_window + + g1_fit_abs, g2_fit_abs, _, _ = _refine_lattice_vectors( + im, + g1_rc=g1_0, + g2_rc=g2_0, + radius_px=refine_radius_px, + refine_gaussian=refine_gaussian, + refine_dft=refine_dft, + refine_all_peaks=refine_all_peaks, + peaks=self.mean_img_peaks, + weights=self.mean_img_weights, + upsample=upsample, + maxfev=gaussian_maxfev, + tr=tr_np, + method=method, + ) + + self.g1_peak_fit.array[r, c, :] = g1_fit_abs + self.g2_peak_fit.array[r, c, :] = g2_fit_abs + + self.g1_array[r, c, 0] = g1_fit_abs[0] + self.g1_array[r, c, 1] = g1_fit_abs[1] + self.g2_array[r, c, 0] = g2_fit_abs[0] + self.g2_array[r, c, 1] = g2_fit_abs[1] + + self.metadata["fit_refine_gaussian"] = refine_gaussian + self.metadata["fit_refine_dft"] = refine_dft + self.metadata["fit_refine_all_peaks"] = refine_all_peaks + self.metadata["fit_refine_radius_px"] = refine_radius_px + self.metadata["fit_upsample"] = upsample + self.metadata["fit_gaussian_maxfev"] = gaussian_maxfev + self.metadata["fit_method"] = method + + # Populate the per-position weight directly from the fitted peak amplitudes, + # mirroring the correlation pipeline where BraggVectors.fit_lattice emits + # mask_weight. create_mask() is therefore optional -- call it only to plot the + # weight map, recompute, or switch to the radial estimator. + self.mask_weight = self._amplitude_mask_weight() + + return self + + def _fit_lattice_vectors_batched( + self, + *, + g1_0: NDArray, + g2_0: NDArray, + refine_gaussian: bool, + refine_radius_px: float, + refine_dft: bool = False, + upsample: int = 16, + device: str, + progressbar: bool, + refine_all_peaks: bool = False, + peaks: NDArray | None = None, + weights: NDArray | None = None, + batch_size: int | None = None, + method: str = "gaussian", + save_to_gpu: bool = True, + ) -> None: + """Batched torch implementation of :meth:`fit_lattice_vectors` (``batch_fit=True``). + + The transform (mask/fill -> ``mode`` -> window -> ``|FFT|`` -> fftshift) and the + peak refinement are evaluated for a stack of scan positions at once on + ``device``. Refinement follows ``method``: ``"cft"`` uses + :func:`_refine_peaks_cft_batched`; ``"gaussian"`` uses + :func:`_refine_peaks_batched` (parabolic, then an isotropic-Gaussian + Levenberg-Marquardt fit when ``refine_gaussian``), optionally followed by + :func:`_refine_peaks_dft_batched` when ``refine_dft``. Results are written into + :attr:`g1_array`/:attr:`g2_array` and :attr:`g1_peak_fit`/:attr:`g2_peak_fit`, + and match the per-position path (see :meth:`fit_lattice_vectors`). + + With ``refine_all_peaks=False`` only the two seed vectors ``g1_0``/``g2_0`` are + refined per position. With ``refine_all_peaks=True`` every peak in ``peaks`` (the + basis detected by :meth:`choose_lattice_vector`, weighted by ``weights``) is + batch-refined across the stack and the basis is recovered per position by the + same intensity-weighted least-squares fit as the per-position all-peaks branch -- + so that path is batched too rather than looping ``scipy.optimize.curve_fit`` over + ``n_positions x n_peaks``. In both cases the amplitude/width/background stored in + :attr:`g1_peak_fit`/:attr:`g2_peak_fit` (columns 2-4, the source of the mask weight) + come from the single-peak refinement at ``g1_0``/``g2_0`` -- the background-subtracted + Gaussian height. The all-peaks least squares only improves the g1/g2 *positions* + (columns 0-1); the raw cepstral value at those positions rides the + central-autocorrelation pedestal and would invert the mask (bright in vacuum), so + it is not used for the weight. + """ + scan_r, scan_c, H, W = self.dataset.array.shape + n_pos = scan_r * scan_c + + mode = self.metadata.get("mode", "linear").lower() + if mode not in ("linear", "log", 'log-min', "gamma"): + raise ValueError("metadata['mode'] must be 'linear', 'log', 'log-min', or 'gamma'") + gamma = float(self.metadata["gamma"]) if mode == "gamma" else None + + dev = torch.device(device) + dt = torch.float32 if (method == "cft" and dev.type != "cpu") else torch.float64 + mask = torch.as_tensor(self.mask_diffraction, dtype=dt, device=dev) + mask_inv = torch.as_tensor(self.mask_diffraction_inv, dtype=dt, device=dev) + use_cache = False + if save_to_gpu and str(dev) != "cpu": + if not hasattr(self, '_gpu_cache') or self._gpu_cache is None: + try: + print(f"Loading dataset to {dev}", end = " ", flush = True) + self._gpu_cache = torch.as_tensor( + np.asarray(self.dataset.array), + dtype = torch.float32, + device=dev, + ) + use_cache = True + except (RuntimeError, torch.cuda.OutOfMemoryError): + print("Out of memory, will read per-batch") + self._gpu_cache = None + use_cache = False + else: + # self._gpu_cache = None + use_cache = True + + if refine_all_peaks: + # Precompute the fixed pieces of the per-position all-peaks fit (these do not + # change across scan positions): the detected peak seeds, their normalized + # weights, and the seed basis used to assign integer (h, k) indices. + peaks_arr = np.asarray(peaks, dtype=float).reshape(-1, 2) + w = np.asarray(weights, dtype=float).reshape(-1) + wsum = float(w.sum()) + w_norm = w / wsum if wsum != 0 else np.full(w.shape, 1.0 / max(1, w.size)) + sqrt_w = np.sqrt(w_norm)[:, None] + A_seed = np.column_stack((g1_0, g2_0)) # (2, 2) + n_pk = peaks_arr.shape[0] + + # Match the correlation chunking heuristic (see BraggVectors._detect_positions). + if batch_size is None: + batch_size = int(min(1024, max(1, 16_000_000 // (H * W)))) + + starts = range(0, n_pos, batch_size) + if progressbar: + try: + from tqdm.auto import tqdm # type: ignore + + bar = tqdm(total=n_pos, desc="fit_lattice_vectors", leave=True) + except Exception: + bar = None + else: + bar = None + + if method not in ("gaussian", "cft"): + raise ValueError("method must be 'gaussian' or 'cft'.") + + def refine(tr: torch.Tensor, ims: torch.Tensor, vec: NDArray) -> torch.Tensor: + if method == "cft": + return _refine_peaks_cft_batched(tr, ims, vec, radius_px=refine_radius_px) + out = _refine_peaks_batched( + ims, vec, radius_px=refine_radius_px, refine_gaussian=refine_gaussian + ) + if refine_dft and upsample > 1: + out[:, :2] = _refine_peaks_dft_batched(tr, out[:, :2], upsample=upsample) + return out + + win_t = torch.as_tensor(self.fft_window, dtype=mask.dtype, device=dev) + + for start in starts: + stop = min(start + batch_size, n_pos) + idxs = [(idx // scan_c, idx % scan_c) for idx in range(start, stop)] + + if use_cache: + rows = [r for r, c in idxs] + cols = [c for r, c in idxs] + dps = self._gpu_cache[rows, cols] + else: + dps = torch.stack( + [ + torch.as_tensor( + np.asarray(self.dataset.array[r, c]), dtype=dt, device=dev + ) + for r, c in idxs + ], + dim=0, + ) + + dpm = dps * mask + mask_inv + if mode == "linear": + tr = dpm + elif mode == "log": + tr = torch.log1p(dpm) + elif mode == "log-min": + dp_min = self.metadata.get("dp_min", 0.0) + log_offset = float(self.metadata.get("log_offset", 0.1)) + tr = torch.log(dpm - dp_min + log_offset) + else: # gamma + tr = dpm.clamp(min=0.0).pow(gamma) + + tr = tr * win_t + ims = torch.fft.fftshift(torch.fft.fft2(tr).abs(), dim=(-2, -1)) + + if refine_all_peaks: + # Batch-refine each detected peak across the stack (one vectorized LM + # solve per peak), then recover the basis per position with the same + # weighted least squares as _refine_lattice_vectors' all-peaks branch. + pts_all = np.empty((len(idxs), n_pk, 2), dtype=float) + for j in range(n_pk): + rj = refine(tr, ims, peaks_arr[j]).cpu().numpy() + pts_all[:, j, :] = rj[:, :2] + # Amplitude/width/background for the mask weight come from the SAME + # single-peak refinement at the g1/g2 seeds as the single-peak branch + # below -- i.e. the background-subtracted Gaussian height. (The all-peaks + # lstsq only improves the g1/g2 *positions*; the raw cepstral value at those + # positions rides the central-autocorrelation pedestal and would invert + # the mask in vacuum.) + g1_fit = refine(tr, ims, g1_0).cpu().numpy() + g2_fit = refine(tr, ims, g2_0).cpu().numpy() + + rows_a = np.fromiter((r for r, _ in idxs), int, len(idxs)) + cols_a = np.fromiter((c for _, c in idxs), int, len(idxs)) + # integer (h, k) indices per position/peak: A_seed is 2x2, so lstsq == solve + ab = np.round(np.einsum("ij,kpj->kpi", np.linalg.inv(A_seed), pts_all)) # (B, n_pk, 2) + M = np.concatenate([ab, np.ones(ab.shape[:2] + (1,))], axis=2) # (B, n_pk, 3) + Mw = M * sqrt_w[None] # sqrt_w: (n_pk, 1) + Pw = pts_all * sqrt_w[None] + uvc = np.linalg.pinv(Mw) @ Pw # (B, 3, 2), batched SVD + g1_b, g2_b = uvc[:, 0], uvc[:, 1] # (B, 2) each + self.g1_peak_fit.array[rows_a, cols_a] = np.concatenate([g1_b, g1_fit[:, 2:]], 1) + self.g2_peak_fit.array[rows_a, cols_a] = np.concatenate([g2_b, g2_fit[:, 2:]], 1) + self.g1_array[rows_a, cols_a] = g1_b + self.g2_array[rows_a, cols_a] = g2_b + else: + g1_np = refine(tr, ims, g1_0).cpu().numpy() + g2_np = refine(tr, ims, g2_0).cpu().numpy() + rows_a = np.fromiter((r for r, _ in idxs), int, len(idxs)) + cols_a = np.fromiter((c for _, c in idxs), int, len(idxs)) + self.g1_peak_fit.array[rows_a, cols_a] = g1_np + self.g2_peak_fit.array[rows_a, cols_a] = g2_np + self.g1_array[rows_a, cols_a] = g1_np[:, :2] + self.g2_array[rows_a, cols_a] = g2_np[:, :2] + + if bar is not None: + bar.update(len(idxs)) + + if bar is not None: + bar.close() + + def _amplitude_mask_weight(self) -> np.ndarray: + """Per-position weight from the fitted g1/g2 peak amplitudes, min-max to ``[0, 1]``. + + The mean of the two fitted lattice-peak amplitudes measures the lattice signal at + each position. It is min-max normalized to ``[0, 1]`` with no contrast windowing; + display contrast is applied later via :meth:`StrainMap.plot_strain`'s + ``mask_range``. Degenerate input (constant / non-finite / amplitudes unavailable) + falls back to uniform full weight. + """ + scan_r = self.dataset.shape[0] + scan_c = self.dataset.shape[1] + if self.g1_peak_fit is None or self.g2_peak_fit is None: + return np.ones((scan_r, scan_c)) + signal = (self.g1_peak_fit.array[:, :, 2] + self.g2_peak_fit.array[:, :, 2]) / 2.0 + lo = np.nanmin(signal) + hi = np.nanmax(signal) + if np.isfinite(lo) and np.isfinite(hi) and hi > lo: + return (signal - lo) / (hi - lo) + return np.ones((scan_r, scan_c)) + + def create_mask( + self, + use_radial_method: bool = False, + exclusion_radius_fraction: float = 0.1, + plot: bool = True, + figsize: tuple[float, float] = (5, 4), + ): + """(Re)compute the per-position weight :attr:`mask_weight` from lattice signal. + + Usually **optional**: :meth:`fit_lattice_vectors` already populates + :attr:`mask_weight` from the fitted peak amplitudes (mirroring the correlation + pipeline, where ``BraggVectors.fit_lattice`` emits ``mask_weight``). Call this + only to plot the weight map, to recompute it, or to switch to the radial + estimator. + + Builds a ``(scan_row, scan_col)`` weight in ``[0, 1]`` measuring the lattice + signal at each position, the analogue of :attr:`BraggVectors.mask_weight`. It is + the default reference weighting handed to :meth:`calculate_strain_map`, so strong, + well-fit positions dominate the reference lattice and weak/vacuum positions are + down-weighted. + + The raw signal is min-max normalized to ``[0, 1]`` with **no** contrast windowing + or smoothing. Display contrast is applied via :meth:`StrainMap.plot_strain`'s + ``mask_range`` argument (e.g. ``mask_range=(0.6, 0.8)``) -- weights at/below + ``low`` render black, at/above ``high`` render full color. (Min-max is sensitive + to a few very bright positions, which compress the bulk toward 0; pick + ``mask_range`` to match where the bulk actually sits -- the weight-map plot here + shows it.) + + Parameters + ---------- + use_radial_method : bool, default=False + If ``True``, weight by the total diffracted intensity *outside* a central + disk (radius ``exclusion_radius_fraction`` of the detector width), excluding + the bright unscattered beam. Use a **small** exclusion (~0.1) so the Bragg + disks are kept: a large exclusion keeps only the far-corner diffuse scatter, + which is *anti*-correlated with crystallinity and inverts the contrast. If + ``False`` (default, recommended), weight by the mean fitted peak amplitude + from :meth:`fit_lattice_vectors` (requires it to have been run), identical to + the weight ``fit_lattice_vectors`` stores automatically. + exclusion_radius_fraction : float, default=0.1 + Central-disk radius (fraction of detector width) excluded by the radial + method. + plot : bool, default=True + If ``True``, show the resulting per-position weight map (the analogue of the + correlation pipeline's :meth:`BraggVectors.fit_lattice` plot). + figsize : tuple of float, default=(5, 4) + Figure size in inches for the weight-map plot. + + Returns + ------- + StrainMapAutocorrelation + ``self``, with :attr:`mask_weight` set. + """ + if not isinstance(self.dataset, (Dataset4d, Dataset4dstem)): + raise ValueError("Dataset must be Dataset4d or Dataset4dstem.") + + scan_r = self.dataset.shape[0] + scan_c = self.dataset.shape[1] + + if use_radial_method: + # center / radius are in DETECTOR coordinates (the last two axes) + center_y = self.dataset.shape[-2] / 2.0 + center_x = self.dataset.shape[-1] / 2.0 + y, x = np.ogrid[:self.dataset.shape[-2], :self.dataset.shape[-1]] + radius_map = np.sqrt((x - center_x)**2 + (y - center_y)**2) + exclusion_radius = exclusion_radius_fraction * self.dataset.shape[-1] + outside_mask = radius_map > exclusion_radius + signal = np.empty(shape=(scan_r, scan_c)) + for r in range(scan_r): + for c in range(scan_c): + dp = self.dataset.array[r, c] + signal[r, c] = np.sum(dp[outside_mask]) + # min-max normalization to [0, 1]; constant/degenerate -> ones + lo = np.nanmin(signal) + hi = np.nanmax(signal) + if np.isfinite(lo) and np.isfinite(hi) and hi > lo: + self.mask_weight = (signal - lo) / (hi - lo) + else: + self.mask_weight = np.ones((scan_r, scan_c)) + else: + if self.g1_peak_fit is None or self.g2_peak_fit is None: + raise RuntimeError("For intensity-based masking, run fit_lattice_vectors() first.") + self.mask_weight = self._amplitude_mask_weight() + + if plot: + fig, ax = plt.subplots(1, 1, figsize=figsize) + handle = ax.imshow(self.mask_weight, cmap="gray", vmin=0.0, vmax=1.0) + ax.set_title("Per-position lattice weight (mask_weight)") + ax.set_xlabel("scan column") + ax.set_ylabel("scan row") + fig.colorbar(handle, ax=ax, fraction=0.046, pad=0.04) + fig.tight_layout() + + return self + + + def calculate_strain_map( + self, + g1_ref: np.ndarray | None = None, + g2_ref: np.ndarray | None = None, + mask: np.ndarray | None = None, + calculation_metric: str = "median", + ) -> StrainMap: + """Build a :class:`StrainMap` from the fitted per-position lattice vectors. + + Mirrors :meth:`BraggVectors.calculate_strain_map` so the downstream strain cells + (``plot_strain``, ``update_reference``, ``estimate_strain_precision``) are + identical for the correlation and cepstral workflows. Because the cepstral + vectors are real-space, ``real_space=True`` is passed and + :func:`~quantem.diffraction.strain._strain_tensor` yields strain matching the + correlation result on the same data. + + Parameters + ---------- + g1_ref : np.ndarray, optional + ``(2,)`` reference for the first lattice vector. Defaults to the median over + the scan inside :class:`StrainMap`. + g2_ref : np.ndarray, optional + ``(2,)`` reference for the second lattice vector. Defaults to the median over + the scan inside :class:`StrainMap`. + mask : np.ndarray, optional + ``(scan_row, scan_col)`` per-position weighting used when computing the + reference lattice. Defaults to :attr:`mask_weight` from :meth:`create_mask` + (the lattice signal strength), so strong, well-fit positions dominate the + reference. + + Returns + ------- + StrainMap + A strain map initialized from the fitted lattice vectors. + """ + if self.g1_array is None or self.g2_array is None: + raise ValueError("Run fit_lattice_vectors() before calculate_strain_map().") + if not isinstance(self.dataset, (Dataset4d, Dataset4dstem)): + raise ValueError("Dataset must be Dataset4d or Dataset4dstem.") + + if mask is None: + mask = self.mask_weight + + ds_sampling = float(self.dataset.sampling[0]) + ds_units = str(self.dataset.units[0]) + + return StrainMap( + g1_array=self.g1_array, + g2_array=self.g2_array, + ds_shape=tuple(self.dataset.shape), + real_space=self.real_space, + g1_ref=g1_ref, + g2_ref=g2_ref, + mask=mask, + ds_sampling=ds_sampling, + ds_units=ds_units, + q_to_r_rotation_ccw_deg = self.metadata['q_to_r_rotation_ccw_deg'], + q_transpose = self.metadata['q_transpose'], + calculation_metric = calculation_metric, + ) + + def plot_lattice_vectors( + self, + subtract_mean: bool = True, + max_shift: float = 1.0, + cmap: str = "PiYG_r", + axsize: tuple[float, float] | None = None, + figsize: tuple[float, float] | None = None, + **imshow_kwargs: Any, + ): + """Plot the four per-position lattice-vector component maps. + + Shows ``g1_row``, ``g1_col``, ``g2_row``, ``g2_col`` from :attr:`g1_array`/ + :attr:`g2_array` as four panels (optionally with the mean subtracted), on a + shared symmetric color range. Positions whose vectors deviate from the + :meth:`choose_lattice_vector` seed by more than ``max_shift`` are masked out, so + bad fits do not blow up the color scale. A quick diagnostic of fit smoothness + before computing strain. + + Parameters + ---------- + subtract_mean : bool, default=True + If ``True``, subtract the (in-range) mean of each component so deviations + from the average lattice are shown. + max_shift : float, default=1.0 + Maximum allowed deviation (pixels) of a vector from the seed before its + position is masked out of the plot and the color-range statistics. + cmap : str, default="PiYG_r" + Diverging colormap; masked positions are drawn black. + axsize : tuple of float, optional + Per-panel size in inches; defaults to ``(4, 4)`` when ``figsize`` is unset. + figsize : tuple of float, optional + Overall figure size in inches; defaults to ``(4 * axsize[0], axsize[1])``. + **imshow_kwargs + Forwarded to :meth:`matplotlib.axes.Axes.imshow`. + + Returns + ------- + tuple + ``(fig, ax)`` with the four-panel figure. + """ + if self.g1_array is None or self.g2_array is None: + raise ValueError("Run fit_lattice_vectors() first to compute g1_array and g2_array.") + if self.g1 is None or self.g2 is None: + raise ValueError("Run choose_lattice_vector() first to set self.g1 and self.g2.") + + im0 = self.g1_array[:, :, 0] + im1 = self.g1_array[:, :, 1] + im2 = self.g2_array[:, :, 0] + im3 = self.g2_array[:, :, 1] + + dg1_0 = im0 - self.g1[0] + dg1_1 = im1 - self.g1[1] + dg2_0 = im2 - self.g2[0] + dg2_1 = im3 - self.g2[1] + + max_shift2 = max_shift * max_shift + mu = (dg1_0 * dg1_0 + dg1_1 * dg1_1) <= max_shift2 + mv = (dg2_0 * dg2_0 + dg2_1 * dg2_1) <= max_shift2 + + if subtract_mean: + if np.any(mu): + im0 = im0 - np.mean(im0[mu]) + im1 = im1 - np.mean(im1[mu]) + else: + im0 = im0 - np.mean(im0) + im1 = im1 - np.mean(im1) + + if np.any(mv): + im2 = im2 - np.mean(im2[mv]) + im3 = im3 - np.mean(im3[mv]) + else: + im2 = im2 - np.mean(im2) + im3 = im3 - np.mean(im3) + + vals = [] + if np.any(mu): + vals.append(np.abs(im0[mu])) + vals.append(np.abs(im1[mu])) + if np.any(mv): + vals.append(np.abs(im2[mv])) + vals.append(np.abs(im3[mv])) + + if vals: + vlim = np.max(np.concatenate(vals)) + else: + vlim = np.max(np.abs(np.stack([im0, im1, im2, im3], axis=0))) + + vmin = -vlim + vmax = vlim + + cm = plt.get_cmap(cmap).copy() + cm.set_bad(color="black") + + m0 = np.ma.array(im0, mask=~mu) + m1 = np.ma.array(im1, mask=~mu) + m2 = np.ma.array(im2, mask=~mv) + m3 = np.ma.array(im3, mask=~mv) + + if axsize is None and figsize is None: + axsize = (4.0, 4.0) + if figsize is None: + figsize = (axsize[0] * 4.0, axsize[1]) + + fig, ax = plt.subplots(1, 4, figsize=figsize) + + ax[0].imshow(m0, cmap=cm, vmin=vmin, vmax=vmax, **imshow_kwargs) + ax[1].imshow(m1, cmap=cm, vmin=vmin, vmax=vmax, **imshow_kwargs) + ax[2].imshow(m2, cmap=cm, vmin=vmin, vmax=vmax, **imshow_kwargs) + ax[3].imshow(m3, cmap=cm, vmin=vmin, vmax=vmax, **imshow_kwargs) + + ax[0].set_title("g1_r") + ax[1].set_title("g1_c") + ax[2].set_title("g2_r") + ax[3].set_title("g2_c") + + for a in ax: + a.set_xticks([]) + a.set_yticks([]) + + return fig, ax + + @property + def reference_qpos(self) -> NDArray | None: + """Detected mean-transform peak positions, center-relative (alias of :attr:`mean_img_peaks`).""" + return self.mean_img_peaks + + @property + def reference_intensity(self) -> NDArray | None: + """Normalized weight of each detected peak (alias of :attr:`mean_img_weights`).""" + return self.mean_img_weights + + @property + def reference_ab(self) -> NDArray | None: + """Integer ``[a, b]`` index of each detected peak (alias of :attr:`mean_img_ab`).""" + return self.mean_img_ab + + +def _nice_length_units(target: float) -> float: + """Round ``target`` to the nearest "nice" scale-bar length (1/2/5 x 10^n).""" + if not np.isfinite(target) or target <= 0: + return 0.0 + exp = np.floor(np.log10(target)) + base = target / (10.0**exp) + if base < 1.5: + nice = 1.0 + elif base < 3.5: + nice = 2.0 + elif base < 7.5: + nice = 5.0 + else: + nice = 10.0 + return nice * (10.0**exp) + + +def _apply_center_crop_limits(ax: Any, shape: tuple[int, int], cropping_factor: float) -> None: + """Zoom ``ax`` to the central ``cropping_factor`` fraction of a ``shape`` image. + + Preserves the existing y-axis direction (inverted for image coordinates). + """ + if cropping_factor >= 1.0: + return + if not (0.0 < cropping_factor <= 1.0): + raise ValueError("cropping_factor must be in (0, 1].") + + H, W = shape + r0 = H // 2 + c0 = W // 2 + half_h = 0.5 * cropping_factor * H + half_w = 0.5 * cropping_factor * W + + ax.set_xlim(c0 - half_w, c0 + half_w) + + y0, y1 = ax.get_ylim() + if y0 > y1: + ax.set_ylim(r0 + half_h, r0 - half_h) + else: + ax.set_ylim(r0 - half_h, r0 + half_h) + + +def _flatten_axes(ax: Any) -> list[Any]: + """Flatten a matplotlib axes container (array/list/tuple) to a flat list of axes.""" + if isinstance(ax, np.ndarray): + return list(ax.ravel()) + if isinstance(ax, (list, tuple)): + out: list[Any] = [] + for a in ax: + out.extend(_flatten_axes(a)) + return out + return [ax] + + +def _raw_vec_to_display(vec_rc: NDArray, *, rotation_ccw_deg: float, transpose: bool) -> NDArray: + """Map a raw-detector ``(row, col)`` vector into the rotated display frame. + + Applies the optional axis transpose, then a counter-clockwise rotation of + ``rotation_ccw_deg``. Inverse of :func:`_display_vec_to_raw`. + """ + v = np.asarray(vec_rc, dtype=float).reshape(2) + dr, dc = v[0], v[1] + + if transpose: + dr, dc = dc, dr + + theta = np.deg2rad(rotation_ccw_deg) + ct = np.cos(theta) + st = np.sin(theta) + + dr2 = ct * dr - st * dc + dc2 = st * dr + ct * dc + return np.array((dr2, dc2), dtype=float) + + +def _display_vec_to_raw(vec_rc: NDArray, *, rotation_ccw_deg: float, transpose: bool) -> NDArray: + """Map a rotated-display ``(row, col)`` vector back to the raw detector frame. + + Inverse of :func:`_raw_vec_to_display`: undo the rotation, then the transpose. + """ + v = np.asarray(vec_rc, dtype=float).reshape(2) + dr, dc = v[0], v[1] + + theta = np.deg2rad(rotation_ccw_deg) + ct = np.cos(theta) + st = np.sin(theta) + + dr2 = ct * dr + st * dc + dc2 = -st * dr + ct * dc + + if transpose: + dr2, dc2 = dc2, dr2 + + return np.array((dr2, dc2), dtype=float) + + +def _plot_lattice_vectors(ax: Any, center_rc: tuple[float, float], g1_rc: NDArray, g2_rc: NDArray) -> None: + """Draw the ``g1`` (red) and ``g2`` (cyan) lattice vectors from ``center_rc`` on ``ax``.""" + r0, c0 = center_rc + + def _draw(vec: NDArray, label: str, color: tuple[float, float, float]) -> None: + dr, dc = vec[0], vec[1] + ax.plot([c0, c0 + dc], [r0, r0 + dr], linewidth=2.75, color=color) + ax.plot([c0 + dc], [r0 + dr], marker="o", markersize=6.0, color=color) + ax.text(c0 + dc, r0 + dr, f" {label}", color=color, fontsize=18, va="center") + + _draw(np.asarray(g1_rc, dtype=float).reshape(2), "g1", (1.0, 0.0, 0.0)) + _draw(np.asarray(g2_rc, dtype=float).reshape(2), "g2", (0.0, 0.7, 1.0)) + +def _plot_peaks(ax: Any, center_rc: tuple[float, float], peaks_plot: NDArray) -> None: + """Mark each detected peak (green dot), offset from ``center_rc``, on ``ax``.""" + r0, c0 = center_rc + + def _draw(vec: NDArray, color: tuple[float, float, float]) -> None: + dr, dc = vec[0], vec[1] + ax.plot([c0 + dc], [r0 + dr], marker="o", markersize=6.0, color=color) + + for p in peaks_plot: + _draw(np.asarray(p, dtype=float).reshape(2), (0.0, 1.0, 0.0)) + +def _overlay_lattice_vectors( + *, + ax: Any, + shape: tuple[int, int], + g1_rc: NDArray, + g2_rc: NDArray, + rot_ccw_deg: float, + q_transpose: bool, + peaks_plot: NDArray | None = None, +) -> None: + """Overlay lattice vectors on the original (and, if present, rotated) transform axes. + + Draws ``g1``/``g2`` (and any ``peaks_plot``) on the first axis in raw coordinates, and + on the second axis (if any) in the rotated display frame. + """ + axs = _flatten_axes(ax) + if not axs: + return + + H, W = shape + center_rc = (H // 2, W // 2) + + _plot_lattice_vectors(axs[0], center_rc, g1_rc, g2_rc) + if peaks_plot is not None: + _plot_peaks(axs[0], center_rc, peaks_plot) + + if len(axs) >= 2: + g1_disp = _raw_vec_to_display(g1_rc, rotation_ccw_deg=rot_ccw_deg, transpose=q_transpose) + g2_disp = _raw_vec_to_display(g2_rc, rotation_ccw_deg=rot_ccw_deg, transpose=q_transpose) + _plot_lattice_vectors(axs[1], center_rc, g1_disp, g2_disp) + + +def _parabolic_vertex_delta(v_m1: float, v_0: float, v_p1: float) -> float: + """Sub-pixel vertex offset (in ``[-1, 1]``) of a parabola through three samples. + + Given values at offsets ``-1, 0, +1``, returns the offset of the parabola's extremum + from the center sample; ``0`` for a degenerate (flat/non-finite) fit. + """ + denom = v_m1 - 2.0 * v_0 + v_p1 + if denom == 0 or not np.isfinite(denom): + return 0.0 + delta = 0.5 * (v_m1 - v_p1) / denom + if not np.isfinite(delta): + return 0.0 + return np.clip(delta, -1.0, 1.0) + + +def _parabolic_peak_rc_amp(im: NDArray, r_guess: float, c_guess: float) -> tuple[float, float, float]: + """3-point parabolic sub-pixel peak near ``(r_guess, c_guess)`` and its amplitude. + + Snaps to the brightest pixel in the 3x3 window around the (absolute-coordinate) + guess, refines row and column independently by a parabolic vertex, and returns + ``(r_sub, c_sub, amp)`` -- the sub-pixel peak position and the image value at the + rounded vertex. Shared by the per-position :func:`_refine_lattice_vectors` and the + batched :meth:`StrainMapAutocorrelation._fit_lattice_vectors_batched` so they agree + exactly. + """ + H, W = im.shape + r0 = int(np.clip(int(np.round(r_guess)), 0, H - 1)) + c0 = int(np.clip(int(np.round(c_guess)), 0, W - 1)) + win = im[max(0, r0 - 1) : min(H, r0 + 2), max(0, c0 - 1) : min(W, c0 + 2)] + if win.size == 0: + return r_guess, c_guess, 0.0 + + ir, ic = np.unravel_index(np.argmax(win), win.shape) + r_peak = max(0, r0 - 1) + ir + c_peak = max(0, c0 - 1) + ic + + if 0 < r_peak < H - 1: + col = im[r_peak - 1 : r_peak + 2, c_peak] + dr = _parabolic_vertex_delta(col[0], col[1], col[2]) + else: + dr = 0.0 + + if 0 < c_peak < W - 1: + row = im[r_peak, c_peak - 1 : c_peak + 2] + dc = _parabolic_vertex_delta(row[0], row[1], row[2]) + else: + dc = 0.0 + + r_sub = r_peak + dr + c_sub = c_peak + dc + r_int = int(np.clip(int(np.round(r_sub)), 0, H - 1)) + c_int = int(np.clip(int(np.round(c_sub)), 0, W - 1)) + return r_sub, c_sub, float(im[r_int, c_int]) + + +def _refine_peak_subpixel( + im: NDArray, + *, + r_guess: float, + c_guess: float, + radius_px: float = 2.0, +) -> tuple[float, float]: + """Refine a peak near ``(r_guess, c_guess)`` to sub-pixel ``(row, col)``. + + Finds the brightest pixel in a ``radius_px`` window, then applies independent + parabolic vertex offsets along each axis. + """ + im = np.asarray(im, dtype=float) + H, W = im.shape + + r0 = int(np.clip(int(np.round(r_guess)), 0, H - 1)) + c0 = int(np.clip(int(np.round(c_guess)), 0, W - 1)) + rad = int(max(0, int(np.ceil(radius_px)))) + + r1 = max(0, r0 - rad) + r2 = min(H, r0 + rad + 1) + c1 = max(0, c0 - rad) + c2 = min(W, c0 + rad + 1) + + win = im[r1:r2, c1:c2] + if win.size == 0: + return r_guess, c_guess + + ir, ic = np.unravel_index(np.argmax(win), win.shape) + r_peak = r1 + ir + c_peak = c1 + ic + + if 0 < r_peak < H - 1: + col = im[r_peak - 1 : r_peak + 2, c_peak] + dr = _parabolic_vertex_delta(col[0], col[1], col[2]) + else: + dr = 0.0 + + if 0 < c_peak < W - 1: + row = im[r_peak, c_peak - 1 : c_peak + 2] + dc = _parabolic_vertex_delta(row[0], row[1], row[2]) + else: + dc = 0.0 + + return r_peak + dr, c_peak + dc + + +def _refine_peak_cft(tr, im, vec, *, radius_px, n_grid=11, n_zoom=4) -> NDArray: + """Single-image version of :func:`_refine_peaks_cft_batched` (numpy). + + Starting from the brightest pixel in the 3x3 window around ``center + vec`` in + ``im``, maximises |FT(tr)| at continuous frequencies on a zooming grid + (``n_grid`` points per axis, ``n_zoom`` levels: step 0.2 -> 0.04 -> 0.008 -> + 0.0016 px for the defaults), constrained to ``+-radius_px`` of the starting pixel. + + Parameters + ---------- + tr : NDArray + Windowed, intensity-scaled pattern ``(H, W)`` whose FFT magnitude is ``im``. + im : NDArray + fftshifted ``|FFT(tr)|``, used for the starting pixel and the background. + vec : NDArray + ``(row, col)`` offset of the peak from the panel center. + radius_px : float + Maximum distance the refined peak may move from the starting pixel. + n_grid, n_zoom : int + Grid points per axis and number of zoom levels. + + Returns + ------- + NDArray + ``(row_off, col_off, amp, 0, bg)``. ``bg`` is the median of the border of the + ``(2*ceil(radius_px)+1)^2`` pixel window in ``im``, and ``amp = |FT| - bg``. + """ + H, W = tr.shape + rc, cc = H // 2, W // 2 + r0 = int(np.clip(round(rc + float(vec[0])), 1, H - 2)) + c0 = int(np.clip(round(cc + float(vec[1])), 1, W - 2)) + ir, ic = np.unravel_index(np.argmax(im[r0 - 1 : r0 + 2, c0 - 1 : c0 + 2]), (3, 3)) + q_r0, q_c0 = r0 - 1 + ir - rc, c0 - 1 + ic - cc + q_r, q_c = float(q_r0), float(q_c0) + nr, nc = np.arange(H), np.arange(W) + t = np.linspace(-1.0, 1.0, n_grid) + half = 1.0 + for _ in range(n_zoom): + kr, kc = q_r + half * t, q_c + half * t + P = np.abs(np.exp((-2j * np.pi / H) * np.outer(kr, nr)) @ tr + @ np.exp((-2j * np.pi / W) * np.outer(nc, kc))) + i, j = np.unravel_index(np.argmax(P), P.shape) + q_r = float(np.clip(kr[i], q_r0 - radius_px, q_r0 + radius_px)) + q_c = float(np.clip(kc[j], q_c0 - radius_px, q_c0 + radius_px)) + half *= 2.0 / (n_grid - 1) + peak = float(P.max()) + rad = int(max(1, np.ceil(radius_px))) + rr = np.clip(np.arange(q_r0 + rc - rad, q_r0 + rc + rad + 1), 0, H - 1) + cw = np.clip(np.arange(q_c0 + cc - rad, q_c0 + cc + rad + 1), 0, W - 1) + win = im[rr[:, None], cw[None, :]] + bg = float(np.median(np.concatenate([win[0], win[-1], win[1:-1, 0], win[1:-1, -1]]))) + return np.array((q_r, q_c, peak - bg, 0.0, bg), dtype=float) + + +def _refine_peak_dft(tr, rc0, *, upsample) -> tuple[float, float]: + """Single-image version of :func:`_refine_peaks_dft_batched` (numpy). + + Evaluates |FT(tr)| on a ``(2*du+1)^2`` grid of spacing ``1/upsample`` + (``du = ceil(1.5*upsample)``, i.e. +-1.5 px) centred on ``rc0``, and refines the + maximum with a 3-point parabola. + + Parameters + ---------- + tr : NDArray + Windowed, intensity-scaled pattern ``(H, W)``. + rc0 : tuple of float + Starting ``(row, col)`` offset from the panel center (e.g. the Gaussian fit). + upsample : int + Upsampling factor; the grid spacing is ``1/upsample`` px. + + Returns + ------- + tuple of float + Refined ``(row, col)`` offset from the panel center. + """ + H, W = tr.shape + up = int(upsample) + du = int(np.ceil(1.5 * up)) + t = np.arange(-du, du + 1) / up + kr, kc = rc0[0] + t, rc0[1] + t + P = np.abs(np.exp((-2j * np.pi / H) * np.outer(kr, np.arange(H))) @ tr + @ np.exp((-2j * np.pi / W) * np.outer(np.arange(W), kc))) + G = P.shape[0] + i, j = np.unravel_index(np.argmax(P), P.shape) + i, j = int(np.clip(i, 1, G - 2)), int(np.clip(j, 1, G - 2)) + di = _parabolic_vertex_delta(P[i - 1, j], P[i, j], P[i + 1, j]) + dj = _parabolic_vertex_delta(P[i, j - 1], P[i, j], P[i, j + 1]) + return float(kr[i] + di / up), float(kc[j] + dj / up) + + +def _refine_lattice_vectors( + im: NDArray, + *, + g1_rc: NDArray, + g2_rc: NDArray, + radius_px: float = 2.0, + refine_gaussian: bool = True, + refine_dft: bool = False, + refine_all_peaks: bool = False, + peaks: NDArray | None = None, + weights: NDArray | None = None, + upsample: int = 16, + maxfev: int = 1000, + threshold_percentile: float = 0.9975, + min_peak_spacing: float = 0, + tr: NDArray | None = None, + method: str = "gaussian", + warn_on_bound: bool = False, +) -> tuple[NDArray, NDArray, NDArray, NDArray, NDArray]: + """Refine the two lattice vectors of a transform panel to subpixel precision. + + Starting from integer-pixel guesses for the ``g1`` and ``g2`` lattice vectors + (expressed as row/column offsets from the panel center), this locates the + corresponding autocorrelation/cepstral peaks and refines them with up to + three successively finer stages: a 3-point parabolic vertex estimate, an + isotropic 2D Gaussian least-squares fit, and DFT upsampling. When + ``refine_all_peaks`` is set, all bright peaks are detected and the lattice + basis is recovered by an intensity-weighted least-squares fit to the full + peak lattice instead of refining only the two seed vectors. + + Parameters + ---------- + im : NDArray + 2D real transform panel (Patterson/cepstral image) to fit. + g1_rc, g2_rc : NDArray + Length-2 initial lattice vectors as ``(row, column)`` offsets relative to + the panel center. + radius_px : float, optional + Half-width (in pixels) of the fitting window used for the Gaussian fit. + Default 2.0. + refine_gaussian : bool, optional + If True (default), refine each peak with an isotropic 2D Gaussian fit. + refine_dft : bool, optional + If True, follow the Gaussian fit with DFT-upsampled refinement + (requires ``upsample > 1``). Default False. + refine_all_peaks : bool, optional + If True, detect every bright peak and solve a weighted least-squares fit + for the lattice basis rather than refining only ``g1_rc`` and ``g2_rc``. + Default False. + peaks : NDArray or None, optional + Precomputed peak positions (row/col offsets from center) to use when + ``refine_all_peaks`` is set; if None they are detected automatically. + weights : NDArray or None, optional + Precomputed peak weights paired with ``peaks``; if None they are derived + from peak amplitudes. + upsample : int, optional + DFT upsampling factor for ``refine_dft``. Default 16. + maxfev : int, optional + Maximum function evaluations for the Gaussian ``curve_fit``. Default 100. + threshold_percentile : float, optional + Fractional intensity percentile (0-1) used as the peak-detection + threshold when auto-detecting peaks. Default 0.9975. + min_peak_spacing : float, optional + Minimum allowed spacing (in pixels) between detected peaks; ``0`` (default) + disables the spacing filter. + warn_on_bound : bool, default=False. + Warn if a Gaussian peak fit ends on its window bound (the peak is likely outside the window). + + Returns + ------- + tuple[NDArray, NDArray, NDArray, NDArray] + ``(g1_result, g2_result, pts, weights)``. ``g1_result`` and ``g2_result`` are + length-5 arrays ``(row_offset, col_offset, amplitude, sigma, background)`` + giving the refined lattice vectors relative to the panel center. ``pts`` + and ``weights`` are the detected peak positions and their normalized + weights when ``refine_all_peaks`` is True, otherwise both are ``None``. + When ``refine_all_peaks`` is True the position comes from the + intensity-weighted all-peaks least squares, but the amplitude/sigma/background + come from the single-peak refinement at the g1/g2 seed -- the + background-subtracted Gaussian height, not the raw cepstral value, which rides + the central pedestal and inverts the mask in vacuum. + """ + from scipy.optimize import curve_fit + + im = np.asarray(im, dtype=float) + if im.ndim != 2: + raise ValueError("im must be 2D.") + + H, W = im.shape + r_center = H // 2 + c_center = W // 2 + + def _fit_gaussian_isotropic( + *, + r0: float, + c0: float, + radius_px: float, + maxfev: int, + ) -> tuple[float, float, float, float, float]: + rad = int(max(1, int(np.ceil(radius_px)))) + r0i = int(np.clip(int(np.round(r0)), 0, H - 1)) + c0i = int(np.clip(int(np.round(c0)), 0, W - 1)) + + r1 = max(0, r0i - rad) + r2 = min(H, r0i + rad + 1) + c1 = max(0, c0i - rad) + c2 = min(W, c0i + rad + 1) + + win = im[r1:r2, c1:c2] + if win.size == 0: + return r0, c0, 0.0, 0.0, 0.0 + + ir, ic = np.unravel_index(np.argmax(win), win.shape) + r_peak = r1 + ir + c_peak = c1 + ic + + bg0 = np.median(win) + amp0 = win[ir, ic] - bg0 + sig0 = max(0.75, radius_px / 2.0) + + rr = np.arange(r1, r2, dtype=float)[:, None] + cc = np.arange(c1, c2, dtype=float)[None, :] + RR = np.broadcast_to(rr, win.shape) + CC = np.broadcast_to(cc, win.shape) + + def _g2( + coords: tuple[NDArray, NDArray], + row: float, + col: float, + amp: float, + sigma: float, + background: float, + ) -> NDArray: + r, c = coords + sig = np.maximum(sigma, 1e-12) + return background + amp * np.exp(-((r - row) ** 2 + (c - col) ** 2) / (2.0 * sig * sig)) + + p0 = (r_peak, c_peak, max(0.0, amp0), sig0, bg0) + + rlo = r1 - 0.5 + rhi = (r2 - 1) + 0.5 + clo = c1 - 0.5 + chi = (c2 - 1) + 0.5 + + bounds_lo = (rlo, clo, 0.0, 0.25, -np.inf) + bounds_hi = (rhi, chi, np.inf, radius_px * 4.0, np.inf) + + try: + popt, _ = curve_fit( + _g2, + (RR.ravel(), CC.ravel()), + win.ravel(), + p0=p0, + bounds=(bounds_lo, bounds_hi), + maxfev=maxfev, + ) + row, col, amp, sig, bg = popt + if not (np.isfinite(row) and np.isfinite(col) and np.isfinite(amp) and np.isfinite(sig) and np.isfinite(bg)): + return r0, c0, p0[2], 0.0, 0.0 + if warn_on_bound and min(row - rlo, rhi - row, col - clo, chi - col) < 1e-3: + warnings.warn( + f"Gaussian peak fit near offset ({r0 - r_center:.2f}, {c0 - c_center:.2f}) " + "stopped at its window bound, so the peak is likely outside the window. " + "Check the g1/g2 guess or increase refine_radius_px.", + UserWarning, stacklevel=4) + return row, col, amp, sig, bg + except Exception: + return r0, c0, p0[2], 0.0, 0.0 + + def _refine_one(vec: NDArray) -> NDArray: + vec = np.asarray(vec, dtype=float).reshape(2) + + if method == "cft": + if tr is None: + raise ValueError("method='cft' needs the windowed pattern `tr`.") + return _refine_peak_cft(tr, im, vec, radius_px=radius_px) + + r_guess = r_center + vec[0] + c_guess = c_center + vec[1] + + r_par, c_par, amp_par = _parabolic_peak_rc_amp(im, r_guess, c_guess) + + if refine_gaussian: + r_fit, c_fit, amp, sig, bg = _fit_gaussian_isotropic( + r0=r_par, + c0=c_par, + radius_px=radius_px, + maxfev=maxfev, + ) + else: + r_fit, c_fit, amp, sig, bg = r_par, c_par, amp_par, 0.0, 0.0 + + if refine_dft and upsample > 1: + if tr is None: + raise ValueError("refine_dft needs the windowed pattern `tr`.") + r_fit, c_fit = _refine_peak_dft( + tr, (r_fit - r_center, c_fit - c_center), upsample=upsample + ) + r_fit, c_fit = r_fit + r_center, c_fit + c_center + + return np.array((r_fit - r_center, c_fit - c_center, amp, sig, bg), dtype=float) + + def _find_initial_peaks_weights(initial_peaks: NDArray) -> tuple[NDArray, NDArray, NDArray]: + if initial_peaks is None: + threshold = np.percentile(im, threshold_percentile*100) + p = np.logical_and.reduce(( + im>np.roll(im,(-1,-1),axis = (0,1)), + im>np.roll(im,(-1, 0),axis = (0,1)), + im>np.roll(im,(-1, 1),axis = (0,1)), + im>np.roll(im,( 0,-1),axis = (0,1)), + im>np.roll(im,( 0, 1),axis = (0,1)), + im>np.roll(im,( 1,-1),axis = (0,1)), + im>np.roll(im,( 1, 0),axis = (0,1)), + im>np.roll(im,( 1, 1),axis = (0,1)), + im>threshold, + ) ) + initial_peaks = np.argwhere(p) + r0 = (r_center, c_center) + initial_peaks = initial_peaks - r0 + if min_peak_spacing > 0: + intensities = im[initial_peaks[:, 0] + r_center, initial_peaks[:, 1] + c_center] + sorted_indices = np.argsort(intensities)[::-1] + initial_peaks_sorted = initial_peaks[sorted_indices] + accepted_peaks = [] + + for peak in initial_peaks_sorted: + if len(accepted_peaks) == 0: + accepted_peaks.append(peak) + else: + distances = np.sqrt(np.sum((np.array(accepted_peaks) - peak)**2, axis=1)) + if np.all(distances >= min_peak_spacing): + accepted_peaks.append(peak) + + initial_peaks = np.array(accepted_peaks) + + peak_sp = np.zeros(initial_peaks.shape) + weights = np.zeros((initial_peaks.shape[0], 1)) + it = 0 + for peak in initial_peaks: + refined_peak = _refine_one(peak) + peak_sp[it, :] = refined_peak[:2] + weights[it] = refined_peak[2] + it += 1 + return peak_sp, weights + + + if refine_all_peaks: + if peaks is None or weights is None: + pts, weights = _find_initial_peaks_weights(None) + else: + pts,_ = _find_initial_peaks_weights(peaks) + + A = np.column_stack((g1_rc, g2_rc)) + ab0_float = np.linalg.lstsq(A, pts.T, rcond=None)[0] + ab0 = (np.round(ab0_float)).T.astype(int) + + weights_norm = weights/weights.sum() + A = np.ones((pts.shape[0], 3)) + A[:,:2] = ab0 + pts_weighted = pts * np.sqrt(weights_norm) + A_weighted = A * np.sqrt(weights_norm) + uvr0 = np.linalg.lstsq(A_weighted, pts_weighted, rcond=None)[0] + g1_refined = uvr0[0,:] + g2_refined = uvr0[1,:] + # Position from the weighted all-peaks lstsq; amplitude/width/background from + # the SAME single-peak refinement at the g1/g2 seeds as the single-peak return + # below -- the background-subtracted Gaussian height (the crystalline order + # parameter), not the raw cepstral value, which rides the central pedestal and + # would invert the mask in vacuum. + g1_amp_fit = _refine_one(g1_rc) + g2_amp_fit = _refine_one(g2_rc) + + return np.array((g1_refined[0], g1_refined[1], g1_amp_fit[2], g1_amp_fit[3], g1_amp_fit[4]), dtype=float), np.array((g2_refined[0], g2_refined[1], g2_amp_fit[2], g2_amp_fit[3], g2_amp_fit[4]), dtype=float), pts, weights, ab0 + + return _refine_one(g1_rc), _refine_one(g2_rc), None, None, None + + +def _refine_peaks_batched( + ims: torch.Tensor, + vec: NDArray, + *, + radius_px: float, + refine_gaussian: bool, + lm_iters: int = 50, +) -> torch.Tensor: + """Batched, vectorized version of the single-image ``_refine_one`` peak refinement. + + Refines one lattice peak (located near ``center + vec``) for every transform in the + stack ``ims`` of shape ``(B, H, W)`` at once. The parabolic sub-pixel step matches + :func:`_refine_lattice_vectors`' ``_parabolic_peak_rc_amp`` exactly; when + ``refine_gaussian`` is set, the per-position ``scipy.optimize.curve_fit`` of an + isotropic 2D Gaussian is replaced by a batched Levenberg-Marquardt solve of the + identical objective and bounds. On clean, well-isolated peaks the two agree to + ~1e-6 px; on noisy, non-Gaussian peaks (e.g. a sinc-like cepstral peak on the sloped + tail of a bright neighbor) the bounded-trf and LM optima can differ by a few 0.01 px + -- amplitude-preserving and well below strain-relevant precision. The batched solve + removes the dominant per-call overhead, so it runs far faster (and scales onto a GPU + via the tensor ``device``). + + Parameters + ---------- + ims : torch.Tensor + Stack of transformed (fftshifted ``|FFT|``) images, shape ``(B, H, W)``, float. + vec : NDArray + Lattice vector ``(row, col)`` relative to the panel center; the peak is sought + near ``(H // 2 + row, W // 2 + col)``. + radius_px : float + Half-width in pixels of the Gaussian-fit window. + refine_gaussian : bool + If ``True``, refine with the batched isotropic-Gaussian LM fit; otherwise return + the parabolic estimate only (``sigma`` and ``background`` set to 0). + lm_iters : int, default=50 + Number of Levenberg-Marquardt iterations. + + Returns + ------- + torch.Tensor + Shape ``(B, 5)``: ``(row_offset, col_offset, amplitude, sigma, background)`` with + ``row_offset``/``col_offset`` measured relative to the panel center (matching + :func:`_refine_lattice_vectors`). + """ + Bn, Hn, Wn = ims.shape + dev, dt = ims.device, ims.dtype + rcent, ccent = Hn // 2, Wn // 2 + bidx = torch.arange(Bn, device=dev) + + # Clamp the seed into the interior so the 3x3 window below is always in-bounds + # (mirrors the clamping in _parabolic_peak_rc_amp; matters only for edge peaks, + # which are interior in normal use but must not crash the batched solve). + r0 = int(np.clip(round(rcent + float(vec[0])), 1, Hn - 2)) + c0 = int(np.clip(round(ccent + float(vec[1])), 1, Wn - 2)) + + # --- 3x3 argmax around the seed (matches _parabolic_peak_rc_amp) --- + win3 = ims[:, r0 - 1 : r0 + 2, c0 - 1 : c0 + 2].reshape(Bn, -1) + am = win3.argmax(1) + r_peak = r0 - 1 + torch.div(am, 3, rounding_mode="floor") + c_peak = c0 - 1 + (am % 3) + + def _parab(vm1: torch.Tensor, v0_: torch.Tensor, vp1: torch.Tensor) -> torch.Tensor: + denom = vm1 - 2.0 * v0_ + vp1 + delta = 0.5 * (vm1 - vp1) / denom + delta = torch.where(torch.isfinite(delta), delta, torch.zeros_like(delta)) + delta = torch.where(denom == 0, torch.zeros_like(delta), delta) + return delta.clamp(-1.0, 1.0) + + v0_ = ims[bidx, r_peak, c_peak] + dr = _parab(ims[bidx, (r_peak - 1).clamp(0, Hn - 1), c_peak], v0_, ims[bidx, (r_peak + 1).clamp(0, Hn - 1), c_peak]) + dr = torch.where((r_peak > 0) & (r_peak < Hn - 1), dr, torch.zeros_like(dr)) + dc = _parab(ims[bidx, r_peak, (c_peak - 1).clamp(0, Wn - 1)], v0_, ims[bidx, r_peak, (c_peak + 1).clamp(0, Wn - 1)]) + dc = torch.where((c_peak > 0) & (c_peak < Wn - 1), dc, torch.zeros_like(dc)) + r_sub = r_peak.to(dt) + dr + c_sub = c_peak.to(dt) + dc + + if not refine_gaussian: + ri = r_sub.round().long().clamp(0, Hn - 1) + ci = c_sub.round().long().clamp(0, Wn - 1) + amp = ims[bidx, ri, ci] + zeros = torch.zeros(Bn, device=dev, dtype=dt) + return torch.stack([r_sub - rcent, c_sub - ccent, amp, zeros, zeros], 1) + + # --- isotropic Gaussian fit over a (2*rad+1)^2 window, batched LM --- + rad = int(max(1, np.ceil(radius_px))) + P = 2 * rad + 1 + r0i = r_sub.round().long() + c0i = c_sub.round().long() + off = torch.arange(-rad, rad + 1, device=dev) + rows_idx = (r0i[:, None, None] + off[None, :, None]).expand(Bn, P, P) + cols_idx = (c0i[:, None, None] + off[None, None, :]).expand(Bn, P, P) + bexp = bidx[:, None, None].expand(Bn, P, P) + win = ims[bexp, rows_idx.clamp(0, Hn - 1), cols_idx.clamp(0, Wn - 1)] + RR = rows_idx.to(dt).reshape(Bn, -1) + CC = cols_idx.to(dt).reshape(Bn, -1) + y = win.reshape(Bn, -1) + + am2 = y.argmax(1) + r_pk2 = (r0i - rad + torch.div(am2, P, rounding_mode="floor")).to(dt) + c_pk2 = (c0i - rad + (am2 % P)).to(dt) + bg0 = y.median(1).values + amp0 = (y.max(1).values - bg0).clamp(min=0.0) + sig0 = torch.full((Bn,), max(0.75, radius_px / 2.0), device=dev, dtype=dt) + p = torch.stack([r_pk2, c_pk2, amp0, sig0, bg0], 1) + + rlo = (r0i - rad).to(dt) - 0.5 + rhi = (r0i + rad).to(dt) + 0.5 + clo = (c0i - rad).to(dt) - 0.5 + chi = (c0i + rad).to(dt) + 0.5 + sig_hi = radius_px * 4.0 + + lam = torch.full((Bn,), 1e-2, device=dev, dtype=dt) + eye5 = torch.eye(5, device=dev, dtype=dt)[None] + + def _sse(p_: torch.Tensor) -> torch.Tensor: + row, col, amp, sig, bg = [p_[:, i : i + 1] for i in range(5)] + sig = sig.clamp(min=1e-9) + E = torch.exp(-((RR - row) ** 2 + (CC - col) ** 2) / (2.0 * sig * sig)) + return ((bg + amp * E - y) ** 2).sum(1) + p_check = p[:, :2].clone() + for it in range(lm_iters): + row, col, amp, sig, bg = [p[:, i : i + 1] for i in range(5)] + sig = sig.clamp(min=1e-9) + d_row = RR - row + d_col = CC - col + E = torch.exp(-(d_row * d_row + d_col * d_col) / (2.0 * sig * sig)) + res = (bg + amp * E) - y + g_row = amp * E * d_row / (sig * sig) + g_col = amp * E * d_col / (sig * sig) + g_amp = E + g_sig = amp * E * ((d_row * d_row + d_col * d_col) / (sig**3)) + g_bg = torch.ones_like(E) + J = torch.stack([g_row, g_col, g_amp, g_sig, g_bg], 2) # (B, N, 5) + JT = J.transpose(1, 2) + JTJ = JT @ J + JTr = JT @ res[..., None] + diag = torch.diagonal(JTJ, dim1=1, dim2=2).clamp(min=1e-12) + A = JTJ + lam[:, None, None] * eye5 * diag[:, None, :] + delta = torch.linalg.solve(A, -JTr)[..., 0] + pn = p + delta + pn[:, 0] = pn[:, 0].clamp(rlo, rhi) + pn[:, 1] = pn[:, 1].clamp(clo, chi) + pn[:, 2] = pn[:, 2].clamp(min=0.0) + pn[:, 3] = pn[:, 3].clamp(0.25, sig_hi) + s_new = _sse(pn) + s_old = res.pow(2).sum(1) + better = s_new < s_old + p = torch.where(better[:, None], pn, p) + lam = torch.where(better, (lam * 0.5).clamp(min=1e-9), (lam * 3.0).clamp(max=1e6)) + if it % 5 == 4: + # stop once no position has moved by more than 1e-6 px over the last 5 iterations + if (p[:, :2] - p_check).abs().amax() < 1e-6: + break + p_check = p[:, :2].clone() + + row, col, amp, sig, bg = [p[:, i] for i in range(5)] + ok = torch.isfinite(p).all(1) + row = torch.where(ok, row, r_sub) + col = torch.where(ok, col, c_sub) + return torch.stack([row - rcent, col - ccent, amp, sig, bg], 1) + +def _refine_peaks_cft_batched( + tr: torch.Tensor, + ims: torch.Tensor, + vec: NDArray, + *, + radius_px: float, + n_grid: int = 11, + n_zoom: int = 4, +) -> torch.Tensor: + """PC-STEM-style peak refinement: maximise the exact continuous FT magnitude. + + ``tr`` is the windowed, intensity-scaled pattern stack (B, H, W) that was FFT'd to + make ``ims`` (fftshifted |FFT|). Starting from the brightest pixel in the 3x3 around + ``center + vec``, |sum tr * exp(-2 pi i q.r / N)| is evaluated on a zooming grid + (step 0.2 -> 0.04 -> 0.008 -> 0.0016 px for the defaults), constrained to + +-radius_px of the starting pixel. Returns (row_off, col_off, amp, 0, bg) like + _refine_peaks_batched, where amp is background-subtracted. + """ + Bn, H, W = tr.shape + dev, dt = tr.device, tr.dtype + rc, cc = H // 2, W // 2 + bidx = torch.arange(Bn, device=dev) + + r0 = int(np.clip(round(rc + float(vec[0])), 1, H - 2)) + c0 = int(np.clip(round(cc + float(vec[1])), 1, W - 2)) + am = ims[:, r0 - 1 : r0 + 2, c0 - 1 : c0 + 2].reshape(Bn, 9).argmax(1) + q_r0 = (r0 - 1 + torch.div(am, 3, rounding_mode="floor") - rc).to(dt) + q_c0 = (c0 - 1 + am % 3 - cc).to(dt) + q_r, q_c = q_r0.clone(), q_c0.clone() + + nr = torch.arange(H, device=dev, dtype=dt) + nc = torch.arange(W, device=dev, dtype=dt) + trc = tr.to(torch.complex128 if dt == torch.float64 else torch.complex64) + t = torch.linspace(-1.0, 1.0, n_grid, device=dev, dtype=dt) + half = 1.0 + for _ in range(n_zoom): + kr = q_r[:, None] + half * t[None, :] # (B, G) + kc = q_c[:, None] + half * t[None, :] + Er = torch.exp((-2j * np.pi / H) * kr[:, :, None] * nr[None, None, :]) # (B, G, H) + Ec = torch.exp((-2j * np.pi / W) * nc[None, :, None] * kc[:, None, :]) # (B, W, G) + P = (Er @ trc @ Ec).abs() # (B, G, G) + idx = P.reshape(Bn, -1).argmax(1) + q_r = torch.maximum(torch.minimum(kr[bidx, idx // n_grid], q_r0 + radius_px), q_r0 - radius_px) + q_c = torch.maximum(torch.minimum(kc[bidx, idx % n_grid], q_c0 + radius_px), q_c0 - radius_px) + half = half * 2.0 / (n_grid - 1) + peak = P.reshape(Bn, -1).max(1).values + + # local background = median of the border of the (2*rad+1)^2 pixel window, + # so mask_weight keeps using a background-subtracted height (see the + # pedestal comment in _fit_lattice_vectors_batched) + rad = int(max(1, np.ceil(radius_px))) + off = torch.arange(-rad, rad + 1, device=dev) + ri = (q_r0 + rc).long()[:, None, None] + off[None, :, None] + ci = (q_c0 + cc).long()[:, None, None] + off[None, None, :] + win = ims[bidx[:, None, None], ri.clamp(0, H - 1), ci.clamp(0, W - 1)] + border = torch.cat([win[:, 0, :], win[:, -1, :], win[:, 1:-1, 0], win[:, 1:-1, -1]], 1) + bg = border.quantile(0.5, dim=1) + + zeros = torch.zeros(Bn, device=dev, dtype=dt) + return torch.stack([q_r, q_c, peak - bg, zeros, bg], 1) + +def _refine_peaks_dft_batched( + tr: torch.Tensor, + rc0: torch.Tensor, + *, + upsample: int, +) -> torch.Tensor: + """Batched DFT-upsampled refinement of a peak position (Guizar-Sicairos style). + + ``tr`` is the windowed, intensity-scaled pattern stack (B, H, W); ``rc0`` (B, 2) is a + starting (row, col) offset from the panel center (e.g. the Gaussian fit). |FT(tr)| is + evaluated on a (2*du+1)^2 grid of spacing 1/upsample centred on ``rc0`` + (du = ceil(1.5*upsample)), i.e. +-1.5 px, and the maximum is refined with a 3-point + parabola. Returns refined (B, 2) offsets from the center. + """ + Bn, H, W = tr.shape + dev, dt = tr.device, tr.dtype + up = int(upsample) + du = int(np.ceil(1.5 * up)) + t = torch.arange(-du, du + 1, device=dev, dtype=dt) / up # (G,) + kr = rc0[:, 0:1].to(dt) + t[None, :] # (B, G) + kc = rc0[:, 1:2].to(dt) + t[None, :] + nr = torch.arange(H, device=dev, dtype=dt) + nc = torch.arange(W, device=dev, dtype=dt) + trc = tr.to(torch.complex128 if dt == torch.float64 else torch.complex64) + Er = torch.exp((-2j * np.pi / H) * kr[:, :, None] * nr[None, None, :]) # (B, G, H) + Ec = torch.exp((-2j * np.pi / W) * nc[None, :, None] * kc[:, None, :]) # (B, W, G) + P = (Er @ trc @ Ec).abs() # (B, G, G) + + G = P.shape[1] + bidx = torch.arange(Bn, device=dev) + idx = P.reshape(Bn, -1).argmax(1) + i = (idx // G).clamp(1, G - 2) + j = (idx % G).clamp(1, G - 2) + + def _parab(vm1, v0, vp1): + denom = vm1 - 2.0 * v0 + vp1 + d = 0.5 * (vm1 - vp1) / denom + return torch.where(torch.isfinite(d) & (denom != 0), d, torch.zeros_like(d)).clamp(-1.0, 1.0) + + di = _parab(P[bidx, i - 1, j], P[bidx, i, j], P[bidx, i + 1, j]) + dj = _parab(P[bidx, i, j - 1], P[bidx, i, j], P[bidx, i, j + 1]) + row = kr[bidx, i] + di / up + col = kc[bidx, j] + dj / up + return torch.stack([row, col], 1) \ No newline at end of file diff --git a/src/quantem/diffraction/strain_visualization.py b/src/quantem/diffraction/strain_visualization.py index f899e168..e141232d 100644 --- a/src/quantem/diffraction/strain_visualization.py +++ b/src/quantem/diffraction/strain_visualization.py @@ -151,6 +151,7 @@ def plot_strain_panels( if figsize is None: figsize = (8, 3) if is_horizontal else (6, 6) + figsize = (max(figsize[0], 4.0), max(figsize[1], 3.0 if is_horizontal else 0.9 * ncols)) if is_horizontal: fig, ax = plt.subplots(1, ncols, figsize=figsize) @@ -301,6 +302,58 @@ def _roi_compose(norm_vals, color_cm): cb_pad = 0.03 cb_min_len = 0.16 + def _fit_cbar_ticks(cbar, norm): + """Keep ticks and their labels inside the bar. + + The locator emits "nice" values that can fall outside [vmin, vmax] for any + range it does not land on exactly -- every range but the shipped defaults. + Matplotlib then draws those by extrapolation past the bar ends, where they + run off the figure on one side and collide with the neighbouring colorbar + on the other. Clip them, then pull the end labels inward so they stop + overhanging the bar. + """ + ticks = [ + float(t) for t in cbar.get_ticks() if norm.vmin - 1e-9 <= t <= norm.vmax + 1e-9 + ] + if ticks: + cbar.set_ticks(ticks) + if cb_orientation == "horizontal": + labels = [lab for lab in cbar.ax.get_xticklabels() if lab.get_text()] + if len(labels) > 1: + labels[0].set_horizontalalignment("left") + labels[-1].set_horizontalalignment("right") + else: + labels = [lab for lab in cbar.ax.get_yticklabels() if lab.get_text()] + if len(labels) > 1: + labels[0].set_verticalalignment("bottom") + labels[-1].set_verticalalignment("top") + + def _fit_panel_gap(margin=1.15): + """Widen ``hspace`` until the inter-panel gap clears the tallest title. + + Stacked panels put each title in the gap above its own panel, so that gap + must be at least as tall as the rendered title. A fixed ``hspace`` cannot do + it: the title height varies with ``figsize`` and ``rotate_title`` (0.028-0.082 + of the figure in practice), needing anywhere from 0.17 to 0.65. Solve + ``n*h + (n-1)*g*h = top - bottom`` for the ``g`` that makes ``g*h >= need``. + """ + renderer = fig.canvas.get_renderer() + heights = [ + a.title.get_window_extent(renderer=renderer) + .transformed(fig.transFigure.inverted()) + .height + for a in ax + if a.title.get_text() + ] + if not heights: + return + need = max(heights) * margin + n = len(ax) + sp = fig.subplotpars + denom = (sp.top - sp.bottom) - need * (n - 1) + if denom > 0: + fig.subplots_adjust(hspace=min(max(need * n / denom, 0.0), 1.2)) + def _finalize_layout(): # set_aspect("equal") only resizes/recenters each panel at draw time, so # get_position() before a draw returns stale boxes -- placing the colorbars @@ -345,14 +398,16 @@ def _finalize_layout(): else: # Top band for the panel titles, right band for the vertical colorbars + labels. - right = 0.55 if need_side_panel else 0.80 + title_gap = 0.15 if arrow_style == "title" else cb_pad + right = (0.55 if need_side_panel else 0.80) - title_gap fig.subplots_adjust(left=0.04, right=right, top=0.92, bottom=0.06, hspace=0.15) _finalize_layout() + _fit_panel_gap() + _finalize_layout() cb_orientation = "vertical" b0 = ax[0].get_position() b2 = ax[n_strain - 1].get_position() - title_gap = 0.15 if arrow_style == "title" else cb_pad cb_x0 = b0.x1 + title_gap strain_cb_pos = [cb_x0, b2.y0, cb_size, b0.y1 - b2.y0] @@ -374,6 +429,7 @@ def _finalize_layout(): cbar1.formatter = FuncFormatter(lambda v, _pos: f"{v:g}%") cbar1.update_ticks() cbar1.ax.tick_params(labelsize=tick_fs) + _fit_cbar_ticks(cbar1, norm_strain) if plot_rotation and rot_cb_pos is not None: cax2 = fig.add_axes(rot_cb_pos) @@ -384,6 +440,7 @@ def _finalize_layout(): cbar2.locator = MaxNLocator(nbins=2) cbar2.update_ticks() cbar2.ax.tick_params(labelsize=tick_fs) + _fit_cbar_ticks(cbar2, norm_rot) def _add_title_arrow(ax, angle_deg, gap_pt=4.0, color="black", fontsize=None): fs = fontsize if fontsize is not None else title_fs diff --git a/src/quantem/diffractive_imaging/origin_models.py b/src/quantem/diffractive_imaging/origin_models.py index cb2e2187..93f90497 100644 --- a/src/quantem/diffractive_imaging/origin_models.py +++ b/src/quantem/diffractive_imaging/origin_models.py @@ -360,9 +360,40 @@ def _estimate_detector_rotation( return rotation_curl def estimate_detector_rotation( - self, rotation_angles_deg: torch.Tensor | NDArray | None = None + self, + rotation_angles_deg: torch.Tensor | NDArray | None = None, + print_result: bool = False, + plot_result: bool = False, + **plot_kwargs, ): - """ """ + """Estimate the detector->scan rotation (and transpose) by curl minimization. + + Sweeps ``rotation_angles_deg`` and, for both the un-transposed and transposed + center-of-mass field, computes the mean absolute curl of the rotated CoM. The + physical orientation makes the CoM field curl-free, so the global minimum over + both branches selects the rotation angle and whether the detector axes are + transposed. Results are stored in :attr:`detector_rotation_deg` and + :attr:`detector_transpose`. + + Parameters + ---------- + rotation_angles_deg : torch.Tensor or NDArray, optional + Angles (degrees) to search. Defaults to ``arange(-89, 90, 1)``. + print_result : bool, default=False + If ``True``, print the minimum curl and angle of each branch and the + chosen orientation. + plot_result : bool, default=False + If ``True``, plot the curl-vs-rotation curves for both branches, marking + each branch minimum. Useful for judging whether the transpose choice is + trustworthy (a clear, deep minimum on one branch) or ambiguous. + **plot_kwargs + Forwarded to :func:`matplotlib.pyplot.subplots` when ``plot_result``. + + Returns + ------- + CenterOfMassOriginModel + ``self``. + """ if rotation_angles_deg is None: rotation_angles_deg = torch.arange(-89, 90, 1, device=self.device, dtype=torch.float) @@ -394,6 +425,31 @@ def estimate_detector_rotation( self._detector_rotation_deg = rotation_angles_deg[ind_min].item() + if print_result or plot_result: + a = rotation_angles_deg.detach().cpu().numpy() + cn = curl_no_transpose.detach().cpu().numpy() + ct = curl_transpose.detach().cpu().numpy() + i_no, i_tr = int(cn.argmin()), int(ct.argmin()) + + if print_result: + chosen = "TRANSPOSE" if self._detector_transpose else "NO transpose" + print(f"no-transpose : min curl {cn[i_no]:.4g} at {a[i_no]:+.2f} deg") + print(f"transpose : min curl {ct[i_tr]:.4g} at {a[i_tr]:+.2f} deg") + print(f"=> chosen: {chosen} at {self._detector_rotation_deg:+.2f} deg") + + if plot_result: + import matplotlib.pyplot as plt + + _fig, ax = plt.subplots(**{"figsize": (6, 4), **plot_kwargs}) + line_no = ax.plot(a, cn, label="no transpose")[0] + line_tr = ax.plot(a, ct, label="transpose")[0] + ax.plot(a[i_no], cn[i_no], "o", color=line_no.get_color()) + ax.plot(a[i_tr], ct[i_tr], "o", color=line_tr.get_color()) + ax.set_xlabel("rotation (deg)") + ax.set_ylabel("mean |curl| of CoM (lower = better)") + ax.set_title("Detector-rotation search") + ax.legend() + return self @property