Adding L1 loss, fixing mask + lattice

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(

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))
 

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,

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