Fitting models clean

src/quantem/core/datastructures/dataset4dstem.py

@@ -1,3 +1,4 @@
+from os import PathLike
 from typing import Any, Self
 
 import matplotlib.pyplot as plt
@@ -97,15 +98,15 @@ def __init__(
         self._virtual_detectors = {}  # Store detector information for regeneration
 
     @classmethod
-    def from_file(cls, file_path: str, file_type: str) -> "Dataset4dstem":
+    def from_file(cls, file_path: str | PathLike, file_type: str | None = None) -> "Dataset4dstem":
         """
         Create a new Dataset4dstem from a file.
 
         Parameters
         ----------
-        file_path : str
+        file_path : str | PathLike
             Path to the data file
-        file_type : str
+        file_type : str | None
             The type of file reader needed. See rosettasciio for supported formats
             https://hyperspy.org/rosettasciio/supported_formats/index.html
 

src/quantem/core/fitting/__init__.py

@@ -0,0 +1,21 @@
+from quantem.core.fitting.background import DCBackground as DCBackground
+from quantem.core.fitting.background import GaussianBackground as GaussianBackground
+from quantem.core.fitting.base import Component as Component
+from quantem.core.fitting.base import Model as Model
+from quantem.core.fitting.base import ModelContext as ModelContext
+from quantem.core.fitting.base import OriginND as OriginND
+from quantem.core.fitting.base import Parameter as Parameter
+from quantem.core.fitting.diffraction import DiskTemplate as DiskTemplate
+from quantem.core.fitting.diffraction import SyntheticDiskLattice as SyntheticDiskLattice
+
+__all__ = [
+    "Component",
+    "DCBackground",
+    "DiskTemplate",
+    "GaussianBackground",
+    "Model",
+    "ModelContext",
+    "OriginND",
+    "Parameter",
+    "SyntheticDiskLattice",
+]

src/quantem/core/fitting/background.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+import torch
+from torch import nn
+
+from quantem.core.fitting.base import OriginND, RenderComponent, RenderContext
+
+
+class DCBackground(RenderComponent):
+    def __init__(
+        self,
+        *,
+        intensity: float | int | Sequence[float | int | None] = 0.0,
+        name: str = "dc_background",
+        constraint_params: dict[str, Any] | None = None,
+    ):
+        """
+        Build a constant background component.
+
+        Notes
+        -----
+        Validity is enforced via hard constraints/parameter bounds. Forward
+        intentionally avoids hard clamps for gradient flow.
+        """
+        super().__init__()
+        self.name = str(name)
+        intensity_init, intensity_lo, intensity_hi = self.parse_bounded_init(
+            intensity, name="intensity"
+        )
+        self.intensity_raw = nn.Parameter(torch.tensor(intensity_init, dtype=torch.float32))
+        bounded_lo = 0.0 if intensity_lo is None else max(float(intensity_lo), 0.0)
+        self.register_parameter_bounds("intensity_raw", bounded_lo, intensity_hi)
+        if constraint_params is not None:
+            self.apply_constraint_params(constraint_params, strict=True)
+        self._enforce_parameter_bounds()
+
+    def forward(self, ctx: RenderContext) -> torch.Tensor:
+        """
+        Render constant background from raw trainable intensity.
+
+        Notes
+        -----
+        Validity is enforced via hard constraints/parameter bounds, not via
+        forward-time hard clamps.
+        """
+        inten = self.intensity_raw.to(device=ctx.device, dtype=ctx.dtype)
+        return torch.ones(ctx.shape, device=ctx.device, dtype=ctx.dtype) * inten
+
+
+class GaussianBackground(RenderComponent):  # TODO this should be N dimensional by default
+    def __init__(
+        self,
+        *,
+        sigma: float | int | Sequence[float | int | None] = (40.0, 5.0, None),
+        intensity: float | int | Sequence[float | int | None] = 0.0,
+        origin: OriginND | None = None,
+        origin_key: str = "origin",
+        name: str = "gaussian_background",
+        constraint_params: dict[str, Any] | None = None,
+    ):
+        """
+        Build a Gaussian background component centered at origin.
+
+        Notes
+        -----
+        ``sigma_raw`` and ``intensity_raw`` validity is enforced via hard
+        constraints/parameter bounds. Forward intentionally avoids hard clamps
+        for gradient flow.
+        """
+        super().__init__()
+        self.name = str(name)
+        self.origin = origin
+        self.origin_key = str(origin_key)
+        sigma_init, sigma_lo, sigma_hi = self.parse_bounded_init(sigma, name="sigma")
+        intensity_init, intensity_lo, intensity_hi = self.parse_bounded_init(
+            intensity, name="intensity"
+        )
+        self.sigma_raw = nn.Parameter(torch.tensor(sigma_init, dtype=torch.float32))
+        sigma_bounded_lo = 1e-6 if sigma_lo is None else max(float(sigma_lo), 1e-6)
+        self.register_parameter_bounds("sigma_raw", sigma_bounded_lo, sigma_hi)
+        self.intensity_raw = nn.Parameter(torch.tensor(intensity_init, dtype=torch.float32))
+        intensity_bounded_lo = 0.0 if intensity_lo is None else max(float(intensity_lo), 0.0)
+        self.register_parameter_bounds("intensity_raw", intensity_bounded_lo, intensity_hi)
+        if constraint_params is not None:
+            self.apply_constraint_params(constraint_params, strict=True)
+        self._enforce_parameter_bounds()
+
+    def set_origin(self, origin: OriginND) -> None:
+        self.origin = origin
+
+    def forward(self, ctx: RenderContext) -> torch.Tensor:
+        """
+        Render Gaussian background from raw trainable parameters.
+
+        Notes
+        -----
+        Validity is enforced via hard constraints/parameter bounds, not via
+        forward-time hard clamps.
+        """
+        if self.origin is None:
+            raise RuntimeError("GaussianBackground requires an OriginND instance.")
+
+        rr = torch.arange(ctx.shape[0], device=ctx.device, dtype=ctx.dtype)[:, None]
+        cc = torch.arange(ctx.shape[1], device=ctx.device, dtype=ctx.dtype)[None, :]
+        r0, c0 = self.origin.coords[0], self.origin.coords[1]
+
+        sigma = self.sigma_raw.to(device=ctx.device, dtype=ctx.dtype)
+        inten = self.intensity_raw.to(device=ctx.device, dtype=ctx.dtype)
+        r2 = (rr - r0) ** 2 + (cc - c0) ** 2
+        return inten * torch.exp(-0.5 * r2 / (sigma * sigma))