Avoid and hide deprecation warnings

src/metatrain/composition/model.py

@@ -343,7 +343,6 @@ def supported_outputs(self) -> Dict[str, ModelOutput]:
 
     def _add_output(self, target_name: str, target_info: TargetInfo) -> None:
         self.outputs[target_name] = ModelOutput(
-            quantity=target_info.quantity,
             unit=target_info.unit,
             sample_kind="atom",
             description=target_info.description,

src/metatrain/composition/tests/test_errors.py

@@ -108,4 +108,4 @@ def test_forward_unknown_output_raises():
 
     systems = read_systems(DATASET_PATH)
     with pytest.raises(ValueError, match="not supported"):
-        model(systems[:1], {"nonexistent": ModelOutput(quantity="energy")})
+        model(systems[:1], {"nonexistent": ModelOutput()})

src/metatrain/experimental/classifier/model.py

@@ -91,7 +91,7 @@ def set_wrapped_model(self, model: ModelInterface) -> None:
 
         # Store capabilities
         outputs = {name: ModelOutput() for name in self.dataset_info.targets.keys()}
-        outputs["feature"] = ModelOutput(quantity="", unit="", sample_kind="system")
+        outputs["feature"] = ModelOutput(unit="", sample_kind="system")
         self.capabilities = ModelCapabilities(
             outputs=outputs,
             atomic_types=old_capabilities.atomic_types,

src/metatrain/experimental/classifier/trainer.py

@@ -202,11 +202,7 @@ def get_lr_schedule(step):
                 # Forward pass - request logits for training
                 outputs = model(
                     systems,
-                    {
-                        target_name_logits: ModelOutput(
-                            quantity="", unit="", sample_kind="system"
-                        )
-                    },
+                    {target_name_logits: ModelOutput(unit="", sample_kind="system")},
                     None,
                 )
 
@@ -252,7 +248,7 @@ def get_lr_schedule(step):
                         systems,
                         {
                             target_name_logits: ModelOutput(
-                                quantity="", unit="", sample_kind="system"
+                                unit="", sample_kind="system"
                             )
                         },
                         None,

src/metatrain/experimental/dpa3/model.py

@@ -214,7 +214,6 @@ def _add_output(self, target_name: str, target: TargetInfo) -> None:
             block.properties for block in target.layout.blocks()
         ]
         self.outputs[target_name] = ModelOutput(
-            quantity=target.quantity,
             unit=target.unit,
             sample_kind="atom",
         )

src/metatrain/experimental/dpa3/tests/test_compile.py

@@ -55,7 +55,7 @@ def test_forward_from_batch_matches_forward():
     # Standard forward
     out = model(
         systems,
-        {"mtt::U0": ModelOutput(quantity="energy", unit="", sample_kind="system")},
+        {"mtt::U0": ModelOutput(unit="", sample_kind="system")},
     )
     # Verify the standard forward works
     assert out["mtt::U0"].block().values.numel() > 0

src/metatrain/experimental/dpa3/tests/test_pretrained.py

@@ -141,9 +141,7 @@ def test_pretrained_forward_matches():
     for s in systems:
         get_system_with_neighbor_lists(s, base.requested_neighbor_lists())
 
-    output_request = {
-        "mtt::U0": ModelOutput(quantity="energy", unit="", sample_kind="system")
-    }
+    output_request = {"mtt::U0": ModelOutput(unit="", sample_kind="system")}
     base_out = base(systems, output_request)
     pretrained_out = pretrained(systems, output_request)
 
@@ -179,9 +177,7 @@ def test_pretrained_checkpoint_roundtrip():
     for s in systems:
         get_system_with_neighbor_lists(s, pretrained.requested_neighbor_lists())
 
-    output_request = {
-        "mtt::U0": ModelOutput(quantity="energy", unit="", sample_kind="system")
-    }
+    output_request = {"mtt::U0": ModelOutput(unit="", sample_kind="system")}
     orig_out = pretrained(systems, output_request)
     reloaded_out = reloaded(systems, output_request)
 

src/metatrain/experimental/dpa3/tests/test_regression.py

@@ -44,7 +44,7 @@ def test_regression_init():
 
     output = model(
         systems,
-        {"mtt::U0": ModelOutput(quantity="energy", unit="", sample_kind="system")},
+        {"mtt::U0": ModelOutput(unit="", sample_kind="system")},
     )
 
     expected_output = torch.tensor(

src/metatrain/experimental/flashmd/model.py

@@ -306,10 +306,8 @@ def requested_neighbor_lists(self) -> List[NeighborListOptions]:
 
     def requested_inputs(self) -> Dict[str, ModelOutput]:
         return {
-            "momentum": ModelOutput(
-                quantity="momentum", unit="(eV*u)^(1/2)", sample_kind="atom"
-            ),
-            "mass": ModelOutput(quantity="mass", unit="u", sample_kind="atom"),
+            "momentum": ModelOutput(unit="(eV*u)^(1/2)", sample_kind="atom"),
+            "mass": ModelOutput(unit="u", sample_kind="atom"),
         }
 
     def forward(
@@ -1256,7 +1254,6 @@ def _add_output(self, target_name: str, target_info: TargetInfo) -> None:
             ] + [len(block.properties.values)]
 
         self.outputs[target_name] = ModelOutput(
-            quantity=target_info.quantity,
             unit=target_info.unit,
             sample_kind="atom",
             description=target_info.description,

src/metatrain/experimental/flashmd/modules/additive.py

@@ -41,7 +41,6 @@ def __init__(self, hypers: PositionAdditiveHypers, dataset_info: DatasetInfo):
                 # skip momenta targets unless `also_momenta` is True
                 continue
             self.outputs[key] = ModelOutput(
-                quantity=value.quantity,
                 unit=value.unit,
                 sample_kind="atom",
                 description=value.description,
@@ -63,7 +62,6 @@ def restart(self, dataset_info: DatasetInfo) -> "PositionAdditive":
                 # skip momenta targets unless `also_momenta` is True
                 continue
             self.outputs[key] = ModelOutput(
-                quantity=value.quantity,
                 unit=value.unit,
                 sample_kind="atom",
                 description=value.description,

src/metatrain/experimental/flashmd/tests/test_functionality.py

@@ -100,8 +100,8 @@ def test_forward():
         system.add_data("momentum", tmap)
 
     outputs = {
-        "position": ModelOutput(quantity="length", unit="angstrom", sample_kind="atom"),
-        "momentum": ModelOutput(quantity="length", unit="angstrom", sample_kind="atom"),
+        "position": ModelOutput(unit="angstrom", sample_kind="atom"),
+        "momentum": ModelOutput(unit="angstrom", sample_kind="atom"),
     }
     result_dict = model(systems, outputs)
 

src/metatrain/experimental/flashmd_symplectic/model.py

@@ -288,10 +288,8 @@ def requested_neighbor_lists(self) -> List[NeighborListOptions]:
 
     def requested_inputs(self) -> Dict[str, ModelOutput]:
         return {
-            "momentum": ModelOutput(
-                quantity="momentum", unit="(eV*u)^(1/2)", sample_kind="atom"
-            ),
-            "mass": ModelOutput(quantity="mass", unit="u", sample_kind="atom"),
+            "momentum": ModelOutput(unit="(eV*u)^(1/2)", sample_kind="atom"),
+            "mass": ModelOutput(unit="u", sample_kind="atom"),
         }
 
     def forward(

src/metatrain/experimental/flashmd_symplectic/tests/test_functionality.py

@@ -100,8 +100,8 @@ def test_forward():
         system.add_data("momentum", tmap)
 
     outputs = {
-        "position": ModelOutput(quantity="length", unit="angstrom", sample_kind="atom"),
-        "momentum": ModelOutput(quantity="length", unit="angstrom", sample_kind="atom"),
+        "position": ModelOutput(unit="angstrom", sample_kind="atom"),
+        "momentum": ModelOutput(unit="angstrom", sample_kind="atom"),
     }
     result_dict = model(systems, outputs)
 

src/metatrain/experimental/mace/model.py

@@ -284,7 +284,6 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None:
         targets = dataset_info.targets
         self.outputs = {
             k: ModelOutput(
-                quantity=targets[k].quantity if k in targets else "",
                 unit=targets[k].unit if k in targets else "",
                 sample_kind="atom",
             )

src/metatrain/experimental/space/model.py

@@ -644,7 +644,6 @@ def _train_dataset_info(self, dataset_info: DatasetInfo) -> DatasetInfo:
 
     def _add_output(self, target_name: str, target_info: TargetInfo) -> None:
         self.outputs[target_name] = ModelOutput(
-            quantity=target_info.quantity,
             unit=target_info.unit,
             sample_kind="atom",
         )

src/metatrain/experimental/space/tests/test_functionality.py

@@ -159,7 +159,7 @@ def test_multiple_targets():
     model = SPACE(hypers, dataset_info)
     system = _make_system(model)
     outputs = {
-        "energy": ModelOutput(quantity="energy", unit="eV", sample_kind="system"),
+        "energy": ModelOutput(unit="eV", sample_kind="system"),
         "dipole": ModelOutput(sample_kind="system"),
         "non_conservative_stress": ModelOutput(sample_kind="system"),
     }

src/metatrain/experimental/space/tests/test_regression.py

@@ -61,7 +61,7 @@ def test_regression_init():
 
     output = model(
         systems,
-        {"mtt::U0": ModelOutput(quantity="energy", unit="", sample_kind="system")},
+        {"mtt::U0": ModelOutput(unit="", sample_kind="system")},
     )
 
     expected_output = torch.tensor(
@@ -158,7 +158,7 @@ def test_regression_train():
     model = torch.jit.script(model)
     output = model(
         systems[:5],
-        {"mtt::U0": ModelOutput(quantity="energy", unit="", sample_kind="system")},
+        {"mtt::U0": ModelOutput(unit="", sample_kind="system")},
     )
 
     expected_output = torch.tensor(
@@ -258,11 +258,7 @@ def test_regression_train_spherical(device):
     ]
     output = model(
         systems,
-        {
-            "mtt::electron_density_basis": ModelOutput(
-                quantity="", unit="", sample_kind="atom"
-            )
-        },
+        {"mtt::electron_density_basis": ModelOutput(unit="", sample_kind="atom")},
     )
 
     expected_output = torch.tensor(

src/metatrain/experimental/space/trainer.py

@@ -58,7 +58,6 @@ def _get_requested_outputs(targets, target_info_dict):
     requested_outputs = {}
     for name, target in targets.items():
         requested_outputs[name] = ModelOutput(
-            quantity=target_info_dict[name].quantity,
             unit=target_info_dict[name].unit,
             sample_kind=target_info_dict[name].sample_kind,
             explicit_gradients=target.block(0).gradients_list(),

src/metatrain/gap/model.py

@@ -73,7 +73,6 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None:
 
         self.outputs = {
             key: ModelOutput(
-                quantity=value.quantity,
                 unit=value.unit,
                 sample_kind="system",
                 description=value.description,

src/metatrain/llpr/model.py

@@ -155,7 +155,6 @@ def set_wrapped_model(self, model: ModelInterface) -> None:
             self.outputs_list.append(name)
             uncertainty_name = _get_uncertainty_name(name)
             additional_capabilities[uncertainty_name] = ModelOutput(
-                quantity=output.quantity,
                 unit=output.unit,
                 sample_kind=output.sample_kind,
                 description=output.description,
@@ -216,13 +215,10 @@ def set_wrapped_model(self, model: ModelInterface) -> None:
             )
             if ensemble_output_name == "mtt::aux::energy_ensemble":
                 ensemble_output_name = "energy_ensemble"
-            explicit_gradients = _ensemble_explicit_gradients(
-                old_capabilities.outputs[name]
-            )
+            explicit_gradients = self._ensemble_explicit_gradients(name)
             if len(explicit_gradients) > 0:
                 self.ensemble_gradient_outputs.append(ensemble_output_name)
             ensemble_outputs[ensemble_output_name] = ModelOutput(
-                quantity=old_capabilities.outputs[name].quantity,
                 unit=old_capabilities.outputs[name].unit,
                 sample_kind=old_capabilities.outputs[name].sample_kind,
                 explicit_gradients=explicit_gradients,
@@ -1146,10 +1142,9 @@ def generate_ensemble(self) -> None:
             if ensemble_name == "mtt::aux::energy_ensemble":
                 ensemble_name = "energy_ensemble"
             new_outputs[ensemble_name] = ModelOutput(
-                quantity=old_outputs[name].quantity,
                 unit=old_outputs[name].unit,
                 sample_kind=old_outputs[name].sample_kind,
-                explicit_gradients=_ensemble_explicit_gradients(old_outputs[name]),
+                explicit_gradients=self._ensemble_explicit_gradients(name),
                 description=f"ensemble of {name}",
             )
         self.capabilities = ModelCapabilities(
@@ -1314,29 +1309,18 @@ def upgrade_checkpoint(cls, checkpoint: Dict) -> Dict:
     def supported_outputs(self) -> Dict[str, ModelOutput]:
         return self.capabilities.outputs
 
+    def _ensemble_explicit_gradients(self, name: str) -> List[str]:
+        """Explicit gradients the ensemble of the ``name`` target is able to produce.
 
-def _ensemble_explicit_gradients(output: ModelOutput) -> List[str]:
-    """Explicit gradients the ensemble of ``output`` is able to produce.
-
-    ``_add_energy_ensemble_gradients`` differentiates one per-system scalar per
-    ensemble member, so what matters is that the output *is* a per-system energy,
-    not what it is called: an energy target may carry any name (``mtt::my_energy``)
-    and any variant (``energy/pbesol``). Selecting on the quantity rather than on the
-    name keeps all of those working.
-
-    Other quantities get nothing: positions/strain gradients of them are not what
-    this computes.
+        Only energy targets are currently supported.
 
-    Note that ``sample_kind`` is deliberately *not* consulted here. In capabilities
-    it describes what the wrapped model is able to produce (PET reports ``"atom"``
-    for its energy even when the target is per-system), not what a given call asks
-    for. Whether the *requested* sample kind is supported is checked in ``forward``,
-    where the request is actually known.
-
-    :param output: the wrapped model's output the ensemble is built from.
-    :return: gradient names for the corresponding ensemble output.
-    """
-    return ["positions", "strain"] if output.quantity == "energy" else []
+        :param name: name of the target the ensemble is built from.
+        :return: gradient names for the corresponding ensemble output.
+        """
+        target = self.dataset_info.targets.get(name)
+        if target is None or target.quantity != "energy":
+            return []
+        return ["positions", "strain"]
 
 
 def _get_uncertainty_name(name: str) -> str:

src/metatrain/pet/model.py

@@ -270,7 +270,7 @@ def requested_neighbor_lists(self) -> List[NeighborListOptions]:
     def requested_inputs(self) -> Dict[str, ModelOutput]:
         if self.system_conditioning is not None:
             return {
-                key: ModelOutput(quantity="", unit="", sample_kind="system")
+                key: ModelOutput(unit="", sample_kind="system")
                 for key in self.system_conditioning.required_data_keys
             }
         return {}
@@ -1051,7 +1051,6 @@ def _add_output(self, target_name: str, target_info: TargetInfo) -> None:
             ] + [len(block.properties.values)]
 
         self.outputs[target_name] = ModelOutput(
-            quantity=target_info.quantity,
             unit=target_info.unit,
             sample_kind="atom",
             description=target_info.description,

src/metatrain/pet/tests/test_regression.py

@@ -60,7 +60,7 @@ def test_regression_init():
 
     output = model(
         systems,
-        {"mtt::U0": ModelOutput(quantity="energy", unit="", sample_kind="system")},
+        {"mtt::U0": ModelOutput(unit="", sample_kind="system")},
     )
 
     expected_output = torch.tensor(
@@ -307,7 +307,7 @@ def test_regression_energy_non_conservative_stress(batch_size):
     outputs = model(
         eval_systems,
         {
-            "energy": ModelOutput(quantity="energy", unit="", sample_kind="system"),
+            "energy": ModelOutput(unit="", sample_kind="system"),
             "non_conservative_stress": ModelOutput(sample_kind="system"),
         },
     )
@@ -426,11 +426,7 @@ def test_regression_train_spherical(device):
     ]
     output = model(
         systems,
-        {
-            "mtt::electron_density_basis": ModelOutput(
-                quantity="", unit="", sample_kind="atom"
-            )
-        },
+        {"mtt::electron_density_basis": ModelOutput(unit="", sample_kind="atom")},
     )
 
     expected_output = torch.tensor(

src/metatrain/scaler/model.py

@@ -304,7 +304,6 @@ def supported_outputs(self) -> Dict[str, ModelOutput]:
 
     def _add_output(self, target_name: str, target_info: TargetInfo) -> None:
         self.outputs[target_name] = ModelOutput(
-            quantity=target_info.quantity,
             unit=target_info.unit,
             sample_kind="atom",
             description=target_info.description,

src/metatrain/scaler/tests/test_scaler.py

@@ -2640,7 +2640,6 @@ def test_scaler_with_neighbor_list_additive_model():
 
     outputs = {
         "energy": ModelOutput(
-            quantity="",
             unit="",
             sample_kind="system",
         )

src/metatrain/soap_bpnn/model.py

@@ -1232,7 +1232,6 @@ def _add_output(self, target_name: str, target: TargetInfo) -> None:
         ]
 
         self.outputs[target_name] = ModelOutput(
-            quantity=target.quantity,
             unit=target.unit,
             sample_kind="atom",
             description=target.description,