pair_e3gnn.cpp uses pair-wise force.

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:
@@ -69,9 +70,13 @@ 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)
@@ -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,12 @@ 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 +197,7 @@ def __init__(
             'energy',
             'forces',
             'stress',
+            'stresses',
             'energies',
         ]
 
@@ -206,15 +222,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 +262,13 @@ 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']
+        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)
 
-        self.results = self.output_to_results(self.model(data))
+        output = self.model(data)
+        self.results = self.output_to_results(output)
 
 
 class SevenNetD3Calculator(SumCalculator):

sevenn/main/sevenn_get_model.py

@@ -64,6 +64,15 @@ def add_args(parser):
         help='Use LAMMPS ML-IAP interface.',
         action='store_true',
     )
+    ag.add_argument(
+        '--atomic_virial',
+        help=(
+            'Serial deploy only: append per-atom virial output '
+            '(inferred_atomic_virial) to TorchScript. This marks model-side '
+            'atomic virial capability for downstream calculator usage.'
+        ),
+        action='store_true',
+    )
 
 
 def run(args):
@@ -78,6 +87,7 @@ def run(args):
     use_cueq = args.enable_cueq
     use_oeq = args.enable_oeq
     use_mliap = args.use_mliap
+    atomic_virial = args.atomic_virial
 
     # Check dependencies
     if use_flash:
@@ -104,9 +114,14 @@ def run(args):
     if use_mliap and get_parallel:
         raise ValueError('Currently, ML-IAP interface does not tested on parallel.')
 
+    if atomic_virial and (use_mliap or get_parallel):
+        raise ValueError('--atomic_virial is only supported for serial deployment.')
+
     # deploy
     if output_prefix is None:
         output_prefix = 'deployed_parallel' if not get_serial else 'deployed_serial'
+        if atomic_virial:
+            output_prefix = 'deployed_model'
 
         if use_mliap:
             output_prefix += '_mliap'
@@ -118,10 +133,24 @@ def run(args):
         checkpoint_path = sevenn.util.pretrained_name_to_path(checkpoint)
 
     if not use_mliap:
-        from sevenn.scripts.deploy import deploy, deploy_parallel
-
-        if get_serial:
-            deploy(checkpoint_path, output_prefix, modal, use_flash=use_flash, use_oeq=use_oeq)  # noqa: E501
+        from sevenn.scripts.deploy import deploy, deploy_parallel, deploy_ts
+
+        if atomic_virial:
+            deploy_ts(
+                checkpoint_path,
+                output_prefix,
+                modal,
+                use_flash=use_flash,
+                use_oeq=use_oeq,
+                atomic_virial=atomic_virial,
+            )
+        elif get_serial:
+            deploy(
+                checkpoint_path,
+                output_prefix,
+                modal,
+                use_flash=use_flash,
+                use_oeq=use_oeq)
         else:
             deploy_parallel(checkpoint_path, output_prefix, modal, use_flash=use_flash, use_oeq=use_oeq)  # noqa: E501
     else:

sevenn/main/sevenn_patch_lammps.py

@@ -36,6 +36,11 @@ def add_args(parser):
         help='Enable OpenEquivariance',
         action='store_true',
     )
+    ag.add_argument(
+        '--atomic_stress',
+        help='Patch pair_e3gnn with atomic-stress enabled source files.',
+        action='store_true',
+    )
     # cxx_standard is detected automatically
 
 
@@ -54,6 +59,12 @@ def run(args):
         d3_support = '0'
         print('  - D3 support disabled')
 
+    atomic_stress = '1' if args.atomic_stress else '0'
+    if args.atomic_stress:
+        print('  - Atomic stress patch enabled')
+    else:
+        print('  - Atomic stress patch disabled')
+
     so_oeq = ''
     if args.enable_oeq:
         try:
@@ -111,6 +122,10 @@ def run(args):
     if args.enable_oeq:
         assert osp.isfile(so_oeq)
         cmd += f' {so_oeq}'
+    else:
+        cmd += ' NONE'
+
+    cmd += f' {atomic_stress}'
 
     res = subprocess.run(cmd.split())
     return res.returncode  # is it meaningless?

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