remove regress_stress/forces and direct_forces properties

src/fairchem/core/common/utils.py

@@ -75,7 +75,13 @@ def decorator(func):
         @wraps(func)
         def cls_method(self, *args, **kwargs):
             f = func
-            if self.regress_forces and not getattr(self, "direct_forces", 0):
+            if hasattr(self, "regress_config"):
+                regress_forces = self.regress_config.forces
+                direct_forces = self.regress_config.direct_forces
+            else:
+                regress_forces = self.regress_forces
+                direct_forces = getattr(self, "direct_forces", 0)
+            if regress_forces and not direct_forces:
                 f = dec(func)
             return f(self, *args, **kwargs)
 

src/fairchem/core/models/base.py

@@ -161,13 +161,6 @@ def tasks(self) -> dict[str, Task]:
         """
         return self._tasks
 
-    @property
-    def direct_forces(self) -> bool:
-        """
-        Whether this model uses direct force prediction.
-        """
-        return getattr(self.backbone, "direct_forces", False)
-
     @property
     def dataset_to_tasks(self) -> dict[str, list]:
         """
@@ -191,8 +184,10 @@ def _validate_task_compatibility(self, task: Task) -> None:
         """
         derivative_properties = ("forces", "stress", "hessian")
 
+        backbone_regress_config = getattr(self.backbone, "regress_config", None)
         if (
-            self.direct_forces
+            backbone_regress_config is not None
+            and backbone_regress_config.direct_forces
             and task.inference_only
             and task.property in derivative_properties
         ):

src/fairchem/core/models/uma/escn_md.py

@@ -512,18 +512,6 @@ def __init__(
         )
         self.register_buffer("coefficient_index", coefficient_index, persistent=False)
 
-    @property  # deprecate this
-    def direct_forces(self) -> bool:
-        return self.regress_config.direct_forces
-
-    @property  # deprecate this
-    def regress_forces(self) -> bool:
-        return self.regress_config.forces
-
-    @property  # deprecate this
-    def regress_stress(self) -> bool:
-        return self.regress_config.stress
-
     def balance_channels(
         self,
         x_message_prime: torch.Tensor,
@@ -903,7 +891,7 @@ def get_default_untrained_tasks(
         stress computation requires energy-conserving force computation.
         """
         # Direct force models can't compute stress via autograd
-        if self.direct_forces:
+        if self.regress_config.direct_forces:
             return []
 
         tasks = []
@@ -1159,14 +1147,6 @@ def __init__(
         backbone.force_block = None
         self.regress_config = backbone.regress_config
 
-    @property
-    def regress_forces(self) -> bool:
-        return self.regress_config.forces
-
-    @property
-    def regress_stress(self) -> bool:
-        return self.regress_config.stress
-
     @conditional_grad(torch.enable_grad())
     def forward(
         self, data: AtomicData, emb: dict[str, torch.Tensor]

src/fairchem/core/models/uma/escn_moe.py

@@ -265,14 +265,6 @@ def __init__(
         self.merged_on_dataset = None
         self.non_merged_dataset_names: list[str] = []
 
-    @property
-    def regress_forces(self) -> bool:
-        return self.regress_config.forces
-
-    @property
-    def regress_stress(self) -> bool:
-        return self.regress_config.stress
-
     @staticmethod
     def _build_expert_mapping(
         dataset_names: list[str] | None,
@@ -424,14 +416,6 @@ def __init__(
         # keep track if this head has been merged or not
         self.merged_on_dataset = None
 
-    @property
-    def regress_forces(self) -> bool:
-        return self.regress_config.forces
-
-    @property
-    def regress_stress(self) -> bool:
-        return self.regress_config.stress
-
     def merge_MOLE_model(self, data):
         self.merged_on_dataset = data.dataset[0]
         self.non_merged_dataset_names = [

src/fairchem/core/units/mlip_unit/predict.py

@@ -196,16 +196,12 @@ def __init__(
         self.assert_on_nans = assert_on_nans
         self._warned_upcast = False
 
-        if self.model.module.direct_forces:
+        if self.model.module.backbone.regress_config.direct_forces:
             logging.warning(
                 "This is a direct-force model. Direct force predictions may lead to "
                 "discontinuities in the potential energy surface and energy conservation errors."
             )
 
-    @property
-    def direct_forces(self) -> bool:
-        return self.model.module.direct_forces
-
     @property
     def dataset_to_tasks(self) -> dict[str, list]:
         return self.model.module.dataset_to_tasks
@@ -478,7 +474,11 @@ def _run_inference(self, data: AtomicData, undo_refs: bool) -> dict:
         """
         Execute model inference.
         """
-        inference_context = torch.no_grad() if self.direct_forces else nullcontext()
+        inference_context = (
+            torch.no_grad()
+            if self.model.module.backbone.regress_config.direct_forces
+            else nullcontext()
+        )
         tf32_context = (
             tf32_context_manager() if self.inference_settings.tf32 else nullcontext()
         )

tests/core/calculate/test_ase_calculator.py

@@ -194,7 +194,7 @@ def test_calculator_setup(all_calculators):
         datasets = list(calc.predictor.dataset_to_tasks.keys())
 
         # all conservative UMA checkpoints should support E/F/S!
-        if not calc.predictor.direct_forces and (
+        if not calc.predictor.model.module.backbone.regress_config.direct_forces and (
             len(datasets) > 1 or (calc.task_name != "omol" and calc.task_name != "odac")
         ):
             print(len(datasets), calc.task_name)

tests/core/models/test_inference_interface.py

@@ -166,13 +166,6 @@ def test_dataset_to_tasks_raises_before_setup(self, mock_hydra_model):
         with pytest.raises(RuntimeError, match="setup_tasks"):
             _ = mock_hydra_model.dataset_to_tasks
 
-    def test_direct_forces_property(self, mock_hydra_model):
-        """Test direct_forces property delegates to backbone."""
-        assert mock_hydra_model.direct_forces is False
-
-        mock_hydra_model.backbone.direct_forces = True
-        assert mock_hydra_model.direct_forces is True
-
 
 class TestBackboneInterface:
     """Tests for backbone interface method implementations."""