Add atomic_virial option

sevenn/_keys.py

@@ -65,6 +65,7 @@
 SCALED_FORCE: Final[str] = 'scaled_force'
 
 PRED_STRESS: Final[str] = 'inferred_stress'
+PRED_ATOMIC_VIRIAL: Final[str] = 'inferred_atomic_virial'
 SCALED_STRESS: Final[str] = 'scaled_stress'
 
 # very general data property for AtomGraphData

sevenn/calculator.py

@@ -40,6 +40,7 @@ def __init__(
         enable_cueq: Optional[bool] = False,
         enable_flash: Optional[bool] = False,
         enable_oeq: Optional[bool] = False,
+        atomic_virial: bool = False,
         sevennet_config: Optional[Dict] = None,  # Not used in logic, just meta info
         **kwargs,
     ) -> None:
@@ -51,7 +52,7 @@ def __init__(
             Name of pretrained models (7net-omni, 7net-mf-ompa, 7net-omat, 7net-l3i5,
             7net-0) or path to the checkpoint, deployed model or the model itself
         file_type: str, default='checkpoint'
-            one of 'checkpoint' | 'torchscript' | 'model_instance'
+            one of 'checkpoint' | 'model_instance'
         device: str | torch.device, default='auto'
             if not given, use CUDA if available
         modal: str | None, default=None
@@ -69,14 +70,18 @@ def __init__(
             if True, use OpenEquivariance to accelerate inference.
         sevennet_config: dict | None, default=None
             Not used, but can be used to carry meta information of this calculator
+        atomic_virial: bool, default=False
+            If True, request per-atom virial output (`stresses`) at runtime.
         """
         super().__init__(**kwargs)
         self.sevennet_config = None
+        self.atomic_virial_requested = atomic_virial
+        self.atomic_virial_from_deploy = False
 
         if isinstance(model, pathlib.PurePath):
             model = str(model)
 
-        allowed_file_types = ['checkpoint', 'torchscript', 'model_instance']
+        allowed_file_types = ['checkpoint', 'model_instance']
         file_type = file_type.lower()
         if file_type not in allowed_file_types:
             raise ValueError(f'file_type not in {allowed_file_types}')
@@ -131,6 +136,7 @@ def __init__(
                 'version': b'',
                 'dtype': b'',
                 'time': b'',
+                'atomic_virial': b'',
             }
             model_loaded = torch.jit.load(
                 model, _extra_files=extra_dict, map_location=self.device
@@ -141,6 +147,9 @@ def __init__(
                 sym_to_num[sym]: i for i, sym in enumerate(chem_symbols.split())
             }
             self.cutoff = float(extra_dict['cutoff'].decode('utf-8'))
+            self.atomic_virial_from_deploy = (
+                extra_dict['atomic_virial'].decode('utf-8') == 'yes'
+            )
 
         elif isinstance(model, AtomGraphSequential):
             if model.type_map is None:
@@ -161,6 +170,13 @@ def __init__(
             self.sevennet_config = sevennet_config
 
         self.model = model_loaded
+        if isinstance(self.model, AtomGraphSequential):
+            force_output = self.model._modules.get('force_output')
+            if force_output is not None:
+                self.atomic_virial_from_deploy = (
+                    self.atomic_virial_from_deploy
+                    or bool(getattr(force_output, 'use_atomic_virial', False))
+                )
 
         self.modal = None
         if isinstance(self.model, AtomGraphSequential):
@@ -182,6 +198,7 @@ def __init__(
             'energy',
             'forces',
             'stress',
+            'stresses',
             'energies',
         ]
 
@@ -206,15 +223,20 @@ def output_to_results(self, output: Dict[str, torch.Tensor]) -> Dict[str, Any]:
             .cpu()
             .numpy()[[0, 1, 2, 4, 5, 3]]  # as voigt notation
         )
-        # Store results
-        return {
+        results: Dict[str, Any] = {
             'free_energy': energy,
             'energy': energy,
             'energies': atomic_energies,
             'forces': forces,
             'stress': stress,
             'num_edges': output[KEY.EDGE_IDX].shape[1],
         }
+        if KEY.PRED_ATOMIC_VIRIAL in output:
+            virial = (
+                output[KEY.PRED_ATOMIC_VIRIAL].detach().cpu().numpy()[:num_atoms, :]
+            )
+            results['stresses'] = virial
+        return results
 
     def calculate(self, atoms=None, properties=None, system_changes=all_changes):
         is_ts_type = isinstance(self.model, torch_script_type)
@@ -241,8 +263,16 @@ def calculate(self, atoms=None, properties=None, system_changes=all_changes):
             data[KEY.EDGE_VEC].requires_grad_(True)  # backward compatibility
             data = data.to_dict()
             del data['data_info']
-
-        self.results = self.output_to_results(self.model(data))
+        elif self.atomic_virial_requested:
+            force_output = self.model._modules.get('force_output')
+            if (
+                force_output is not None
+                and hasattr(force_output, 'use_atomic_virial')
+            ):
+                setattr(force_output, 'use_atomic_virial', True)
+
+        output = self.model(data)
+        self.results = self.output_to_results(output)
 
 
 class SevenNetD3Calculator(SumCalculator):

sevenn/nn/force_output.py

@@ -149,7 +149,9 @@ def __init__(
         data_key_energy: str = KEY.PRED_TOTAL_ENERGY,
         data_key_force: str = KEY.PRED_FORCE,
         data_key_stress: str = KEY.PRED_STRESS,
+        data_key_atomic_virial: str = KEY.PRED_ATOMIC_VIRIAL,
         data_key_cell_volume: str = KEY.CELL_VOLUME,
+        use_atomic_virial: bool = False,
     ) -> None:
 
         super().__init__()
@@ -158,7 +160,9 @@ def __init__(
         self.key_energy = data_key_energy
         self.key_force = data_key_force
         self.key_stress = data_key_stress
+        self.key_atomic_virial = data_key_atomic_virial
         self.key_cell_volume = data_key_cell_volume
+        self.use_atomic_virial = use_atomic_virial
         self._is_batch_data = True
 
     def get_grad_key(self) -> str:
@@ -206,6 +210,8 @@ def forward(self, data: AtomGraphDataType) -> AtomGraphDataType:
             _s = torch.zeros(tot_num, 6, dtype=fij.dtype, device=fij.device)
             _edge_dst6 = broadcast(edge_idx[1], _virial, 0)
             _s.scatter_reduce_(0, _edge_dst6, _virial, reduce='sum')
+            if self.use_atomic_virial:
+                data[self.key_atomic_virial] = torch.neg(_s)
 
             if self._is_batch_data:
                 batch = data[KEY.BATCH]  # for deploy, must be defined first