diff --git a/src/quantem/core/fitting/diffraction.py b/src/quantem/core/fitting/diffraction.py index 2ac2d0f9..62997511 100644 --- a/src/quantem/core/fitting/diffraction.py +++ b/src/quantem/core/fitting/diffraction.py @@ -530,6 +530,7 @@ def __init__( center_intensity_0: float | Sequence[float] | None = None, exclude_indices: Iterable[tuple[int, int]] | None = None, boundary_px: float = 0.0, + min_frac_inside_mask: float | None = None, origin: OriginND | None = None, origin_key: str = "origin", constraint_params: dict[str, Any] | None = None, @@ -586,6 +587,12 @@ def __init__( self.u_max = int(u_max) self.v_max = int(v_max) self.boundary_px = float(boundary_px) + # Drop a lattice disk from the render/fit when less than this fraction + # of its (circular) template patch falls inside ctx.mask -- i.e. the disk + # has too little illuminated data to constrain it. None disables it. + self.min_frac_inside_mask = ( + None if min_frac_inside_mask is None else float(min_frac_inside_mask) + ) if max_intensity_order is None: max_intensity_order = 1 if bool(per_disk_slopes) else 0 @@ -748,6 +755,36 @@ def enforce_hard_constraints(self, ctx: RenderContext) -> None: super().enforce_hard_constraints(ctx) + def _mask_keep( + self, ctx: RenderContext, centers_r: torch.Tensor, centers_c: torch.Tensor + ) -> torch.Tensor | None: + """ + Per-disk keep mask from ``min_frac_inside_mask``. + + Returns a boolean tensor shaped like ``centers_r`` that is True where at + least ``min_frac_inside_mask`` of the disk's circular template patch lands + on a True pixel of ``ctx.mask``; returns None when the filter is disabled + or no mask is set. Off-frame patch pixels count as outside the mask. + """ + if self.min_frac_inside_mask is None or ctx.mask is None: + return None + mask = ctx.mask + h, w = int(mask.shape[0]), int(mask.shape[1]) + dr = cast(torch.Tensor, self.disk.dr).to(device=ctx.device) + dc = cast(torch.Tensor, self.disk.dc).to(device=ctx.device) + tv = self.disk.patch_values().detach() + supp = tv > 0.5 * tv.max().clamp(min=1e-12) + drs = dr[supp] + dcs = dc[supp] + if drs.numel() == 0: + return None + rr = (centers_r.reshape(-1)[:, None] + drs[None, :]).round().to(torch.long) + cc = (centers_c.reshape(-1)[:, None] + dcs[None, :]).round().to(torch.long) + in_frame = (rr >= 0) & (rr < h) & (cc >= 0) & (cc < w) + mval = mask[rr.clamp(0, h - 1), cc.clamp(0, w - 1)].to(torch.bool) & in_frame + frac = mval.to(torch.float32).mean(dim=1) + return (frac >= self.min_frac_inside_mask).reshape(centers_r.shape) + def forward(self, ctx: RenderContext) -> torch.Tensor: if self.origin is None: raise RuntimeError("SyntheticDiskLattice requires an OriginND instance.") @@ -767,6 +804,9 @@ def forward(self, ctx: RenderContext) -> torch.Tensor: b = torch.as_tensor(self.boundary_px, device=ctx.device, dtype=ctx.dtype) keep = (centers_r >= b) & (centers_r <= (ctx.shape[0] - 1) - b) keep = keep & (centers_c >= b) & (centers_c <= (ctx.shape[1] - 1) - b) + mk = self._mask_keep(ctx, centers_r, centers_c) + if mk is not None: + keep = keep & mk if not torch.any(keep): return out @@ -869,6 +909,9 @@ def forward_batched( bb = torch.as_tensor(self.boundary_px, device=ctx.device, dtype=ctx.dtype) keep = (r0_kb >= bb) & (r0_kb <= (ctx.shape[0] - 1) - bb) keep = keep & (c0_kb >= bb) & (c0_kb <= (ctx.shape[1] - 1) - bb) + mk = self._mask_keep(ctx, r0_kb, c0_kb) + if mk is not None: + keep = keep & mk keep_f = keep.to(dtype=ctx.dtype) active_order = int( diff --git a/src/quantem/diffraction/model_fitting.py b/src/quantem/diffraction/model_fitting.py index 5545a4cc..768925b7 100644 --- a/src/quantem/diffraction/model_fitting.py +++ b/src/quantem/diffraction/model_fitting.py @@ -306,6 +306,9 @@ def get_overlay_coordinates(self) -> tuple[np.ndarray, np.ndarray]: ) keep = (centers_r >= b) & (centers_r <= (self.ctx.shape[0] - 1) - b) keep = keep & (centers_c >= b) & (centers_c <= (self.ctx.shape[1] - 1) - b) + mk = component._mask_keep(self.ctx, centers_r, centers_c) + if mk is not None: + keep = keep & mk if torch.any(keep): rc = torch.stack((centers_r[keep], centers_c[keep]), dim=1) centers.append(rc.detach().cpu().numpy().astype(np.float32, copy=False)) @@ -907,6 +910,8 @@ def fit_individual_diffraction_pattern_batched( per_sample_loss = ((pred_mod - tgt_mod) ** 2).mean(dim=(1, 2)) elif isinstance(loss_fn, LogMSELoss): per_sample_loss = ((torch.log1p(pred) - torch.log1p(targets)) ** 2).mean(dim=(1, 2)) + elif isinstance(loss_fn, torch.nn.L1Loss): + per_sample_loss = diff2.abs().mean(dim=(1, 2)) else: per_sample_loss = (diff2 * diff2).mean(dim=(1, 2)) diff --git a/src/quantem/diffraction/strain.py b/src/quantem/diffraction/strain.py index e4b9a70c..c99e1431 100644 --- a/src/quantem/diffraction/strain.py +++ b/src/quantem/diffraction/strain.py @@ -304,7 +304,8 @@ def plot_strain_roi( def plot_strain( self, - rotation_angle: float = 20.0, + rotation_angle_deg: float = 0.0, + transpose: bool = False, 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), @@ -324,9 +325,15 @@ def plot_strain( Parameters ---------- - rotation_angle : float, default=20.0 + 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 + 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) @@ -359,12 +366,23 @@ def plot_strain( tuple ``(fig, ax)`` from :func:`plot_strain_panels`. """ - e_uu, e_vv, e_uv = self.rotate_strain(rotation_angle) + e_rr = self.e_rr.array + e_cc = self.e_cc.array + e_rc = self.e_rc.array + phi = self.phi.array + if transpose: + # 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) return plot_strain_panels( e_uu, e_vv, e_uv, - self.phi.array, + phi, self.mask, self.u_ref, self.v_ref, diff --git a/src/quantem/diffractive_imaging/origin_models.py b/src/quantem/diffractive_imaging/origin_models.py index b74b2f27..7353ff30 100644 --- a/src/quantem/diffractive_imaging/origin_models.py +++ b/src/quantem/diffractive_imaging/origin_models.py @@ -309,9 +309,40 @@ def _estimate_detector_rotation( return rotation_curl def estimate_detector_rotation( - self, rotation_angles_deg: torch.Tensor | NDArray | None = None + self, + rotation_angles_deg: torch.Tensor | NDArray | None = None, + print_result: bool = False, + plot_result: bool = False, + **plot_kwargs, ): - """ """ + """Estimate the detector->scan rotation (and transpose) by curl minimization. + + Sweeps ``rotation_angles_deg`` and, for both the un-transposed and transposed + center-of-mass field, computes the mean absolute curl of the rotated CoM. The + physical orientation makes the CoM field curl-free, so the global minimum over + both branches selects the rotation angle and whether the detector axes are + transposed. Results are stored in :attr:`detector_rotation_deg` and + :attr:`detector_transpose`. + + Parameters + ---------- + rotation_angles_deg : torch.Tensor or NDArray, optional + Angles (degrees) to search. Defaults to ``arange(-89, 90, 1)``. + print_result : bool, default=False + If ``True``, print the minimum curl and angle of each branch and the + chosen orientation. + plot_result : bool, default=False + If ``True``, plot the curl-vs-rotation curves for both branches, marking + each branch minimum. Useful for judging whether the transpose choice is + trustworthy (a clear, deep minimum on one branch) or ambiguous. + **plot_kwargs + Forwarded to :func:`matplotlib.pyplot.subplots` when ``plot_result``. + + Returns + ------- + CenterOfMassOriginModel + ``self``. + """ if rotation_angles_deg is None: rotation_angles_deg = torch.arange(-89, 90, 1, device=self.device, dtype=torch.float) @@ -343,6 +374,31 @@ def estimate_detector_rotation( self._detector_rotation_deg = rotation_angles_deg[ind_min].item() + if print_result or plot_result: + a = rotation_angles_deg.detach().cpu().numpy() + cn = curl_no_transpose.detach().cpu().numpy() + ct = curl_transpose.detach().cpu().numpy() + i_no, i_tr = int(cn.argmin()), int(ct.argmin()) + + if print_result: + chosen = "TRANSPOSE" if self._detector_transpose else "NO transpose" + print(f"no-transpose : min curl {cn[i_no]:.4g} at {a[i_no]:+.2f} deg") + print(f"transpose : min curl {ct[i_tr]:.4g} at {a[i_tr]:+.2f} deg") + print(f"=> chosen: {chosen} at {self._detector_rotation_deg:+.2f} deg") + + if plot_result: + import matplotlib.pyplot as plt + + _fig, ax = plt.subplots(**{"figsize": (6, 4), **plot_kwargs}) + line_no = ax.plot(a, cn, label="no transpose")[0] + line_tr = ax.plot(a, ct, label="transpose")[0] + ax.plot(a[i_no], cn[i_no], "o", color=line_no.get_color()) + ax.plot(a[i_tr], ct[i_tr], "o", color=line_tr.get_color()) + ax.set_xlabel("rotation (deg)") + ax.set_ylabel("mean |curl| of CoM (lower = better)") + ax.set_title("Detector-rotation search") + ax.legend() + return self @property