From 7f2be6f4720fad629e637a63aa68ee06fff4185b Mon Sep 17 00:00:00 2001 From: mitis1 Date: Mon, 27 Jul 2026 23:19:09 -0700 Subject: [PATCH 1/7] fixing visualize, dilation plot, strain panel labels, and model fitting constraints --- src/quantem/core/fitting/base.py | 4 +- src/quantem/diffraction/model_fitting.py | 13 +++--- .../model_fitting_visualizations.py | 42 ++++++++++++++++++- src/quantem/diffraction/strain.py | 18 +++----- .../diffraction/strain_visualization.py | 35 ++++++++-------- 5 files changed, 73 insertions(+), 39 deletions(-) diff --git a/src/quantem/core/fitting/base.py b/src/quantem/core/fitting/base.py index c28b33637..b0a1f620a 100644 --- a/src/quantem/core/fitting/base.py +++ b/src/quantem/core/fitting/base.py @@ -802,8 +802,8 @@ def __init__( def forward(self, pred, target): eps = 1 - pred_modified = (pred - pred.min().detach() + eps) ** self.gamma - target_modified = (target - target.min().detach() + eps) ** self.gamma + 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 diff --git a/src/quantem/diffraction/model_fitting.py b/src/quantem/diffraction/model_fitting.py index 768925b71..3dfcdc0b7 100644 --- a/src/quantem/diffraction/model_fitting.py +++ b/src/quantem/diffraction/model_fitting.py @@ -903,8 +903,8 @@ def fit_individual_diffraction_pattern_batched( if isinstance(loss_fn, SqrtMSELoss): gamma = float(loss_fn.gamma) eps = 1.0 - pred_min = pred.amin(dim=(1, 2), keepdim=True).detach() - tgt_min = targets.amin(dim=(1, 2), keepdim=True).detach() + pred_min = pred.amin(dim=(1, 2), keepdim=True) + tgt_min = targets.amin(dim=(1, 2), keepdim=True) pred_mod = (pred - pred_min + eps) ** gamma tgt_mod = (targets - tgt_min + eps) ** gamma per_sample_loss = ((pred_mod - tgt_mod) ** 2).mean(dim=(1, 2)) @@ -1664,11 +1664,12 @@ def apply_hard_constraints( disk_intensity_frozen = "disk.intensity_raw" in skip_keys # DiskTemplate composite hard constraints - if self.disk is not None and not disk_template_frozen: + if self.disk is not None: template = stacked.get("disk.template_raw") intensity = stacked.get("disk.intensity_raw") cfg = self.disk.constraint_config - if template is not None and intensity is not None: + 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)): @@ -1677,12 +1678,14 @@ def apply_hard_constraints( 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 bool(self.disk.hard_constraints.get("force_positive", False)): + 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" diff --git a/src/quantem/diffraction/model_fitting_visualizations.py b/src/quantem/diffraction/model_fitting_visualizations.py index c87fefafe..6bc35a09d 100644 --- a/src/quantem/diffraction/model_fitting_visualizations.py +++ b/src/quantem/diffraction/model_fitting_visualizations.py @@ -6,6 +6,11 @@ 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 @@ -343,8 +348,41 @@ def plot_model( 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()))) + # 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( diff --git a/src/quantem/diffraction/strain.py b/src/quantem/diffraction/strain.py index c99e14313..03605e9b4 100644 --- a/src/quantem/diffraction/strain.py +++ b/src/quantem/diffraction/strain.py @@ -92,13 +92,6 @@ def __init__( self.ds_sampling = 1.0 if ds_sampling is None else ds_sampling self.ds_units = "pixels" if ds_units is None else ds_units - # Per-position weighting / ROI in [0, 1]. The mask producers - # (BraggVectors.fit_lattice, StrainMapAutocorrelation.create_mask) already emit a - # [0, 1] weight, so a well-formed mask is taken as-is: re-normalizing it here - # would collide with that scaling -- a near-constant mask (e.g. the radial - # cepstral weight) would be squashed to ~0 and blank the strain display. Only a - # mask that falls outside [0, 1] (e.g. a raw-intensity ROI) is rescaled, and a - # constant / empty / all-NaN mask falls back to uniform full weight. m = np.ones(ds_shape[:2], dtype=float) if mask is None else np.asarray(mask, dtype=float) m_lo = np.nanmin(m) m_hi = np.nanmax(m) @@ -212,6 +205,8 @@ def plot_strain_roi( plot_rotation: bool = True, cmap_strain: str = "RdBu_r", cmap_rotation: str = "PiYG", + strain_range_percent: tuple[float, float] | None = None, + rotation_range_degrees: tuple[float, float] | None = None, rotate_strain: bool = False, rotate_title: bool = False, plot_dilation: bool = False, @@ -283,8 +278,8 @@ def plot_strain_roi( self.ds_shape, ds_sampling=self.ds_sampling, ds_units=self.ds_units, - strain_range_percent=(-smax, smax), - rotation_range_degrees=(-rmax, rmax), + strain_range_percent=(-smax, smax) if strain_range_percent is None else strain_range_percent, + rotation_range_degrees=(-rmax, rmax) if rotation_range_degrees is None else rotation_range_degrees, roi=inside, plot_rotation=plot_rotation, cmap_strain=cmap_strain, @@ -304,8 +299,7 @@ def plot_strain_roi( def plot_strain( self, - rotation_angle_deg: float = 0.0, - transpose: bool = False, + rotation_angle: float = 0.0, strain_range_percent: tuple[float, float] = (-3.0, 3.0), rotation_range_degrees: tuple[float, float] = (-2.0, 2.0), mask_range: tuple[float, float] = (0.0, 1.0), @@ -854,7 +848,7 @@ def _strain_tensor( # const = -1 is the reciprocal-space (nanobeam) shear/rotation convention. Both # modalities reduce strain_trans to F.T above, so the convention is shared. - const = -1 + const = 1 if real_space else -1 e_rr = strain_trans[:, :, 0, 0] - 1 e_cc = strain_trans[:, :, 1, 1] - 1 e_rc = strain_trans[:, :, 1, 0] * 0.5 * const + strain_trans[:, :, 0, 1] * 0.5 * const diff --git a/src/quantem/diffraction/strain_visualization.py b/src/quantem/diffraction/strain_visualization.py index e6cc61139..11cdeaf10 100644 --- a/src/quantem/diffraction/strain_visualization.py +++ b/src/quantem/diffraction/strain_visualization.py @@ -124,7 +124,7 @@ def _roi_compose(norm_vals, color_cm): etot_disp = _roi_compose(norm_strain(etot_pct), cm_strain) if rotate_strain: etot_disp = etot_disp.transpose(1,0,2) - ax[0].imshow(euu_disp * mask[:, :, np.newaxis]) + ax[0].imshow(etot_disp * mask[:, :, np.newaxis]) ax[1].imshow(euv_disp * mask[:, :, np.newaxis]) else: ax[0].imshow(euu_disp * mask[:, :, np.newaxis]) @@ -137,23 +137,22 @@ def _roi_compose(norm_vals, color_cm): title_fs = 16 * fs_scale tick_fs = 12 * fs_scale title_val = 'vertical' if rotate_title else 'horizontal' - if panel_titles is None and not plot_dilation: - panel_titles = ( - r"$\epsilon_{uu}$ $\updownarrow$", - r"$\epsilon_{vv}$ $\leftrightarrow$", - r"$\epsilon_{uv}$ $\nwarrow\!\!\!\!\!\!\!\!\searrow$", - ) - ax[0].set_title(panel_titles[0], fontsize=title_fs, rotation=title_val) - ax[1].set_title(panel_titles[1], fontsize=title_fs, rotation=title_val) - ax[2].set_title(panel_titles[2], fontsize=title_fs, rotation=title_val) - if plot_dilation and panel_titles is None: - panel_titles = ( - r"$\epsilon_{uu} + \epsilon_{vv}$", - r"$\epsilon_{uv}$ $\nwarrow\!\!\!\!\!\!\!\!\!\:\searrow$", - "" - ) - ax[0].set_title(panel_titles[0], fontsize=title_fs, rotation=title_val) - ax[1].set_title(panel_titles[1], fontsize=title_fs, rotation=title_val) + if panel_titles is None: + if plot_dilation: + panel_titles = ( + r"$\epsilon_{uu} + \epsilon_{vv}$", + r"$\epsilon_{uv}$ $\nwarrow\!\!\!\!\!\!\!\!\!\:\searrow$", + "", + ) + else: + panel_titles = ( + r"$\epsilon_{uu}$ $\updownarrow$", + r"$\epsilon_{vv}$ $\leftrightarrow$", + r"$\epsilon_{uv}$ $\nwarrow\!\!\!\!\!\!\!\!\searrow$", + ) + # apply to the strain panels whether panel_titles was defaulted or passed in + for i in range(n_strain): + ax[i].set_title(panel_titles[i], fontsize=title_fs, rotation=title_val) if plot_rotation: norm_rot = Normalize(vmin=rotation_range_degrees[0], vmax=rotation_range_degrees[1]) From 40c296d7101afb7debe106ac6e51a2df9be8f5d5 Mon Sep 17 00:00:00 2001 From: mitis1 Date: Tue, 28 Jul 2026 16:59:46 -0700 Subject: [PATCH 2/7] initial strain rotation bug fixes and renaming some variables --- src/quantem/diffraction/strain.py | 16 +++++----- .../diffraction/strain_visualization.py | 29 ++++++++++++++++--- 2 files changed, 34 insertions(+), 11 deletions(-) diff --git a/src/quantem/diffraction/strain.py b/src/quantem/diffraction/strain.py index 03605e9b4..19ebae436 100644 --- a/src/quantem/diffraction/strain.py +++ b/src/quantem/diffraction/strain.py @@ -207,7 +207,7 @@ def plot_strain_roi( cmap_rotation: str = "PiYG", strain_range_percent: tuple[float, float] | None = None, rotation_range_degrees: tuple[float, float] | None = None, - rotate_strain: bool = False, + transpose_image: bool = False, rotate_title: bool = False, plot_dilation: bool = False, layout: str = "horizontal", @@ -285,7 +285,7 @@ def plot_strain_roi( cmap_strain=cmap_strain, cmap_rotation=cmap_rotation, layout=layout, - rotate_strain = rotate_strain, + transpose_image = transpose_image, rotate_title = rotate_title, plot_dilation = plot_dilation, figsize=figsize, @@ -309,7 +309,8 @@ def plot_strain( cmap_strain: str = "RdBu_r", cmap_rotation: str = "PiYG", layout: str = "horizontal", - rotate_strain: bool = False, + transpose_image: bool = False, + transpose_strain: bool = False, rotate_title: bool = False, plot_dilation: bool = False, figsize: tuple[float, float] | None = None, @@ -322,7 +323,7 @@ def plot_strain( rotation_angle_deg : float, default=0.0 Angle (degrees) by which the strain tensor is rotated into the display frame before plotting. - transpose : bool, default=False + transpose_strain : bool, default=False If ``True``, transpose the detector (row/col) axes before rotating, matching the DPC convention (see :func:`~quantem.diffraction.strain_autocorrelation._raw_vec_to_display`): @@ -364,14 +365,14 @@ def plot_strain( e_cc = self.e_cc.array e_rc = self.e_rc.array phi = self.phi.array - if transpose: + if transpose_strain: # Detector-axis transpose, applied BEFORE the rotation to match the DPC # convention shared across quantem (see _raw_vec_to_display): swapping the # (row, col) axes swaps the normal strains, keeps the shear unchanged, and # reverses the sense of the rotation field. e_rr, e_cc = e_cc, e_rr phi = -phi - e_uu, e_vv, e_uv = _rotate_strain_tensor(e_rr, e_cc, e_rc, rotation_angle_deg) + e_uu, e_vv, e_uv = _rotate_strain_tensor(e_rr, e_cc, e_rc, rotation_angle) return plot_strain_panels( e_uu, e_vv, @@ -392,10 +393,11 @@ def plot_strain( cmap_strain=cmap_strain, cmap_rotation=cmap_rotation, layout=layout, - rotate_strain = rotate_strain, + transpose_image = transpose_image, rotate_title = rotate_title, plot_dilation = plot_dilation, figsize=figsize, + strain_rotation_angle=rotation_angle, **kwargs, ) diff --git a/src/quantem/diffraction/strain_visualization.py b/src/quantem/diffraction/strain_visualization.py index 11cdeaf10..c17bf8a56 100644 --- a/src/quantem/diffraction/strain_visualization.py +++ b/src/quantem/diffraction/strain_visualization.py @@ -31,11 +31,12 @@ def plot_strain_panels( cmap_strain: str = "RdBu_r", cmap_rotation: str = "PiYG", layout: str = "horizontal", - rotate_strain: bool = False, + transpose_image: bool = False, rotate_title: bool = False, plot_dilation: bool = False, figsize: tuple[float, float] | None = None, panel_titles: tuple[str, str, str] | None = None, + strain_rotation_angle: float = 0.0, **kwargs, ): """Render strain (e_uu, e_vv, e_uv) and rotation panels. @@ -113,7 +114,7 @@ def _roi_compose(norm_vals, color_cm): evv_disp = _roi_compose(norm_strain(evv_pct), cm_strain) euv_disp = _roi_compose(norm_strain(euv_pct), cm_strain) - if rotate_strain: + if transpose_image: euu_disp = euu_disp.transpose(1,0,2) evv_disp = evv_disp.transpose(1,0,2) euv_disp = euv_disp.transpose(1,0,2) @@ -122,7 +123,7 @@ def _roi_compose(norm_vals, color_cm): if plot_dilation: etot_pct = (e_uu + e_vv) * 100 etot_disp = _roi_compose(norm_strain(etot_pct), cm_strain) - if rotate_strain: + if transpose_image: etot_disp = etot_disp.transpose(1,0,2) ax[0].imshow(etot_disp * mask[:, :, np.newaxis]) ax[1].imshow(euv_disp * mask[:, :, np.newaxis]) @@ -131,6 +132,18 @@ def _roi_compose(norm_vals, color_cm): ax[1].imshow(evv_disp * mask[:, :, np.newaxis]) ax[2].imshow(euv_disp * mask[:, :, np.newaxis]) + + def _add_title_arrow(ax, angle_deg, x=0.80, y=1.06, length=0.045, + color="black", lw=1.5): + a = np.deg2rad(angle_deg) + dx, dy = length * np.cos(a), length * np.sin(a) + ax.annotate( + "", + xy=(x + dx, y + dy), xytext=(x - dx, y - dy), + xycoords="axes fraction", + annotation_clip=False, + arrowprops=dict(arrowstyle="<->", color=color, lw=lw), + ) ref_dim = figsize[1] if is_horizontal else figsize[0] fs_threshold = 3.0 fs_scale = min(1.0, max(0.5, ref_dim / fs_threshold)) @@ -144,20 +157,28 @@ def _roi_compose(norm_vals, color_cm): r"$\epsilon_{uv}$ $\nwarrow\!\!\!\!\!\!\!\!\!\:\searrow$", "", ) + title_arrow_angles = (None, -45 + strain_rotation_angle, None) else: panel_titles = ( r"$\epsilon_{uu}$ $\updownarrow$", r"$\epsilon_{vv}$ $\leftrightarrow$", r"$\epsilon_{uv}$ $\nwarrow\!\!\!\!\!\!\!\!\searrow$", ) + title_arrow_angles = (90 + strain_rotation_angle, 0 + strain_rotation_angle, -45 + strain_rotation_angle) + else: + title_arrow_angles = (None, None, None) + # apply to the strain panels whether panel_titles was defaulted or passed in for i in range(n_strain): ax[i].set_title(panel_titles[i], fontsize=title_fs, rotation=title_val) + angle = title_arrow_angles[i] + if angle is not None: + _add_title_arrow(ax[i], angle, color="black") if plot_rotation: norm_rot = Normalize(vmin=rotation_range_degrees[0], vmax=rotation_range_degrees[1]) rot_disp = _roi_compose(norm_rot(rot_deg), cm_rot) - if rotate_strain: rot_disp = rot_disp.transpose(1,0,2) + if transpose_image: rot_disp = rot_disp.transpose(1,0,2) ax[-1].imshow(rot_disp * mask[:, :, np.newaxis]) ax[-1].set_title(r"Rotation $\circlearrowleft$", fontsize=title_fs, rotation=title_val) From d5812bd6bfff17a8f9da5c7f8c5461191f6b78a6 Mon Sep 17 00:00:00 2001 From: mitis1 Date: Wed, 29 Jul 2026 13:50:18 -0700 Subject: [PATCH 3/7] initial rotation implementation --- src/quantem/diffraction/bragg_vectors.py | 55 ++++- src/quantem/diffraction/model_fitting.py | 38 ++++ src/quantem/diffraction/strain.py | 118 +++++++++-- .../diffraction/strain_autocorrelation.py | 2 + .../diffraction/strain_visualization.py | 193 ++++++++++++------ 5 files changed, 325 insertions(+), 81 deletions(-) diff --git a/src/quantem/diffraction/bragg_vectors.py b/src/quantem/diffraction/bragg_vectors.py index 7b571e04e..344910ba4 100644 --- a/src/quantem/diffraction/bragg_vectors.py +++ b/src/quantem/diffraction/bragg_vectors.py @@ -441,7 +441,8 @@ def detect_disks( detector dimensions. progressbar : bool, default=True If ``True``, show a tqdm progress bar over the full-scan detection. - + save_to_gpu: bool, default = True, + Transfers dataset to gpu for faster calculation and less cpu load. Returns ------- Vector @@ -924,6 +925,8 @@ def calculate_strain_map( u_ref: np.ndarray | None = None, v_ref: np.ndarray | None = None, mask: np.ndarray | None = None, + q_to_r_rotation_ccw_deg: float | None = None, + q_transpose: bool | None = None, ) -> StrainMap: """Build a :class:`StrainMap` from the fitted per-position lattice vectors. @@ -954,8 +957,52 @@ def calculate_strain_map( if mask is None: mask = self.mask_weight - ds_sampling = float(self.dataset.sampling[0]) - ds_units = str(self.dataset.units[0]) + ds_units = None + ds_sampling = None + if hasattr(self.dataset, 'units'): + if isinstance(self.dataset.units, (tuple, list)): + ds_units = str(self.dataset.units[0]) + else: + ds_units = str(self.dataset.units) + if hasattr(self.dataset, 'sampling'): + if isinstance(self.dataset.sampling, (tuple, list, np.ndarray)): + ds_sampling = float(self.dataset.sampling[0]) + else: + ds_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( u_array=self.u_array, @@ -967,6 +1014,8 @@ def calculate_strain_map( mask=mask, ds_sampling=ds_sampling, ds_units=ds_units, + q_to_r_rotation_ccw_deg = q_to_r_rotation_ccw_deg, + q_transpose = q_transpose, ) # ---- visualization ---- diff --git a/src/quantem/diffraction/model_fitting.py b/src/quantem/diffraction/model_fitting.py index 3dfcdc0b7..5e23135e8 100644 --- a/src/quantem/diffraction/model_fitting.py +++ b/src/quantem/diffraction/model_fitting.py @@ -1125,6 +1125,8 @@ def calculate_strain_map( u_ref: np.ndarray | None = None, v_ref: np.ndarray | None = None, mask: np.ndarray | None = None, + q_to_r_rotation_ccw_deg: float | None = None, + q_transpose: bool | None = None, ) -> StrainMap: """Build a :class:`StrainMap` from the fitted per-position lattice vectors. @@ -1172,6 +1174,40 @@ def calculate_strain_map( 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( u_array = self.u_array, @@ -1183,6 +1219,8 @@ def calculate_strain_map( 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, ) @property diff --git a/src/quantem/diffraction/strain.py b/src/quantem/diffraction/strain.py index 19ebae436..0c000e364 100644 --- a/src/quantem/diffraction/strain.py +++ b/src/quantem/diffraction/strain.py @@ -4,6 +4,7 @@ import numpy as np from numpy.lib.stride_tricks import sliding_window_view +from numpy.typing import NDArray from quantem.core.datastructures.dataset2d import Dataset2d from quantem.core.io.serialize import AutoSerialize @@ -81,10 +82,18 @@ def __init__( mask: np.ndarray | None = None, ds_sampling: float | None = None, ds_units: str | None = None, + q_to_r_rotation_ccw_deg: float = 0.0, + q_transpose: bool = False, ): super().__init__() self.u_array = u_array self.v_array = v_array + + self.q_to_r_rotation_ccw_deg = q_to_r_rotation_ccw_deg + self.q_transpose = q_transpose + + self.u_array = _raw_vec_to_display(self.u_array, rotation_ccw_deg = q_to_r_rotation_ccw_deg, transpose=q_transpose) + self.v_array = _raw_vec_to_display(self.v_array, rotation_ccw_deg = q_to_r_rotation_ccw_deg, transpose=q_transpose) self.ds_shape = ds_shape self.real_space = real_space @@ -102,8 +111,12 @@ def __init__( self.mask = m # user-supplied reference vectors persist across re-fits (None = use median) - self._u_ref_fixed = None if u_ref is None else np.asarray(u_ref, dtype=float) - self._v_ref_fixed = None if v_ref is None else np.asarray(v_ref, dtype=float) + self._u_ref_fixed = None if u_ref is None else _raw_vec_to_display(np.asarray(u_ref, dtype=float), + rotation_ccw_deg=q_to_r_rotation_ccw_deg, + transpose=q_transpose) + self._v_ref_fixed = None if v_ref is None else _raw_vec_to_display(np.asarray(v_ref, dtype=float), + rotation_ccw_deg=q_to_r_rotation_ccw_deg, + transpose=q_transpose) self.u_ref = None self.v_ref = None @@ -117,6 +130,7 @@ def update_reference( u_ref: np.ndarray | None = None, v_ref: np.ndarray | None = None, plot_strain_roi: bool = False, + define_in_rotated_frame: bool = False, **plot_kwargs, ) -> "StrainMap": """(Re)compute the reference lattice and strain tensor maps. @@ -141,6 +155,8 @@ def update_reference( If ``True``, show the recomputed strain via :meth:`plot_strain_roi` (color-scaled to the ROI) so the chosen reference region can be checked for flatness. + define_in_rotated_frame: bool, default = False + If ''True'' means the u_ref and v_ref passed into the function is defined in the rotated detector frame **plot_kwargs Forwarded to :meth:`plot_strain_roi` when ``plot_strain_roi=True``. @@ -152,14 +168,26 @@ def update_reference( u_med, v_med = _reference_lattice(self.u_array, self.v_array, self.mask, strain_mask) if u_ref is not None: - self.u_ref = np.asarray(u_ref, dtype=float) + if define_in_rotated_frame: + self.u_ref = np.asarray(u_ref, dtype=float) + else: + self.u_ref = _raw_vec_to_display( + np.asarray(u_ref, dtype=float), + rotation_ccw_deg=self.q_to_r_rotation_ccw_deg, + transpose=self.q_transpose) elif self._u_ref_fixed is not None: self.u_ref = self._u_ref_fixed else: self.u_ref = u_med if v_ref is not None: - self.v_ref = np.asarray(v_ref, dtype=float) + if define_in_rotated_frame: + self.v_ref = np.asarray(v_ref, dtype=float) + else: + self.v_ref = _raw_vec_to_display( + np.asarray(v_ref, dtype=float), + rotation_ccw_deg=self.q_to_r_rotation_ccw_deg, + transpose=self.q_transpose) elif self._v_ref_fixed is not None: self.v_ref = self._v_ref_fixed else: @@ -211,6 +239,7 @@ def plot_strain_roi( rotate_title: bool = False, plot_dilation: bool = False, layout: str = "horizontal", + arrow_style: str = "title", figsize: tuple[float, float] | None = None, **kwargs, ): @@ -236,8 +265,21 @@ def plot_strain_roi( Colormap for the strain panels. cmap_rotation : str, default="PiYG" Colormap for the rotation panel. + strain_range_percent : tuple of float, default=(-3.0, 3.0) + Symmetric color range for the strain panels, in percent. + rotation_range_degrees : tuple of float, default=(-2.0, 2.0) + Symmetric color range for the rotation panel, in degrees. + transpose_image: bool, default = False + If ''True'' transpose the real space image before plotting strain + rotate_title: bool, default = False + If ''True'', rotates panel titles by 90 degrees. + plot_dilation: bool, default = False + If ''True'' plots euu + evv, and euv instead of euu, evv, euv layout : {"horizontal", "vertical"}, default="horizontal" Panel arrangement. + arrow_style: str, default="title" + Plots the directional arrows along with the strain titles. + Alternatively can be "legend" where it plots it on the side figsize : tuple of float, optional Figure size in inches; if omitted it is derived from the layout. **kwargs @@ -249,6 +291,10 @@ def plot_strain_roi( tuple ``(fig, ax)`` from :func:`plot_strain_panels`. """ + + if arrow_style not in ("title", "legend"): + raise ValueError("arrow_style must be 'title' or 'legend'") + roi_src = self.mask if strain_mask is None else strain_mask e_rr, e_cc, e_rc, phi = ( self.e_rr.array, @@ -294,6 +340,7 @@ def plot_strain_roi( r"$\epsilon_{cc}$ $\leftrightarrow$", r"$\epsilon_{rc}$ $\nwarrow\!\!\!\!\!\!\!\!\!\:\searrow$", ), + arrow_style = arrow_style, **kwargs, ) @@ -308,11 +355,12 @@ def plot_strain( plot_scalebar: bool = False, cmap_strain: str = "RdBu_r", cmap_rotation: str = "PiYG", - layout: str = "horizontal", transpose_image: bool = False, transpose_strain: bool = False, rotate_title: bool = False, plot_dilation: bool = False, + layout: str = "horizontal", + arrow_style: str = "title", figsize: tuple[float, float] | None = None, **kwargs, ): @@ -320,15 +368,9 @@ def plot_strain( Parameters ---------- - rotation_angle_deg : float, default=0.0 + rotation_angle : float, default=0.0 Angle (degrees) by which the strain tensor is rotated into the display frame before plotting. - transpose_strain : bool, default=False - If ``True``, transpose the detector (row/col) axes before rotating, - matching the DPC convention (see - :func:`~quantem.diffraction.strain_autocorrelation._raw_vec_to_display`): - transpose first, then rotate. This swaps the normal strain components, - leaves the shear unchanged, and reverses the sign of the rotation field. strain_range_percent : tuple of float, default=(-3.0, 3.0) Symmetric color range for the strain panels, in percent. rotation_range_degrees : tuple of float, default=(-2.0, 2.0) @@ -348,10 +390,25 @@ def plot_strain( Colormap for the strain panels. cmap_rotation : str, default="PiYG" Colormap for the rotation panel. + transpose_image: bool, default = False + If ''True'' transpose the real space image before plotting strain + transpose_strain : bool, default=False + If ``True``, transpose the detector (row/col) axes before rotating, + matching the DPC convention (see + :func:`~quantem.diffraction.strain_autocorrelation._raw_vec_to_display`): + transpose first, then rotate. This swaps the normal strain components, + leaves the shear unchanged, and reverses the sign of the rotation field. + rotate_title: bool, default = False + If ''True'', rotates panel titles by 90 degrees. + plot_dilation: bool, default = False + If ''True'' plots euu + evv, and euv instead of euu, evv, euv layout : {"horizontal", "vertical"}, default="horizontal" Panel arrangement. figsize : tuple of float, optional Figure size in inches; if omitted it is derived from the layout. + arrow_style: str, default="title" + Plots the directional arrows along with the strain titles. + Alternatively can be "legend" where it plots it on the side **kwargs Forwarded to :func:`~quantem.diffraction.strain_visualization.plot_strain_panels`. @@ -361,6 +418,9 @@ def plot_strain( tuple ``(fig, ax)`` from :func:`plot_strain_panels`. """ + 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 @@ -398,6 +458,7 @@ def plot_strain( plot_dilation = plot_dilation, figsize=figsize, strain_rotation_angle=rotation_angle, + arrow_style = arrow_style, **kwargs, ) @@ -863,6 +924,7 @@ def _rotate_strain_tensor( e_cc: np.ndarray, e_rc: np.ndarray, rotation_angle: float, + real_space=False ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Rotate a 2D strain tensor by ``rotation_angle`` (degrees). @@ -876,7 +938,8 @@ def _rotate_strain_tensor( Row-column (shear) strain component. rotation_angle : float Frame rotation angle, in degrees. - + real_space: bool + Tells whether vector is defined in real or reciprocal space Returns ------- tuple of np.ndarray @@ -885,7 +948,28 @@ def _rotate_strain_tensor( angle = np.deg2rad(rotation_angle) c = np.cos(angle) s = np.sin(angle) - e_uu = e_rr * (c * c) + 2.0 * e_rc * (c * s) + e_cc * (s * s) - e_vv = e_rr * (s * s) - 2.0 * e_rc * (c * s) + e_cc * (c * c) - e_uv = (e_cc - e_rr) * (c * s) + e_rc * (c * c - s * s) - return e_uu, e_vv, e_uv \ No newline at end of file + sign = -1.0 if real_space else 1.0 + e_uu = e_rr * (c * c) + sign * 2.0 * e_rc * (c * s) + e_cc * (s * s) + e_vv = e_rr * (s * s) - sign * 2.0 * e_rc * (c * s) + e_cc * (c * c) + e_uv = sign * (e_cc - e_rr) * (c * s) + e_rc * (c * c - s * s) + return e_uu, e_vv, e_uv + +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) + 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.stack((dr2, dc2), axis=-1) \ No newline at end of file diff --git a/src/quantem/diffraction/strain_autocorrelation.py b/src/quantem/diffraction/strain_autocorrelation.py index 0c9a8d14b..5d939f9fb 100644 --- a/src/quantem/diffraction/strain_autocorrelation.py +++ b/src/quantem/diffraction/strain_autocorrelation.py @@ -1275,6 +1275,8 @@ def calculate_strain_map( 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'], ) def plot_lattice_vectors( diff --git a/src/quantem/diffraction/strain_visualization.py b/src/quantem/diffraction/strain_visualization.py index c17bf8a56..8ad0888f7 100644 --- a/src/quantem/diffraction/strain_visualization.py +++ b/src/quantem/diffraction/strain_visualization.py @@ -37,6 +37,7 @@ def plot_strain_panels( figsize: tuple[float, float] | None = None, panel_titles: tuple[str, str, str] | None = None, strain_rotation_angle: float = 0.0, + arrow_style: str = "title", **kwargs, ): """Render strain (e_uu, e_vv, e_uv) and rotation panels. @@ -133,17 +134,6 @@ def _roi_compose(norm_vals, color_cm): ax[2].imshow(euv_disp * mask[:, :, np.newaxis]) - def _add_title_arrow(ax, angle_deg, x=0.80, y=1.06, length=0.045, - color="black", lw=1.5): - a = np.deg2rad(angle_deg) - dx, dy = length * np.cos(a), length * np.sin(a) - ax.annotate( - "", - xy=(x + dx, y + dy), xytext=(x - dx, y - dy), - xycoords="axes fraction", - annotation_clip=False, - arrowprops=dict(arrowstyle="<->", color=color, lw=lw), - ) ref_dim = figsize[1] if is_horizontal else figsize[0] fs_threshold = 3.0 fs_scale = min(1.0, max(0.5, ref_dim / fs_threshold)) @@ -154,39 +144,38 @@ def _add_title_arrow(ax, angle_deg, x=0.80, y=1.06, length=0.045, if plot_dilation: panel_titles = ( r"$\epsilon_{uu} + \epsilon_{vv}$", - r"$\epsilon_{uv}$ $\nwarrow\!\!\!\!\!\!\!\!\!\:\searrow$", + r"$\epsilon_{uv}$", "", ) title_arrow_angles = (None, -45 + strain_rotation_angle, None) else: panel_titles = ( - r"$\epsilon_{uu}$ $\updownarrow$", - r"$\epsilon_{vv}$ $\leftrightarrow$", - r"$\epsilon_{uv}$ $\nwarrow\!\!\!\!\!\!\!\!\searrow$", + r"$\epsilon_{uu}$", + r"$\epsilon_{vv}$", + r"$\epsilon_{uv}$", ) title_arrow_angles = (90 + strain_rotation_angle, 0 + strain_rotation_angle, -45 + strain_rotation_angle) else: title_arrow_angles = (None, None, None) - # apply to the strain panels whether panel_titles was defaulted or passed in - for i in range(n_strain): - ax[i].set_title(panel_titles[i], fontsize=title_fs, rotation=title_val) - angle = title_arrow_angles[i] - if angle is not None: - _add_title_arrow(ax[i], angle, color="black") if plot_rotation: norm_rot = Normalize(vmin=rotation_range_degrees[0], vmax=rotation_range_degrees[1]) rot_disp = _roi_compose(norm_rot(rot_deg), cm_rot) if transpose_image: rot_disp = rot_disp.transpose(1,0,2) ax[-1].imshow(rot_disp * mask[:, :, np.newaxis]) - ax[-1].set_title(r"Rotation $\circlearrowleft$", fontsize=title_fs, rotation=title_val) + if arrow_style == "title": + ax[-1].set_title(r"$\phi$ $\circlearrowleft$", fontsize=title_fs, rotation=title_val) + else: + ax[-1].set_title(r"$\phi$", fontsize=title_fs, rotation=title_val) + for a in ax: a.set_xticks([]) a.set_yticks([]) a.set_facecolor("black") a.set_aspect("equal") + a.set_anchor("W" if not is_horizontal else "C") if plot_scalebar: scalebar_kwargs = {} @@ -236,12 +225,13 @@ def _finalize_layout(): except AttributeError: # matplotlib < 3.5 fig.canvas.draw() + need_side_panel = plot_gvecs or arrow_style == "legend" if is_horizontal: # Reserve a bottom band wide enough for the colorbar + its tick labels and # title (fontsize 16) and a right band for the rotation-panel gap; widen the # right band when the g-vector compass is drawn in it. These keep the figure # usable when saved "as is" (no bbox_inches='tight'). - right = 0.78 if plot_gvecs else 0.93 + right = 0.72 if need_side_panel else 0.93 fig.subplots_adjust(left=0.04, right=right, top=0.88, bottom=0.24, wspace=0.05) if plot_rotation: # nudge the rotation panel right for a visual gap from the strain panels; @@ -269,20 +259,23 @@ def _finalize_layout(): else: # Top band for the panel titles, right band for the vertical colorbars + labels. - fig.subplots_adjust(left=0.04, right=0.80, top=0.92, bottom=0.06, hspace=0.15) + right = 0.55 if need_side_panel else 0.80 + fig.subplots_adjust(left=0.04, right=right, top=0.92, bottom=0.06, hspace=0.15) _finalize_layout() cb_orientation = "vertical" b0 = ax[0].get_position() b2 = ax[n_strain - 1].get_position() - strain_cb_pos = [b0.x1 + cb_pad, b2.y0, cb_size, b0.y1 - b2.y0] + 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] if plot_rotation: b3 = ax[-1].get_position() rot_cb_h = max(b3.y1 - b3.y0, cb_min_len) rot_cb_cy = 0.5 * (b3.y0 + b3.y1) rot_cb_y0 = min(max(rot_cb_cy - 0.5 * rot_cb_h, 0.0), 0.99 - rot_cb_h) - rot_cb_pos = [b0.x1 + cb_pad, rot_cb_y0, cb_size, rot_cb_h] + rot_cb_pos = [cb_x0, rot_cb_y0, cb_size, rot_cb_h] last_pos = b3 else: rot_cb_pos = None @@ -306,51 +299,129 @@ def _finalize_layout(): cbar2.update_ticks() cbar2.ax.tick_params(labelsize=tick_fs) - if plot_gvecs: - if u_ref is None or v_ref is None: - print("Warning: u_ref and v_ref not found. Call fit_strain() first.") - return fig, ax - - # The compass goes in the reserved margin beside the last panel; clamp its - # right edge to 0.99 so it never spills off the figure when saved "as is". - if is_horizontal: - ref_left = last_pos.x1 + 0.005 - ref_width = min(last_pos.width, 0.99 - ref_left) - ref_ax = fig.add_axes([ref_left, last_pos.y0, ref_width, last_pos.height]) - else: - ref_left = min(last_pos.x1 + 0.18, 0.74) - ref_width = min(last_pos.width, 0.99 - ref_left) - ref_ax = fig.add_axes([ref_left, last_pos.y0, ref_width, last_pos.height]) + 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 + try: + ax.figure.draw_without_rendering() + except AttributeError: # matplotlib < 3.5 + ax.figure.canvas.draw() + renderer = ax.figure.canvas.get_renderer() + bbox_ax = ax.title.get_window_extent(renderer=renderer).transformed(ax.transAxes.inverted()) + y = 0.5 * (bbox_ax.y0 + bbox_ax.y1) + ax.annotate( + "\u2194", + xy=(bbox_ax.x1, y), xycoords=ax.transAxes, + xytext=(gap_pt + fs / 2.0, 0), textcoords="offset points", + ha="center", va="center", + rotation=angle_deg, rotation_mode="anchor", + fontsize=fs, color=color, + annotation_clip=False, + ) + + def _add_arrow_legend(fig, x0, y_top, entries, plot_rotation, fontsize, color="black"): + fig_w_in, fig_h_in = figsize + row_h_in = fontsize * 1.6 / 72.0 + box_w_in = 1.6 + n_rows = len(entries) + 1 + (2 if plot_rotation else 0) + row_h = row_h_in / fig_h_in + box_h = row_h * n_rows + box_w = min(box_w_in / fig_w_in, 0.99 - x0) + leg_ax = fig.add_axes([x0, y_top - box_h, box_w, box_h]) + leg_ax.set_xlim(0, 1) + leg_ax.set_ylim(0, 1) + leg_ax.axis("off") + + dy = 1.0 / n_rows + y = 1.0 - dy / 2 + leg_ax.text(0.0, y, "Strain", fontsize=fontsize, fontweight="bold", ha="left", va="center") + for label, angle_deg in entries: + y -= dy + leg_ax.text(0.15, y, "\u2194", rotation=angle_deg, rotation_mode="anchor", + ha="center", va="center", fontsize=fontsize, color=color) + leg_ax.text(0.32, y, label, fontsize=fontsize, ha="left", va="center") + + if plot_rotation: + y -= dy + leg_ax.text(0.0, y, "Rotation", fontsize=fontsize, fontweight="bold", ha="left", va="center") + y -= dy + leg_ax.text(0.15, y, "\u21ba", fontsize=fontsize, ha="center", va="center") + leg_ax.text(0.32, y, r"$\phi$", fontsize=fontsize, ha="left", va="center") + return box_h + + for i in range(n_strain): + ax[i].set_title(panel_titles[i], fontsize=title_fs, rotation=title_val) + angle = title_arrow_angles[i] + if arrow_style == "title" and angle is not None: + _add_title_arrow(ax[i], angle, color="black") + + if is_horizontal: + _finalize_layout() + renderer = fig.canvas.get_renderer() + panel_edge = last_pos.x1 + if plot_rotation: + title_edge = ax[-1].title.get_window_extent(renderer=renderer).transformed(fig.transFigure.inverted()).x1 + panel_edge = max(panel_edge, title_edge) + margin_x0 = panel_edge + 0.03 + else: + _finalize_layout() + renderer = fig.canvas.get_renderer() + margin_x0 = cax1.get_tightbbox(renderer).transformed(fig.transFigure.inverted()).x1 + 0.02 + if plot_rotation and rot_cb_pos is not None: + rot_edge = cax2.get_tightbbox(renderer).transformed(fig.transFigure.inverted()).x1 + margin_x0 = max(margin_x0, rot_edge + 0.02) + top_bound = 0.88 if is_horizontal else 0.92 + bottom_bound = 0.24 if is_horizontal else 0.06 + center_y = 0.5 * (top_bound + bottom_bound) + + entries = [] + leg_h = 0.0 + if arrow_style == "legend": + entries = [(panel_titles[i], title_arrow_angles[i]) for i in range(n_strain) + if title_arrow_angles[i] is not None] + n_rows = len(entries) + 1 + (2 if plot_rotation else 0) + leg_h = (title_fs * 1.6 / 72.0 / figsize[1]) * n_rows + + show_gvecs = plot_gvecs and u_ref is not None and v_ref is not None + if plot_gvecs and not show_gvecs: + print("Warning: u_ref and v_ref not found. Call fit_strain() first.") + fig_aspect = figsize[0] / figsize[1] + gvec_w = min(0.99 - margin_x0, 0.15) if show_gvecs else 0.0 + gvec_h = gvec_w * fig_aspect if show_gvecs else 0.0 + + gap = 0.03 if (leg_h > 0 and gvec_h > 0) else 0.0 + total_needed = leg_h + gap + gvec_h + available_span = top_bound - bottom_bound + side_scale = min(1.0, available_span / total_needed) if total_needed > 0 else 1.0 + leg_h *= side_scale + gvec_w *= side_scale + gvec_h *= side_scale + legend_fontsize = title_fs * side_scale + + y_top = center_y + (leg_h + gap * side_scale + gvec_h) / 2.0 + + if leg_h > 0: + _add_arrow_legend(fig, margin_x0, y_top, entries, plot_rotation=plot_rotation, + fontsize=legend_fontsize, color="black") + y_top -= leg_h + gap * side_scale + + if show_gvecs: + ref_ax = fig.add_axes([margin_x0, y_top - gvec_h, gvec_w, gvec_h]) ref_ax.set_xlim(-1.5, 1.5) ref_ax.set_ylim(-1.5, 1.5) ref_ax.set_aspect("equal") ref_ax.axis("off") u_norm = u_ref / np.linalg.norm(u_ref) v_norm = v_ref / np.linalg.norm(v_ref) - u_row, u_col = u_norm v_row, v_col = v_norm arrow_props_ref = dict(arrowstyle="->", lw=3, mutation_scale=25) - - u_arrow = FancyArrowPatch( - (0, 0), (u_col, -u_row), - color="darkred", **arrow_props_ref - ) - ref_ax.add_patch(u_arrow) - - v_arrow = FancyArrowPatch( - (0, 0), (v_col, -v_row), - color="darkblue", **arrow_props_ref - ) - ref_ax.add_patch(v_arrow) - ref_ax.text(u_col * 1.3, -u_row * 1.3, r"$\mathbf{g}_{1}$", - fontsize=14, fontweight="bold", color="darkred", - ha="center", va="center") - - ref_ax.text(v_col * 1.3, -v_row * 1.3, r"$\mathbf{g}_{2}$", - fontsize=14, fontweight="bold", color="darkblue", - ha="center", va="center") + ref_ax.add_patch(FancyArrowPatch((0, 0), (u_col, -u_row), color="darkred", **arrow_props_ref)) + ref_ax.add_patch(FancyArrowPatch((0, 0), (v_col, -v_row), color="darkblue", **arrow_props_ref)) + ref_ax.text(u_col * 1.3, -u_row * 1.3, r"$\mathbf{g}_{1}$", fontsize=14, fontweight="bold", + color="darkred", ha="center", va="center") + ref_ax.text(v_col * 1.3, -v_row * 1.3, r"$\mathbf{g}_{2}$", fontsize=14, fontweight="bold", + color="darkblue", ha="center", va="center") return fig, ax From 5e660145776f90634fdeff21caac0a585980f3bd Mon Sep 17 00:00:00 2001 From: mitis1 Date: Wed, 29 Jul 2026 22:13:01 -0700 Subject: [PATCH 4/7] fixing subpixel dft fit for cepstral, batched mask implementation for model fitting, and optimizer batched implementation --- src/quantem/diffraction/model_fitting.py | 230 +++++++++++++++--- .../diffraction/strain_autocorrelation.py | 21 +- .../diffraction/strain_visualization.py | 2 +- 3 files changed, 209 insertions(+), 44 deletions(-) diff --git a/src/quantem/diffraction/model_fitting.py b/src/quantem/diffraction/model_fitting.py index 5e23135e8..c5f940a4e 100644 --- a/src/quantem/diffraction/model_fitting.py +++ b/src/quantem/diffraction/model_fitting.py @@ -20,6 +20,7 @@ ) 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 @@ -830,6 +831,12 @@ def fit_individual_diffraction_pattern_batched( 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) @@ -868,11 +875,12 @@ def fit_individual_diffraction_pattern_batched( stacked = plan.build_stacked_params(B) - adam_state: dict[str, dict[str, torch.Tensor]] = { - name: { - "m": torch.zeros_like(p.detach()), - "v": torch.zeros_like(p.detach()), - } + 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() } @@ -903,17 +911,37 @@ def fit_individual_diffraction_pattern_batched( if isinstance(loss_fn, SqrtMSELoss): gamma = float(loss_fn.gamma) eps = 1.0 - pred_min = pred.amin(dim=(1, 2), keepdim=True) - tgt_min = targets.amin(dim=(1, 2), keepdim=True) + if fidelity_mask is not None: + pred_min = pred.masked_fill(~fidelity_mask, float("inf")).amin(dim=(1, 2), keepdim=True) + tgt_min = targets.masked_fill(~fidelity_mask, float("inf")).amin(dim=(1, 2), keepdim=True) + else: + pred_min = pred.amin(dim=(1, 2), keepdim=True) + tgt_min = targets.amin(dim=(1, 2), keepdim=True) pred_mod = (pred - pred_min + eps) ** gamma tgt_mod = (targets - tgt_min + eps) ** gamma - per_sample_loss = ((pred_mod - tgt_mod) ** 2).mean(dim=(1, 2)) + sq = (pred_mod - tgt_mod) ** 2 + if fidelity_mask is not None: + per_sample_loss = (sq * fidelity_mask).sum(dim=(1, 2)) / fidelity_valid_count + else: + per_sample_loss = sq.mean(dim=(1, 2)) elif isinstance(loss_fn, LogMSELoss): - per_sample_loss = ((torch.log1p(pred) - torch.log1p(targets)) ** 2).mean(dim=(1, 2)) + sq = (torch.log1p(pred) - torch.log1p(targets)) ** 2 + if fidelity_mask is not None: + per_sample_loss = (sq * fidelity_mask).sum(dim=(1, 2)) / fidelity_valid_count + else: + per_sample_loss = sq.mean(dim=(1, 2)) elif isinstance(loss_fn, torch.nn.L1Loss): - per_sample_loss = diff2.abs().mean(dim=(1, 2)) + ad = diff2.abs() + if fidelity_mask is not None: + per_sample_loss = (ad * fidelity_mask).sum(dim=(1, 2)) / fidelity_valid_count + else: + per_sample_loss = ad.mean(dim=(1, 2)) else: - per_sample_loss = (diff2 * diff2).mean(dim=(1, 2)) + sq = diff2 * diff2 + if fidelity_mask is not None: + per_sample_loss = (sq * fidelity_mask).sum(dim=(1, 2)) / fidelity_valid_count + else: + per_sample_loss = sq.mean(dim=(1, 2)) # Add per-sample soft constraint loss. if float(constraint_weight) != 0.0: @@ -940,7 +968,10 @@ def fit_individual_diffraction_pattern_batched( t = step + 1 current_lrs = plan.lr_at_step(t, int(n_steps)) - _adam_step_inplace(stacked, tuple(grads_list), adam_state, current_lrs, chunk_trainable, t) + _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) @@ -983,6 +1014,7 @@ def get_individual_uv_vectors(self) -> "ModelDiffraction": if pos_state is None: self.u_array[r,c,:] = None self.v_array[r,c,:] = None + continue for key in pos_state.keys(): key_parts = key.split('.') if(key_parts[-1] == 'u_row'): @@ -1264,52 +1296,156 @@ def _lr_for_component( 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.setdefault("type", "adam") + d.setdefault("lr", component.DEFAULT_LR) + return OptimizerParams.parse_dict(d) + + +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], - chunk_trainable: dict[str, torch.Tensor], # NEW parameter + hparams: dict[str, dict[str, Any]], + chunk_trainable: dict[str, torch.Tensor], t: int, - beta1: float = 0.9, - beta2: float = 0.999, - eps: float = 1e-8, + decoupled_wd: bool = False, ) -> None: """ - In-place Adam step with per-sample trainability support. - + 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 + 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). """ - bias1 = 1.0 - beta1 ** t - bias2 = 1.0 - beta2 ** t - 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_masked = g * mask_expanded - - st["m"].mul_(beta1).add_(g_masked, alpha=1.0 - beta1) - st["v"].mul_(beta2).addcmul_(g_masked, g_masked, value=1.0 - beta2) - else: - st["m"].mul_(beta1).add_(g, alpha=1.0 - beta1) - st["v"].mul_(beta2).addcmul_(g, g, value=1.0 - beta2) - - m_hat = st["m"] / bias1 - v_hat = st["v"] / bias2 - + + 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.""" @@ -1320,6 +1456,8 @@ def __init__(self) -> 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] = [] @@ -1483,6 +1621,28 @@ def from_model( # 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) diff --git a/src/quantem/diffraction/strain_autocorrelation.py b/src/quantem/diffraction/strain_autocorrelation.py index 5d939f9fb..5e63e69cb 100644 --- a/src/quantem/diffraction/strain_autocorrelation.py +++ b/src/quantem/diffraction/strain_autocorrelation.py @@ -984,8 +984,8 @@ def _fit_lattice_vectors_batched( self._gpu_cache = None use_cache = False else: - self._gpu_cache = None - use_cache = False + # 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 @@ -1275,7 +1275,7 @@ def calculate_strain_map( 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_to_r_rotation_ccw_deg = self.metadata['q_to_r_rotation_ccw_deg'], q_transpose = self.metadata['q_transpose'], ) @@ -1698,8 +1698,13 @@ def _refine_peak_subpixel_dft( F = np.fft.fft2(np.fft.fftshift(im)) up = upsample - du = int(np.fix(np.ceil(1.5 * up))) - patch = np.abs(dft_upsample(F, up=up, shift=(r0, c0))) + H, W = im.shape + du = int(np.floor(np.ceil(1.5 * up) / 2.0)) + off_r = -(-H // 2) # ceil(H/2) + off_c = -(-W // 2) # ceil(W/2) + + shift = (du + up * (r0 - off_r), du + up * (c0 - off_c)) + patch = np.abs(dft_upsample(F, up=up, shift=shift)) patch = np.asarray(patch, dtype=float) i0, j0 = np.unravel_index(np.argmax(patch), patch.shape) @@ -1715,9 +1720,9 @@ def _refine_peak_subpixel_dft( dj = _parabolic_vertex_delta(row[0], row[1], row[2]) else: dj = 0.0 - M, N = im.shape - dr = ((float(i0) - du + di)) / up - dc = ((float(j0) - du + dj)) / up + + dr = ((du - float(i0) - di)) / up + dc = ((du - float(j0) - dj)) / up return r0 + dr, c0 + dc diff --git a/src/quantem/diffraction/strain_visualization.py b/src/quantem/diffraction/strain_visualization.py index 8ad0888f7..819891cb7 100644 --- a/src/quantem/diffraction/strain_visualization.py +++ b/src/quantem/diffraction/strain_visualization.py @@ -154,7 +154,7 @@ def _roi_compose(norm_vals, color_cm): r"$\epsilon_{vv}$", r"$\epsilon_{uv}$", ) - title_arrow_angles = (90 + strain_rotation_angle, 0 + strain_rotation_angle, -45 + strain_rotation_angle) + title_arrow_angles = (0 + strain_rotation_angle, 90 + strain_rotation_angle, -45 + strain_rotation_angle) else: title_arrow_angles = (None, None, None) From 957694fe1b3ffc0bdf02a1949bf311ffd94b1083 Mon Sep 17 00:00:00 2001 From: mitis1 Date: Thu, 30 Jul 2026 22:46:59 -0700 Subject: [PATCH 5/7] adding colin background subtraction and fixing merge commits in model fitting and strain mapping --- src/quantem/core/fitting/diffraction.py | 297 +++++++++++++++++- src/quantem/diffraction/bragg_vectors.py | 136 +++++++- src/quantem/diffraction/disk_detection.py | 63 +++- src/quantem/diffraction/model_fitting.py | 112 +++++-- .../diffraction/strain_autocorrelation.py | 24 +- 5 files changed, 565 insertions(+), 67 deletions(-) diff --git a/src/quantem/core/fitting/diffraction.py b/src/quantem/core/fitting/diffraction.py index 62997511c..1c8c15983 100644 --- a/src/quantem/core/fitting/diffraction.py +++ b/src/quantem/core/fitting/diffraction.py @@ -48,6 +48,143 @@ def put(rr: torch.Tensor, cc: torch.Tensor, ww: torch.Tensor) -> None: 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], *, @@ -125,6 +262,8 @@ def __init__( 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, ): @@ -136,6 +275,12 @@ def __init__( 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 ------- @@ -154,6 +299,11 @@ def __init__( 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( @@ -209,6 +359,8 @@ def from_array( 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, @@ -221,6 +373,8 @@ def from_array( origin=origin, origin_key=origin_key, intensity=intensity, + render_sigma=render_sigma, + render_kernel=render_kernel, constraint_params=constraint_params, constraint_config=constraint_config, @@ -246,7 +400,14 @@ def add_patch( 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) - _splat_patch(out, r0=r0, c0=c0, patch_vals=vals, dr=dr, dc=dc, scale=scale) + 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: """ @@ -285,6 +446,11 @@ def forward_batched( 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 ) @@ -531,6 +697,10 @@ def __init__( 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, @@ -553,6 +723,13 @@ def __init__( 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 @@ -593,6 +770,23 @@ def __init__( 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 @@ -728,18 +922,44 @@ def __init__( 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 ----- - Positivity is enforced as a hard projection after optimizer steps. - The forward path intentionally avoids clamp-based dead gradients. - Only ``i0_raw`` is projected; slope terms remain unconstrained. + ``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)): @@ -752,6 +972,10 @@ def enforce_hard_constraints(self, ctx: RenderContext) -> 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) @@ -785,6 +1009,28 @@ def _mask_keep( 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 torch.clamp(inten, min=0.0) + + def constraint_loss( + self, ctx: RenderContext, params: dict[str, object] | None = None + ) -> torch.Tensor: + if self.slope_l2_weight <= 0.0 or self.ir is None or self.ic is None: + return torch.zeros((), device=ctx.device, dtype=ctx.dtype) + ir = self.ir.to(device=ctx.device, dtype=ctx.dtype) + ic = self.ic.to(device=ctx.device, dtype=ctx.dtype) + return self.slope_l2_weight * torch.mean(ir * ir + ic * ic) + def forward(self, ctx: RenderContext) -> torch.Tensor: if self.origin is None: raise RuntimeError("SyntheticDiskLattice requires an OriginND instance.") @@ -848,6 +1094,8 @@ def forward(self, ctx: RenderContext) -> torch.Tensor: 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) @@ -855,15 +1103,23 @@ def forward(self, ctx: RenderContext) -> torch.Tensor: 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) - _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) - ) + 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( @@ -949,6 +1205,8 @@ def forward_batched( + 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 @@ -956,6 +1214,11 @@ def forward_batched( 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, @@ -964,7 +1227,13 @@ def forward_batched( 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('disk.') and param.requires_grad: + 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 {} diff --git a/src/quantem/diffraction/bragg_vectors.py b/src/quantem/diffraction/bragg_vectors.py index 344910ba4..5c7bfdcd8 100644 --- a/src/quantem/diffraction/bragg_vectors.py +++ b/src/quantem/diffraction/bragg_vectors.py @@ -211,7 +211,7 @@ def make_template_synthetic( radius: float | None = None, edge: float = 1.0, center: tuple[float, float] | None = None, - subtract_mean: bool = True, + subtract_mean: bool = False, ) -> "BraggVectors": """Build the template from a synthetic soft-edged disk. @@ -227,9 +227,10 @@ def make_template_synthetic( center : tuple of float, optional ``(row, col)`` disk center; defaults to the detector center ``(H // 2, W // 2)``. - subtract_mean : bool, default=True - If ``True``, make the template zero-sum — a band-pass kernel that - suppresses uniform background in the correlation. + subtract_mean : bool, default=False + If ``True``, make the template zero-sum. The default keeps the + unit-sum positive template, so correlation values stay positive and + roughly measure the probe-weighted counts under each peak. Returns ------- @@ -257,7 +258,7 @@ def make_template_synthetic( def make_template_from_data( self, roi: NDArray | None = None, - subtract_mean: bool = True, + subtract_mean: bool = False, center: tuple[float, float] | None = None, ) -> "BraggVectors": """Build the template by averaging diffraction patterns from the data. @@ -269,9 +270,10 @@ def make_template_from_data( ideally a vacuum / single-disk region so the unscattered probe is isolated. ``None`` (default) averages the whole scan (the mean diffraction pattern). - subtract_mean : bool, default=True - If ``True``, make the template zero-sum — a band-pass kernel that - suppresses uniform background in the correlation. + subtract_mean : bool, default=False + If ``True``, make the template zero-sum. The default keeps the + unit-sum positive template, so correlation values stay positive and + roughly measure the probe-weighted counts under each peak. center : tuple of float, optional ``(row, col)`` probe center rolled to the origin; defaults to the probe's intensity centroid. @@ -318,7 +320,7 @@ def make_template_from_probe( self, probe: NDArray | torch.Tensor, center: tuple[float, float] | None = None, - subtract_mean: bool = True, + subtract_mean: bool = False, ) -> "BraggVectors": """Build the template from an explicit probe image (e.g. a measured vacuum probe). @@ -329,9 +331,10 @@ def make_template_from_probe( center : tuple of float, optional ``(row, col)`` probe center rolled to the origin; defaults to the probe's intensity centroid. - subtract_mean : bool, default=True - If ``True``, make the template zero-sum — a band-pass kernel that - suppresses uniform background in the correlation. + subtract_mean : bool, default=False + If ``True``, make the template zero-sum. The default keeps the + unit-sum positive template, so correlation values stay positive and + roughly measure the probe-weighted counts under each peak. Returns ------- @@ -366,7 +369,9 @@ def template(self) -> np.ndarray | None: return None return torch.fft.fftshift(self._template).detach().cpu().numpy() - def correlation_map(self, row: int, col: int) -> np.ndarray: + def correlation_map( + self, row: int, col: int, background_sigma: float | str | None = None + ) -> np.ndarray: """Cross-correlation map of one diffraction pattern with the template (numpy). Peaks in the returned map sit at absolute disk positions (no fftshift @@ -389,7 +394,9 @@ def correlation_map(self, row: int, col: int) -> np.ndarray: dp = torch.as_tensor( np.asarray(self.dataset.array[row, col]), dtype=torch.float, device=self.device ) - corr, _ = cross_correlation(dp, self._template_ft) + corr, _ = cross_correlation( + dp, self._template_ft, self._resolve_background_sigma(background_sigma) + ) return corr.detach().cpu().numpy() def detect_disks( @@ -402,6 +409,7 @@ def detect_disks( subpixel: str = "upsample", upsample_factor: int = 16, max_num_peaks: int = 1000, + background_sigma: float | str | None = None, batch_size: int | None = None, progressbar: bool = True, save_to_gpu: bool = True, @@ -459,6 +467,7 @@ def detect_disks( subpixel=subpixel, upsample_factor=upsample_factor, max_num_peaks=max_num_peaks, + background_sigma=self._resolve_background_sigma(background_sigma), ) if save_to_gpu and self.device != "cpu": @@ -498,6 +507,78 @@ def detect_disks( self.compute_bvm() return peaks + def correct_peak_origins( + self, + origins: NDArray, + origin_ref: NDArray | tuple[float, float] | None = None, + *, + inplace: bool = False, + ) -> "BraggVectors": + """Shift the detected peak coordinates so all positions share one origin. + + Subtracts each scan position's measured diffraction origin (e.g. the + plane-fitted center-of-mass of the central beam -- the descan) from its + peaks and adds back a common reference ``origin_ref``, so the peak + coordinates from every position live in a single detector frame. The + diffraction data itself is untouched: this calibrates the measurements, + not the images. The Bragg vector map is recomputed from the corrected + peaks. + + By default a corrected *copy* of the workflow is returned and ``self`` + keeps the raw detections, so calling this repeatedly (e.g. re-running a + notebook cell) never double-applies the shift. + + Parameters + ---------- + origins : np.ndarray + ``(scan_row, scan_col, 2)`` per-position diffraction origins in + detector pixels (row, col). + 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 + 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. + """ + 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]) + 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}.") + if origin_ref is None: + origin_ref = origins.mean(axis=(0, 1)) + origin_ref = np.asarray(origin_ref, dtype=float).reshape(2) + + if inplace: + bv = self + else: + bv = type(self)(dataset=self.dataset, device=self.device, _token=type(self)._token) + bv._template = self._template + bv._template_ft = self._template_ft + bv.metadata = _copy.deepcopy(self.metadata) + bv.peaks = self.peaks.copy() + + peaks = bv.peaks + # rowwise transform on the flat peak table: one shift per scan cell, + # repeated per detected peak (cells and shifts share raster order) + flat = peaks.flatten() + counts = np.asarray(peaks.row_counts(), dtype=int) + shifts = np.repeat(origin_ref[None, :] - origins.reshape(-1, 2), counts, axis=0) + flat[:, :2] += shifts + peaks.set_flattened(flat) + + bv.metadata["origin_correction"] = { + "origin_ref": (float(origin_ref[0]), float(origin_ref[1])), + } + bv.compute_bvm() + return bv + def compute_bvm(self, sampling: float = 1.0) -> Dataset2d: """Accumulate all detected peaks into a Bragg vector map (intensity histogram). @@ -1166,6 +1247,7 @@ def show_detection( subpixel: str = "upsample", upsample_factor: int = 16, max_num_peaks: int = 1000, + background_sigma: float | str | None = None, image: np.ndarray | None = None, peak_radius: float = 6.0, marker_radius: float | None = None, @@ -1237,6 +1319,7 @@ def show_detection( subpixel=subpixel, upsample_factor=upsample_factor, max_num_peaks=max_num_peaks, + background_sigma=background_sigma, progressbar=False, ) if image is None: @@ -1295,6 +1378,31 @@ def peak_histogram(self, *, returnfig: bool = False, **kwargs): # ---- helpers ---- + def _resolve_background_sigma( + self, background_sigma: float | str | None + ) -> float | None: + """Resolve the ``background_sigma`` argument to a value in pixels. + + ``"auto"`` (the default everywhere) maps to twice the central-beam + radius: wide enough that the disk-scale correlation peaks pass + untouched, narrow enough to remove the zero-sum template's negative + moat around a bright unscattered beam -- which otherwise pushes weak + disk peaks below zero, where the correlation clamp erases them before + peak finding. Pass ``None`` to disable the background subtraction or a + float to set the scale explicitly. + """ + if background_sigma is None: + return None + if isinstance(background_sigma, str): + if background_sigma != "auto": + raise ValueError("background_sigma must be a float, None, or 'auto'.") + radius = self.metadata.get("template", {}).get("radius") + if radius is None: + dp_mean = np.asarray(self.dataset.dp_mean.array) + _, radius = estimate_central_beam(dp_mean) + return 2.0 * float(radius) + return float(background_sigma) + def _set_template( self, probe: torch.Tensor, diff --git a/src/quantem/diffraction/disk_detection.py b/src/quantem/diffraction/disk_detection.py index 4af9534f8..e857a3e9a 100644 --- a/src/quantem/diffraction/disk_detection.py +++ b/src/quantem/diffraction/disk_detection.py @@ -237,9 +237,51 @@ def template_fourier(template: torch.Tensor) -> torch.Tensor: return torch.conj(torch.fft.fft2(template)) +def _background_highpass( + shape: tuple[int, int], + background_sigma: float, + device, + rfft: bool = False, +) -> torch.Tensor: + """Fourier-domain high-pass ``1 - G`` removing correlation background. + + ``G`` is the transform of a real-space Gaussian of standard deviation + ``background_sigma`` pixels, so multiplying the Fourier product by ``1 - G`` + subtracts a Gaussian-smoothed copy of the correlation map -- the slowly + varying background (the zero-sum template's negative moat around the bright + central beam) that otherwise pushes weak disk peaks below zero, where the + ``relu`` clamp erases them before peak finding. + + Parameters + ---------- + shape : tuple of int + ``(H, W)`` detector shape. + background_sigma : float + Real-space standard deviation of the subtracted background, in pixels. + device + Torch device for the filter tensor. + rfft : bool, default=False + If ``True``, return the ``(H, W // 2 + 1)`` half-plane filter for + ``rfft2`` products instead of the full ``(H, W)`` filter. + + Returns + ------- + torch.Tensor + The ``1 - G`` filter in the requested Fourier layout. + """ + H, W = int(shape[0]), int(shape[1]) + qr = torch.fft.fftfreq(H, device=device, dtype=torch.float)[:, None] + if rfft: + qc = torch.fft.rfftfreq(W, device=device, dtype=torch.float)[None, :] + else: + qc = torch.fft.fftfreq(W, device=device, dtype=torch.float)[None, :] + g = torch.exp(-2.0 * (torch.pi**2) * (float(background_sigma) ** 2) * (qr**2 + qc**2)) + return 1.0 - g + def cross_correlation( dp: torch.Tensor, template_ft: torch.Tensor, + background_sigma: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Cross-correlate a diffraction pattern with a template. @@ -261,6 +303,8 @@ def cross_correlation( """ dp = torch.as_tensor(dp) m = torch.fft.fft2(dp) * template_ft + if background_sigma is not None and background_sigma > 0: + m = m * _background_highpass(m.shape[-2:], background_sigma, m.device) corr_map = torch.clamp(torch.fft.ifft2(m).real, min=0.0) return corr_map, m @@ -268,6 +312,7 @@ def cross_correlation( def cross_correlation_batch( dps: torch.Tensor, template_ft: torch.Tensor, + background_sigma: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Cross-correlate a stack of diffraction patterns with one template. @@ -291,11 +336,17 @@ def cross_correlation_batch( """ dps = torch.as_tensor(dps) m = torch.fft.fft2(dps) * template_ft + if background_sigma is not None and background_sigma > 0: + m = m * _background_highpass(m.shape[-2:], background_sigma, m.device) corr_map = torch.clamp(torch.fft.ifft2(m).real, min=0.0) return corr_map, m -def _corr_map_rfft(dps: torch.Tensor, template_ft: torch.Tensor) -> torch.Tensor: +def _corr_map_rfft( + dps: torch.Tensor, + template_ft: torch.Tensor, + background_sigma: float | None = None, +) -> torch.Tensor: """Real-FFT correlation map(s), used when no Fourier product is needed downstream. For real ``dps`` and a real template the Fourier product is conjugate-symmetric, @@ -318,6 +369,8 @@ def _corr_map_rfft(dps: torch.Tensor, template_ft: torch.Tensor) -> torch.Tensor dps = torch.as_tensor(dps) H, W = dps.shape[-2], dps.shape[-1] prod = torch.fft.rfft2(dps) * template_ft[..., : W // 2 + 1] + if background_sigma is not None and background_sigma > 0: + prod = prod * _background_highpass((H, W), background_sigma, prod.device, rfft=True) corr_map = torch.fft.irfft2(prod, s=(H, W)) return torch.clamp(corr_map, min=0.0) @@ -332,6 +385,7 @@ def detect_disks( subpixel: str = "upsample", upsample_factor: int = 16, max_num_peaks: int = 1000, + background_sigma: float | None = None, ) -> np.ndarray: """Detect Bragg disks in one diffraction pattern by template matching. @@ -366,7 +420,7 @@ def detect_disks( if subpixel not in SUBPIXEL_MODES: raise ValueError(f"subpixel must be in {SUBPIXEL_MODES}, got {subpixel!r}") - corr_map, m = cross_correlation(dp, template_ft) + corr_map, m = cross_correlation(dp, template_ft, background_sigma) peaks = _local_maxima(corr_map, edge_boundary) peaks = _filter_maxima(peaks, min_abs_intensity, min_spacing, max_num_peaks) @@ -393,6 +447,7 @@ def detect_disks_batch( subpixel: str = "upsample", upsample_factor: int = 16, max_num_peaks: int = 1000, + background_sigma: float | None = None, ) -> list[np.ndarray]: """Detect Bragg disks across a stack of diffraction patterns (batched). @@ -437,9 +492,9 @@ def detect_disks_batch( raise ValueError(f"subpixel must be in {SUBPIXEL_MODES}, got {subpixel!r}") if subpixel == "upsample": - corr_map, m = cross_correlation_batch(dps, template_ft) + corr_map, m = cross_correlation_batch(dps, template_ft, background_sigma) else: - corr_map = _corr_map_rfft(dps, template_ft) + corr_map = _corr_map_rfft(dps, template_ft, background_sigma) m = None peaks_all, bidx, counts = _detect_peaks_batched( diff --git a/src/quantem/diffraction/model_fitting.py b/src/quantem/diffraction/model_fitting.py index c5f940a4e..12f2cd1a7 100644 --- a/src/quantem/diffraction/model_fitting.py +++ b/src/quantem/diffraction/model_fitting.py @@ -14,9 +14,11 @@ from quantem.core.fitting.base import ( AdditiveRenderModel, FitBase, + LogMSELoss, OriginND, RenderComponent, RenderContext, + SqrtMSELoss, ) from quantem.core.fitting.diffraction import DiskTemplate, SyntheticDiskLattice from quantem.core.io.serialize import AutoSerialize @@ -325,6 +327,7 @@ def preprocess( self, *, align: bool = False, + origins: Any = None, edge_blend: float = 8.0, upsample_factor: int = 32, max_shift: float | None = None, @@ -393,6 +396,29 @@ def preprocess( 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 @@ -606,6 +632,11 @@ def reset( 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: @@ -626,6 +657,7 @@ def fit_individual_diffraction_pattern( 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, @@ -645,6 +677,7 @@ def fit_individual_diffraction_pattern( 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, @@ -727,6 +760,7 @@ def fit_individual_diffraction_pattern_batched( 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, @@ -901,47 +935,57 @@ def fit_individual_diffraction_pattern_batched( 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) + 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) - # Per-sample fidelity loss summed → scalar with per-sample grads - diff2 = (pred.float() - targets.float()) - # Match SqrtMSELoss behavior approximately when loss_fn is SqrtMSELoss: - # gamma-power transform of (x - min(x) + 1), per-sample independently. - from quantem.core.fitting.base import SqrtMSELoss, LogMSELoss - if isinstance(loss_fn, SqrtMSELoss): - gamma = float(loss_fn.gamma) - eps = 1.0 + + if is_sqrt: if fidelity_mask is not None: pred_min = pred.masked_fill(~fidelity_mask, float("inf")).amin(dim=(1, 2), keepdim=True) - tgt_min = targets.masked_fill(~fidelity_mask, float("inf")).amin(dim=(1, 2), keepdim=True) else: pred_min = pred.amin(dim=(1, 2), keepdim=True) - tgt_min = targets.amin(dim=(1, 2), keepdim=True) - pred_mod = (pred - pred_min + eps) ** gamma - tgt_mod = (targets - tgt_min + eps) ** gamma - sq = (pred_mod - tgt_mod) ** 2 - if fidelity_mask is not None: - per_sample_loss = (sq * fidelity_mask).sum(dim=(1, 2)) / fidelity_valid_count - else: - per_sample_loss = sq.mean(dim=(1, 2)) - elif isinstance(loss_fn, LogMSELoss): - sq = (torch.log1p(pred) - torch.log1p(targets)) ** 2 - if fidelity_mask is not None: - per_sample_loss = (sq * fidelity_mask).sum(dim=(1, 2)) / fidelity_valid_count - else: - per_sample_loss = sq.mean(dim=(1, 2)) - elif isinstance(loss_fn, torch.nn.L1Loss): - ad = diff2.abs() - if fidelity_mask is not None: - per_sample_loss = (ad * fidelity_mask).sum(dim=(1, 2)) / fidelity_valid_count - else: - per_sample_loss = ad.mean(dim=(1, 2)) + pred_mod = (pred - pred_min + sqrt_eps) ** sqrt_gamma + elif is_log: + pred_mod = torch.log1p(pred) else: - sq = diff2 * diff2 - if fidelity_mask is not None: - per_sample_loss = (sq * fidelity_mask).sum(dim=(1, 2)) / fidelity_valid_count - else: - per_sample_loss = sq.mean(dim=(1, 2)) + pred_mod = pred + + diff_mod = pred_mod - tgt_mod + val = diff_mod.abs() if is_l1 else diff_mod * diff_mod + 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: diff --git a/src/quantem/diffraction/strain_autocorrelation.py b/src/quantem/diffraction/strain_autocorrelation.py index 5e63e69cb..1dbb316d7 100644 --- a/src/quantem/diffraction/strain_autocorrelation.py +++ b/src/quantem/diffraction/strain_autocorrelation.py @@ -228,9 +228,31 @@ def diffraction_mask( 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: + import warnings + + 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( - distance_transform_edt(np.logical_not(mask_init)) / edge_blend, + edge_distance / edge_blend, 0.0, 1.0, ) From a42311ba2bfb5386523c262115e060006f9dded7 Mon Sep 17 00:00:00 2001 From: mitis1 Date: Thu, 30 Jul 2026 23:23:29 -0700 Subject: [PATCH 6/7] bug fixes from previous commit --- src/quantem/core/fitting/diffraction.py | 2 +- src/quantem/diffraction/bragg_vectors.py | 1 + src/quantem/diffraction/model_fitting.py | 38 +++++++++++++++++-- .../diffraction/strain_visualization.py | 4 +- 4 files changed, 39 insertions(+), 6 deletions(-) diff --git a/src/quantem/core/fitting/diffraction.py b/src/quantem/core/fitting/diffraction.py index 1c8c15983..f1929c83c 100644 --- a/src/quantem/core/fitting/diffraction.py +++ b/src/quantem/core/fitting/diffraction.py @@ -1020,7 +1020,7 @@ def _apply_positivity(self, inten: torch.Tensor) -> torch.Tensor: if self.intensity_softplus_beta is not None: beta = self.intensity_softplus_beta return F.softplus(beta * inten) / beta - return torch.clamp(inten, min=0.0) + return (inten) def constraint_loss( self, ctx: RenderContext, params: dict[str, object] | None = None diff --git a/src/quantem/diffraction/bragg_vectors.py b/src/quantem/diffraction/bragg_vectors.py index 5c7bfdcd8..24d7b2877 100644 --- a/src/quantem/diffraction/bragg_vectors.py +++ b/src/quantem/diffraction/bragg_vectors.py @@ -1,5 +1,6 @@ from __future__ import annotations +import copy as _copy from pathlib import Path from typing import Any, Literal, Sequence, Union diff --git a/src/quantem/diffraction/model_fitting.py b/src/quantem/diffraction/model_fitting.py index 12f2cd1a7..0b7349873 100644 --- a/src/quantem/diffraction/model_fitting.py +++ b/src/quantem/diffraction/model_fitting.py @@ -1897,10 +1897,32 @@ def apply_hard_constraints( if self.lats is not None: for lat_name, lat in zip(self.lat_names, self.lats): - for pname, (lo, hi) in lat.parameter_bounds.items(): - key = f"{lat_name}.{pname}" - if key in stacked and key not in skip_keys: - self._clamp_bounds_inplace(stacked[key], lo, hi) + 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 @@ -2189,6 +2211,14 @@ def batched_constraint_loss( 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): + w = float(getattr(lat, "slope_l2_weight", 0.0)) + if w <= 0.0: + continue + ir = stacked.get(f"{lat_name}.ir") + ic = stacked.get(f"{lat_name}.ic") + if ir is not None and ic is not None: + out = out + w * (ir * ir + ic * ic).mean(dim=1) return out def resolve_component_keys(self, components: Any) -> list[str]: diff --git a/src/quantem/diffraction/strain_visualization.py b/src/quantem/diffraction/strain_visualization.py index 819891cb7..c4b3258d2 100644 --- a/src/quantem/diffraction/strain_visualization.py +++ b/src/quantem/diffraction/strain_visualization.py @@ -154,7 +154,9 @@ def _roi_compose(norm_vals, color_cm): r"$\epsilon_{vv}$", r"$\epsilon_{uv}$", ) - title_arrow_angles = (0 + strain_rotation_angle, 90 + strain_rotation_angle, -45 + strain_rotation_angle) + title_arrow_angles = (90 + strain_rotation_angle, 0 + strain_rotation_angle, -45 + strain_rotation_angle) + if transpose_image: + title_arrow_angles = (0 + strain_rotation_angle, 90 + strain_rotation_angle, 45 + strain_rotation_angle) else: title_arrow_angles = (None, None, None) From c2a9d73c7d9e0325a62907092fb96da7d597d4f9 Mon Sep 17 00:00:00 2001 From: mitis1 Date: Sat, 8 Aug 2026 18:06:46 -0700 Subject: [PATCH 7/7] initial data consistency implementation --- src/quantem/core/fitting/base.py | 10 +- src/quantem/core/fitting/diffraction.py | 85 +++++++++++++-- src/quantem/diffraction/model_fitting.py | 132 +++++++++++++++++++++-- 3 files changed, 205 insertions(+), 22 deletions(-) diff --git a/src/quantem/core/fitting/base.py b/src/quantem/core/fitting/base.py index b0a1f620a..6050cd96b 100644 --- a/src/quantem/core/fitting/base.py +++ b/src/quantem/core/fitting/base.py @@ -452,11 +452,15 @@ def apply_hard_constraints(self, ctx: RenderContext) -> None: component = cast(RenderComponent, module) component.enforce_hard_constraints(ctx) - def total_constraint_loss(self, ctx: RenderContext) -> torch.Tensor: + 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) - loss = loss + component.constraint_loss(ctx) + 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, @@ -1057,7 +1061,7 @@ def _constraint_loss( ) -> 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) + return self.model.total_constraint_loss(self.ctx, **kwargs) def set_component_trainable( self, diff --git a/src/quantem/core/fitting/diffraction.py b/src/quantem/core/fitting/diffraction.py index f1929c83c..cd01e6f1a 100644 --- a/src/quantem/core/fitting/diffraction.py +++ b/src/quantem/core/fitting/diffraction.py @@ -671,6 +671,10 @@ 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, @@ -1022,15 +1026,6 @@ def _apply_positivity(self, inten: torch.Tensor) -> torch.Tensor: return F.softplus(beta * inten) / beta return (inten) - def constraint_loss( - self, ctx: RenderContext, params: dict[str, object] | None = None - ) -> torch.Tensor: - if self.slope_l2_weight <= 0.0 or self.ir is None or self.ic is None: - return torch.zeros((), device=ctx.device, dtype=ctx.dtype) - ir = self.ir.to(device=ctx.device, dtype=ctx.dtype) - ic = self.ic.to(device=ctx.device, dtype=ctx.dtype) - return self.slope_l2_weight * torch.mean(ir * ir + ic * ic) - def forward(self, ctx: RenderContext) -> torch.Tensor: if self.origin is None: raise RuntimeError("SyntheticDiskLattice requires an OriginND instance.") @@ -1237,4 +1232,74 @@ def get_optimization_parameters(self) -> dict[str, list[torch.nn.Parameter]]: params.append(param) if not params: return {} - return {'default': params} \ No newline at end of file + 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: + cfg = self.effective_soft_constraints(cast(dict[str, object] | None, params)) + tv_weight = float(cfg.get("consistency_weight", 0.0)) + + 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/diffraction/model_fitting.py b/src/quantem/diffraction/model_fitting.py index 0b7349873..70249a352 100644 --- a/src/quantem/diffraction/model_fitting.py +++ b/src/quantem/diffraction/model_fitting.py @@ -725,6 +725,7 @@ def fit_individual_diffraction_pattern( 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), @@ -734,6 +735,7 @@ def fit_individual_diffraction_pattern( scheduler_params=scheduler_params, progress=False, run_key=f"individual_{r}_{c}", + neighbor_target=nt_single, ) s_fit = self._get_model_state_dict_copy() @@ -909,6 +911,10 @@ def fit_individual_diffraction_pattern_batched( 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())} @@ -989,7 +995,7 @@ def fit_individual_diffraction_pattern_batched( # Add per-sample soft constraint loss. if float(constraint_weight) != 0.0: - constraint_per_sample = plan.batched_constraint_loss(ctx, stacked) + 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() @@ -1325,6 +1331,108 @@ def _resolve_rows_cols_for_batched( 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, @@ -2203,8 +2311,8 @@ def batched_constraint_loss( self, ctx: RenderContext, stacked: dict[str, torch.Tensor], + neighbor_target=None ) -> torch.Tensor: - """Per-sample soft-constraint loss summed across components.""" 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: @@ -2212,13 +2320,19 @@ def batched_constraint_loss( 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): - w = float(getattr(lat, "slope_l2_weight", 0.0)) - if w <= 0.0: - continue - ir = stacked.get(f"{lat_name}.ir") - ic = stacked.get(f"{lat_name}.ic") - if ir is not None and ic is not None: - out = out + w * (ir * ir + ic * ic).mean(dim=1) + 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]: