diff --git a/.gitignore b/.gitignore index d87d85c8..a66c803d 100644 --- a/.gitignore +++ b/.gitignore @@ -167,6 +167,9 @@ cython_debug/ # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ +# BasedPyRight Config +pyrightconfig.json + # Ruff stuff: .ruff_cache/ diff --git a/src/quantem/__init__.py b/src/quantem/__init__.py index b9aa54b4..ba70f629 100644 --- a/src/quantem/__init__.py +++ b/src/quantem/__init__.py @@ -1,4 +1,5 @@ from pkgutil import extend_path + __path__ = extend_path(__path__, __name__) from importlib.metadata import version diff --git a/src/quantem/core/ml/constraints.py b/src/quantem/core/ml/constraints.py index 553b0611..590da4cb 100644 --- a/src/quantem/core/ml/constraints.py +++ b/src/quantem/core/ml/constraints.py @@ -1,13 +1,22 @@ from abc import ABC, abstractmethod from copy import deepcopy from dataclasses import dataclass -from typing import Any, Self +from typing import Any, Generic, Self, TypeVar import numpy as np import torch from numpy.typing import NDArray +@dataclass +class BaseContext(ABC): + """ + Constraints should contain a context object that contains all necessary data for the constraints to be applied. + """ + pass + +T_ctx = TypeVar("T_ctx", bound=BaseContext) + @dataclass(slots=False) class Constraints(ABC): """ @@ -47,7 +56,7 @@ def __str__(self) -> str: ) -class BaseConstraints(ABC): +class BaseConstraints(ABC, Generic[T_ctx]): """ Base class for constraints. """ @@ -86,14 +95,14 @@ def constraints(self, constraints: Constraints | dict[str, Any]): # --- Required methods tha tneeds to implemented in subclasses --- @abstractmethod - def apply_hard_constraints(self, *args, **kwargs) -> torch.Tensor: + def apply_hard_constraints(self, pred: torch.Tensor) -> torch.Tensor: """ Apply hard constraints to the model. """ raise NotImplementedError @abstractmethod - def apply_soft_constraints(self, *args, **kwargs) -> torch.Tensor: + def apply_soft_constraints(self, ctx: T_ctx) -> torch.Tensor: """ Apply soft constraints to the model. """ diff --git a/src/quantem/core/ml/models/__init__.py b/src/quantem/core/ml/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/quantem/core/ml/models/kplanes.py b/src/quantem/core/ml/models/kplanes.py new file mode 100644 index 00000000..cdf55261 --- /dev/null +++ b/src/quantem/core/ml/models/kplanes.py @@ -0,0 +1,732 @@ +""" +Tensor Decomposition Methods for INR-based reconstructions +""" + +import itertools +from typing import Callable, Optional, Sequence + +# import tinycudann as tcnn +import torch +import torch.nn.functional as F +from torch import nn + +from .model_base import PPLR, TensorDecompositionModel +from .so3params import SO3ParamQuat, SO3ParamR9SVD + +""" +K-planes utility functions +""" + + +def grid_sample_wrapper( + grid: torch.Tensor, coords: torch.Tensor, align_corners: bool = True +) -> torch.Tensor: + """ + Performs bilinear interpolation on a grid at given coordinates. + + Args: + grid: Grid tensor of shape [B, C, H, W] or [C, H, W] + coords: Coordinate tensor of shape [B, N, 2] or [N, 2] + align_corners: Whether to align corners + + Returns: + Interpolated values of shape [B, N, C] or [N, C] + """ + grid_dim = coords.shape[-1] + + if grid.dim() == grid_dim + 1: + # no batch dimension present, need to add it + grid = grid.unsqueeze(0) + if coords.dim() == 2: + coords = coords.unsqueeze(0) + + if grid_dim == 2 or grid_dim == 3: + grid_sampler = F.grid_sample + else: + raise NotImplementedError( + f"Grid-sample was called with {grid_dim}D data but is only " + f"implemented for 2 and 3D data." + ) + + coords = coords.view([coords.shape[0]] + [1] * (grid_dim - 1) + list(coords.shape[1:])) + B, feature_dim = grid.shape[:2] + n = coords.shape[-2] + interp = grid_sampler( + grid, # [B, feature_dim, reso, ...] + coords, # [B, 1, ..., n, grid_dim] + align_corners=align_corners, + mode="bilinear", + padding_mode="border", + ) + interp = interp.view(B, feature_dim, n).transpose(-1, -2) # [B, n, feature_dim] + interp = interp.squeeze() # [B?, n, feature_dim?] + return interp + + +def init_planes( + in_dim: int, + out_dim: int, + resolution: Sequence[int], + init_range: tuple = (0.1, 0.5), +) -> nn.ParameterList: + """Create the set of 2D planes for a k-plane decomposition. + + For in_dim=3 (spatial), this creates 3 planes: XY, XZ, YZ. + For in_dim=4 (spatial + time), this creates 6 planes: XY, XZ, XT, YZ, YT, ZT. + Time planes (those involving axis 3) are initialized to 1 so they start + as identity multipliers. + + Args: + in_dim: Dimensionality of the input coordinates (3 or 4). + out_dim: Number of feature channels per plane. + resolution: Resolution along each axis, e.g. [128, 128, 128]. + init_range: (a, b) for uniform initialization of spatial planes. + + Returns: + nn.ParameterList of plane parameters, each of shape [1, out_dim, res_j, res_i]. + """ + assert len(resolution) == in_dim + # All pairs of axes + axis_pairs = list(itertools.combinations(range(in_dim), 2)) + planes = nn.ParameterList() + a, b = init_range + for pair in axis_pairs: + # grid_sample expects (N, C, H, W) — so resolution is reversed + shape = [1, out_dim] + [resolution[ax] for ax in reversed(pair)] + param = nn.Parameter(torch.empty(*shape)) + # Time planes init to 1; spatial planes init uniform + if in_dim == 4 and 3 in pair: + nn.init.ones_(param) + else: + nn.init.uniform_(param, a=a, b=b) + planes.append(param) + return planes + + +def query_planes( + pts: torch.Tensor, + planes: nn.ParameterList, + in_dim: int, +) -> float: + """Query the k-plane representation at a batch of points. + + Projects each point onto every axis-pair plane, bilinearly interpolates, + and returns the element-wise product across all planes. + + Args: + pts: (B, in_dim) coordinates in [-1, 1]. + planes: The ParameterList from init_planes. + in_dim: 3 or 4. + + Returns: + (B, out_dim) features. + """ + axis_pairs = list(itertools.combinations(range(in_dim), 2)) + result = 1.0 + for plane_param, pair in zip(planes, axis_pairs): + # Extract the 2D coords for this plane + coords_2d = pts[..., list(pair)] # (B, 2) + coords_2d = coords_2d.view(1, -1, 1, 2) # (1, B, 1, 2) for grid_sample + # grid_sample: input (N,C,H,W), grid (N, H_out, W_out, 2) + sampled = F.grid_sample( + plane_param, # (1, C, H, W) + coords_2d, # (1, B, 1, 2) + align_corners=True, + mode="bilinear", + padding_mode="border", + ) # -> (1, C, B, 1) + sampled = sampled.squeeze(0).squeeze(-1).T # (B, C) + result = result * sampled + return result # pyright: ignore[reportReturnType] + + +def interpolate_ms_features( + pts: torch.Tensor, + ms_grids: nn.ParameterList, +) -> torch.Tensor: + mat_mode = [[0, 1], [0, 2], [1, 2]] + coord_plane = torch.stack( + [ + pts[:, mat_mode[0]], + pts[:, mat_mode[1]], + pts[:, mat_mode[2]], + ] + ).view(3, -1, 1, 2) + + per_scale = [] + for plane_coef in ms_grids: + C = plane_coef.shape[1] + feats = F.grid_sample( + plane_coef, coord_plane, align_corners=True, mode="bilinear", padding_mode="border" + ).reshape(3, C, -1) + fused = feats[0] * feats[1] * feats[2] + per_scale.append(fused.T) + + return torch.cat(per_scale, dim=-1) + + +class KPlanes(PPLR, TensorDecompositionModel): + """ + K-Planes model adapted from Fridovich-Keil et al., https://arxiv.org/abs/2301.10241 + """ + def __init__( + self, + # Grid parameters + grid_dimensions: int = 2, + input_coords_dims: int = 3, + M_features: int = 32, + resolution: Sequence[int] = (200, 200, 200), + multiscale_res_multipliers: Optional[Sequence[int]] = None, + concat_features: bool = True, + density_activation: Callable = lambda x: F.softplus( + x - 1 + ), # Keep playing around with this and trunc_exp + # Hybrid MLP parameters + use_hybrid_mlp: bool = False, + hybrid_hidden_dim: int = 64, + hybrid_num_layers: int = 2, + ): + """ + Assume coords are [-1, 1] in each dimension. + """ + super().__init__() + self._td_type = "kplanes" + self.grid_dimensions = grid_dimensions + self.input_coords_dims = input_coords_dims + self.M_features = M_features + self.resolution = resolution + self.multiscale_res_multipliers = multiscale_res_multipliers or [1] + self.concat_features = concat_features + self.density_activation = density_activation + + self.grids = nn.ParameterList() + self.feature_dim = 0 + for res_mult in self.multiscale_res_multipliers: + scaled_res = [int(r * res_mult) for r in self.resolution] + plane = nn.Parameter(torch.empty(3, self.M_features, scaled_res[1], scaled_res[0])) + nn.init.uniform_(plane, 0.1, 0.5) + self.grids.append(plane) + self.feature_dim += self.M_features + + # Network head + if use_hybrid_mlp: + hybrid_hidden_dim = int(hybrid_hidden_dim) + hybrid_num_layers = int(hybrid_num_layers) + if hybrid_hidden_dim <= 0: + raise ValueError(f"hybrid_hidden_dim must be >= 1, got {hybrid_hidden_dim}") + if hybrid_num_layers <= 0: + raise ValueError(f"hybrid_num_layers must be >= 1, got {hybrid_num_layers}") + + factory = {} # add dtype/device kwargs here if needed + layers = [] + in_dim = self.feature_dim + for _ in range(hybrid_num_layers): + lin = nn.Linear(in_dim, hybrid_hidden_dim, **factory) + nn.init.kaiming_uniform_(lin.weight, a=0.0, nonlinearity="relu") + nn.init.zeros_(lin.bias) + layers.append(lin) + layers.append(nn.ReLU(inplace=True)) + in_dim = hybrid_hidden_dim + + out = nn.Linear(in_dim, 1, bias=True, **factory) + nn.init.normal_(out.weight, std=0.01) + nn.init.zeros_(out.bias) + layers.append(out) + self.sigma_net = nn.Sequential(*layers) + + def get_densities(self, coords: torch.Tensor): + """Computes and returns densities""" + pts = coords.reshape(-1, 3) + features = interpolate_ms_features( + pts=pts, + ms_grids=self.grids, + ) + density_before_activation = self.sigma_net(features) + density = self.density_activation(density_before_activation) + return density + + def forward(self, pts: torch.Tensor): + return self.get_densities(pts) + + def get_params(self) -> dict[str, list[torch.nn.Parameter]]: + return { + "grids": list(self.grids.parameters()), + "sigma_net": list(self.sigma_net.parameters()), + } + + @property + def param_keys(self) -> list[str]: + return ["grids", "sigma_net"] + + @property + def td_type(self) -> str: + return self._td_type + + @td_type.setter + def td_type(self, td_type: str): + if not isinstance(td_type, str): + raise TypeError("td_type must be a string") + self._td_type = td_type + + @property + def tilted(self) -> bool: + return False + + @tilted.setter + def tilted(self, tilted: bool): + if not isinstance(tilted, bool): + raise TypeError("tilted must be a boolean") + self._tilted = tilted + + @property + def grids(self) -> torch.nn.ParameterList: + return self._grids + + @grids.setter + def grids(self, grids: torch.nn.ParameterList): + if not isinstance(grids, torch.nn.ParameterList): + raise TypeError("Grids must be a ParameterList") + self._grids = grids + + @property + def resolution(self) -> list[int]: + return self._resolution + + @resolution.setter + def resolution(self, resolution: Sequence[int]): + if not isinstance(resolution, Sequence): + raise TypeError("Resolution must be a sequence") + self._resolution = list(resolution) + + +# --------------------------------------------------------------------------- +# KPlanesTILTED +# --------------------------------------------------------------------------- + + +def interpolate_ms_features_tilted( + pts: torch.Tensor, # (B, 3) + ms_grids: nn.ParameterList, # each grid: (3*T, C, H, W) + rotation_matrices: torch.Tensor, # (T, 3, 3) +) -> torch.Tensor: + """ + Fully-vectorized multi-scale, multi-rotation K-Planes feature interpolation. + Returns features of shape (B, C * T * num_scales). + """ + T = rotation_matrices.shape[0] + B = pts.shape[0] + + # (T, B, 3) — rotate all points by all rotations at once + rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts) + + # Build (T, 3, B, 2) coords for planes XY, ZX, YZ in one shot. + # index_select is faster and cleaner than advanced indexing with python lists. + # Plane axis layout: XY=(0,1), ZX=(2,0), YZ=(1,2) + idx = torch.tensor([[0, 1], [2, 0], [1, 2]], device=pts.device) # (3, 2) + # rotated: (T, B, 3) -> gather along last dim with idx (3, 2) + # Result: (T, 3, B, 2) + coords = ( + rotated.unsqueeze(1).expand(T, 3, B, 3).gather(-1, idx.view(1, 3, 1, 2).expand(T, 3, B, 2)) + ) + + # Flatten (T, 3) -> 3*T so it matches grid's first dim, and add the H_out=1 axis + coord_tensor = coords.reshape(3 * T, B, 1, 2) # (3T, B, 1, 2) + + per_scale_features = [] + for plane_coef in ms_grids: + # plane_coef: (3T, C, H, W) + C = plane_coef.shape[1] + + sampled = F.grid_sample( + plane_coef, + coord_tensor, + align_corners=True, + mode="bilinear", + padding_mode="border", + ) # (3T, C, B, 1) + + # (3T, C, B) -> (T, 3, C, B) -> Hadamard across the "3" dim -> (T, C, B) + sampled = sampled.squeeze(-1).view(T, 3, C, B).prod(dim=1) + + # (T, C, B) -> (B, T, C) -> (B, T*C) to concatenate rotations along feature dim + per_scale_features.append(sampled.permute(2, 0, 1).reshape(B, T * C)) + + # Concatenate across scales -> (B, T * C * num_scales) + return torch.cat(per_scale_features, dim=-1) + + +class KPlanesTILTED(KPlanes): + """ + K-Planes with T learned SO(3) rotations (TILTED). Adapted from Yi et al., https://arxiv.org/abs/2308.15461 + + Inherits KPlanes for the sigma_net, density_activation, and get_params + interface. Overrides: + * __init__ – replaces the axis-aligned grids with (3*T)-plane + grids and adds SO3Param. + * get_densities – calls the TILTED interpolation instead. + * get_params – adds "so3" key so callers can set a separate lr. + * param_keys – updated list. + + Two-phase optimization + ---------------------- + Phase 1: instantiate with small `M_features` (and optionally smaller + `resolution` / fewer scales) — the bottleneck model. Train it + until τ converges, then call `extract_tau_state()`. + Phase 2: instantiate at full capacity, call `load_tau_state(M_bneck)` + to seed the rotations, then train normally. + + Parameters + ---------- + M_features : int + Feature channels *per transform per scale*. Total feature_dim will + be M_features * T * len(multiscale_res_multipliers). + T : int + Number of learned rotations (TILTED-T in the paper; 4 or 8 recommended). + Must match between phase 1 and phase 2 when doing two-phase transfer. + tau_init : str + "random" (paper default) or "identity". + Irrelevant if you're calling load_tau_state() right after __init__. + All other args are forwarded to KPlanes. + """ + + def __init__( + self, + # Grid parameters + input_coords_dims: int = 3, + M_features: int = 32, + resolution: Sequence[int] = (200, 200, 200), + multiscale_res_multipliers: Optional[Sequence[float]] = None, + density_activation: Callable = lambda x: F.softplus(x - 1), + # TILTED parameters + T: int = 4, + tau_init: str = "random", + # Hybrid MLP parameters + use_hybrid_mlp: bool = False, + hybrid_hidden_dim: int = 64, + hybrid_num_layers: int = 2, + so3_param_type: str = "r9svd", + ): + self._td_type = "tilted" + if input_coords_dims != 3: + raise NotImplementedError("KPlanesTILTED is implemented for 3D only.") + if T < 1: + raise ValueError(f"T must be >= 1, got {T}") + + multiscale_res_multipliers = list(multiscale_res_multipliers or [1]) + num_scales = len(multiscale_res_multipliers) + + # Total feature dim seen by the MLP head. + # Each scale contributes M_features * T channels. + feature_dim = M_features * T * num_scales + + # Call KPlanes.__init__ with grid_dimensions=2 so it builds sigma_net + # correctly; we immediately replace self.grids below. + super().__init__( + grid_dimensions=2, + input_coords_dims=3, + M_features=M_features, + resolution=resolution, + multiscale_res_multipliers=multiscale_res_multipliers, + concat_features=True, + density_activation=density_activation, + use_hybrid_mlp=use_hybrid_mlp, + hybrid_hidden_dim=hybrid_hidden_dim, + hybrid_num_layers=hybrid_num_layers, + ) + + self.T = T + + # ---- Rebuild grids: (3*T, M_features, H, W) per scale ---- + self.grids = nn.ParameterList() + for res_mult in multiscale_res_multipliers: + scaled_res = [int(r * res_mult) for r in resolution] + plane = nn.Parameter(torch.empty(3 * T, M_features, scaled_res[1], scaled_res[0])) + nn.init.uniform_(plane, 0.1, 0.5) + self.grids.append(plane) + + # ---- Rebuild sigma_net with the correct feature_dim ---- + # KPlanes built sigma_net with self.feature_dim (= M * num_scales), + # which is wrong for T > 1. Rebuild here. + self.feature_dim = feature_dim + self._build_sigma_net(use_hybrid_mlp, hybrid_hidden_dim, hybrid_num_layers) + + # ---- Learnable rotations ---- + self.set_so3_param_type(so3_param_type, init=tau_init) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _build_sigma_net( + self, + use_hybrid_mlp: bool, + hybrid_hidden_dim: int, + hybrid_num_layers: int, + ) -> None: + """Rebuild sigma_net for self.feature_dim (called after grids are set).""" + if use_hybrid_mlp: + layers = [] + in_dim = self.feature_dim + for _ in range(hybrid_num_layers): + lin = nn.Linear(in_dim, hybrid_hidden_dim) + nn.init.kaiming_uniform_(lin.weight, a=0.0, nonlinearity="relu") + nn.init.zeros_(lin.bias) + layers.append(lin) + layers.append(nn.ReLU(inplace=True)) + in_dim = hybrid_hidden_dim + out = nn.Linear(in_dim, 1, bias=True) + nn.init.normal_(out.weight, std=0.01) + nn.init.zeros_(out.bias) + layers.append(out) + self.sigma_net = nn.Sequential(*layers) + else: + # Single-linear "explicit" decoder. Small init -> density ~ 0 initially. + self.sigma_net = nn.Linear(self.feature_dim, 1, bias=True) + nn.init.normal_(self.sigma_net.weight, std=0.01) + nn.init.zeros_(self.sigma_net.bias) + + # ------------------------------------------------------------------ + # Core forward + # ------------------------------------------------------------------ + + def get_densities(self, coords: torch.Tensor) -> torch.Tensor: + pts = coords.reshape(-1, 3) + R = self.so3.as_matrix() # (T, 3, 3) + features = interpolate_ms_features_tilted( + pts=pts, + ms_grids=self.grids, + rotation_matrices=R, + ) + density_before_activation = self.sigma_net(features) + return self.density_activation(density_before_activation) + + def forward(self, pts: torch.Tensor) -> torch.Tensor: + return self.get_densities(pts) + + # ------------------------------------------------------------------ + # Parameter groups + # ------------------------------------------------------------------ + + def get_params(self) -> dict[str, list[nn.Parameter]]: + return { + "grids": list(self.grids.parameters()), + "sigma_net": list(self.sigma_net.parameters()), + "so3": list(self.so3.parameters()), + } + + @property + def param_keys(self) -> list[str]: + return ["grids", "sigma_net", "so3"] + + # ------------------------------------------------------------------ + # Two-phase transfer: extract / load learned rotations + # ------------------------------------------------------------------ + + def extract_tau_state(self) -> torch.Tensor: + """ + Returns the current raw R^9 matrices (detached copy) so they can be + used to initialise a phase-2 model via `load_tau_state`. + + Returns + ------- + torch.Tensor of shape (T, 3, 3) + """ + return self.so3.M.detach().cpu().clone() + + def load_tau_state(self, M: torch.Tensor) -> None: + """ + Load pre-trained rotation matrices (e.g. from a bottleneck phase-1 model). + + No orthogonalization is needed — SO3Param.as_matrix() projects to SO(3) + via SVD on every forward pass. + + Parameters + ---------- + M : torch.Tensor of shape (T, 3, 3) + Raw unconstrained matrices from `extract_tau_state()`. + """ + if M.shape != self.so3.M.shape: + raise ValueError( + f"Shape mismatch: got {M.shape}, expected {self.so3.M.shape}. " + f"Make sure T matches between phase 1 and phase 2." + ) + with torch.no_grad(): + self.so3.M.copy_(M.to(self.so3.M.device)) + + # ------------------------------------------------------------------ + # Pretty print + # ------------------------------------------------------------------ + + def extra_repr(self) -> str: + return ( + f"T={self.T}, " + f"M_features={self.M_features}, " + f"feature_dim={self.feature_dim}, " + f"num_scales={len(self.multiscale_res_multipliers)}" + ) + + def set_so3_param_type(self, so3_param_type: str, init: str = "rand") -> None: + """ + Set the SO3 parameterization type. + + Parameters + ---------- + so3_param_type : str + SO3 parameterization type ("quat" or "r9svd"). + """ + if so3_param_type == "r9svd": + self.so3 = SO3ParamR9SVD(self.T, init=init) + elif so3_param_type == "quat": + self.so3 = SO3ParamQuat(self.T, init=init) + else: + raise ValueError(f"Invalid SO3 parameterization type: {so3_param_type}") + + @property + def tilted(self) -> bool: + return True + + +# CP Decomp for Warmup SO3 rotations + + +def interpolate_ms_features_cp_tilted( + pts: torch.Tensor, # (B, 3) + ms_grids: nn.ParameterList, # each grid: (3*T, C, L) — 1D lines + rotation_matrices: torch.Tensor, # (T, 3, 3) +) -> torch.Tensor: + """ + CP (vector outer product) version of TILTED interpolation. + Returns features of shape (B, C * T * num_scales). + """ + T = rotation_matrices.shape[0] + B = pts.shape[0] + + # Rotate all points by all rotations: (T, B, 3) + rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts) + + per_scale_features = [] + for line_coef in ms_grids: + # line_coef: (3T, C, L) — three 1D feature lines per transform (x, y, z) + C, _ = line_coef.shape[1], line_coef.shape[2] + + # For each transform t, we need three 1D samples: at x_t, y_t, z_t. + # Lay them out as (3T, B) coords, matching line_coef's first dim. + # Axis order per transform: x, y, z. + coords_1d = rotated.reshape(T, B, 3).permute(0, 2, 1).reshape(3 * T, B) + # coords_1d: (3T, B), each row is samples along one axis for one transform + + # grid_sample wants 4D input for 2D sampling, or we can use 1D via a + # (3T, C, 1, L) reshape and pass 2D coords with y fixed at 0. + # Simpler: use F.grid_sample with a 4D trick, or just do manual linear interp. + # Here's the grid_sample way: + line_coef_4d = line_coef.unsqueeze(2) # (3T, C, 1, L) + # grid: need (3T, Hout=1, Wout=B, 2), with x = coord, y = 0 + grid = torch.stack( + [ + coords_1d, # x + torch.zeros_like(coords_1d), # y + ], + dim=-1, + ).unsqueeze(1) # (3T, 1, B, 2) + + sampled = F.grid_sample( + line_coef_4d, + grid, + align_corners=True, + mode="bilinear", + padding_mode="border", + ).squeeze(2) # (3T, C, B) + + # Hadamard across the 3 axes per transform: (T, 3, C, B) -> (T, C, B) + sampled = sampled.view(T, 3, C, B).prod(dim=1) + + # (T, C, B) -> (B, T*C) + per_scale_features.append(sampled.permute(2, 0, 1).reshape(B, T * C)) + + return torch.cat(per_scale_features, dim=-1) + + +class CPTilted(PPLR, TensorDecompositionModel): + """ + CP decomposition with TILTED rotations — the true bottleneck model for + phase 1. Rank-1-per-channel feature representation. Adapted from Yi et al., https://arxiv.org/abs/2308.15461 + + Shares the SO3Param and sigma_net design with KPlanesTILTED so you can + lift τ directly across: cp_model.extract_tau_state() -> + kplanes_model.load_tau_state(). + """ + + def __init__( + self, + C: int = 4, # channels per transform per scale + resolution: Sequence[int] = (128, 128, 128), + multiscale_res_multipliers: Optional[Sequence[int]] = None, + T: int = 4, + tau_init: str = "random", + density_activation: Callable = lambda x: F.softplus(x - 1), + so3_param_type: str = "r9svd", + ): + super().__init__() + self._td_type = "cp_tilted" + self.T = T + self.C = C + self.multiscale_res_multipliers = list(multiscale_res_multipliers or [1]) + self.density_activation = density_activation + + # 1D feature lines, one per axis per transform per scale. + # Shape per scale: (3*T, C, L). We use max(resolution) for L; if your + # scene is strongly anisotropic use 3 separate lines per axis. + self.grids = nn.ParameterList() + for mult in self.multiscale_res_multipliers: + L = int(max(resolution) * mult) + line = nn.Parameter(torch.empty(3 * T, C, L)) + nn.init.uniform_(line, 0.1, 0.5) + self.grids.append(line) + + self.feature_dim = C * T * len(self.multiscale_res_multipliers) + + # Same minimal single-linear decoder as your KPlanesTILTED default. + self.sigma_net = nn.Linear(self.feature_dim, 1, bias=True) + nn.init.normal_(self.sigma_net.weight, std=0.01) + nn.init.zeros_(self.sigma_net.bias) + + if so3_param_type == "r9svd": + self.so3 = SO3ParamR9SVD(T, init=tau_init) + elif so3_param_type == "quat": + self.so3 = SO3ParamQuat(T, init=tau_init) + else: + raise ValueError(f"Unknown SO3 param type: {so3_param_type}") + + def get_densities(self, coords: torch.Tensor) -> torch.Tensor: + pts = coords.reshape(-1, 3) + R = self.so3.as_matrix() + features = interpolate_ms_features_cp_tilted(pts, self.grids, R) + return self.density_activation(self.sigma_net(features)) + + def forward(self, pts): + return self.get_densities(pts) + + def get_params(self): + return { + "grids": list(self.grids.parameters()), + "sigma_net": list(self.sigma_net.parameters()), + "so3": list(self.so3.parameters()), + } + + @property + def param_keys(self): + return ["grids", "sigma_net", "so3"] + + @property + def td_type(self) -> str: + return self._td_type + + def extract_tau_state(self) -> torch.Tensor: + return self.so3.M.detach().clone() + + @property + def tilted(self) -> bool: + return True + + +KPlanesType = KPlanes | KPlanesTILTED | CPTilted diff --git a/src/quantem/core/ml/models/model_base.py b/src/quantem/core/ml/models/model_base.py new file mode 100644 index 00000000..60c1c4f6 --- /dev/null +++ b/src/quantem/core/ml/models/model_base.py @@ -0,0 +1,50 @@ +from abc import ABC, abstractmethod +from typing import Dict + +import torch.nn as nn + + +class PPLR(ABC): + """ + Abstract base class for models that require multi-parameter optimization. + """ + + @abstractmethod + def get_params(self) -> Dict[str, list[nn.Parameter]]: + """ + Return a dictionary of parameters grouped by key. + + For example if your nn.Module has multiple optimizable parameter groups, + you can return a dictionary with the keys "grids" and "sigma_net" + (KPlanes example). + """ + pass + + @property + @abstractmethod + def param_keys(self) -> list[str]: + """List of available parameter-group keys.""" + pass + + +class TensorDecompositionModel(nn.Module, ABC): + """ + Base class for factored tensor-decomposition models. + + Subclasses must set ``td_type`` as a normal attribute in ``__init__``. + """ + + td_type: str + + +class PlanarDecompositionModel(TensorDecompositionModel, PPLR): + """ + Planar factored-grid models: K-Planes, K-Planes-TILTED, HexPlane, tri-planes. + + Subclasses must set ``grids``, ``tilted``, and ``resolution`` as normal + attributes in ``__init__``. + """ + + grids: nn.ParameterList + tilted: bool + resolution: list[int] diff --git a/src/quantem/core/ml/models/so3params.py b/src/quantem/core/ml/models/so3params.py new file mode 100644 index 00000000..cd55d0ac --- /dev/null +++ b/src/quantem/core/ml/models/so3params.py @@ -0,0 +1,192 @@ +import math +from typing import Literal + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# --- Tilted KPlanes --- + +# --------------------------------------------------------------------------- +# SO(3) quaternion parameter module +# --------------------------------------------------------------------------- + + +class SO3ParamQuat(nn.Module): + """ + Stores T unit quaternions as learnable parameters in R^4 and normalises + them on every call to `as_matrix()`. + + Quaternion convention: [x, y, z, w] (scalar-last, same as scipy). + + Initialisation + -------------- + "random" – uniform sampling over SO(3) via Shoemake's method. + "identity" – all rotations start as the identity (good for fine-tuning). + """ + + def __init__(self, T: int, init: str = "random"): + super().__init__() + if T < 1: + raise ValueError(f"T must be >= 1, got {T}") + quats = self._init_quaternions(T, init) # (T, 4) + self.quats = nn.Parameter(quats) + + @staticmethod + def quat_to_rotmat(q: torch.Tensor) -> torch.Tensor: + """Unit quaternion (..., 4) [x, y, z, w] -> rotation matrix (..., 3, 3). + Assumes q is already normalized.""" + x, y, z, w = q.unbind(dim=-1) + xx, yy, zz = x * x, y * y, z * z + xy, xz, yz = x * y, x * z, y * z + wx, wy, wz = w * x, w * y, w * z + R = torch.stack( + [ + 1 - 2 * (yy + zz), 2 * (xy - wz), 2 * (xz + wy), + 2 * (xy + wz), 1 - 2 * (xx + zz), 2 * (yz - wx), + 2 * (xz - wy), 2 * (yz + wx), 1 - 2 * (xx + yy), + ], + dim=-1, + ).reshape(*q.shape[:-1], 3, 3) + return R + + @staticmethod + def rotmat_to_quat(R: torch.Tensor) -> torch.Tensor: + """Rotation matrix (..., 3, 3) -> unit quaternion (..., 4) [x, y, z, w]. + + Shepperd's method: build the four candidate quaternions, each dividing + by a different diagonal combination, then per-element pick the branch + with the largest denominator so we never divide by a near-zero number. + The naive trace-only formula blows up when trace ~ -1 (180deg rotations). + """ + m00, m01, m02 = R[..., 0, 0], R[..., 0, 1], R[..., 0, 2] + m10, m11, m12 = R[..., 1, 0], R[..., 1, 1], R[..., 1, 2] + m20, m21, m22 = R[..., 2, 0], R[..., 2, 1], R[..., 2, 2] + + # 4 * (component^2) for w, x, y, z respectively; these sum to 4. + t = torch.stack( + [ + 1.0 + m00 + m11 + m22, # 4 w^2 + 1.0 + m00 - m11 - m22, # 4 x^2 + 1.0 - m00 + m11 - m22, # 4 y^2 + 1.0 - m00 - m11 + m22, # 4 z^2 + ], + dim=-1, + ) # (..., 4) + + eps = torch.finfo(R.dtype).eps + S = 2.0 * torch.sqrt(t.clamp_min(eps)) # S[k] = 4 * |component_k| + S0, S1, S2, S3 = S.unbind(-1) + + # each candidate in [x, y, z, w] order + cand_w = torch.stack([(m21 - m12) / S0, (m02 - m20) / S0, (m10 - m01) / S0, 0.25 * S0], dim=-1) + cand_x = torch.stack([0.25 * S1, (m01 + m10) / S1, (m02 + m20) / S1, (m21 - m12) / S1], dim=-1) + cand_y = torch.stack([(m01 + m10) / S2, 0.25 * S2, (m12 + m21) / S2, (m02 - m20) / S2], dim=-1) + cand_z = torch.stack([(m02 + m20) / S3, (m12 + m21) / S3, 0.25 * S3, (m10 - m01) / S3], dim=-1) + + cands = torch.stack([cand_w, cand_x, cand_y, cand_z], dim=-2) # (..., 4, 4) + idx = t.argmax(dim=-1) # (...,) + idx = idx[..., None, None].expand(*idx.shape, 1, 4) # (..., 1, 4) + q = cands.gather(-2, idx).squeeze(-2) # (..., 4) + return F.normalize(q, p=2, dim=-1) + + def as_matrix(self) -> torch.Tensor: + return self.quat_to_rotmat(self.normalized()) + + @classmethod + def from_matrix(cls, R: torch.Tensor) -> "SO3ParamQuat": + """Initialize a bank close to the given rotations R (T, 3, 3).""" + obj = cls(R.shape[0], init="identity") + with torch.no_grad(): + obj.quats.copy_(cls.rotmat_to_quat(R)) + return obj + + def extra_repr(self) -> str: + return f"T={self.quats.shape[0]}" + + # ------------------------------------------------------------------ + # Initialisers + # ------------------------------------------------------------------ + + @staticmethod + def _shoemake_sample(T: int) -> torch.Tensor: + """Uniform SO(3) sampling via Shoemake (1992). Returns (T, 4) [x,y,z,w].""" + u = torch.rand(T, 3) + sqrt1_u0 = torch.sqrt(1.0 - u[:, 0]) + sqrt_u0 = torch.sqrt(u[:, 0]) + two_pi = 2.0 * math.pi + x = sqrt1_u0 * torch.sin(two_pi * u[:, 1]) + y = sqrt1_u0 * torch.cos(two_pi * u[:, 1]) + z = sqrt_u0 * torch.sin(two_pi * u[:, 2]) + w = sqrt_u0 * torch.cos(two_pi * u[:, 2]) + return torch.stack([x, y, z, w], dim=-1) # (T, 4) + + @staticmethod + def _identity(T: int) -> torch.Tensor: + """All-identity rotations: [0,0,0,1] * T.""" + q = torch.zeros(T, 4) + q[:, 3] = 1.0 + return q + + @classmethod + def _init_quaternions(cls, T: int, init: str) -> torch.Tensor: + if init == "random": + return cls._shoemake_sample(T) + elif init == "identity": + return cls._identity(T) + else: + raise ValueError(f"Unknown init '{init}'; choose 'random' or 'identity'.") + + # ------------------------------------------------------------------ + # Forward helpers + # ------------------------------------------------------------------ + + def normalized(self) -> torch.Tensor: + """Returns (T, 4) unit quaternions.""" + return F.normalize(self.quats, p=2, dim=-1) + + +class SO3ParamR9SVD(nn.Module): + """ + SO(3) rotation bank using R9+SVD parameterization. + Each rotation is stored as an unconstrained 3x3 matrix M, + projected to SO(3) via SVD+(M) = U diag(1,1,det(UVt)) Vt. + + Based on Rene Geist et al., 2024: https://arxiv.org/abs/2404.11735v1 + """ + + def __init__(self, T: int, init: Literal["random", "identity"] = "random"): + super().__init__() + if init == "random": + M = torch.eye(3).unsqueeze(0).repeat(T, 1, 1) + 0.1 * torch.randn(T, 3, 3) + elif init == "identity": + M = torch.eye(3).unsqueeze(0).repeat(T, 1, 1) + else: + raise ValueError(f"Unknown init '{init}'") + self.M = nn.Parameter(M) + + @staticmethod + def rotmat_to_r9(R: torch.Tensor) -> torch.Tensor: + """Rotation matrix (..., 3, 3) -> R9. Identity embedding: any R in SO(3) + is a fixed point of the SVD projection, so this just stores R directly.""" + return R + + @staticmethod + def r9_to_rotmat(M: torch.Tensor) -> torch.Tensor: + """R9 (..., 3, 3) -> nearest SO(3) matrix via SVD+.""" + U, _, Vh = torch.linalg.svd(M) + d = torch.det(U @ Vh) + diag = torch.ones(*M.shape[:-2], 3, device=M.device, dtype=M.dtype) + diag[..., 2] = d + return U @ (diag.unsqueeze(-1) * Vh) + + def as_matrix(self) -> torch.Tensor: + return self.r9_to_rotmat(self.M) + + @classmethod + def from_matrix(cls, R: torch.Tensor) -> "SO3ParamR9SVD": + """Initialize a bank close to given rotations R (T, 3, 3).""" + obj = cls(R.shape[0], init="identity") + with torch.no_grad(): + obj.M.copy_(cls.rotmat_to_r9(R)) + return obj \ No newline at end of file diff --git a/src/quantem/core/ml/optimizer_mixin.py b/src/quantem/core/ml/optimizer_mixin.py index 4ea263c0..ece31dff 100644 --- a/src/quantem/core/ml/optimizer_mixin.py +++ b/src/quantem/core/ml/optimizer_mixin.py @@ -1,7 +1,7 @@ import textwrap from abc import abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, Generator, Iterator, Literal, Sequence +from typing import TYPE_CHECKING, Any, Literal from quantem.core import config @@ -219,7 +219,7 @@ def parse_dict(cls, d: dict): raise ValueError(f"Unknown optimizer type: {name.lower()}") -OptimizerType = ( +OptimizerParamsType = ( OptimizerParams.Adam | OptimizerParams.AdamW | OptimizerParams.SGD @@ -514,7 +514,7 @@ def parse_dict(cls, d: dict): raise ValueError(f"Unknown scheduler type: {name}") -SchedulerType = ( +SchedulerParamsType = ( SchedulerParams.Plateau | SchedulerParams.Exponential | SchedulerParams.Cyclic @@ -531,13 +531,16 @@ class OptimizerMixin: """ DEFAULT_OPTIMIZER_TYPE = "adamw" + DEFAULT_OPTIMIZER_KEY = "default" def __init__(self): """Initialize the optimizer mixin.""" self._optimizer = None self._scheduler = None - self._optimizer_params: OptimizerType = OptimizerParams.NoneOptimizer() - self._scheduler_params: SchedulerType = SchedulerParams.NoneScheduler() + self._optimizer_params: dict[str, OptimizerParamsType] = { + self.DEFAULT_OPTIMIZER_KEY: OptimizerParams.NoneOptimizer() + } + self._scheduler_params: SchedulerParamsType = SchedulerParams.NoneScheduler() # Don't call super().__init__() in mixin classes to avoid MRO issues @property @@ -551,83 +554,141 @@ def scheduler(self) -> "torch.optim.lr_scheduler.LRScheduler | None": return self._scheduler @property - def optimizer_params(self) -> OptimizerType: + def optimizer_params(self) -> dict[str, OptimizerParamsType]: """Get the optimizer parameters.""" return self._optimizer_params @optimizer_params.setter - def optimizer_params(self, params: OptimizerType | dict): - """Set the optimizer parameters.""" - if isinstance(params, dict): - params = OptimizerParams.parse_dict(d=params) - if not isinstance(params, OptimizerType): - raise TypeError(f"optimizer parameters must be a OptimizerType, got {type(params)}") - self._optimizer_params = params + def optimizer_params( + self, params: OptimizerParamsType | dict[str, OptimizerParamsType] | dict[str, Any] + ) -> None: + self._optimizer_params = self._normalize_optimizer_params(params) + + def _normalize_optimizer_params( + self, params: OptimizerParamsType | dict[str, Any] + ) -> dict[str, OptimizerParamsType]: + """Normalize input to dict[str, OptimizerParamsType]. Subclasses can override to validate keys.""" + # Single optimizer, already an OptimizerParamsType + if isinstance(params, OptimizerParamsType): + return {self.DEFAULT_OPTIMIZER_KEY: params} + if not isinstance(params, dict): + raise TypeError(f"optimizer_params must be OptimizerParamsType or dict, got {type(params)}") + # Single optimizer as dict shorthand, e.g. {"name": "adam", "lr": 1e-3} + if self._is_single_optimizer_dict(params): + return {self.DEFAULT_OPTIMIZER_KEY: OptimizerParams.parse_dict(d=params)} + # dict-of-OptimizerParamsType form (PPLR) + return { + k: v if isinstance(v, OptimizerParamsType) else OptimizerParams.parse_dict(d=v) + for k, v in params.items() + } + + @staticmethod + def _is_single_optimizer_dict(d: dict) -> bool: + return "type" in d or "name" in d @property - def scheduler_params(self) -> SchedulerType: + def scheduler_params(self) -> SchedulerParamsType: """Get the scheduler parameters.""" return self._scheduler_params @scheduler_params.setter - def scheduler_params(self, params: SchedulerType | dict): + def scheduler_params(self, params: SchedulerParamsType | dict): """Set the scheduler parameters.""" if isinstance(params, dict): params = SchedulerParams.parse_dict(d=params) - if not isinstance(params, SchedulerType): - raise TypeError(f"scheduler parameters must be a SchedulerType, got {type(params)}") + if not isinstance(params, SchedulerParamsType): + raise TypeError(f"scheduler parameters must be a SchedulerParamsType, got {type(params)}") self._scheduler_params = params @abstractmethod def get_optimization_parameters( self, - ) -> "torch.Tensor | Sequence[torch.Tensor] | Iterator[torch.Tensor]": + ) -> "dict[str, list[torch.Tensor]]": """ - Get the parameters that should be optimized for this model. - This could be replaced with just module.parameters(), but this allows for flexibility - in the future to allow for per parameter LRs. + Get the parameters that should be optimized for this model, grouped by name. + + Returns a mapping ``{group_key: [tensors]}``. The group keys MUST match the keys of + ``optimizer_params`` (the common single-group case uses ``DEFAULT_OPTIMIZER_KEY``). + ``set_optimizer`` joins each group to its optimizer spec by key and bakes the per-group + hyperparameters (``spec.params()``) into the torch param group — implementations return + only the tensors, NOT pre-baked hyperparameters. Return ``{}`` when there is nothing + to optimize. """ raise NotImplementedError("Subclasses must implement get_optimization_parameters") - def set_optimizer(self, opt_params: OptimizerType | dict | None = None) -> None: + def set_optimizer(self, opt_params: OptimizerParamsType | dict | None = None) -> None: """ - Set the optimizer for this model. - Currently supports single LR for all parameters, TODO allow for per parameter LRs by - updating get_optimization_parameters to return a list of parameters and their LRs. + Set the optimizer for this model, supporting per-parameter-group learning rates (PPLR). + + ``optimizer_params`` is a ``dict[str, OptimizerParamsType]`` keyed by parameter group. Each + group's spec is joined by key to the tensors returned by ``get_optimization_parameters()`` + and its hyperparameters are baked into the corresponding torch param group here. All + groups must use the same optimizer class. If every group is a ``NoneOptimizer`` (or there + are no groups), the optimizer is removed. """ if opt_params is not None: self.optimizer_params = opt_params - if not self._optimizer_params: - self._optimizer = None - return - - if isinstance(self._optimizer_params, OptimizerParams.NoneOptimizer): + # Single canonical "disable" path: drop NoneOptimizer sentinels and, if nothing is left, + # remove the optimizer. Done BEFORE get_optimization_parameters() because some models + # (e.g. the dataset model) raise / return nothing when there is nothing to optimize. + specs = { + key: spec + for key, spec in self.optimizer_params.items() + if not isinstance(spec, OptimizerParams.NoneOptimizer) + } + if not specs: self.remove_optimizer() return - params = self.get_optimization_parameters() - if isinstance(params, torch.Tensor): - params = [params] - elif isinstance(params, Generator): - params = list(params) + # All groups must agree on the optimizer class. + spec_list = list(specs.values()) + for spec in spec_list[1:]: + if type(spec) is not type(spec_list[0]): + raise ValueError( + f"All parameter groups must use the same optimizer type, " + f"got {type(spec_list[0]).__name__} and {type(spec).__name__}" + ) + + # Join specs to param groups by key; bake each group's hyperparameters here. + groups = self.get_optimization_parameters() # dict[str, list[tensor]] + if set(groups) != set(specs): + raise ValueError( + f"optimizer_params keys {set(specs)} do not match parameter group keys " + f"{set(groups)} from {type(self).__name__}.get_optimization_parameters()" + ) - # Ensure parameters require gradients - for p in params: - p.requires_grad_(True) + param_groups = [] + for key, tensors in groups.items(): + for p in tensors: + p.requires_grad_(True) + param_groups.append({"params": tensors, **specs[key].params()}) + self._optimizer = self._build_optimizer(spec_list[0], param_groups) - match self._optimizer_params: + def _build_optimizer(self, opt_params, param_groups) -> "torch.optim.Optimizer": + """Construct the torch optimizer for ``opt_params`` over pre-baked ``param_groups``. + + ``param_groups`` already carry their per-group hyperparameters (see ``set_optimizer``), + so each group's ``lr`` etc. overrides the optimizer-level default. ``NoneOptimizer`` must + have been filtered out by the caller. + """ + match opt_params: case OptimizerParams.Adam(): - self._optimizer = torch.optim.Adam(params, **self._optimizer_params.params()) + return torch.optim.Adam(param_groups) case OptimizerParams.AdamW(): - self._optimizer = torch.optim.AdamW(params, **self._optimizer_params.params()) + return torch.optim.AdamW(param_groups) case OptimizerParams.SGD(): - self._optimizer = torch.optim.SGD(params, **self._optimizer_params.params()) + return torch.optim.SGD(param_groups) + case OptimizerParams.NoneOptimizer(): + raise ValueError( + "NoneOptimizer must be filtered out before _build_optimizer; " + "set_optimizer should have short-circuited to remove_optimizer()." + ) case _: - raise NotImplementedError(f"Unknown optimizer type: {self._optimizer_params}") + raise NotImplementedError(f"Unknown optimizer type: {opt_params}") def set_scheduler( - self, scheduler_params: SchedulerType | dict | None = None, num_iter: int | None = None + self, scheduler_params: SchedulerParamsType | dict | None = None, num_iter: int | None = None ) -> None: """Set the scheduler for this model.""" if scheduler_params is not None: @@ -638,7 +699,10 @@ def set_scheduler( return optimizer = self._optimizer - base_LR = optimizer.param_groups[0]["lr"] + # Schedulers scale every torch param group proportionally off its own initial_lr; this + # scalar only seeds scheduler config (e.g. min_lr, cyclic bounds). Use the max group LR + # as the representative (collapses to group 0 in the single-group case). + base_LR = max(pg["lr"] for pg in optimizer.param_groups) params = self._scheduler_params.params(base_LR, num_iter=num_iter) match self.scheduler_params: case SchedulerParams.NoneScheduler(): @@ -703,7 +767,7 @@ def get_current_lr(self) -> float: def remove_optimizer(self) -> None: """Remove the optimizer and scheduler.""" self._optimizer = None - self._optimizer_params = OptimizerParams.NoneOptimizer() + self._optimizer_params = {self.DEFAULT_OPTIMIZER_KEY: OptimizerParams.NoneOptimizer()} self._scheduler = None self._scheduler_params = SchedulerParams.NoneScheduler() @@ -721,56 +785,44 @@ def reconnect_optimizer_to_parameters(self) -> None: if self._optimizer is None: return - current_params = self.get_optimization_parameters() - if isinstance(current_params, torch.Tensor): - current_params = [current_params] - elif isinstance(current_params, Generator): - current_params = list(current_params) - - optimizable_params = [ - p for p in current_params if isinstance(p, torch.Tensor) and p.is_leaf - ] - - if not optimizable_params: - print( - f"souldn't be getting here! No optimizable parameters found for {self.__class__.__name__}, removing optimizer" - ) + new_groups = self.get_optimization_parameters() + if not new_groups: + print(f"No optimizable parameters for {type(self).__name__}, removing optimizer") self.remove_optimizer() return - for p in optimizable_params: - p.requires_grad_(True) + # Ensure leaf params with grad + for tensors in new_groups.values(): + for p in tensors: + if not p.is_leaf: + raise ValueError("Non-leaf tensor in param group; build groups from leaves") + p.requires_grad_(True) - # Preserve optimizer state and param_group settings - old_state = self._optimizer.state.copy() - current_param_group = self._optimizer.param_groups[0].copy() + old_state = dict(self._optimizer.state) + old_hyperparams = [ + {k: v for k, v in pg.items() if k != "params"} for pg in self._optimizer.param_groups + ] - # Reconnect to new parameters self._optimizer.param_groups.clear() - self._optimizer.add_param_group({"params": optimizable_params}) + for tensors in new_groups.values(): + self._optimizer.add_param_group({"params": tensors}) + + # Restore per-group hyperparameters by index + for new_pg, old_pg in zip(self._optimizer.param_groups, old_hyperparams): + new_pg.update(old_pg) - # Update state mapping and move tensors to correct device + # Remap state for tensors that survived new_state = {} - device = optimizable_params[0].device - for i, old_param in enumerate(old_state.keys()): - if i < len(optimizable_params): - new_param = optimizable_params[i] - new_state[new_param] = {} - for key, value in old_state[old_param].items(): - if isinstance(value, torch.Tensor): - new_state[new_param][key] = value.to(device) - else: - new_state[new_param][key] = value + for new_pg in self._optimizer.param_groups: + for new_param in new_pg["params"]: + if new_param in old_state: + new_state[new_param] = { + k: (v.to(new_param.device) if isinstance(v, torch.Tensor) else v) + for k, v in old_state[new_param].items() + } self._optimizer.state.clear() self._optimizer.state.update(new_state) - # Restore param_group settings (LR, betas, etc.) but keep new parameters - self._optimizer.param_groups[0].update( - {k: v for k, v in current_param_group.items() if k != "params"} - ) - - # Reconnect scheduler - if self._scheduler is not None and self._optimizer is not None: + if self._scheduler is not None: self._scheduler.optimizer = self._optimizer - return diff --git a/src/quantem/core/utils/tomography_utils.py b/src/quantem/core/utils/tomography_utils.py index 9f5df102..006f8cfc 100644 --- a/src/quantem/core/utils/tomography_utils.py +++ b/src/quantem/core/utils/tomography_utils.py @@ -168,7 +168,6 @@ def fourier_binning(img, crop_size): center = np.array(img.shape) // 2 fft_img = np.fft.fftshift(np.fft.fft2(img)) - cropped_fft = fft_img[ center[0] - crop_size[0] // 2 : center[0] + crop_size[0] // 2, center[1] - crop_size[1] // 2 : center[1] + crop_size[1] // 2, diff --git a/src/quantem/core/visualization/visualization.py b/src/quantem/core/visualization/visualization.py index 77f647b4..e252caa6 100644 --- a/src/quantem/core/visualization/visualization.py +++ b/src/quantem/core/visualization/visualization.py @@ -40,10 +40,7 @@ # There might be a cleaner way to do this, but better to have it here than in the functions NormInputCell: TypeAlias = NormalizationConfig | ShowParams.Norm | dict | str Show2dNormInput: TypeAlias = ( - NormInputCell - | None - | Sequence[NormInputCell] - | Sequence[Sequence[NormInputCell]] + NormInputCell | None | Sequence[NormInputCell] | Sequence[Sequence[NormInputCell]] ) ScalebarInputCell: TypeAlias = ScalebarConfig | ShowParams.Scalebar | dict | bool | None @@ -57,7 +54,7 @@ | Sequence[Sequence[ScalebarInputCell]] ) -CmapType: TypeAlias = str | colors.Colormap +CmapType: TypeAlias = str | colors.Colormap def _show_2d_array( diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py index ef55c92e..751c9245 100644 --- a/src/quantem/diffractive_imaging/dataset_models.py +++ b/src/quantem/diffractive_imaging/dataset_models.py @@ -1,4 +1,6 @@ +import warnings from abc import abstractmethod +from dataclasses import replace from pathlib import Path from typing import Any, Literal, Self @@ -11,7 +13,7 @@ from quantem.core.datastructures.dataset3d import Dataset3d from quantem.core.datastructures.dataset4dstem import Dataset4dstem from quantem.core.io.serialize import AutoSerialize -from quantem.core.ml.optimizer_mixin import OptimizerMixin +from quantem.core.ml.optimizer_mixin import OptimizerMixin, OptimizerParams from quantem.core.utils.utils import electron_wavelength_angstrom, tqdmnd from quantem.core.utils.validators import ( validate_array, @@ -32,7 +34,7 @@ class PtychographyDatasetBase(AutoSerialize, OptimizerMixin, torch.nn.Module): _token = object() _patch_indices: torch.Tensor - # TODO update optimizers and such to allow for different lrs for different parameters + # TODO make this a PPLR so different lrs can be used for different parameters DEFAULT_LRS = { "descan": 1e-3, "scan_positions": 1e-3, @@ -95,18 +97,47 @@ def __init__( self._constraints = {} self._probe_energy = None - def get_optimization_parameters(self): - """Get the combined descan and scan position parameters for optimization.""" - params = [] + def get_optimization_parameters(self) -> "dict[str, list[torch.Tensor]]": + """Descan and scan-position parameters as separate PPLR groups. + + Returns one group per *learnable* parameter set; ``{}`` when neither is learnable + (``set_optimizer`` then short-circuits to removing the optimizer). + """ + groups: dict[str, list[torch.Tensor]] = {} if self.learn_descan: - params.append(self._descan_shifts) + groups["descan"] = [self._descan_shifts] if self.learn_scan_positions: - params.append(self._scan_positions_px) - if len(params) == 0: - raise RuntimeError( - "No parameters to optimize for dataset: learn_descan and learn_scan_positions are both False" - ) - return params + groups["scan_positions"] = [self._scan_positions_px] + return groups + + def _normalize_optimizer_params(self, params): + """Broadcast a single optimizer spec to the learnable descan/scan_position groups. + + A single ``OptimizerParamsType`` / single-optimizer dict (normalized to the ``"default"`` key) + is fanned out to whichever groups are currently learnable, so the common single-LR caller + keeps working. An explicit PPLR dict (keyed by ``descan``/``scan_positions``) passes through. + """ + norm = super()._normalize_optimizer_params(params) + if set(norm) == {self.DEFAULT_OPTIMIZER_KEY}: + spec = norm[self.DEFAULT_OPTIMIZER_KEY] + learnable = [ + key + for key, on in ( + ("descan", self.learn_descan), + ("scan_positions", self.learn_scan_positions), + ) + if on + ] + if not learnable and not isinstance(spec, OptimizerParams.NoneOptimizer): + warnings.warn( + f"{type(self).__name__}: an optimizer was requested but nothing is " + "learnable (both learn_descan and learn_scan_positions are False); " + "the optimizer will be removed. Enable learn_descan and/or " + "learn_scan_positions to optimize.", + stacklevel=2, + ) + return {key: replace(spec) for key in learnable} if learnable else {} + return norm def to(self, *args, **kwargs): """Move all relevant tensors to a different device.""" diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py index 1c6fdab9..d311d8dd 100644 --- a/src/quantem/diffractive_imaging/object_models.py +++ b/src/quantem/diffractive_imaging/object_models.py @@ -14,7 +14,11 @@ from quantem.core.io.serialize import AutoSerialize from quantem.core.ml.blocks import reset_weights from quantem.core.ml.loss_functions import get_loss_module -from quantem.core.ml.optimizer_mixin import OptimizerMixin, OptimizerType, SchedulerType +from quantem.core.ml.optimizer_mixin import ( + OptimizerMixin, + OptimizerParamsType, + SchedulerParamsType, +) from quantem.core.utils.rng import RNGMixin from quantem.core.utils.validators import ( validate_arr_gt, @@ -194,7 +198,7 @@ def obj(self): @property @abstractmethod - def params(self): + def params(self) -> list[nn.Parameter]: raise NotImplementedError() @abstractmethod @@ -232,16 +236,12 @@ def to(self, *args, **kwargs): def name(self) -> str: raise NotImplementedError() - def get_optimization_parameters(self): - """Get the parameters that should be optimized for this model.""" - try: - params = self.params - if params is None: - return [] - return params - except NotImplementedError: - # This happens when params is not implemented yet in abstract base - return [] + def get_optimization_parameters(self) -> "dict[str, list[torch.Tensor]]": + """Get the parameters that should be optimized for this model, keyed by group.""" + params = self.params + if params is None: + return {} + return {self.DEFAULT_OPTIMIZER_KEY: list(params)} def _propagate_array( self, array: "torch.Tensor", propagator_array: "torch.Tensor" @@ -633,9 +633,9 @@ def num_slices(self) -> int: return self._obj.shape[0] @property - def params(self): + def params(self) -> list[nn.Parameter]: """optimization parameters""" - return self._obj + return [self._obj] @property def initial_obj(self): @@ -1025,9 +1025,9 @@ def to(self, *args, **kwargs): return self @property - def params(self): + def params(self) -> list[nn.Parameter]: """optimization parameters""" - return self.model.parameters() + return list(self.model.parameters()) def reset(self): """Reset the object model to its initial or pre-trained state""" @@ -1050,8 +1050,8 @@ def pretrain( pretrain_target: torch.Tensor | None = None, reset: bool = False, num_iters: int = 100, - optimizer_params: dict | OptimizerType | None = None, - scheduler_params: dict | SchedulerType | None = None, + optimizer_params: dict | OptimizerParamsType | None = None, + scheduler_params: dict | SchedulerParamsType | None = None, loss_fn: Callable | str = "l2", apply_constraints: bool = False, show: bool = True, diff --git a/src/quantem/diffractive_imaging/probe_models.py b/src/quantem/diffractive_imaging/probe_models.py index f47ea1f7..af638b85 100644 --- a/src/quantem/diffractive_imaging/probe_models.py +++ b/src/quantem/diffractive_imaging/probe_models.py @@ -15,7 +15,11 @@ from quantem.core.io.serialize import AutoSerialize from quantem.core.ml.blocks import reset_weights from quantem.core.ml.loss_functions import get_loss_module -from quantem.core.ml.optimizer_mixin import OptimizerMixin, OptimizerType, SchedulerType +from quantem.core.ml.optimizer_mixin import ( + OptimizerMixin, + OptimizerParamsType, + SchedulerParamsType, +) from quantem.core.utils.rng import RNGMixin from quantem.core.utils.utils import electron_wavelength_angstrom, to_numpy from quantem.core.utils.validators import ( @@ -94,16 +98,12 @@ def __init__( if roi_shape is not None: self.roi_shape = roi_shape - def get_optimization_parameters(self): - """Get the parameters that should be optimized for this model.""" - try: - params = self.params - if params is None: - return [] - return params - except NotImplementedError: - # This happens when params is not implemented yet in abstract base - return [] + def get_optimization_parameters(self) -> "dict[str, list[torch.Tensor]]": + """Get the parameters that should be optimized for this model, keyed by group.""" + params = self.params + if params is None: + return {} + return {self.DEFAULT_OPTIMIZER_KEY: list(params)} @property def learn_probe_tilt(self) -> bool: @@ -1359,8 +1359,8 @@ def pretrain( pretrain_target: torch.Tensor | None = None, reset: bool = False, num_iters: int = 100, - optimizer_params: dict | OptimizerType | None = None, - scheduler_params: dict | SchedulerType | None = None, + optimizer_params: dict | OptimizerParamsType | None = None, + scheduler_params: dict | SchedulerParamsType | None = None, loss_fn: Callable | str = "l2", apply_constraints: bool = False, show: bool = True, diff --git a/src/quantem/diffractive_imaging/ptychography_opt.py b/src/quantem/diffractive_imaging/ptychography_opt.py index c18150e0..ff1b3a2c 100644 --- a/src/quantem/diffractive_imaging/ptychography_opt.py +++ b/src/quantem/diffractive_imaging/ptychography_opt.py @@ -4,9 +4,9 @@ from quantem.core import config from quantem.core.ml.optimizer_mixin import ( OptimizerParams, - OptimizerType, + OptimizerParamsType, SchedulerParams, - SchedulerType, + SchedulerParamsType, ) from quantem.diffractive_imaging.ptychography_base import PtychographyBase @@ -23,7 +23,7 @@ class PtychographyOpt(PtychographyBase): """ OPTIMIZABLE_VALS = ["object", "probe", "dataset"] - DEFAULT_OPTIMIZER_TYPE: OptimizerType = OptimizerParams.Adam() + DEFAULT_OPTIMIZER_TYPE: OptimizerParamsType = OptimizerParams.Adam() def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -42,7 +42,7 @@ def _get_default_lr(self, key: str) -> float: # region --- explicit properties and setters --- @property - def optimizer_params(self) -> dict[str, OptimizerType]: + def optimizer_params(self) -> dict[str, OptimizerParamsType | dict[str, OptimizerParamsType]]: return { key: params for key, params in [ @@ -56,7 +56,7 @@ def optimizer_params(self) -> dict[str, OptimizerType]: @optimizer_params.setter def optimizer_params(self, d: dict) -> None: """ - Takes a dictionary mapping optimizable keys to either an ``OptimizerType`` + Takes a dictionary mapping optimizable keys to either an ``OptimizerParamsType`` dataclass or a plain dict (with optional ``"name"``/``"type"`` and ``"lr"`` keys). Missing ``"name"`` / ``"lr"`` are filled from ``DEFAULT_OPTIMIZER_TYPE`` and ``_get_default_lr`` respectively. @@ -71,7 +71,7 @@ def optimizer_params(self, d: dict) -> None: d = {k: {} for k in d} for k, v in d.items(): - if isinstance(v, OptimizerType): + if isinstance(v, OptimizerParamsType): pass # already a dataclass, pass through elif isinstance(v, dict): if not v: @@ -82,7 +82,7 @@ def optimizer_params(self, d: dict) -> None: if "lr" not in v: v["lr"] = self._get_default_lr(k) else: - raise TypeError(f"Expected OptimizerType or dict for key '{k}', got {type(v)}") + raise TypeError(f"Expected OptimizerParamsType or dict for key '{k}', got {type(v)}") if k == "object": self.obj_model.optimizer_params = v @@ -131,7 +131,7 @@ def remove_optimizer(self, key: str) -> None: self.dset.remove_optimizer() @property - def scheduler_params(self) -> dict[str, SchedulerType]: + def scheduler_params(self) -> dict[str, SchedulerParamsType]: """Returns the parameters used to set the schedulers.""" return { "object": self.obj_model.scheduler_params, @@ -142,7 +142,7 @@ def scheduler_params(self) -> dict[str, SchedulerType]: @scheduler_params.setter def scheduler_params(self, d: dict) -> None: """ - Takes a dictionary mapping optimizable keys to either a ``SchedulerType`` + Takes a dictionary mapping optimizable keys to either a ``SchedulerParamsType`` dataclass or a plain dict. Keys not present in ``d`` are set to ``SchedulerParams.NoneScheduler()`` (disables scheduling for that model). @@ -178,7 +178,7 @@ def schedulers(self) -> dict[str, "torch.optim.lr_scheduler._LRScheduler"]: schedulers["dataset"] = self.dset.scheduler return schedulers - def set_schedulers(self, params: dict[str, SchedulerType], num_iter: int | None = None): + def set_schedulers(self, params: dict[str, SchedulerParamsType], num_iter: int | None = None): """Set schedulers for each model.""" for key, scheduler_params in params.items(): if key not in self.OPTIMIZABLE_VALS: diff --git a/src/quantem/imaging/drift.py b/src/quantem/imaging/drift.py index 424e18e6..ba29e417 100644 --- a/src/quantem/imaging/drift.py +++ b/src/quantem/imaging/drift.py @@ -1,9 +1,9 @@ +import warnings from collections.abc import Sequence from typing import List, Optional, Union import matplotlib.pyplot as plt import numpy as np -import warnings from numpy.typing import NDArray from scipy.interpolate import interp1d from scipy.ndimage import distance_transform_edt, gaussian_filter diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index f6343e4e..4838395d 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -232,11 +232,13 @@ def from_data( # --- Optimization Parameters --- - def get_optimization_parameters(self) -> list[nn.Parameter]: - """ - Get the parameters that should be optimized for this model. + def get_optimization_parameters(self) -> dict[str, list[torch.Tensor]]: + """Single param group keyed by DEFAULT_OPTIMIZER_KEY. + + Hyperparameters are baked by ``set_optimizer``, not here — return only the tensors, + matching the ``dict[str, list[tensor]]`` contract the object models use. """ - return list(self.parameters()) + return {self.DEFAULT_OPTIMIZER_KEY: list(self.parameters())} # --- Forward pass --- @abstractmethod @@ -643,6 +645,32 @@ def to(self, device: torch.device | str): self.device = device self.reconnect_optimizer_to_parameters() + # --- Save learned parameters --- + + def save_parameters(self, path: str): + """ + Saves the learned parameters to a file. + """ + torch.save( + { + "z1": self._z1_params.detach().cpu(), + "z3": self._z3_params.detach().cpu(), + "shifts": self._shifts_params.detach().cpu(), + }, + path, + ) + + def load_parameters(self, path: str): + """ + Loads the learned parameters from a file. + """ + data = torch.load(path) + self._z1_params = nn.Parameter(data["z1"]).to(self.device) + self._z3_params = nn.Parameter(data["z3"]).to(self.device) + self._shifts_params = nn.Parameter(data["shifts"]).to(self.device) + if self.optimizer is not None: + self.reconnect_optimizer_to_parameters() + class TomographyINRPretrainDataset(Dataset): """ diff --git a/src/quantem/tomography/object_models.py b/src/quantem/tomography/object_models.py index c5eb2406..c6b85b64 100644 --- a/src/quantem/tomography/object_models.py +++ b/src/quantem/tomography/object_models.py @@ -1,7 +1,7 @@ from abc import abstractmethod from copy import deepcopy from dataclasses import dataclass -from typing import Any, Callable, Generator +from typing import Any, Callable, Generator, Optional, cast import numpy as np import torch @@ -13,9 +13,11 @@ from quantem.core.ml.constraints import BaseConstraints, Constraints from quantem.core.ml.ddp import DDPMixin from quantem.core.ml.loss_functions import get_loss_module +from quantem.core.ml.models.model_base import PlanarDecompositionModel from quantem.core.ml.optimizer_mixin import OptimizerMixin from quantem.core.utils.rng import RNGMixin from quantem.tomography.dataset_models import TomographyINRPretrainDataset +from quantem.tomography.tomography_context import ReconstructionContext class ObjConstraintParams: @@ -99,13 +101,42 @@ class ObjINRConstraints(Constraints): sparsity: float = 0.0 _name: str = "obj_inr" - soft_constraint_keys = ["tv_vol", "sparsity"] + soft_constraint_keys = ["tv_vol", "tv_plane", "sparsity"] + hard_constraint_keys = ["positivity", "shrinkage"] + + @dataclass + class ObjTensorDecompConstraints(Constraints): + """ + Constraints for a tensor decomposition object representation. + + Attributes + ---------- + positivity : bool + If ``True``, enforces non-negative values in the reconstruction. + shrinkage : float + Shrinkage regularization strength; pushes values toward zero. + tv_vol : float + Total variation regularization weight for the 3-D volume. + soft_constraint_keys : list[str] + Constraint fields penalized softly during optimization. + hard_constraint_keys : list[str] + Constraint fields enforced strictly during optimization. + """ + + positivity: bool = True + shrinkage: float = 0.0 + tv_vol: float = 0.0 + tv_plane: float = 0.0 + sparsity: float = 0.0 + _name: str = "obj_tensor_decomp" + + soft_constraint_keys = ["tv_vol", "tv_plane", "sparsity"] hard_constraint_keys = ["positivity", "shrinkage"] @classmethod def parse_dict( cls, d: dict - ) -> "ObjConstraintParams.ObjPixelatedConstraints | ObjConstraintParams.ObjINRConstraints": + ) -> "ObjConstraintParams.ObjPixelatedConstraints | ObjConstraintParams.ObjINRConstraints | ObjConstraintParams.ObjTensorDecompConstraints": """ Instantiate an object constraint dataclass from a configuration dictionary. @@ -153,15 +184,26 @@ def parse_dict( return ObjConstraintParams.ObjPixelatedConstraints(**d) elif name == "obj_inr": return ObjConstraintParams.ObjINRConstraints(**d) + elif name == "obj_tensor_decomp": + return ObjConstraintParams.ObjTensorDecompConstraints(**d) else: raise ValueError(f"Unknown object constraint type: {name.lower()}") ObjConstraintsType = ( - ObjConstraintParams.ObjPixelatedConstraints | ObjConstraintParams.ObjINRConstraints + ObjConstraintParams.ObjPixelatedConstraints + | ObjConstraintParams.ObjINRConstraints + | ObjConstraintParams.ObjTensorDecompConstraints ) +def _unwrap(model: nn.Module | nn.parallel.DistributedDataParallel) -> PlanarDecompositionModel: + """Unwrap a DistributedDataParallel model to get the underlying module ONLY for tensor decomposition models.""" + if isinstance(model, nn.parallel.DistributedDataParallel): + return cast(PlanarDecompositionModel, model.module) + return cast(PlanarDecompositionModel, model) + + class ObjectBase(AutoSerialize, nn.Module, RNGMixin, OptimizerMixin): DEFAULT_LRS = { "object": 8e-6, @@ -190,77 +232,78 @@ def __init__( # --- Instantiation ---- - # --- Properties --- - @property - def shape(self) -> tuple[int, int, int]: - """ - Shape of the object (x, y, z). - """ - return self._shape - - @shape.setter - def shape(self, new_shape: tuple[int, int, int]): - self._shape = new_shape - - @property - def obj(self) -> torch.Tensor: - """ - Returns the object, should be implemented in subclasses. - """ - raise NotImplementedError - - @property - def model(self) -> nn.Module: - """ - Returns the model, should be implemented in subclasses. - """ - raise NotImplementedError - - @abstractmethod - def dtype(self) -> torch.dtype: - """ - Returns the dtype of the object. - """ - raise NotImplementedError - - @abstractmethod - def forward(self, *args, **kwargs) -> torch.Tensor: - """ - Forward pass, should be implemented in subclasses. Note for any nn.Module this is - a required method. - """ - raise NotImplementedError - - @abstractmethod - def reset(self) -> None: - """ - Reset the object, should be implemented in subclasses. - """ - raise NotImplementedError - - @property - def params(self) -> Generator[torch.nn.Parameter, None, None]: - """ - Get the parameters that should be optimized for this model. - - Should be implemented in subclasses. - """ - raise NotImplementedError - - # --- Helper Functions --- - def get_optimization_parameters(self) -> list[nn.Parameter]: - """ - Get the parameters that should be optimized for this model. - """ - return list(self.params) - - @abstractmethod # Each subclass should implement this. - def to(self, *args, **kwargs): - """ - Move the object to a device - """ - - raise NotImplementedError + # --- Properties --- + @property + def shape(self) -> tuple[int, int, int]: + """ + Shape of the object (x, y, z). + """ + return self._shape + + @shape.setter + def shape(self, new_shape: tuple[int, int, int]): + self._shape = new_shape + + @property + def obj(self) -> torch.Tensor: + """ + Returns the object, should be implemented in subclasses. + """ + raise NotImplementedError + + @property + def model(self) -> nn.Module: + """ + Returns the model, should be implemented in subclasses. + """ + raise NotImplementedError + + @property + def dtype(self) -> torch.dtype: + """ + Returns the dtype of the object. + """ + raise NotImplementedError + + @abstractmethod + def forward(self, coords: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Forward pass, should be implemented in subclasses. Note for any nn.Module this is + a required method. + """ + raise NotImplementedError + + @abstractmethod + def reset(self) -> None: + """ + Reset the object, should be implemented in subclasses. + """ + raise NotImplementedError + + @property + def params(self) -> Generator[torch.nn.Parameter, None, None]: + """ + Get the parameters that should be optimized for this model. + + Should be implemented in subclasses. + """ + raise NotImplementedError + + # --- Helper Functions --- + def get_optimization_parameters(self) -> "dict[str, list[torch.Tensor]]": + """Default: a single param group keyed by DEFAULT_OPTIMIZER_KEY. + + Hyperparameters are baked by ``set_optimizer``, not here — return only the tensors. + """ + return {self.DEFAULT_OPTIMIZER_KEY: list(self.params)} + + @abstractmethod # Each subclass should implement this. + def to(self, device: str | torch.device): + """ + Move the object to a device + """ + + raise NotImplementedError class ObjectConstraints(BaseConstraints, ObjectBase): # TODO: Ask Arthur why we still need this @@ -268,7 +311,7 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @abstractmethod - def get_tv_loss(self, **kwargs) -> torch.Tensor: + def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor: """ Get the TV loss for the object model. Must be implemented in each subclass. """ @@ -342,17 +385,9 @@ def obj(self, obj: torch.Tensor): def obj_view(self) -> np.ndarray: return self.obj.cpu().unsqueeze(0).numpy() - @property - def shape(self) -> tuple[int, int, int]: - return self._shape - - @shape.setter - def shape(self, shape: tuple[int, int, int]): - self._shape = shape - - @property - def soft_loss(self) -> torch.Tensor: - return self.apply_soft_constraints(self._obj) + # @property + # def soft_loss(self) -> torch.Tensor: + # return self.apply_soft_constraints(self._obj) @property def name(self) -> str: @@ -384,30 +419,32 @@ def apply_hard_constraints( # TODO: Need to implement the other hard constraints: Fourier Filter and Circular Mask. return obj2 - def apply_soft_constraints(self, obj: torch.Tensor) -> torch.Tensor: - soft_loss = torch.tensor(0.0, device=obj.device, dtype=obj.dtype, requires_grad=True) + def apply_soft_constraints(self, ctx: ReconstructionContext) -> torch.Tensor: + assert ctx.obj is not None, "ObjectPixelated requires ctx.obj to be set" + soft_loss = torch.tensor( + 0.0, device=ctx.obj.device, dtype=ctx.obj.dtype, requires_grad=True + ) if self.constraints.tv_vol > 0: - tv_loss = self.get_tv_loss( - obj.unsqueeze(0).unsqueeze(0), tv_weight=self.constraints.tv_vol - ) + tv_loss = self.get_tv_loss(ctx) soft_loss += tv_loss return soft_loss # --- Forward method --- - def forward(self, dummy_input=None) -> torch.Tensor: + def forward(self, coords=None) -> torch.Tensor: return self.obj # --- Defining the TV loss --- - def get_tv_loss(self, obj: torch.Tensor, tv_weight: float = 1e-3) -> torch.Tensor: # pyright: ignore[reportIncompatibleMethodOverride] -> get_tv_loss has different arguments depending on the object. - tv_d = torch.pow(obj[:, :, 1:, :, :] - obj[:, :, :-1, :, :], 2).sum() - tv_h = torch.pow(obj[:, :, :, 1:, :] - obj[:, :, :, :-1, :], 2).sum() - tv_w = torch.pow(obj[:, :, :, :, 1:] - obj[:, :, :, :, :-1], 2).sum() + def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor: + assert ctx.obj is not None, "ObjectPixelated requires ctx.obj to be set" + tv_d = torch.pow(ctx.obj[:, :, 1:, :, :] - ctx.obj[:, :, :-1, :, :], 2).sum() + tv_h = torch.pow(ctx.obj[:, :, :, 1:, :] - ctx.obj[:, :, :, :-1, :], 2).sum() + tv_w = torch.pow(ctx.obj[:, :, :, :, 1:] - ctx.obj[:, :, :, :, :-1], 2).sum() tv_loss = tv_d + tv_h + tv_w - return tv_loss * tv_weight / (torch.prod(torch.tensor(obj.shape))) + return tv_loss * self.constraints.tv_vol / (torch.prod(torch.tensor(ctx.obj.shape))) # --- Helper Functions --- - def to(self, device: str | torch.device): # pyright: ignore[reportIncompatibleMethodOverride] -> better to do this device change + def to(self, device: str | torch.device): if isinstance(device, str): device = torch.device(device) self._device = device @@ -435,7 +472,7 @@ def __init__( ) self._pretrain_losses = [] self._pretrain_lrs = [] - self.constraints: ObjConstraintsType = self.DEFAULT_CONSTRAINTS.copy() + self.constraints: ObjConstraintParams.ObjINRConstraints = self.DEFAULT_CONSTRAINTS.copy() # Register the network submodule (important: real nn.Module attribute) if model is not None: self.setup_distributed(device=device) @@ -499,40 +536,23 @@ def obj_view(self) -> np.ndarray: def apply_soft_constraints( self, - coords: torch.Tensor, - pred: torch.Tensor, + ctx: ReconstructionContext, ) -> torch.Tensor: - soft_loss = torch.tensor(0.0, device=pred.device) + soft_loss = torch.tensor(0.0, device=ctx.coords.device) if self.constraints.tv_vol > 0: - num_tv_samples = min(10_000, coords.shape[0]) - tv_indices = torch.randperm(coords.shape[0], device=coords.device)[:num_tv_samples] - - tv_coords = coords[tv_indices].detach().requires_grad_(True) - tv_densities_recomputed = self.model(tv_coords) - if isinstance(tv_densities_recomputed, tuple): - tv_densities_recomputed = tv_densities_recomputed[0] - - # Ensure shape is [num_samples, num_channels] - if tv_densities_recomputed.dim() == 1: - tv_densities_recomputed = tv_densities_recomputed.unsqueeze(-1) - - # Compute gradients for each channel - grad_outputs = torch.autograd.grad( - outputs=tv_densities_recomputed, - inputs=tv_coords, - grad_outputs=torch.ones_like(tv_densities_recomputed), - create_graph=True, - )[0] # Shape: [num_samples, coord_dim] - - # Compute TV loss - gradient magnitude per sample - grad_norm = torch.norm(grad_outputs, dim=1) # Shape: [num_samples] - soft_loss += self.constraints.tv_vol * grad_norm.mean() + assert ctx.coords is not None, ( + "coords must be provided for INR object model to compute the TV loss" + ) + soft_loss += self.get_tv_loss(ctx) if ( isinstance(self.constraints, ObjConstraintParams.ObjINRConstraints) and self.constraints.sparsity > 0 ): # NOTE: For the linter, I must make this :) - sparsity_loss = self.constraints.sparsity * torch.norm(pred, p=1) + assert ctx.pred is not None, ( + "pred must be provided for INR object model to compute the sparsity loss" + ) + sparsity_loss = self.constraints.sparsity * torch.norm(ctx.pred, p=1) soft_loss += sparsity_loss return soft_loss @@ -549,14 +569,42 @@ def apply_hard_constraints(self, pred: torch.Tensor) -> torch.Tensor: return pred + # --- Define get_tv_loss --- + + def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor: + """ + Compute the total variation loss for the INR model. + """ + assert ctx.coords is not None, "coords must be provided for INR object model" + num_tv_samples = min(10_000, ctx.coords.shape[0]) + tv_indices = torch.randperm(ctx.coords.shape[0], device=ctx.coords.device)[:num_tv_samples] + + tv_coords = ctx.coords[tv_indices].detach().requires_grad_(True) + tv_densities_recomputed = self.model(tv_coords) + if isinstance(tv_densities_recomputed, tuple): + tv_densities_recomputed = tv_densities_recomputed[0] + + # Ensure shape is [num_samples, num_channels] + if tv_densities_recomputed.dim() == 1: + tv_densities_recomputed = tv_densities_recomputed.unsqueeze(-1) + + # Compute gradients for each channel + grad_outputs = torch.autograd.grad( + outputs=tv_densities_recomputed, + inputs=tv_coords, + grad_outputs=torch.ones_like(tv_densities_recomputed), + create_graph=True, + )[0] # Shape: [num_samples, coord_dim] + + # Compute TV loss - gradient magnitude per sample + grad_norm = torch.norm(grad_outputs, dim=1) # Shape: [num_samples] + return self.constraints.tv_vol * grad_norm.mean() + # --- Optimization Parameters --- @property def params(self) -> Generator[torch.nn.Parameter, None, None]: return self.model.parameters() # type: ignore[attr-defined] - def get_optimization_parameters(self) -> list[nn.Parameter]: - return list(self.params) - # Pretraining @property def pretrained_weights(self) -> dict[str, torch.Tensor]: @@ -587,14 +635,6 @@ def dtype(self) -> torch.dtype: # TODO: This is a temporary solution to get the dtype of the object. return torch.float32 - @property - def shape(self) -> tuple[int, int, int]: - return self._shape - - @shape.setter - def shape(self, shape: tuple[int, int, int]): - self._shape = shape - # --- Helper Functions --- def rebuild_model(self): self._model = self.distribute_model(self._model) @@ -609,8 +649,10 @@ def reset(self): # --- Forward Method --- - def forward(self, coords: torch.Tensor) -> torch.Tensor: + def forward(self, coords: Optional[torch.Tensor] = None) -> torch.Tensor: """forward pass for the INR model""" + assert coords is not None, "ObjectINR.forward requires coords" + all_densities = self.model(coords) if all_densities.dim() > 1: @@ -808,43 +850,206 @@ def create_volume(self, return_vol: bool = False): self._obj = pred_full.detach().cpu() - def get_tv_loss( # pyright: ignore[reportIncompatibleMethodOverride] + def to(self, device: str | torch.device): # pyright: ignore[reportIncompatibleMethodOverride] -> better to do this device change + if isinstance(device, str): + device = torch.device(device) + self._device = device + if self.world_size == 1: + self._model = self._model.to(device) + elif not isinstance(self._model, torch.nn.parallel.DistributedDataParallel): + self.distribute_model(self._model) + self.reconnect_optimizer_to_parameters() + + +class ObjectTensorDecomp(ObjectINR): + DEFAULT_CONSTRAINTS = ObjConstraintParams.ObjTensorDecompConstraints() + + def __init__( self, - coords: torch.Tensor, - ) -> torch.Tensor: - tv_loss = torch.tensor(0.0, device=coords.device) + shape: tuple[int, int, int], + device: str = "cpu", + rng: np.random.Generator | int | None = None, + model: nn.Module | None = None, + _token: object | None = None, + ): + super().__init__( + shape=shape, + device=device, + rng=rng, + _token=self._token, + ) + self._pretrain_losses = [] + self._pretrain_lrs = [] + self.constraints: ObjConstraintParams.ObjTensorDecompConstraints = ( + self.DEFAULT_CONSTRAINTS.copy() + ) + # Register the network submodule (important: real nn.Module attribute) + if model is not None: + self.setup_distributed(device=device) + self._model = self.distribute_model(model) - num_tv_samples = min(10000, coords.shape[0]) - tv_indices = torch.randperm(coords.shape[0], device=coords.device)[:num_tv_samples] + @classmethod + def from_model( + cls, + model: nn.Module, + shape: tuple[int, int, int], + device: str = "cpu", + rng: np.random.Generator | int | None = None, + ): + obj_model = cls( + shape=shape, + device=device, + rng=rng, + model=model, # ✅ build/register in __init__ + ) - tv_coords = coords[tv_indices].detach().requires_grad_(True) + obj_model.setup_distributed(device=device) + obj_model.to(device) + return obj_model - tv_densities_recomputed = self.forward(tv_coords) + # --- Constraints --- - if tv_densities_recomputed.dim() > 1: - tv_densities_recomputed = tv_densities_recomputed.squeeze(-1) + def apply_soft_constraints(self, ctx: ReconstructionContext) -> torch.Tensor: + soft_loss = torch.tensor( + 0.0, device=ctx.pred.device if ctx.pred is not None else self.device + ) + if self.constraints.tv_vol > 0: + assert ctx.coords is not None, "Coordinates must be provided for TV loss" + assert ctx.pred is not None, "Prediction must be provided for TV loss" + soft_loss += self.get_tv_loss(ctx) - grad_outputs = torch.autograd.grad( - outputs=tv_densities_recomputed, - inputs=tv_coords, - grad_outputs=torch.ones_like(tv_densities_recomputed), - create_graph=True, - )[0] + if self.constraints.sparsity > 0: # NOTE: For the linter, I must make this :) + assert ctx.all_densities is not None, ( + "All densities must be provided for sparsity loss" + ) + sparsity_loss = self.constraints.sparsity * ctx.all_densities.abs().mean() + soft_loss += sparsity_loss - grad_norm = torch.norm(grad_outputs, dim=1) + return soft_loss + + # TV Losses + + def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor: + """ + Gets the summed total variational loss for the tensor decomposition model. - tv_loss += self.constraints.tv_vol * grad_norm.mean() + _get_plane_tv_loss: Total-variation across the planes. + _get_volume_tv_loss: Isotropic volume TV + """ + assert ctx.coords is not None, "Coordinates must be provided for TV loss" + assert ctx.pred is not None, "Prediction must be provided for TV loss" + tv_loss = torch.tensor(0.0, device=ctx.pred.device) + tv_loss += self._get_plane_tv_loss() + tv_loss += self.get_volume_tv_loss(ctx.coords) return tv_loss - def to(self, device: str | torch.device): # pyright: ignore[reportIncompatibleMethodOverride] -> better to do this device change - if isinstance(device, str): - device = torch.device(device) - self._device = device - if self.world_size == 1: - self._model = self._model.to(device) - elif not isinstance(self._model, torch.nn.parallel.DistributedDataParallel): - self.distribute_model(self._model) - self.reconnect_optimizer_to_parameters() + def _get_plane_tv_loss(self) -> torch.Tensor: + """ + Gets the total-variation across the planes. + """ + is_tilted = self.model.tilted + per_level = [] + + model = _unwrap(self.model) + for p in model.grids: + # p: (3*T, C, H, W) for TILTED, (3, C, H, W) for KPlanes + dh = (p[:, :, 1:, :] - p[:, :, :-1, :]).pow(2).mean(dim=(1, 2, 3)) + dw = (p[:, :, :, 1:] - p[:, :, :, :-1]).pow(2).mean(dim=(1, 2, 3)) + per_plane = dh + dw # (3*T,) or (3,) + + if is_tilted: + T = self.model.T + per_rotation = per_plane.view(T, 3).sum(dim=1) # sum 3 planes per rotation + level_tv = per_rotation.mean() # avg across rotations + else: + level_tv = per_plane.sum() + + per_level.append(level_tv) + + return self.constraints.tv_plane * torch.stack(per_level).sum() + + def get_volume_tv_loss(self, coords: torch.Tensor) -> torch.Tensor: + """ + Isotropic volume TV via finite differences. Same form as the autograd + version (L1 of gradient L2-norm) but avoids double-backward, so it + works for KPlanesTILTED, CPTilted, and anything else. + """ + num_tv_samples = min(10_000, coords.shape[0]) + tv_indices = torch.randperm(coords.shape[0], device=coords.device)[:num_tv_samples] + tv_coords = coords[tv_indices] # (N, 3) + + model = _unwrap(self.model) + h = 2.0 / min(model.resolution) + + pred = model(tv_coords) + if isinstance(pred, tuple): + pred = pred[0] + if pred.dim() == 1: + pred = pred.unsqueeze(-1) # (N, 1) + + grads = [] + for axis in range(3): + offset = torch.zeros(3, device=tv_coords.device) + offset[axis] = h + shifted_pred = self.model(tv_coords + offset) + if isinstance(shifted_pred, tuple): + shifted_pred = shifted_pred[0] + if shifted_pred.dim() == 1: + shifted_pred = shifted_pred.unsqueeze(-1) + grads.append((shifted_pred - pred) / h) # (N, 1) + + grad_stack = torch.stack(grads, dim=-1) # (N, C, 3) + grad_norm = torch.norm(grad_stack, dim=-1) # (N, C) + + return self.constraints.tv_vol * grad_norm.mean() + + def apply_hard_constraints(self, pred: torch.Tensor) -> torch.Tensor: + """ + Apply hard constraints to the predicted values of the INR model. + """ + + if self.constraints.positivity: + pred = torch.clamp(pred, min=0.0, max=None) + if self.constraints.shrinkage: + pred = torch.max(pred - self.constraints.shrinkage, torch.zeros_like(pred)) + + return pred + + # --- Optimization Parameters --- + @property + def params(self) -> Generator[torch.nn.Parameter, None, None]: + """ + Returns the optimization parameters, here we also check if PPLR is used and return the appropriate parameters. + """ + + return self.model.parameters() # type: ignore[attr-defined] + + def get_optimization_parameters(self) -> "dict[str, list[torch.Tensor]]": + """PPLR: per-key param groups (hyperparameters are baked by set_optimizer).""" + model = _unwrap(self.model) + return {key: list(model.get_params()[key]) for key in model.param_keys} + + def _normalize_optimizer_params(self, params): + """ObjectTensorDecomp requires a dict matching model.param_keys.""" + if not isinstance(params, dict) or self._is_single_optimizer_dict(params): + raise TypeError( + f"ObjectTensorDecomp requires dict[str, OptimizerParamsType] keyed by " + f"param_keys; got {type(params)}" + ) + model = _unwrap(self.model) + expected = set(model.param_keys) + got = set(params.keys()) + if got != expected: + raise ValueError( + f"optimizer_params keys must match model.param_keys: " + f"got {got}, expected {expected}" + ) + return super()._normalize_optimizer_params(params) + + def pretrain(self) -> None: + raise NotImplementedError( + "Tensor decomposition pretraining is not usually required, and for TILTED there is a two-phase warmup approach." + ) -ObjectModelType = ObjectPixelated | ObjectINR +ObjectModelType = ObjectPixelated | ObjectINR | ObjectTensorDecomp diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py index 907ad994..8976272e 100644 --- a/src/quantem/tomography/tomography.py +++ b/src/quantem/tomography/tomography.py @@ -10,6 +10,7 @@ from quantem.core.io.serialize import load as autoserialize_load from quantem.core.ml.loss_functions import get_loss_module +from quantem.core.ml.models.kplanes import CPTilted from quantem.core.utils.filter import gaussian_filter_2d_stack, gaussian_kernel_1d from quantem.core.utils.tomography_utils import torch_phase_cross_correlation from quantem.tomography.dataset_models import ( @@ -25,9 +26,11 @@ ObjConstraintsType, ObjectINR, ObjectPixelated, + ObjectTensorDecomp, ) from quantem.tomography.radon.radon import iradon_torch, radon_torch from quantem.tomography.tomography_base import TomographyBase +from quantem.tomography.tomography_context import ReconstructionContext from quantem.tomography.tomography_opt import TomographyOpt @@ -41,7 +44,7 @@ class Tomography(TomographyOpt, TomographyBase): def from_models( cls, dset: DatasetModelType, - obj_model: ObjectINR, + obj_model: ObjectINR | ObjectTensorDecomp, logger: LoggerTomography | None = None, device: str = "cuda", verbose: int | bool = True, @@ -179,7 +182,9 @@ def reconstruct( consistency_loss = torch.tensor(0.0, device=self.device) total_loss = torch.tensor(0.0, device=self.device) epoch_soft_constraint_loss = torch.tensor(0.0, device=self.device) - if isinstance(self.obj_model, ObjectINR): + if isinstance(self.obj_model, ObjectINR) or isinstance( + self.obj_model, ObjectTensorDecomp + ): self.obj_model.model.train() else: raise NotImplementedError( @@ -201,7 +206,7 @@ def reconstruct( with torch.autocast( device_type=self.device.type, dtype=torch.bfloat16, - enabled=True, + enabled=False, ): all_coords = self.dset.get_coords(batch, N, curr_num_samples_per_ray) @@ -214,9 +219,14 @@ def reconstruct( ) pred = integrated_densities.float() - soft_constraints_loss = 0.0 - if self.num_epochs > 0: - soft_constraints_loss = self.obj_model.apply_soft_constraints(all_coords, pred) + + soft_constraints_loss = self.obj_model.apply_soft_constraints( + ctx=ReconstructionContext( + coords=all_coords, + pred=pred, + all_densities=all_densities, + ) + ) target = batch["target_value"].to(self.device, non_blocking=True).float() @@ -235,6 +245,25 @@ def reconstruct( total_loss += batch_loss.detach() consistency_loss += batch_consistency_loss.detach() + if isinstance(self.obj_model.model, CPTilted): + if a0 == 0: + prev_R = self.obj_model.model.so3.as_matrix().detach().clone() + elif (a0 + 1) % 20 == 0: + R_now = self.obj_model.model.so3.as_matrix().detach() + # Cumulative angular change per rotation over the last 20 iters. + # trace(R_prev^T R_now) = 1 + 2*cos(theta), so theta = acos((trace - 1) / 2). + rel_trace = torch.einsum("tij,tij->t", prev_R, R_now) + angle = torch.acos(((rel_trace - 1) / 2).clamp(-1, 1)) # (T,) radians + angle_deg = torch.rad2deg(angle) + per_tau_str = ", ".join(f"{a:.2f}°" for a in angle_deg.tolist()) + print( + f"iter {a0}: 20-iter τ change " + f"max={angle_deg.max().item():.2f}°, " + f"mean={angle_deg.mean().item():.2f}°, " + f"per-τ=[{per_tau_str}]" + ) + prev_R = R_now.clone() + if self.world_size > 1: dist.all_reduce(total_loss, dist.ReduceOp.AVG) dist.all_reduce(consistency_loss, dist.ReduceOp.AVG) diff --git a/src/quantem/tomography/tomography_base.py b/src/quantem/tomography/tomography_base.py index 58909bdf..419b5ede 100644 --- a/src/quantem/tomography/tomography_base.py +++ b/src/quantem/tomography/tomography_base.py @@ -11,6 +11,7 @@ ObjConstraintsType, ObjectINR, ObjectModelType, + ObjectTensorDecomp, ) @@ -51,7 +52,7 @@ def __init__( self._val_losses: list[float] = [] self._lrs: dict[str, list] = {} # DDP Initialization - if isinstance(obj_model, ObjectINR): + if isinstance(obj_model, ObjectINR) or isinstance(obj_model, ObjectTensorDecomp): self.setup_distributed(device=device) if self.global_rank == 0: print("Setting up DDP for obj_model") diff --git a/src/quantem/tomography/tomography_context.py b/src/quantem/tomography/tomography_context.py new file mode 100644 index 00000000..d322287c --- /dev/null +++ b/src/quantem/tomography/tomography_context.py @@ -0,0 +1,30 @@ +from dataclasses import dataclass +from typing import Optional +from quantem.core.ml.constraints import BaseContext + +import torch + + +@dataclass +class ReconstructionContext(BaseContext): + """ + Handles all reconstruction parameters to be passed into object models. + + Subclasses will pick whatever parameter they need + - Pixelated reads ".volume" + - INR reads ".coords" and recomputes via the model. + - TensorDecomp reads ".coords" and ".pred" (and ".all densities") + + Variable descriptions: + - volume: Reconstructed object (volume). + - coords: Used for INR reconstructions to provide the coordinates to the model. + - pred: Predicted values per coordinate position from the model. + - all_densities: Integrated densities per ray from the model. + - obj: Object model (INR, TensorDecomp, etc.). + """ + + volume: Optional[torch.Tensor] = None + coords: Optional[torch.Tensor] = None + pred: Optional[torch.Tensor] = None + all_densities: Optional[torch.Tensor] = None + obj: Optional[torch.Tensor] = None diff --git a/src/quantem/tomography/tomography_opt.py b/src/quantem/tomography/tomography_opt.py index 16c846ab..cf990379 100644 --- a/src/quantem/tomography/tomography_opt.py +++ b/src/quantem/tomography/tomography_opt.py @@ -2,7 +2,11 @@ import torch -from quantem.core.ml.optimizer_mixin import OptimizerParams, OptimizerType, SchedulerType +from quantem.core.ml.optimizer_mixin import ( + OptimizerParams, + OptimizerParamsType, + SchedulerParamsType, +) from quantem.tomography.tomography_base import TomographyBase @@ -12,7 +16,7 @@ class TomographyOpt(TomographyBase): """ OPTIMIZABLE_VALS = ["object", "pose"] - DEFAULT_OPTIMIZER_TYPE = "adam" + DEFAULT_OPTIMIZER_TYPE: OptimizerParamsType = OptimizerParams.Adam() def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -27,7 +31,7 @@ def _get_default_lr(self, key: str) -> float: raise ValueError(f"Unknown optimization key: {key}") @property - def optimizer_params(self) -> dict[str, OptimizerType]: + def optimizer_params(self) -> dict[str, OptimizerParamsType | dict[str, OptimizerParamsType]]: return { key: params for key, params in [ @@ -38,7 +42,7 @@ def optimizer_params(self) -> dict[str, OptimizerType]: } @optimizer_params.setter - def optimizer_params(self, d: dict[str, OptimizerType] | dict[str, dict]): + def optimizer_params(self, d: dict[str, OptimizerParamsType] | dict[str, dict]): """Set the optimizer parameters.""" if isinstance(d, (tuple, list)): d = {k: {} for k in d} @@ -52,8 +56,8 @@ def optimizer_params(self, d: dict[str, OptimizerType] | dict[str, dict]): if k not in targets: raise ValueError(f"Unknown optimization key: {k}") - if not isinstance(v, OptimizerType): - v = OptimizerParams.parse_dict(v) + # if not isinstance(v, OptimizerParamsType): + # v = OptimizerParams.parse_dict(v) targets[k].optimizer_params = v @@ -100,7 +104,7 @@ def remove_optimizer(self, key: str): raise ValueError(f"Unknown optimization key: {key}") @property - def scheduler_params(self) -> dict[str, SchedulerType]: + def scheduler_params(self) -> dict[str, SchedulerParamsType]: """Returns the parameters used to set the schedulers.""" return { "object": self.obj_model.scheduler_params, @@ -136,7 +140,7 @@ def schedulers(self) -> dict[str, torch.optim.lr_scheduler._LRScheduler]: return schedulers def set_schedulers( - self, params: Mapping[str, SchedulerType | dict], num_iter: int | None = None + self, params: Mapping[str, SchedulerParamsType | dict], num_iter: int | None = None ): for key, scheduler_params in params.items(): if key == "object": diff --git a/tests/datastructures/test_dataset3d_show.py b/tests/datastructures/test_dataset3d_show.py index 27c18f41..7938c66f 100644 --- a/tests/datastructures/test_dataset3d_show.py +++ b/tests/datastructures/test_dataset3d_show.py @@ -31,16 +31,19 @@ def extract_frame_indices_from_figure(fig): class TestShowInputValidation: """Test that invalid inputs raise clear errors.""" - @pytest.mark.parametrize("kwargs,match", [ - ({"step": 0}, "cannot be zero"), - ({"start": 100}, "out of bounds"), - ({"start": -100}, "out of bounds"), - ({"start": 5, "end": 5}, "No frames to display"), - ({"ncols": 0}, "ncols must be >= 1"), - ({"ncols": -1}, "ncols must be >= 1"), - ({"max": 0}, "max must be >= 1"), - ({"max": -1}, "max must be >= 1"), - ]) + @pytest.mark.parametrize( + "kwargs,match", + [ + ({"step": 0}, "cannot be zero"), + ({"start": 100}, "out of bounds"), + ({"start": -100}, "out of bounds"), + ({"start": 5, "end": 5}, "No frames to display"), + ({"ncols": 0}, "ncols must be >= 1"), + ({"ncols": -1}, "ncols must be >= 1"), + ({"max": 0}, "max must be >= 1"), + ({"max": -1}, "max must be >= 1"), + ], + ) def test_raises_value_error(self, dataset_with_10_frames, kwargs, match): with pytest.raises(ValueError, match=match): dataset_with_10_frames.show(**kwargs) @@ -49,39 +52,44 @@ def test_raises_value_error(self, dataset_with_10_frames, kwargs, match): class TestShowFrameSelection: """Test frame selection with start, end, step, max combinations.""" - @pytest.mark.parametrize("kwargs,expected_indices", [ - ({}, list(range(20))), - ({"max": 5}, [0, 1, 2, 3, 4]), - ({"max": None}, list(range(100))), - ({"start": 90}, list(range(90, 100))), - ({"start": 95, "max": 3}, [95, 96, 97]), - ]) + @pytest.mark.parametrize( + "kwargs,expected_indices", + [ + ({}, list(range(20))), + ({"max": 5}, [0, 1, 2, 3, 4]), + ({"max": None}, list(range(100))), + ({"start": 90}, list(range(90, 100))), + ({"start": 95, "max": 3}, [95, 96, 97]), + ], + ) def test_large_dataset(self, dataset_with_100_frames, kwargs, expected_indices): fig, _ = dataset_with_100_frames.show(returnfig=True, **kwargs) assert extract_frame_indices_from_figure(fig) == expected_indices plt.close(fig) - @pytest.mark.parametrize("kwargs,expected_indices", [ - # Default shows all frames (< max) - ({}, list(range(10))), - # Start and end - ({"start": 5}, [5, 6, 7, 8, 9]), - ({"end": 5}, [0, 1, 2, 3, 4]), - # Step - ({"step": 2}, [0, 2, 4, 6, 8]), - ({"step": 3}, [0, 3, 6, 9]), - ({"start": 2, "end": 8, "step": 2}, [2, 4, 6]), - # Negative start index - ({"start": -1, "max": 1}, [9]), - ({"start": -3, "max": 2}, [7, 8]), - # Negative step (reverse order) - ({"start": 9, "step": -1}, [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]), - ({"start": 9, "end": 4, "step": -1}, [9, 8, 7, 6, 5]), - ({"start": 9, "step": -2}, [9, 7, 5, 3, 1]), - ({"start": 9, "step": -1, "max": 3}, [9, 8, 7]), - ]) + @pytest.mark.parametrize( + "kwargs,expected_indices", + [ + # Default shows all frames (< max) + ({}, list(range(10))), + # Start and end + ({"start": 5}, [5, 6, 7, 8, 9]), + ({"end": 5}, [0, 1, 2, 3, 4]), + # Step + ({"step": 2}, [0, 2, 4, 6, 8]), + ({"step": 3}, [0, 3, 6, 9]), + ({"start": 2, "end": 8, "step": 2}, [2, 4, 6]), + # Negative start index + ({"start": -1, "max": 1}, [9]), + ({"start": -3, "max": 2}, [7, 8]), + # Negative step (reverse order) + ({"start": 9, "step": -1}, [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]), + ({"start": 9, "end": 4, "step": -1}, [9, 8, 7, 6, 5]), + ({"start": 9, "step": -2}, [9, 7, 5, 3, 1]), + ({"start": 9, "step": -1, "max": 3}, [9, 8, 7]), + ], + ) def test_small_dataset(self, dataset_with_10_frames, kwargs, expected_indices): fig, _ = dataset_with_10_frames.show(returnfig=True, **kwargs) assert extract_frame_indices_from_figure(fig) == expected_indices plt.close(fig) - diff --git a/tests/ml/test_optimizermixin.py b/tests/ml/test_optimizermixin.py index d829af41..d6347ce6 100644 --- a/tests/ml/test_optimizermixin.py +++ b/tests/ml/test_optimizermixin.py @@ -348,3 +348,106 @@ def test_parse_invalid_name_type_raises(self): def test_parse_default_name_is_none(self): result = SchedulerParams.parse_dict({}) assert isinstance(result, SchedulerParams.NoneScheduler) + + +# ─── OptimizerMixin.set_optimizer / PPLR behavior ─────────────────────────── + +from quantem.core import config # noqa: E402 +from quantem.core.ml.optimizer_mixin import OptimizerMixin # noqa: E402 + +torch = pytest.importorskip("torch") if config.get("has_torch") else None +requires_torch = pytest.mark.skipif( + not config.get("has_torch"), reason="requires torch" +) + + +def _param(value=1.0): + return torch.nn.Parameter(torch.tensor([value])) + + +class _FakeModel(OptimizerMixin): + """Minimal concrete OptimizerMixin for exercising set_optimizer/reset_optimizer. + + ``groups`` is the dict[str, list[Parameter]] returned by get_optimization_parameters. + Pass ``raise_on_params=True`` to prove the disable path short-circuits before the call. + """ + + def __init__(self, groups, raise_on_params=False): + super().__init__() + self._groups = groups + self._raise_on_params = raise_on_params + + def get_optimization_parameters(self): + if self._raise_on_params: + raise AssertionError("get_optimization_parameters should not have been called") + return self._groups + + +@requires_torch +class TestSetOptimizer: + def test_single_optimizer_build(self): + model = _FakeModel({"default": [_param()]}) + model.set_optimizer(OptimizerParams.Adam(lr=1e-3)) + assert model.has_optimizer() + assert isinstance(model.optimizer, torch.optim.Adam) + assert len(model.optimizer.param_groups) == 1 + assert model.optimizer.param_groups[0]["lr"] == 1e-3 + + def test_single_optimizer_dict_shorthand(self): + model = _FakeModel({"default": [_param()]}) + model.set_optimizer({"name": "sgd", "lr": 5e-2, "momentum": 0.9}) + assert isinstance(model.optimizer, torch.optim.SGD) + assert model.optimizer.param_groups[0]["lr"] == 5e-2 + assert model.optimizer.param_groups[0]["momentum"] == 0.9 + + def test_none_optimizer_removes_without_touching_params(self): + # raise_on_params proves the disable path short-circuits before get_optimization_parameters + model = _FakeModel({"default": [_param()]}, raise_on_params=True) + model.set_optimizer(OptimizerParams.NoneOptimizer()) + assert model.optimizer is None + assert not model.has_optimizer() + assert model.optimizer_params == { + OptimizerMixin.DEFAULT_OPTIMIZER_KEY: OptimizerParams.NoneOptimizer() + } + + def test_multi_group_pplr_applies_per_group_lr(self): + model = _FakeModel({"descan": [_param()], "scan_positions": [_param(2.0)]}) + model.set_optimizer( + {"descan": OptimizerParams.SGD(lr=1e-2), "scan_positions": OptimizerParams.SGD(lr=1e-3)} + ) + groups = model.optimizer.param_groups + assert len(groups) == 2 + # key order is preserved (descan, then scan_positions) + assert groups[0]["lr"] == 1e-2 + assert groups[1]["lr"] == 1e-3 + + def test_key_mismatch_raises(self): + model = _FakeModel({"default": [_param()]}) + with pytest.raises(ValueError, match="do not match"): + model.set_optimizer( + {"descan": OptimizerParams.SGD(lr=1e-2), "scan_positions": OptimizerParams.SGD()} + ) + + def test_mixed_optimizer_classes_raises(self): + model = _FakeModel({"a": [_param()], "b": [_param()]}) + with pytest.raises(ValueError, match="same optimizer type"): + model.set_optimizer( + {"a": OptimizerParams.Adam(lr=1e-3), "b": OptimizerParams.SGD(lr=1e-3)} + ) + + def test_reset_optimizer_on_unconfigured_model_is_noop(self): + model = _FakeModel({"default": [_param()]}) + # fresh model defaults to {"default": NoneOptimizer()} + model.reset_optimizer() + assert model.optimizer is None + assert not model.has_optimizer() + + def test_set_scheduler_base_lr_uses_max_group_lr(self): + model = _FakeModel({"descan": [_param()], "scan_positions": [_param(2.0)]}) + model.set_optimizer( + {"descan": OptimizerParams.SGD(lr=1e-2), "scan_positions": OptimizerParams.SGD(lr=1e-3)} + ) + model.set_scheduler(SchedulerParams.Plateau(), num_iter=10) + assert model.scheduler is not None + # Plateau min_lr defaults to base_LR / 20, with base_LR = max group lr (1e-2) + assert model.scheduler.min_lrs[0] == pytest.approx(1e-2 / 20) diff --git a/tests/tomography/conftest.py b/tests/tomography/conftest.py new file mode 100644 index 00000000..06957095 --- /dev/null +++ b/tests/tomography/conftest.py @@ -0,0 +1,72 @@ +"""Shared fixtures and markers for the tomography test suite. + +The suite is split into three tiers (see the plan / individual modules): + +* CPU, always-on -- radon, utils, object/dataset models, optimizer-param wiring. +* CPU, slow -- conventional SIRT/FBP reconstruction (``--runslow``). +* GPU, slow -- INR / KPlanes reconstruction (``requires_gpu`` + ``--runslow``). + +``torch`` and ``scikit-image`` are core dependencies of quantem, so they are always +importable; ``requires_torch`` is kept only for parity with the existing test style. The +meaningful gate is ``requires_gpu`` (CI runs CPU-only) combined with ``@pytest.mark.slow``. +""" + +import numpy as np +import pytest +import torch + +from quantem.core import config +from quantem.tomography.utils import rot_ZXZ + +# --- Markers --------------------------------------------------------------- +requires_torch = pytest.mark.skipif(not config.get("has_torch"), reason="requires torch") +requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") + + +@pytest.fixture +def torch_device() -> str: + """Device for torch-only tests. + + ``cuda:0`` when available, else ``cpu``. Note: object models that go through + ``setup_distributed`` must be built on ``cuda`` when CUDA is present (the CPU path is + only valid when no CUDA device exists), so torch-only construction tests follow this + fixture rather than hard-coding ``"cpu"``. + """ + return "cuda:0" if torch.cuda.is_available() else "cpu" + + +# --- Synthetic data -------------------------------------------------------- +def make_tilt_series(volume: torch.Tensor, angles: np.ndarray) -> np.ndarray: + """Project a volume into a tilt series with quantem's own forward model. + + Mirrors ``tomography_00_generate_phantom``: rotate by ``rot_ZXZ`` (Euler ZXZ, tilt on + the X axis) then sum along the beam axis. Using this rather than ``radon_torch`` keeps + the synthetic data consistent with the geometry the reconstructors assume. + """ + projections = [] + vol = volume.unsqueeze(0) # (1, Z, Y, X) + for angle in angles: + rotated = rot_ZXZ(vol, 0.0, float(angle), 0.0, device="cpu") + projections.append(rotated[0].sum(0)) + return torch.stack(projections).numpy().astype(np.float32) + + +@pytest.fixture(scope="module") +def phantom_volume() -> np.ndarray: + """Small deterministic (32, 32, 32) phantom with a couple of solid blocks.""" + vol = np.zeros((32, 32, 32), dtype=np.float32) + vol[8:24, 10:20, 12:22] = 1.0 + vol[18:26, 6:12, 16:24] = 0.6 + return vol + + +@pytest.fixture(scope="module") +def tilt_angles() -> np.ndarray: + """Nine tilt angles spanning -70..70 degrees.""" + return np.linspace(-70, 70, 9).astype(np.float32) + + +@pytest.fixture(scope="module") +def tilt_series(phantom_volume: np.ndarray, tilt_angles: np.ndarray) -> np.ndarray: + """Synthetic tilt series, shape (n_angles, 32, 32).""" + return make_tilt_series(torch.from_numpy(phantom_volume), tilt_angles) diff --git a/tests/tomography/test_dataset_models.py b/tests/tomography/test_dataset_models.py new file mode 100644 index 00000000..78ded427 --- /dev/null +++ b/tests/tomography/test_dataset_models.py @@ -0,0 +1,106 @@ +"""Tests for ``quantem.tomography.dataset_models``. + +Covers constraint parsing, the pixelated dataset (validation, normalisation, tilt-angle +convention, pose-parameter materialisation) and the INR / pretrain datasets. +""" + +import numpy as np +import pytest +import torch + +from quantem.tomography.dataset_models import ( + DatasetConstraintParams, + DatasetValue, + TomographyINRDataset, + TomographyINRPretrainDataset, + TomographyPixDataset, +) + +from .conftest import requires_torch + + +class TestDatasetConstraintParse: + def test_parse_base_by_name(self): + c = DatasetConstraintParams.parse_dict({"name": "base_tomography_dataset", "tv_zs": 0.1}) + assert isinstance(c, DatasetConstraintParams.BaseTomographyDatasetConstraints) + assert c.tv_zs == 0.1 + + def test_parse_base_by_type_key(self): + c = DatasetConstraintParams.parse_dict( + {"type": "base_tomography_dataset", "tv_shifts": 0.2} + ) + assert c.tv_shifts == 0.2 + + def test_unknown_name_raises(self): + with pytest.raises(ValueError): + DatasetConstraintParams.parse_dict({"name": "nope"}) + + +def _stack(nang=5, n=12, seed=0): + rng = np.random.default_rng(seed) + return (rng.random((nang, n, n)) * 10).astype(np.float32) + + +class TestTomographyPixDataset: + def test_wrong_projection_axis_raises(self): + # projections must live on axis 0 (i.e. fewer than the image dims). + bad = np.zeros((20, 5, 5), dtype=np.float32) + with pytest.raises(ValueError): + TomographyPixDataset.from_data(bad, np.linspace(-60, 60, 20).astype(np.float32)) + + def test_tilt_angles_are_negated(self): + angles = np.linspace(-40, 60, 5).astype(np.float32) + d = TomographyPixDataset.from_data(_stack(), angles) + np.testing.assert_allclose(d.tilt_angles.numpy(), -angles, atol=1e-5) + + def test_normalised_by_95th_quantile(self): + d = TomographyPixDataset.from_data(_stack(), np.linspace(-60, 60, 5).astype(np.float32)) + q95 = torch.quantile(d.tilt_stack, 0.95) + assert torch.isclose(q95, torch.tensor(1.0), atol=1e-4) + + def test_reference_idx_and_learnable_tilts(self): + # negated angles -> [40, 15, -10, -35, -60]; smallest |angle| is index 2. + angles = np.linspace(-40, 60, 5).astype(np.float32) + d = TomographyPixDataset.from_data(_stack(), angles) + assert d.reference_tilt_idx == 2 + assert d.learnable_tilts == 4 + + def test_forward_returns_dataset_value(self): + angles = np.linspace(-40, 60, 5).astype(np.float32) + d = TomographyPixDataset.from_data(_stack(nang=5, n=12), angles) + out = d.forward(0) + assert isinstance(out, DatasetValue) + assert out.target.shape == (12, 12) + assert out.tilt_angle == pytest.approx(float(-angles[0])) + + def test_to_materialises_pose_parameters(self): + d = TomographyPixDataset.from_data(_stack(), np.linspace(-60, 60, 5).astype(np.float32)) + d.to("cpu") + assert isinstance(d.z1_params, torch.nn.Parameter) + assert d.shifts_params.shape == (d.learnable_tilts, 2) + + +@requires_torch +class TestTomographyINRDataset: + def test_len_is_projections_times_pixels(self): + d = TomographyINRDataset.from_data( + _stack(nang=5, n=12), np.linspace(-60, 60, 5, dtype="f4") + ) + assert len(d) == 5 * 12 * 12 + + def test_getitem_keys(self): + d = TomographyINRDataset.from_data( + _stack(nang=5, n=12), np.linspace(-60, 60, 5, dtype="f4") + ) + item = d[0] + assert {"phi", "pixel_i", "pixel_j", "projection_idx", "target_value"} <= set(item.keys()) + + +class TestTomographyINRPretrainDataset: + def test_len_and_getitem(self): + vol = torch.rand(1, 8, 8, 8) + ds = TomographyINRPretrainDataset(pretrain_target=vol) + assert len(ds) == 8**3 + item = ds[0] + assert set(item.keys()) == {"coords", "target"} + assert item["coords"].shape == (3,) diff --git a/tests/tomography/test_object_models.py b/tests/tomography/test_object_models.py new file mode 100644 index 00000000..f9f1922c --- /dev/null +++ b/tests/tomography/test_object_models.py @@ -0,0 +1,159 @@ +"""Tests for ``quantem.tomography.object_models``. + +The constraint-parsing and ``ObjectPixelated`` tests are pure CPU. The INR / tensor-decomp +construction tests are ``requires_torch`` and follow the ``torch_device`` fixture (they must +be built on CUDA when CUDA is present; see conftest). +""" + +import numpy as np +import pytest +import torch + +from quantem.tomography.object_models import ( + ObjConstraintParams, + ObjectBase, + ObjectINR, + ObjectPixelated, + ObjectTensorDecomp, +) +from quantem.tomography.tomography_context import ReconstructionContext + +from .conftest import requires_torch + + +class TestObjConstraintParse: + def test_parse_pixelated_by_name(self): + c = ObjConstraintParams.parse_dict({"name": "obj_pixelated", "tv_vol": 0.01}) + assert isinstance(c, ObjConstraintParams.ObjPixelatedConstraints) + assert c.tv_vol == 0.01 + + def test_parse_inr_by_type_key(self): + c = ObjConstraintParams.parse_dict({"type": "obj_inr", "sparsity": 0.05}) + assert isinstance(c, ObjConstraintParams.ObjINRConstraints) + assert c.sparsity == 0.05 + + def test_parse_tensor_decomp(self): + c = ObjConstraintParams.parse_dict({"name": "obj_tensor_decomp", "tv_plane": 0.1}) + assert isinstance(c, ObjConstraintParams.ObjTensorDecompConstraints) + assert c.tv_plane == 0.1 + + def test_missing_name_raises(self): + with pytest.raises(ValueError): + ObjConstraintParams.parse_dict({"tv_vol": 0.1}) + + def test_unknown_name_raises(self): + with pytest.raises(ValueError): + ObjConstraintParams.parse_dict({"name": "obj_nope"}) + + def test_constraint_key_partitions(self): + c = ObjConstraintParams.ObjPixelatedConstraints() + assert "positivity" in c.hard_constraint_keys + assert "tv_vol" in c.soft_constraint_keys + + +class TestObjectPixelatedConstruction: + def test_from_uniform_is_zeros(self): + obj = ObjectPixelated.from_uniform(shape=(8, 8, 8), device="cpu") + assert obj.shape == (8, 8, 8) + assert torch.allclose(obj.obj, torch.zeros(8, 8, 8)) + assert obj.obj_type == "pixelated" + + def test_from_array_numpy(self): + arr = np.random.default_rng(0).random((6, 6, 6)).astype(np.float32) + obj = ObjectPixelated.from_array(arr, device="cpu") + assert obj.shape == (6, 6, 6) + assert torch.allclose(obj.obj, torch.from_numpy(arr)) + assert obj.dtype == torch.float32 + + def test_from_array_torch_is_copied(self): + t = torch.ones(4, 4, 4) + obj = ObjectPixelated.from_array(t, device="cpu") + t += 5.0 + assert torch.allclose(obj.obj, torch.ones(4, 4, 4)) # original copy untouched + + def test_obj_view_shape(self): + obj = ObjectPixelated.from_uniform(shape=(5, 6, 7), device="cpu") + assert obj.obj_view.shape == (1, 5, 6, 7) + + def test_forward_returns_obj(self): + obj = ObjectPixelated.from_array(torch.full((4, 4, 4), 2.0), device="cpu") + assert torch.allclose(obj.forward(), obj.obj) + + +class TestObjectPixelatedConstraints: + def test_positivity_clamps_negatives(self): + obj = ObjectPixelated.from_array(torch.full((4, 4, 4), -1.0), device="cpu") + obj.constraints.positivity = True + assert torch.all(obj.obj >= 0.0) + + def test_positivity_off_keeps_negatives(self): + obj = ObjectPixelated.from_array(torch.full((4, 4, 4), -1.0), device="cpu") + obj.constraints.positivity = False + assert torch.all(obj.obj < 0.0) + + def test_shrinkage_subtracts_then_floors(self): + obj = ObjectPixelated.from_array(torch.full((4, 4, 4), 1.0), device="cpu") + obj.constraints.positivity = False + obj.constraints.shrinkage = 0.25 + assert torch.allclose(obj.obj, torch.full((4, 4, 4), 0.75)) + + def test_tv_loss_scales_with_weight(self): + ctx = ReconstructionContext(obj=torch.rand(1, 1, 8, 8, 8)) + obj = ObjectPixelated.from_uniform(shape=(8, 8, 8), device="cpu") + obj.constraints.tv_vol = 1.0 + loss1 = obj.get_tv_loss(ctx) + obj.constraints.tv_vol = 2.0 + loss2 = obj.get_tv_loss(ctx) + assert torch.isclose(loss2, 2.0 * loss1) + + def test_soft_constraint_zero_when_tv_off(self): + ctx = ReconstructionContext(obj=torch.rand(1, 1, 8, 8, 8)) + obj = ObjectPixelated.from_uniform(shape=(8, 8, 8), device="cpu") + obj.constraints.tv_vol = 0.0 + assert float(obj.apply_soft_constraints(ctx).detach()) == 0.0 + + +class TestFactoryGuard: + def test_objectbase_requires_token(self): + with pytest.raises(RuntimeError): + ObjectBase(shape=(4, 4, 4)) + + +@requires_torch +class TestObjectINR: + def test_from_model_builds(self, torch_device): + from quantem.core.ml.inr import HSiren + + model = HSiren(in_features=3, out_features=1, hidden_layers=1, hidden_features=8) + obj = ObjectINR.from_model(model, shape=(16, 16, 16), device=torch_device) + assert obj.shape == (16, 16, 16) + assert obj.model is not None + + def test_optimization_parameters_single_group(self, torch_device): + from quantem.core.ml.inr import HSiren + + model = HSiren(in_features=3, out_features=1, hidden_layers=1, hidden_features=8) + obj = ObjectINR.from_model(model, shape=(16, 16, 16), device=torch_device) + groups = obj.get_optimization_parameters() + assert list(groups.keys()) == ["default"] + assert len(groups["default"]) > 0 + + +@requires_torch +class TestObjectTensorDecomp: + def _model(self, n=16): + from quantem.core.ml.models.kplanes import KPlanesTILTED + + return KPlanesTILTED( + M_features=2, resolution=(n, n, n), multiscale_res_multipliers=[1], T=2 + ) + + def test_pplr_optimization_parameter_keys(self, torch_device): + obj = ObjectTensorDecomp.from_model(self._model(), shape=(16, 16, 16), device=torch_device) + keys = set(obj.get_optimization_parameters().keys()) + assert keys == {"grids", "sigma_net", "so3"} + + def test_pretrain_not_implemented(self, torch_device): + obj = ObjectTensorDecomp.from_model(self._model(), shape=(16, 16, 16), device=torch_device) + with pytest.raises(NotImplementedError): + obj.pretrain() diff --git a/tests/tomography/test_radon.py b/tests/tomography/test_radon.py new file mode 100644 index 00000000..4f61837a --- /dev/null +++ b/tests/tomography/test_radon.py @@ -0,0 +1,111 @@ +"""Tests for the pure-torch Radon transform (``quantem.tomography.radon.radon``). + +All CPU, deterministic. Cross-checked against scikit-image where a ground truth helps. +""" + +import numpy as np +import pytest +import torch + +from quantem.tomography.radon.radon import ( + get_fourier_filter_torch, + iradon_torch, + radon_torch, +) + + +def _disk(n: int, cy: int, cx: int, r: int) -> torch.Tensor: + yy, xx = np.mgrid[0:n, 0:n] + return torch.from_numpy((((yy - cy) ** 2 + (xx - cx) ** 2) < r**2).astype(np.float32)) + + +class TestRadonShapes: + def test_2d_input_returns_angles_by_pixels(self): + img = _disk(64, 32, 32, 10) + theta = torch.linspace(0, 180, 30) + sino = radon_torch(img, theta=theta) + assert sino.shape == (30, 64) + + def test_batched_input_returns_batch_angles_pixels(self): + imgs = torch.stack([_disk(48, 24, 20, 8), _disk(48, 24, 28, 8)]) + theta = torch.linspace(0, 180, 20) + sino = radon_torch(imgs, theta=theta) + assert sino.shape == (2, 20, 48) + + def test_default_theta_is_180_angles(self): + sino = radon_torch(_disk(32, 16, 16, 6)) + assert sino.shape == (180, 32) + + def test_iradon_shapes(self): + sino = radon_torch(_disk(40, 20, 20, 8), theta=torch.linspace(0, 180, 25)) + rec = iradon_torch(sino, theta=torch.linspace(0, 180, 25)) + assert rec.shape == (40, 40) + + def test_iradon_output_size_override(self): + sino = radon_torch(_disk(40, 20, 20, 8), theta=torch.linspace(0, 180, 25)) + rec = iradon_torch(sino, theta=torch.linspace(0, 180, 25), output_size=32) + assert rec.shape == (32, 32) + + +class TestFourierFilter: + def test_even_size_ok(self): + f = get_fourier_filter_torch(64, "ramp") + assert f.shape == (1, 64) + + def test_odd_size_raises(self): + with pytest.raises(ValueError): + get_fourier_filter_torch(63, "ramp") + + def test_unknown_filter_raises(self): + with pytest.raises(ValueError): + get_fourier_filter_torch(64, "not-a-filter") + + def test_none_filter_is_all_ones(self): + f = get_fourier_filter_torch(64, None) + assert torch.allclose(f, torch.ones_like(f)) + + @pytest.mark.parametrize("name", ["ramp", "shepp-logan", "cosine", "hamming", "hann"]) + def test_named_filters_run(self, name): + f = get_fourier_filter_torch(64, name) + assert f.shape == (1, 64) + assert torch.isfinite(f).all() + + +class TestRadonBehaviour: + def test_circular_mask_zeros_corners(self): + """The forward transform masks to the inscribed circle, so corner mass is dropped.""" + img = torch.ones(32, 32) + full = img.sum() + sino = radon_torch(img, theta=torch.tensor([0.0])) + # A single 0-degree projection sums columns; total equals the masked mass < full. + assert sino.sum() < full + + def test_iradon_circle_zeros_outside(self): + sino = radon_torch(_disk(48, 24, 24, 10), theta=torch.linspace(0, 180, 30)) + rec = iradon_torch(sino, theta=torch.linspace(0, 180, 30), circle=True) + n = rec.shape[0] + yy, xx = np.mgrid[0:n, 0:n] + outside = ((yy - n // 2) ** 2 + (xx - n // 2) ** 2) > (n // 2) ** 2 + assert torch.allclose(rec[outside], torch.zeros(int(outside.sum()))) + + def test_roundtrip_recovers_structure(self): + disk = _disk(64, 32, 24, 9) + theta = torch.linspace(0, 180, 60) + rec = iradon_torch(radon_torch(disk, theta=theta), theta=theta, filter_name="ramp") + corr = np.corrcoef(disk.numpy().ravel(), rec.numpy().ravel())[0, 1] + assert corr > 0.9 + + +class TestRadonVsSkimage: + """Loose cross-check against scikit-image's reference implementation.""" + + def test_forward_matches_skimage(self): + sk = pytest.importorskip("skimage.transform") + n = 64 + disk = _disk(n, n // 2, n // 2, 12) + theta = np.linspace(0.0, 180.0, 45, endpoint=False).astype(np.float32) + ours = radon_torch(disk, theta=torch.from_numpy(theta)).numpy() # (A, N) + ref = sk.radon(disk.numpy(), theta=theta, circle=True).T # skimage: (N, A) -> (A, N) + # Different interpolation conventions; require strong agreement, not equality. + corr = np.corrcoef(ours.ravel(), ref.ravel())[0, 1] + assert corr > 0.95 diff --git a/tests/tomography/test_tomography_conventional.py b/tests/tomography/test_tomography_conventional.py new file mode 100644 index 00000000..a7cfb5e0 --- /dev/null +++ b/tests/tomography/test_tomography_conventional.py @@ -0,0 +1,72 @@ +"""End-to-end conventional (SIRT / FBP) reconstruction. + +CPU, deterministic, but marked ``slow`` because it runs a (tiny) iterative reconstruction. +Reconstructions are capped at a few iterations: the suite checks behaviour and wiring (loss +decreases, output stays physical) rather than convergence quality, so spatial agreement with +the phantom is only a loose lower bound. +""" + +import numpy as np +import pytest + +from quantem.tomography.dataset_models import TomographyPixDataset +from quantem.tomography.object_models import ObjConstraintParams, ObjectPixelated +from quantem.tomography.tomography import TomographyConventional + +pytestmark = pytest.mark.slow + + +def _build(tilt_series, tilt_angles, n): + dset = TomographyPixDataset.from_data( + tilt_series, tilt_angles, learn_shift=False, learn_tilt_axis=False + ) + obj = ObjectPixelated.from_uniform(shape=(n, n, n), device="cpu") + return TomographyConventional.from_models( + dset=dset, obj_model=obj, device="cpu", verbose=False + ) + + +class TestSIRT: + def test_loss_decreases_and_output_physical(self, phantom_volume, tilt_series, tilt_angles): + n = phantom_volume.shape[0] + tomo = _build(tilt_series, tilt_angles, n) + tomo.reconstruct( + num_iter=4, + mode="sirt", + obj_constraints=ObjConstraintParams.ObjPixelatedConstraints(positivity=True), + ) + losses = tomo.epoch_losses + assert tomo.num_epochs == 4 + assert losses[-1] < losses[0] + rec = tomo.obj_model.obj.detach().cpu().numpy() + assert np.isfinite(rec).all() + assert rec.min() >= 0.0 # positivity + + def test_recon_correlates_with_phantom(self, phantom_volume, tilt_series, tilt_angles): + n = phantom_volume.shape[0] + tomo = _build(tilt_series, tilt_angles, n) + tomo.reconstruct( + num_iter=4, + mode="sirt", + obj_constraints=ObjConstraintParams.ObjPixelatedConstraints(positivity=True), + ) + rec = tomo.obj_model.obj.detach().cpu().numpy() + corr = np.corrcoef(rec.ravel(), phantom_volume.ravel())[0, 1] + assert corr > 0.15 # loose: only a handful of iterations + + def test_obj_constraints_accepts_dict(self, phantom_volume, tilt_series, tilt_angles): + n = phantom_volume.shape[0] + tomo = _build(tilt_series, tilt_angles, n) + tomo.reconstruct(num_iter=2, mode="sirt", obj_constraints={"name": "obj_pixelated"}) + assert tomo.num_epochs == 2 + + +class TestFBP: + def test_fbp_runs_single_epoch(self, phantom_volume, tilt_series, tilt_angles): + n = phantom_volume.shape[0] + tomo = _build(tilt_series, tilt_angles, n) + tomo.reconstruct(num_iter=5, mode="fbp") + # FBP breaks after the first epoch regardless of num_iter. + assert tomo.num_epochs == 1 + rec = tomo.obj_model.obj.detach().cpu().numpy() + assert np.isfinite(rec).all() diff --git a/tests/tomography/test_tomography_inr.py b/tests/tomography/test_tomography_inr.py new file mode 100644 index 00000000..fbd5f1a1 --- /dev/null +++ b/tests/tomography/test_tomography_inr.py @@ -0,0 +1,96 @@ +"""End-to-end INR / KPlanes (tensor-decomposition) reconstruction. + +These exercise the full learned-reconstruction path (model + pose optimisation, autocast, +spawned dataloader workers), so they are gated behind ``requires_gpu`` and ``slow`` and only +run locally with ``--runslow``. Reconstructions are capped at 4 iterations; the assertion is +loss-decreases plus finite output, not convergence quality. + +The ``num_workers=2`` is required, not incidental: ``setup_dataloader`` hard-codes +``multiprocessing_context="spawn"``, which is invalid with ``num_workers=0``. +""" + +import numpy as np +import pytest +import torch + +from quantem.tomography.dataset_models import TomographyINRDataset +from quantem.tomography.object_models import ObjectTensorDecomp +from quantem.tomography.tomography import Tomography +from quantem.tomography.tomography_lite import TomographyLiteConv, TomographyLiteINR + +from .conftest import make_tilt_series, requires_gpu + +pytestmark = [requires_gpu, pytest.mark.slow] + +DEVICE = "cuda:0" + + +@pytest.fixture(scope="module") +def small_phantom(): + vol = torch.zeros(1, 24, 24, 24) + vol[0, 6:18, 6:14, 8:16] = 1.0 + angles = np.linspace(-60, 60, 7).astype(np.float32) + series = make_tilt_series(vol[0], angles) + return series, angles + + +class TestLiteINR: + def test_reconstruct_reduces_loss(self, small_phantom): + series, angles = small_phantom + tomo = TomographyLiteINR.from_dataset( + tilt_series=series, tilt_angles=angles, device=DEVICE + ) + tomo.reconstruct(num_iter=4, num_workers=2, batch_size=256, learn_pose=True) + losses = tomo.epoch_losses + assert len(losses) == 4 + assert losses[-1] < losses[0] + view = tomo.obj_model.obj_view + assert view.shape == (1, 24, 24, 24) + assert np.isfinite(view).all() + + +class TestLiteConv: + def test_smoke(self, small_phantom): + series, angles = small_phantom + tomo = TomographyLiteConv.from_dataset( + tilt_series=series, tilt_angles=angles, device=DEVICE + ) + tomo.reconstruct(num_iter=3, mode="sirt") + assert tomo.num_epochs == 3 + assert np.isfinite(tomo.obj_model.obj.detach().cpu().numpy()).all() + + +class TestKPlanes: + def test_pplr_reconstruct_reduces_loss(self, small_phantom): + from quantem.core.ml.models.kplanes import KPlanesTILTED + from quantem.core.ml.optimizer_mixin import OptimizerParams, SchedulerParams + + series, angles = small_phantom + n = series.shape[1] + model = KPlanesTILTED( + M_features=2, resolution=(n, n, n), multiscale_res_multipliers=[1], T=2 + ) + obj = ObjectTensorDecomp.from_model(model, shape=(n, n, n), device=DEVICE) + dset = TomographyINRDataset.from_data(series, angles) + tomo = Tomography.from_models(dset=dset, obj_model=obj, device=DEVICE, verbose=False) + tomo.reconstruct( + optimizer_params={ + "object": { + "grids": OptimizerParams.Adam(lr=1e-2), + "sigma_net": OptimizerParams.Adam(lr=1e-3), + "so3": OptimizerParams.Adam(lr=1e-2), + }, + "pose": OptimizerParams.Adam(lr=1e-2), + }, + scheduler_params={ + "object": SchedulerParams.CosineAnnealing(T_max=4), + "pose": SchedulerParams.CosineAnnealing(T_max=4), + }, + num_iter=4, + batch_size=256, + num_samples_per_ray=20, + num_workers=2, + ) + losses = tomo.epoch_losses + assert len(losses) == 4 + assert losses[-1] < losses[0] diff --git a/tests/tomography/test_tomography_opt.py b/tests/tomography/test_tomography_opt.py new file mode 100644 index 00000000..4ec56a13 --- /dev/null +++ b/tests/tomography/test_tomography_opt.py @@ -0,0 +1,143 @@ +"""Tests for the tomography optimizer / scheduler wiring (``TomographyOpt``). + +This is the surface the PPLR ``OptimizerParamsType`` / ``SchedulerParamsType`` refactor +touched. In particular, ``test_set_optimizers_builds_object_and_pose`` and the PPLR test +regression-guard the pose-optimizer path: ``TomographyDatasetBase.get_optimization_parameters`` +must return a ``dict[str, list[tensor]]`` (it previously returned a ``list`` and crashed +``set_optimizer`` with ``TypeError: unhashable type: 'dict'``). + +Construction only -- no forward passes -- so these run on CPU under CI. +""" + +import numpy as np +import pytest + +from quantem.core.ml.optimizer_mixin import OptimizerParams, SchedulerParams +from quantem.tomography.dataset_models import TomographyINRDataset +from quantem.tomography.object_models import ObjectINR, ObjectTensorDecomp +from quantem.tomography.tomography import Tomography + +from .conftest import requires_torch + + +def _tilts(nang=5, n=12): + rng = np.random.default_rng(0) + angles = np.linspace(-60, 60, nang).astype(np.float32) + stack = rng.random((nang, n, n)).astype(np.float32) + return stack, angles + + +def _inr_tomography(device): + from quantem.core.ml.inr import HSiren + + model = HSiren(in_features=3, out_features=1, hidden_layers=1, hidden_features=8) + obj = ObjectINR.from_model(model, shape=(16, 16, 16), device=device) + stack, angles = _tilts() + dset = TomographyINRDataset.from_data(stack, angles) + return Tomography.from_models(dset=dset, obj_model=obj, device=device, verbose=False) + + +def _td_tomography(device): + from quantem.core.ml.models.kplanes import KPlanesTILTED + + model = KPlanesTILTED( + M_features=2, resolution=(16, 16, 16), multiscale_res_multipliers=[1], T=2 + ) + obj = ObjectTensorDecomp.from_model(model, shape=(16, 16, 16), device=device) + stack, angles = _tilts() + dset = TomographyINRDataset.from_data(stack, angles) + return Tomography.from_models(dset=dset, obj_model=obj, device=device, verbose=False) + + +@requires_torch +class TestOptimizerParams: + def test_setter_getter_roundtrip(self, torch_device): + tomo = _inr_tomography(torch_device) + tomo.optimizer_params = { + "object": OptimizerParams.Adam(lr=1e-3), + "pose": OptimizerParams.Adam(lr=1e-2), + } + assert set(tomo.optimizer_params.keys()) == {"object", "pose"} + + def test_unknown_key_raises(self, torch_device): + tomo = _inr_tomography(torch_device) + with pytest.raises(ValueError): + tomo.optimizer_params = {"banana": OptimizerParams.Adam(lr=1e-3)} + + def test_set_optimizers_builds_object_and_pose(self, torch_device): + """Regression guard: the pose path must not raise (see module docstring).""" + tomo = _inr_tomography(torch_device) + tomo.optimizer_params = { + "object": OptimizerParams.Adam(lr=1e-3), + "pose": OptimizerParams.Adam(lr=1e-2), + } + tomo.set_optimizers() + assert set(tomo.optimizers.keys()) == {"object", "pose"} + + def test_current_lrs(self, torch_device): + tomo = _inr_tomography(torch_device) + tomo.optimizer_params = { + "object": OptimizerParams.Adam(lr=1e-3), + "pose": OptimizerParams.Adam(lr=1e-2), + } + tomo.set_optimizers() + lrs = tomo.get_current_lrs() + assert set(lrs.keys()) == {"object", "pose"} + assert lrs["object"] == pytest.approx(1e-3) + assert lrs["pose"] == pytest.approx(1e-2) + + def test_remove_optimizer(self, torch_device): + tomo = _inr_tomography(torch_device) + tomo.optimizer_params = { + "object": OptimizerParams.Adam(lr=1e-3), + "pose": OptimizerParams.Adam(lr=1e-2), + } + tomo.set_optimizers() + tomo.remove_optimizer("object") + assert "object" not in tomo.optimizers + assert "pose" in tomo.optimizers + + def test_pplr_object_groups(self, torch_device): + """Per-parameter LR: object optimizer carries one torch param group per key.""" + tomo = _td_tomography(torch_device) + tomo.optimizer_params = { + "object": { + "grids": OptimizerParams.Adam(lr=1e-2), + "sigma_net": OptimizerParams.Adam(lr=1e-3), + "so3": OptimizerParams.Adam(lr=1e-2), + }, + "pose": OptimizerParams.Adam(lr=1e-2), + } + tomo.set_optimizers() + assert len(tomo.optimizers["object"].param_groups) == 3 + assert "pose" in tomo.optimizers + + +@requires_torch +class TestSchedulerParams: + def test_scheduler_setter_getter(self, torch_device): + tomo = _inr_tomography(torch_device) + tomo.scheduler_params = { + "object": SchedulerParams.CosineAnnealing(T_max=10), + "pose": SchedulerParams.CosineAnnealing(T_max=10), + } + assert set(tomo.scheduler_params.keys()) == {"object", "pose"} + + def test_set_schedulers_builds(self, torch_device): + tomo = _inr_tomography(torch_device) + tomo.optimizer_params = { + "object": OptimizerParams.Adam(lr=1e-3), + "pose": OptimizerParams.Adam(lr=1e-2), + } + tomo.set_optimizers() + tomo.scheduler_params = { + "object": SchedulerParams.CosineAnnealing(T_max=10), + "pose": SchedulerParams.CosineAnnealing(T_max=10), + } + tomo.set_schedulers(tomo.scheduler_params, num_iter=10) + assert set(tomo.schedulers.keys()) == {"object", "pose"} + + def test_bad_scheduler_type_raises(self, torch_device): + tomo = _inr_tomography(torch_device) + with pytest.raises(TypeError): + tomo.obj_model.scheduler_params = 123 diff --git a/tests/tomography/test_utils.py b/tests/tomography/test_utils.py new file mode 100644 index 00000000..1fa6a47c --- /dev/null +++ b/tests/tomography/test_utils.py @@ -0,0 +1,78 @@ +"""Tests for ``quantem.tomography.utils``: 1D total-variation loss and the +differentiable ZXZ rotation operators. All CPU.""" + +import pytest +import torch + +from quantem.tomography.utils import ( + differentiable_rotz_vectorized, + rot_ZXZ, + tv_loss_1d, +) + + +class TestTVLoss1D: + def test_constant_input_is_zero(self): + assert tv_loss_1d(torch.ones(10)) == 0.0 + + def test_known_value_mean(self): + # diffs are [1, 1, 1], abs-mean = 1.0 + x = torch.tensor([0.0, 1.0, 2.0, 3.0]) + assert torch.isclose(tv_loss_1d(x, reduction="mean"), torch.tensor(1.0)) + + def test_known_value_sum(self): + x = torch.tensor([0.0, 1.0, 2.0, 3.0]) + assert torch.isclose(tv_loss_1d(x, reduction="sum"), torch.tensor(3.0)) + + def test_reduction_none_shape(self): + x = torch.zeros(2, 5) + out = tv_loss_1d(x, reduction="none") + assert out.shape == (2, 4) + + def test_bad_reduction_raises(self): + with pytest.raises(ValueError): + tv_loss_1d(torch.zeros(4), reduction="median") + + +def _block_volume(n: int = 16) -> torch.Tensor: + """(1, n, n, n) volume with an off-centre block so rotations are detectable.""" + vol = torch.zeros(1, n, n, n) + vol[0, 4:12, 4:10, 5:11] = 1.0 + return vol + + +class TestRotations: + def test_zero_rotation_is_identity(self): + vol = _block_volume() + out = rot_ZXZ(vol, 0.0, 0.0, 0.0, device="cpu") + assert torch.max(torch.abs(out - vol)) < 1e-4 + + def test_rotation_preserves_mass(self): + vol = _block_volume() + rotated = rot_ZXZ(vol, 0.0, 30.0, 0.0, device="cpu") + rel_err = abs(float(rotated.sum()) - float(vol.sum())) / float(vol.sum()) + assert rel_err < 0.02 + + def test_accepts_python_float_and_tensor_angle(self): + vol = _block_volume() + out_float = rot_ZXZ(vol, 0.0, 25.0, 0.0, device="cpu") + out_tensor = rot_ZXZ( + vol, + torch.tensor(0.0), + torch.tensor(25.0), + torch.tensor(0.0), + device="cpu", + ) + assert torch.allclose(out_float, out_tensor, atol=1e-5) + + def test_rotation_changes_volume(self): + vol = _block_volume() + rotated = rot_ZXZ(vol, 0.0, 90.0, 0.0, device="cpu") + assert torch.max(torch.abs(rotated - vol)) > 0.1 + + def test_gradient_flows_through_rotation(self): + vol = _block_volume().requires_grad_(True) + out = differentiable_rotz_vectorized(vol, torch.tensor(20.0)) + out.sum().backward() + assert vol.grad is not None + assert torch.isfinite(vol.grad).all()