Openequivariance implementation

CHANGELOG.md

@@ -1,6 +1,14 @@
 # Changelog
 All notable changes to this project will be documented in this file.
 
+## [0.12.2.dev]
+### Added
+- Support OpenEquivariance
+
+### Changed
+- **[Breaking]** Rename optional dependency group `mliap` into `mliap12` (reflecting its CUDA 12.x dependency).
+- Add `cueq13` and `mliap13` optional dependency groups for CUDA 13.x.
+
 ## [0.12.1]
 ### Fixed
 - FlashTP with LAMMPS parallel in torch
@@ -13,7 +21,7 @@ All notable changes to this project will be documented in this file.
 - all CLI argument name of enabling flashTP is now 'enable_flash', not 'enable_flashTP'
 
 ### Added
-- TorchSim interface adnd docs
+- TorchSim interface and docs
 
 ## [0.12.0]
 ### Added

docs/source/user_guide/accelerator.md

@@ -3,14 +3,14 @@
 This document describes available accelerator integrations in SevenNet and their installation guide.
 
 :::{caution}
-We do not support CuEquivariance for [LAMMPS: Torch](./lammps_mliap.md). You must use [LAMMPS: ML-IAP](./lammps_torch.md) for CuEquivariance.
+We do not support CuEquivariance for [LAMMPS: Torch](./lammps_torch.md). You must use [LAMMPS: ML-IAP](./lammps_mliap.md) for CuEquivariance.
 :::
 
-[CuEquivariance](https://github.com/NVIDIA/cuEquivariance) and [FlashTP](https://openreview.net/forum?id=wiQe95BPaB) provide acceleration for both SevenNet training and inference. (For speed, check the section 2.7 of [SevenNet-Omni paper](https://arxiv.org/abs/2510.11241))
+[FlashTP](https://openreview.net/forum?id=wiQe95BPaB), [CuEquivariance](https://github.com/NVIDIA/cuEquivariance), and [OpenEquivariance](https://github.com/PASSIONLab/OpenEquivariance) provide acceleration for both SevenNet training and inference.
 
 :::{tip}
-For small systems, FlashTP with [LAMMPS: Torch](./lammps_mliap.md) shows performance advantage over cuEquivariance with [LAMMPS: ML-IAP](./lammps_torch.md).
-A performance crossover occurs at around 10³ atoms, beyond which cuEquivariance becomes more efficient.
+For small systems, FlashTP with [LAMMPS: Torch](./lammps_torch.md) shows performance advantage over cuEquivariance with [LAMMPS: ML-IAP](./lammps_mliap.md).
+A performance crossover occurs at around 10³ atoms, beyond which cuEquivariance becomes more efficient.  (For more informattion, check the section 2.7 of [SevenNet-Omni paper](https://arxiv.org/abs/2510.11241))
 
 FlashTP with [LAMMPS: Torch](./lammps_mliap.md) is generally faster than FlashTP with [LAMMPS: ML-IAP](./lammps_mliap.md).
 :::
@@ -24,9 +24,16 @@ CuEquivariance is an NVIDIA Python library designed to facilitate the constructi
 - cuEquivariance >= 0.6.1
 
 ### Installation
-After installation of SevenNet, install cuEquivariance following their guideline.
-To use cuEquivariance with SevenNet, you need to install `cuequivariance`, `cuequivariance-torch`, and `cuequivariance-ops-torch-cu{12, 13}` (depending on your CUDA version)
-[cuEquivariance](https://github.com/NVIDIA/cuEquivariance).
+After installation of SevenNet, install cuEquivariance via the pip extras:
+
+```bash
+pip install sevenn[cueq12]  # For CUDA 12.x
+pip install sevenn[cueq13]  # For CUDA 13.x
+```
+
+This will install `cuequivariance`, `cuequivariance-torch`, and the corresponding `cuequivariance-ops-torch-cu{12,13}`.
+
+Alternatively, you can install cuEquivariance manually following [their guideline](https://github.com/NVIDIA/cuEquivariance).
 
 :::{note}
 Some GeForce GPUs do not support `pynvml`,
@@ -79,8 +86,33 @@ True
 
 For more information, see [FlashTP](https://github.com/SNU-ARC/flashTP).
 
+## [OpenEquivariance (experimental)](https://github.com/PASSIONLab/OpenEquivariance)
+
+:::{caution}
+OpenEquivariance support in SevenNet is currently experimental.
+:::
+
+OpenEquivariance (oEq) is a library for acceleration of equivariant tensor products. Unlike cuEquivariance, its kernels are compiled once per irreps configuration, eliminating runtime recompilation overhead during training on diverse datasets.
+
+### Requirements
+- Python >= 3.10
+- openequivariance
+
+### Installation
+```bash
+pip install sevenn[oeq]
+```
+
+Check your installation:
+```bash
+python -c 'from sevenn.nn.oeq_helper import is_oeq_available; print(is_oeq_available())'
+True
+```
+
+For more information, see [OpenEquivariance](https://github.com/PASSIONLab/OpenEquivariance).
+
 ## Usage
-After the installation, you can leverage the accelerator with appropriate flag (`--enable_cueq`) or (`--enable_flash`) options
+After the installation, you can leverage the accelerator with appropriate flag (`--enable_cueq`), (`--enable_flash`), or (`--enable_oeq`) options
 
 - [Training](./cli.md#sevenn-train)
 - [ASE Calculator](./ase_calculator.md)

docs/source/user_guide/ase_calculator.md

@@ -14,11 +14,11 @@ from sevenn.calculator import SevenNetD3Calculator
 calc = SevenNetD3Calculator(model='7net-0', device='cuda')
 ```
 
-Use enable_cueq or enable_flash to use cuEquivariance or flashTP for faster inference.
-For more information about cuEq and flashTP, follow [here](./accelerator.md)
+Use enable_cueq, enable_flash, or enable_oeq to use cuEquivariance, flashTP, or OpenEquivariance for faster inference.
+For more information about accelerators, follow [here](./accelerator.md)
 ```python
 from sevenn.calculator import SevenNetCalculator
-calc = SevenNetCalculator(model='7net-0', enable_cueq=True) # or enable_flash=True
+calc = SevenNetCalculator(model='7net-0', enable_cueq=True) # or enable_flash=True or enable_oeq=True
 ```
 
 If you encounter the error `CUDA is not installed or nvcc is not available`, please ensure the `nvcc` compiler is available. Currently, CPU + D3 is not supported.

docs/source/user_guide/cli.md

@@ -44,10 +44,10 @@ We support multi-GPU training using PyTorch DDP (distributed data parallel) with
 torchrun --standalone --nnodes {number of nodes} --nproc_per_node {number of GPUs} --no_python sevenn input.yaml -d
 ```
 
-Enable cuEquivariance or flashTP while training by adding the --enable_cueq or --enable_flash.
+Enable cuEquivariance, flashTP, or OpenEquivariance while training by adding the --enable_cueq, --enable_flash, or --enable_oeq.
 
 ```bash
-sevenn train input.yaml -s --enable_cueq # or --enable_flash
+sevenn train input.yaml -s --enable_cueq # or --enable_flash or --enable_oeq
 ```
 
 Please note that `batch_size` in `input.yaml` refers to the per-GPU batch size.
@@ -79,19 +79,19 @@ See {doc}`accelerator` for installation of accelerators.
 sevenn get_model \
     {pretrained_name or checkpoint_path} \
     {--use_mliap} \  # For LAMMPS ML-IAP use.
-    {--enable_flash | --enable_cueq} \  # For accelerators.
+    {--enable_flash | --enable_cueq | --enable_oeq} \  # For accelerators.
     {--modal {task_name}} \  # Required when using multi-fidelity model
     {--get_parallel}  # For parallel MD simulations
 ```
 
 If `--use_mliap` is not set (TorchScript version), it will create `deployed_serial.pt` or a directory containing several `deployed_parallel_*.pt` files for parallel execution.
 They can be used as a LAMMPS potential with the `e3gnn` or `e3gnn/parallel` pair_style in LAMMPS.
-In this case, only `--enable_flash` is available.
+In this case, `--enable_flash` and `--enable_oeq` are available.
 Check {doc}`lammps_torch` for installation and lammps script in this use case.
 
 If `--use_mliap` is set, it will create `deployed_serial_mliap.pt`.
 The file can be used with the `mliap` pair_style in LAMMPS.
-In this case, both `--enable_cueq` and `--enable_flash` are available, but you cannot specify both at the same time.
+In this case, `--enable_cueq`, `--enable_flash`, and `--enable_oeq` are available, but you cannot specify multiple at the same time.
 Parallel execution is not tested in this case.
 Check {doc}`lammps_mliap` for installation and lammps script in this use case.
 

docs/source/user_guide/lammps_mliap.md

@@ -6,13 +6,15 @@ Currently the parallel implementation of LAMMPS/ML-IAP is not tested.
 
 ## Requirements
 - cython == 3.0.11
-- cupy-cuda12x
+- cupy-cuda12x or cupy-cuda13x
 - flashTP (optional, follow [here](accelerator.md#flashtp))
 - cuEquivariance (optional, follow [here](accelerator.md#cuequivariance))
+- OpenEquivariance (optional, follow [here](accelerator.md#openequivariance))
 
 Install via:
 ```bash
-pip install sevenn[mliap]
+pip install sevenn[mliap12]  # For CUDA 12.x
+pip install sevenn[mliap13]  # For CUDA 13.x
 ```
 
 ## Build
@@ -119,7 +121,7 @@ Please check [sevenn graph_build](./cli.md#sevenn-graph-build) for detail.
 An ML-IAP potential checkpoint can be deployed using ``sevenn get_model`` command with ``--use_mliap`` flag.
 - By default, output file name will be ``deployed_serial_mliap.pt``.
   (You can customize the output file name using ``--output_prefix`` flag.)
-- You can accelerate the inference with ``--enable_cueq`` or ``--enable_flash`` flag:
+- You can accelerate the inference with ``--enable_cueq``, ``--enable_flash``, or ``--enable_oeq`` flag:
 ```bash
 sevenn get_model \
     {pretrained_name or checkpoint_path} \

docs/source/user_guide/lammps_torch.md

@@ -16,13 +16,15 @@ Ensure the LAMMPS version is `stable_2Aug2023_update3`. You can easily switch th
 
 ```bash
 git clone https://github.com/lammps/lammps.git lammps_sevenn --branch stable_2Aug2023_update3 --depth=1
-sevenn patch_lammps ./lammps_sevenn {--enable_flash} {--d3}
+sevenn patch_lammps ./lammps_sevenn {--enable_flash} {--enable_oeq} {--d3}
 ```
 You can refer to `sevenn/pair_e3gnn/patch_lammps.sh` for details of the patch process.
 
 
 :::{tip}
 (Optional) Add `--enable_flash` option to accelerate SevenNet for LAMMPS using flashTP. You must preinstall [flashTP](accelerator.md#flashtp) before building LAMMPS with flashTP.
+
+(Optional) Add `--enable_oeq` option to accelerate SevenNet for LAMMPS using OpenEquivariance. You must preinstall [OpenEquivariance](accelerator.md#openequivariance) before building LAMMPS with OpenEquivariance.
 :::
 
 ### (Optional) Build with GPU-D3 pair style

docs/source/user_guide/torchsim.md

@@ -24,6 +24,12 @@ from sevenn.torchsim import SevenNetModel
 model = SevenNetModel(model="7net-omni", modal="mpa")
 ```
 
+You can enable accelerators (cuEquivariance, flashTP, or OpenEquivariance) via the corresponding flags. For more information about accelerators, follow [here](./accelerator.md).
+```python
+model = SevenNetModel(model="7net-omni", modal="mpa", enable_oeq=True)
+# or enable_cueq=True or enable_flash=True
+```
+
 The `device` parameter defaults to `'auto'` (CUDA if available, otherwise CPU).
 
 ### Batched MD

pyproject.toml

@@ -36,7 +36,10 @@ dependencies = [
 test = ["pytest", "pytest-cov>=5", "ipython"]
 torchsim = ["torch-sim-atomistic>=0.5.2; python_version >= '3.12'"]
 cueq12 = ["cuequivariance>=0.6.0; python_version >= '3.10'", "cuequivariance-torch>=0.6.0; python_version >= '3.10'", "cuequivariance-ops-torch-cu12; python_version >= '3.10'"]
-mliap = ["cython==3.0.11", "cupy-cuda12x"]
+cueq13 = ["cuequivariance>=0.7.0; python_version >= '3.10'", "cuequivariance-torch>=0.7.0; python_version >= '3.10'", "cuequivariance-ops-torch-cu13; python_version >= '3.10'"]
+oeq = ["openequivariance"]
+mliap12 = ["cython==3.0.11", "cupy-cuda12x"]
+mliap13 = ["cython==3.0.11", "cupy-cuda13x"]
 
 [project.scripts]
 sevenn = "sevenn.main.sevenn:main"

sevenn/_const.py

@@ -131,6 +131,7 @@ def error_record_condition(x):
     KEY._NORMALIZE_SPH: True,
     KEY.USE_FLASH_TP: False,
     KEY.CUEQUIVARIANCE_CONFIG: {},
+    KEY.USE_OEQ: False,
 }
 
 
@@ -178,6 +179,7 @@ def error_record_condition(x):
     KEY._NORMALIZE_SPH: bool,
     KEY.USE_FLASH_TP: bool,
     KEY.CUEQUIVARIANCE_CONFIG: dict,
+    KEY.USE_OEQ: bool,
 }
 
 

sevenn/_keys.py

@@ -222,6 +222,7 @@
 
 USE_FLASH_TP = 'use_flash_tp'
 CUEQUIVARIANCE_CONFIG = 'cuequivariance_config'
+USE_OEQ = 'use_oeq'
 
 _NORMALIZE_SPH = '_normalize_sph'
 OPTIMIZE_BY_REDUCE = 'optimize_by_reduce'

sevenn/calculator.py

@@ -39,6 +39,7 @@ def __init__(
         modal: Optional[str] = None,
         enable_cueq: Optional[bool] = False,
         enable_flash: Optional[bool] = False,
+        enable_oeq: Optional[bool] = False,
         sevennet_config: Optional[Dict] = None,  # Not used in logic, just meta info
         **kwargs,
     ) -> None:
@@ -64,6 +65,8 @@ def __init__(
             if True, use cuEquivariant to accelerate inference.
         enable_flash: bool, default=None (use the checkpoint's backend)
             if True, use FlashTP to accelerate inference.
+        enable_oeq: bool, default=False
+            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
         """
@@ -80,17 +83,21 @@ def __init__(
 
         enable_cueq = os.getenv('SEVENNET_ENABLE_CUEQ') == '1' or enable_cueq
         enable_flash = os.getenv('SEVENNET_ENABLE_FLASH') == '1' or enable_flash
-        print('cueq')
-        print(enable_cueq)
-        print('flash')
-        print(enable_flash)
+        enable_oeq = os.getenv('SEVENNET_ENABLE_OEQ') == '1' or enable_oeq
 
         if enable_cueq and file_type in ['model_instance', 'torchscript']:
             warnings.warn(
                 'file_type should be checkpoint to enable cueq. cueq set to False'
             )
             enable_cueq = False
 
+        # TODO: not verified this line
+        if enable_oeq and file_type in ['model_instance', 'torchscript']:
+            warnings.warn(
+                'file_type should be checkpoint to enable oeq. oeq set to False'
+            )
+            enable_oeq = False
+
         if isinstance(device, str):  # TODO: do we really need this?
             if device == 'auto':
                 self.device = torch.device(
@@ -105,7 +112,7 @@ def __init__(
             cp = util.load_checkpoint(model)
 
             model_loaded = cp.build_model(
-                enable_cueq=enable_cueq, enable_flash=enable_flash
+                enable_cueq=enable_cueq, enable_flash=enable_flash, enable_oeq=enable_oeq  # noqa: E501
             )
             model_loaded.set_is_batch_data(False)
 
@@ -244,12 +251,16 @@ def __init__(
         model: Union[str, pathlib.PurePath, AtomGraphSequential] = '7net-0',
         file_type: str = 'checkpoint',
         device: Union[torch.device, str] = 'auto',
-        sevennet_config: Optional[Any] = None,  # hold meta information
+        modal: Optional[str] = None,
+        enable_cueq: Optional[bool] = False,
+        enable_flash: Optional[bool] = False,
+        enable_oeq: Optional[bool] = False,
+        sevennet_config: Optional[Any] = None,
         damping_type: str = 'damp_bj',
         functional_name: str = 'pbe',
         vdw_cutoff: float = 9000,  # au^2, 0.52917726 angstrom = 1 au
         cn_cutoff: float = 1600,  # au^2, 0.52917726 angstrom = 1 au
-        **kwargs,
+        **kwargs,  # pass extra kwargs to both calculators
     ) -> None:
         """Initialize SevenNetD3Calculator. CUDA required.
 
@@ -263,21 +274,27 @@ def __init__(
         device: str | torch.device, default='auto'
             if not given, use CUDA if available
         modal: str | None, default=None
-            modal (fidelity) if given model is multi-modal model. for 7net-mf-ompa,
-            for 7net-omni, use one of 'mpa', 'omat24', 'matpes_pbe', 'matpes_r2scan',
-            'mp_r2scan', 'oc20', 'oc22', 'odac23', 'omol25_low', 'omol25_high',
-            'spice', 'qcml', 'pet_mad'
-            it should be one of 'mpa' (MPtrj + sAlex) or 'omat24' (OMat24)
+            modal (fidelity) if given model is multi-modal model.
+            for 7net-mf-ompa, it should be one of 'mpa' or 'omat24'.
+            for 7net-omni, use one of 'mpa', 'omat24', 'matpes_pbe',
+            'matpes_r2scan', 'mp_r2scan', 'oc20', 'oc22', 'odac23',
+            'omol25_low', 'omol25_high', 'spice', 'qcml', 'pet_mad'
         enable_cueq: bool, default=False
-            if True, use cuEquivariant to accelerate inference.
+            if True, use cuEquivariance backend.
+        enable_flash: bool, default=False
+            if True, use flashTP backend.
+        enable_oeq: bool, default=False
+            if True, use OpenEquivariance backend.
+        sevennet_config: dict | None, default=None
+            hold meta information
         damping_type: str, default='damp_bj'
             Damping type of D3, one of 'damp_bj' | 'damp_zero'
         functional_name: str, default='pbe'
             Target functional name of D3 parameters.
         vdw_cutoff: float, default=9000
-            vdw cutoff of D3 calculator in au
+            vdw cutoff of D3 calculator in au^2
         cn_cutoff: float, default=1600
-            cn cutoff of D3 calculator in au
+            cn cutoff of D3 calculator in au^2
         """
         d3_calc = D3Calculator(
             damping_type=damping_type,
@@ -291,6 +308,10 @@ def __init__(
             model=model,
             file_type=file_type,
             device=device,
+            modal=modal,
+            enable_cueq=enable_cueq,
+            enable_flash=enable_flash,
+            enable_oeq=enable_oeq,
             sevennet_config=sevennet_config,
             **kwargs,
         )

sevenn/checkpoint.py

@@ -312,6 +312,7 @@ def build_model(
         *,
         enable_cueq: Optional[bool] = None,
         enable_flash: Optional[bool] = None,
+        enable_oeq: Optional[bool] = None,
         _flash_lammps: bool = False,
     ) -> AtomGraphSequential:
         """
@@ -328,11 +329,19 @@ def build_model(
         cp_using_flash = self.config.get(KEY.USE_FLASH_TP, False)
         enable_flash = cp_using_flash if enable_flash is None else enable_flash
 
+        cp_using_oeq = self.config.get(KEY.USE_OEQ, False)
+        enable_oeq = cp_using_oeq if enable_oeq is None else enable_oeq
+
+        if sum([enable_cueq, enable_flash, enable_oeq]) > 1:
+            raise ValueError('Only one TP accelerator can be enabled.')
+
         assert not _flash_lammps or enable_flash
         cfg_new = self.config
         cfg_new['_flash_lammps'] = _flash_lammps
+        cfg_new[KEY.USE_OEQ] = enable_oeq
 
-        if (cp_using_cueq, cp_using_flash) == (enable_cueq, enable_flash):
+        if (cp_using_cueq, cp_using_flash, cp_using_oeq) \
+                == (enable_cueq, enable_flash, enable_oeq):
             # backend not given, or checkpoint backend is same as requested
             model = build_E3_equivariant_model(cfg_new)
             state_dict = compat.patch_state_dict_if_old(
@@ -347,6 +356,7 @@ def build_model(
 
             cfg_new[KEY.CUEQUIVARIANCE_CONFIG] = {'use': enable_cueq}
             cfg_new[KEY.USE_FLASH_TP] = enable_flash
+            cfg_new[KEY.USE_OEQ] = enable_oeq
             model = build_E3_equivariant_model(cfg_new)
             stct_src = compat.patch_state_dict_if_old(
                 self.model_state_dict, self.config, model