Tomography INR Speed-ups and quantem-cuda API Built-in

src/quantem/core/ml/inr.py

@@ -111,18 +111,20 @@ def _build(self) -> None:
         self.net = nn.Sequential(*net_list)
 
         if self.winner_initialization:
-            if type(self.winner_initialization) is int:
-                rng = torch.Generator()
-                rng.manual_seed(self.winner_initialization)
-            else:
-                rng = torch.Generator()
-                rng.manual_seed(42)
+            seed = self.winner_initialization if type(self.winner_initialization) is int else 42
+            rng = torch.Generator()
+            rng.manual_seed(seed)
+            # torch.randn_like ignores generators, so the noise must come from
+            # torch.randn with the seeded generator -- otherwise the "winner" seed
+            # silently has no effect and the perturbation is not reproducible.
             with torch.no_grad():
-                self.net[0].linear.weight += (  # type: ignore[reportAttributeAccessIssue]
-                    torch.randn_like(self.net[0].linear.weight) * 5 / self.first_omega_0  # type:ignore
-                )
-                self.net[1].linear.weight += (  # type: ignore[reportAttributeAccessIssue]
-                    torch.randn_like(self.net[1].linear.weight) * 0.1 / self.hidden_omega_0  # type:ignore
+                w0 = self.net[0].linear.weight  # type: ignore[reportAttributeAccessIssue]
+                w0 += torch.randn(w0.shape, generator=rng, dtype=w0.dtype) * 5 / self.first_omega_0
+                w1 = self.net[1].linear.weight  # type: ignore[reportAttributeAccessIssue]
+                w1 += (
+                    torch.randn(w1.shape, generator=rng, dtype=w1.dtype)
+                    * 0.1
+                    / self.hidden_omega_0
                 )
 
     def forward(self, coords: torch.Tensor) -> torch.Tensor:

src/quantem/core/ml/models/kplanes.py

@@ -169,6 +169,7 @@ class KPlanes(PPLR, TensorDecompositionModel):
     """
     K-Planes model adapted from Fridovich-Keil et al., https://arxiv.org/abs/2301.10241
     """
+
     def __init__(
         self,
         # Grid parameters
@@ -199,6 +200,15 @@ def __init__(
         self.concat_features = concat_features
         self.density_activation = density_activation
 
+        # All three planes share one (3, C, res[1], res[0]) tensor, which ignores
+        # res[2]: an anisotropic resolution would silently give the XZ/YZ planes the
+        # wrong grid along z. Refuse it rather than misallocate.
+        if len(set(self.resolution)) != 1:
+            raise ValueError(
+                f"KPlanes currently requires an isotropic resolution, got {self.resolution}; "
+                "the plane grids are allocated as (res[1], res[0]) for all three axis pairs."
+            )
+
         self.grids = nn.ParameterList()
         self.feature_dim = 0
         for res_mult in self.multiscale_res_multipliers:
@@ -320,14 +330,18 @@ def interpolate_ms_features_tilted(
     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))
-    )
+    # Plane axis layout: XY=(0,1), ZX=(2,0), YZ=(1,2). Stacking views avoids the
+    # per-call index-tensor allocation (a host-to-device copy) and the
+    # (T, 3, B, 3) expand+gather this used to do.
+    x, y, z = rotated.unbind(-1)  # each (T, B)
+    coords = torch.stack(
+        (
+            torch.stack((x, y), dim=-1),
+            torch.stack((z, x), dim=-1),
+            torch.stack((y, z), dim=-1),
+        ),
+        dim=1,
+    )  # (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)

src/quantem/core/ml/optimizer_mixin.py

@@ -294,14 +294,15 @@ class Plateau:
         _name: str = "plateau"
 
         def params(self, base_LR: float, num_iter: int | None = None) -> dict:
-            if self.min_lr is None:
-                self.min_lr = self.min_lr_factor * base_LR
+            # Derived values stay local: params() must not mutate the dataclass, or a
+            # shared instance bakes in the first optimizer's base LR for every later one.
+            min_lr = self.min_lr if self.min_lr is not None else self.min_lr_factor * base_LR
             return {
                 "mode": self.mode,
                 "factor": self.factor,
                 "patience": self.patience,
                 "threshold": self.threshold,
-                "min_lr": self.min_lr,
+                "min_lr": min_lr,
                 "cooldown": self.cooldown,
             }
 
@@ -332,13 +333,12 @@ def params(self, base_LR: float, num_iter: int | None = None) -> dict:
             if effective_num_iter is None:
                 raise ValueError("num_iter must be set if num_iter is not provided")
 
-            self.num_iter = effective_num_iter
-
+            gamma = self.gamma
             if self.factor is not None:
-                self.gamma = self.factor ** (1.0 / effective_num_iter)
+                gamma = self.factor ** (1.0 / effective_num_iter)
 
             return {
-                "gamma": self.gamma,
+                "gamma": gamma,
             }
 
     @dataclass
@@ -385,13 +385,11 @@ class Cyclic:
         _name: str = "cyclic"
 
         def params(self, base_LR: float, num_iter: int | None = None) -> dict:
-            if self.base_lr is None:
-                self.base_lr = self.base_lr_factor * base_LR
-            if self.max_lr is None:
-                self.max_lr = self.max_lr_factor * base_LR
+            base_lr = self.base_lr if self.base_lr is not None else self.base_lr_factor * base_LR
+            max_lr = self.max_lr if self.max_lr is not None else self.max_lr_factor * base_LR
             return {
-                "base_lr": self.base_lr,
-                "max_lr": self.max_lr,
+                "base_lr": base_lr,
+                "max_lr": max_lr,
                 "step_size_up": self.step_size_up,
                 "step_size_down": self.step_size_down,
                 "mode": self.mode,
@@ -426,12 +424,11 @@ def params(self, base_LR: float, num_iter: int | None = None) -> dict:
                 raise ValueError(
                     "total_iters must be set if num_iter is not provided"
                 )  # Should never be reached
-            if self.total_iters is None:
-                self.total_iters = num_iter
+            total_iters = self.total_iters if self.total_iters is not None else num_iter
             return {
                 "start_factor": self.start_factor,
                 "end_factor": self.end_factor,
-                "total_iters": self.total_iters,
+                "total_iters": total_iters,
             }
 
     @dataclass
@@ -459,10 +456,9 @@ def params(self, base_LR: float, num_iter: int | None = None) -> dict:
                 raise ValueError(
                     "T_max must be set if num_iter is not provided"
                 )  # Should never be reached
-            if self.T_max is None:
-                self.T_max = num_iter
+            T_max = self.T_max if self.T_max is not None else num_iter
             return {
-                "T_max": self.T_max,
+                "T_max": T_max,
                 "eta_min": self.eta_min,
             }
 
@@ -572,7 +568,9 @@ def _normalize_optimizer_params(
         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)}")
+            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)}
@@ -597,7 +595,9 @@ def scheduler_params(self, params: SchedulerParamsType | dict):
         if isinstance(params, dict):
             params = SchedulerParams.parse_dict(d=params)
         if not isinstance(params, SchedulerParamsType):
-            raise TypeError(f"scheduler parameters must be a SchedulerParamsType, got {type(params)}")
+            raise TypeError(
+                f"scheduler parameters must be a SchedulerParamsType, got {type(params)}"
+            )
         self._scheduler_params = params
 
     @abstractmethod
@@ -688,7 +688,9 @@ def _build_optimizer(self, opt_params, param_groups) -> "torch.optim.Optimizer":
                 raise NotImplementedError(f"Unknown optimizer type: {opt_params}")
 
     def set_scheduler(
-        self, scheduler_params: SchedulerParamsType | 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: