Added correct typehints for scaler and composition weights

Author: pfebrer

src/metatrain/pet/hypers.py

@@ -5,13 +5,11 @@
 
 from typing_extensions import TypedDict
 
-from metatrain.utils.hypers import (
- CompositionWeightsDict,
- ScalingWeightsDict,
- init_with_defaults,
-)
+from metatrain.utils.additive import FixedCompositionWeights
+from metatrain.utils.hypers import init_with_defaults
 from metatrain.utils.long_range import LongRangeHypers
 from metatrain.utils.loss import LossSpecification
+from metatrain.utils.scaler import FixedScalerWeights
 
 from .modules.finetuning import FinetuneHypers
 
@@ -117,10 +115,10 @@ class PETTrainerHypers(TypedDict):
  """Interval to save checkpoints."""
  scale_targets: bool = True
  """Normalize targets to unit std during training."""
- fixed_composition_weights: CompositionWeightsDict = {}
+ fixed_composition_weights: FixedCompositionWeights = {}
  """Weights for atomic contributions."""
- fixed_scaling_weights: ScalingWeightsDict = {}
-
+ fixed_scaling_weights: FixedScalerWeights = {}
+ """Weights for target scaling."""
  per_structure_targets: list[str] = []
  """Targets to calculate per-structure losses."""
  num_workers: Optional[int] = None

src/metatrain/pet/modules/finetuning.py

@@ -54,7 +54,7 @@ class LoRaFinetuneHypers(TypedDict):
 
 class HeadsFinetuneHypers(TypedDict):
  """Hyperparameters for heads finetuning of PET models.
- 
+
  Freezes all model parameters except for the prediction heads
  and last layers.
  """

src/metatrain/pet/tests/test_functionality.py

@@ -3,16 +3,17 @@
 import torch
 from metatomic.torch import ModelOutput, System
 from omegaconf import OmegaConf
+from pydantic import ValidationError
 
-from metatrain.pet import PET
+from metatrain.pet import PET, Trainer
 from metatrain.pet.modules.transformer import AttentionBlock
-from metatrain.utils.architectures import check_architecture_options
 from metatrain.utils.data import DatasetInfo
 from metatrain.utils.data.target_info import (
  get_energy_target_info,
  get_generic_target_info,
 )
 from metatrain.utils.neighbor_lists import get_system_with_neighbor_lists
+from metatrain.utils.pydantic import validate_architecture_options
 
 from . import DEFAULT_HYPERS, MODEL_HYPERS
 
@@ -365,16 +366,16 @@ def test_fixed_composition_weights():
  }
  }
  hypers = OmegaConf.create(hypers)
- check_architecture_options(name="pet", options=OmegaConf.to_container(hypers))
+ validate_architecture_options(OmegaConf.to_container(hypers), PET, Trainer)
 
 
 def test_fixed_composition_weights_error():
  """Test that only input of type Dict[str, Dict[int, float]] are allowed."""
  hypers = DEFAULT_HYPERS.copy()
  hypers["training"]["fixed_composition_weights"] = {"energy": {"H": 300.0}}
  hypers = OmegaConf.create(hypers)
- with pytest.raises(ValueError, match=r"'H' does not match '\^\[0-9\]\+\$'"):
- check_architecture_options(name="pet", options=OmegaConf.to_container(hypers))
+ with pytest.raises(ValidationError, match=r"Input should be a valid integer"):
+ validate_architecture_options(OmegaConf.to_container(hypers), PET, Trainer)
 
 
 @pytest.mark.parametrize("per_atom", [True, False])

src/metatrain/utils/additive/__init__.py

@@ -1,3 +1,3 @@
-from .composition import CompositionModel # noqa: F401
+from .composition import CompositionModel, FixedCompositionWeights # noqa: F401
 from .remove import get_remove_additive_transform, remove_additive # noqa: F401
 from .zbl import ZBL # noqa: F401

src/metatrain/utils/additive/_base_composition.py

@@ -12,6 +12,9 @@
 from metatomic.torch import ModelOutput, System
 
 
+FixedCompositionWeights = dict[str, dict[int, float]]
+
+
 class BaseCompositionModel(torch.nn.Module):
  """
  Fits a composition model for a dict of targets.
@@ -270,7 +273,7 @@ def accumulate(
 
  def fit(
  self,
- fixed_weights: Optional[Dict[str, Dict[int, float]]] = None,
+ fixed_weights: Optional[FixedCompositionWeights] = None,
  targets_to_fit: Optional[List[str]] = None,
  ) -> None:
  """

src/metatrain/utils/additive/composition.py

@@ -17,7 +17,11 @@
 from ..data import DatasetInfo, TargetInfo, unpack_batch
 from ..jsonschema import validate
 from ..transfer import batch_to
-from ._base_composition import BaseCompositionModel, _include_key
+from ._base_composition import (
+ BaseCompositionModel,
+ FixedCompositionWeights,
+ _include_key,
+)
 from .remove import remove_additive
 
 
@@ -175,7 +179,7 @@ def train_model(
  additive_models: List[torch.nn.Module],
  batch_size: int,
  is_distributed: bool,
- fixed_weights: Optional[Dict[str, Dict[int, float]]] = None,
+ fixed_weights: Optional[FixedCompositionWeights] = None,
  ) -> None:
  """
  Train the composition model on the provided training data in the ``datasets``.

src/metatrain/utils/hypers.py

@@ -25,46 +25,3 @@ def get_hypers_annotation(module: type) -> Type[dict]:
  :return: The TypedDict annotation of the 'hypers' parameter.
  """
  return inspect.signature(module).parameters["hypers"].annotation
-
-
-# Generic dictionary types to use as annotations in hypers
-ScalingWeightsDict = dict
-# class ScalingWeightsDict(dict):
-# """Dictionary type for scaling weights."""
-
-# @classmethod
-# def get_json_schema(cls) -> dict:
-# return {
-# "type": "object",
-# "patternProperties": {
-# "^.*$": {
-# "anyOf": [
-# {
-# "type": "object",
-# "propertyNames": {"pattern": "^[0-9]+$"},
-# "additionalProperties": {"type": "number"},
-# },
-# {"type": "number"},
-# ]
-# }
-# },
-# "additionalProperties": False,
-# }
-
-CompositionWeightsDict = dict
-# class CompositionWeightsDict(dict):
-# """Dictionary type for composition weights."""
-
-# @classmethod
-# def get_json_schema(cls) -> dict:
-# return {
-# "type": "object",
-# "patternProperties": {
-# "^.*$": {
-# "type": "object",
-# "propertyNames": {"pattern": "^[0-9]+$"},
-# "additionalProperties": {"type": "number"},
-# }
-# },
-# "additionalProperties": False,
-# }

src/metatrain/utils/scaler/__init__.py

@@ -1,2 +1,2 @@
 from .remove import get_remove_scale_transform, remove_scale # noqa: F401
-from .scaler import Scaler # noqa: F401
+from .scaler import FixedScalerWeights, Scaler # noqa: F401

src/metatrain/utils/scaler/_base_scaler.py

@@ -11,6 +11,9 @@
 from metatomic.torch import System
 
 
+FixedScalerWeights = dict[str, Union[float, dict[int, float]]]
+
+
 class BaseScaler(torch.nn.Module):
  """
  Fits a scaler for a dict of targets. Scales are computed as the per-property (and
@@ -236,7 +239,7 @@ def accumulate(
 
  def fit(
  self,
- fixed_weights: Optional[Dict[str, Union[float, Dict[int, float]]]] = None,
+ fixed_weights: Optional[FixedScalerWeights] = None,
  targets_to_fit: Optional[List[str]] = None,
  ) -> None:
  """

src/metatrain/utils/scaler/scaler.py

@@ -17,7 +17,7 @@
 from ..data import DatasetInfo, TargetInfo, unpack_batch
 from ..jsonschema import validate
 from ..transfer import batch_to
-from ._base_scaler import BaseScaler
+from ._base_scaler import BaseScaler, FixedScalerWeights
 
 
 class Scaler(torch.nn.Module):
@@ -139,7 +139,7 @@ def train_model(
  additive_models: List[torch.nn.Module],
  batch_size: int,
  is_distributed: bool,
- fixed_weights: Optional[Dict[str, Union[float, Dict[int, float]]]] = None,
+ fixed_weights: Optional[FixedScalerWeights] = None,
  ) -> None:
  """
  Placeholder docs.