Per-parameter Learning Rate (PPLR) Implementation, Tensor Decomposition Models, and Tomography Updates

.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/
 

src/quantem/__init__.py

@@ -1,4 +1,5 @@
 from pkgutil import extend_path
+
 __path__ = extend_path(__path__, __name__)
 
 from importlib.metadata import version

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.
         """