aimnet 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

aimnet/calculators/model_registry.py

@@ -1,6 +1,5 @@
 import logging
 import os
-from typing import Dict, Optional
 
 import click
 import requests
@@ -9,7 +8,7 @@ import yaml
 logging.basicConfig(level=logging.INFO)
 
 
-def load_model_registry(registry_file: Optional[str] = None) -> Dict[str, str]:
+def load_model_registry(registry_file: str | None = None) -> dict[str, str]:
     registry_file = registry_file or os.path.join(os.path.dirname(__file__), "model_registry.yaml")
     with open(os.path.join(os.path.dirname(__file__), "model_registry.yaml")) as f:
         return yaml.load(f, Loader=yaml.SafeLoader)
@@ -43,9 +42,7 @@ def _maybe_download_asset(file: str, url: str) -> str:
 
 def get_model_path(s: str) -> str:
     # direct file path
-    if os.path.isfile(s):
-        print("Found model file:", s)
-    else:
+    if not os.path.isfile(s):
         s = get_registry_model_path(s)
     return s
 

aimnet/calculators/model_registry.yaml

@@ -1,33 +1,71 @@
 # map file name to url
 models:
     aimnet2_wb97m_d3_0:
-        file: aimnet2_wb97m_d3_0.jpt
-        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_wb97m_d3_0.jpt
+        file: aimnet2_wb97m_d3_0.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2/aimnet2_wb97m_d3_0.pt
     aimnet2_wb97m_d3_1:
-        file: aimnet2_wb97m_d3_1.jpt
-        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_wb97m_d3_1.jpt
+        file: aimnet2_wb97m_d3_1.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2/aimnet2_wb97m_d3_1.pt
     aimnet2_wb97m_d3_2:
-        file: aimnet2_wb97m_d3_2.jpt
-        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_wb97m_d3_2.jpt
+        file: aimnet2_wb97m_d3_2.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2/aimnet2_wb97m_d3_2.pt
     aimnet2_wb97m_d3_3:
-        file: aimnet2_wb97m_d3_3.jpt
-        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_wb97m_d3_3.jpt
+        file: aimnet2_wb97m_d3_3.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2/aimnet2_wb97m_d3_3.pt
     aimnet2_b973c_d3_0:
-        file: aimnet2_b973c_d3_0.jpt
-        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_b973c_d3_0.jpt
+        file: aimnet2_b973c_d3_0.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2/aimnet2_b973c_d3_0.pt
     aimnet2_b973c_d3_1:
-        file: aimnet2_b973c_d3_1.jpt
-        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_b973c_d3_1.jpt
+        file: aimnet2_b973c_d3_1.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2/aimnet2_b973c_d3_1.pt
     aimnet2_b973c_d3_2:
-        file: aimnet2_b973c_d3_2.jpt
-        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_b973c_d3_2.jpt
+        file: aimnet2_b973c_d3_2.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2/aimnet2_b973c_d3_2.pt
     aimnet2_b973c_d3_3:
-        file: aimnet2_b973c_d3_3.jpt
-        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_b973c_d3_3.jpt
+        file: aimnet2_b973c_d3_3.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2/aimnet2_b973c_d3_3.pt
+    aimnet2_b973c_2025_d3_0:
+        file: aimnet2_2025_b973c_d3_0.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2/aimnet2_2025_b973c_d3_0.pt
+    aimnet2_b973c_2025_d3_1:
+        file: aimnet2_2025_b973c_d3_1.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2/aimnet2_2025_b973c_d3_1.pt
+    aimnet2_b973c_2025_d3_2:
+        file: aimnet2_2025_b973c_d3_2.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2/aimnet2_2025_b973c_d3_2.pt
+    aimnet2_b973c_2025_d3_3:
+        file: aimnet2_2025_b973c_d3_3.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2/aimnet2_2025_b973c_d3_3.pt
+    aimnet2nse_0:
+        file: aimnet2nse_wb97m_0.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2NSE/aimnet2nse_wb97m_0.pt
+    aimnet2nse_1:
+        file: aimnet2nse_wb97m_1.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2NSE/aimnet2nse_wb97m_1.pt
+    aimnet2nse_2:
+        file: aimnet2nse_wb97m_2.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2NSE/aimnet2nse_wb97m_2.pt
+    aimnet2nse_3:
+        file: aimnet2nse_wb97m_3.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2NSE/aimnet2nse_wb97m_3.pt
+    aimnet2-pd_0:
+        file: aimnet2-pd_0.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2Pd/aimnet2-pd_0.pt
+    aimnet2-pd_1:
+        file: aimnet2-pd_1.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2Pd/aimnet2-pd_1.pt
+    aimnet2-pd_2:
+        file: aimnet2-pd_2.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2Pd/aimnet2-pd_2.pt
+    aimnet2-pd_3:
+        file: aimnet2-pd_3.pt
+        url: https://storage.googleapis.com/aimnetcentral/aimnet2v2/AIMNet2Pd/aimnet2-pd_3.pt
 
 # map model alias to file name
 aliases:
     aimnet2: aimnet2_wb97m_d3_0
     aimnet2_wb97m: aimnet2_wb97m_d3_0
     aimnet2_b973c: aimnet2_b973c_d3_0
-    aimnet2_qr: aimnet2_qr_v0
+    aimnet2nse: aimnet2nse_0
+    aimnet2pd: aimnet2-pd_0
+    aimnet2_2025: aimnet2_b973c_2025_d3_0

aimnet/cli.py

@@ -1,8 +1,8 @@
+import sys
+
 import click
 
-from .train.calc_sae import calc_sae
-from .train.pt2jpt import jitcompile
-from .train.train import train
+from .calculators.model_registry import clear_assets
 
 
 @click.group()
@@ -10,9 +10,65 @@ def cli():
     """AIMNet2 command line tool"""
 
 
-cli.add_command(train, name="train")
-cli.add_command(jitcompile, name="jitcompile")
-cli.add_command(calc_sae, name="calc_sae")
+# Always available commands
+cli.add_command(clear_assets, name="clear_model_cache")
+
+
+# Register convert command (doesn't need heavy training dependencies)
+try:
+    from .models.convert import convert_legacy_jpt
+
+    cli.add_command(convert_legacy_jpt, name="convert")
+except ImportError:
+
+    @cli.command(name="convert")
+    def convert_stub():
+        """Convert legacy JIT model to new format (requires aimnet[train])"""
+        click.echo(
+            "Training dependencies not installed.\nInstall with: pip install aimnet[train]",
+            err=True,
+        )
+        sys.exit(1)
+
+
+# Try to lazily register training commands (requires ignite, etc.)
+try:
+    from .train.calc_sae import calc_sae
+    from .train.export_model import export_model
+    from .train.train import train
+
+    cli.add_command(train, name="train")
+    cli.add_command(export_model, name="export")
+    cli.add_command(calc_sae, name="calc_sae")
+except ImportError:
+    # If training dependencies are not available, register stub commands with helpful error messages
+
+    @cli.command(name="train")
+    def train_stub():
+        """Train AIMNet2 models (requires aimnet[train])"""
+        click.echo(
+            "Training dependencies not installed.\nInstall with: pip install aimnet[train]",
+            err=True,
+        )
+        sys.exit(1)
+
+    @cli.command(name="export")
+    def export_stub():
+        """Export trained model to distributable format (requires aimnet[train])"""
+        click.echo(
+            "Training dependencies not installed.\nInstall with: pip install aimnet[train]",
+            err=True,
+        )
+        sys.exit(1)
+
+    @cli.command(name="calc_sae")
+    def calc_sae_stub():
+        """Calculate SAE (requires aimnet[train])"""
+        click.echo(
+            "Training dependencies not installed.\nInstall with: pip install aimnet[train]",
+            err=True,
+        )
+        sys.exit(1)
 
 
 if __name__ == "__main__":

aimnet/config.py

@@ -1,6 +1,7 @@
 import os
+from collections.abc import Callable, Iterator
 from importlib import import_module
-from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
+from typing import Any
 
 import yaml
 from jinja2 import Template
@@ -27,7 +28,7 @@ def get_module(name: str) -> Callable:
     return func  # type: ignore[no-any-return]
 
 
-def get_init_module(name: str, args: Optional[List] = None, kwargs: Optional[Dict] = None) -> Callable:
+def get_init_module(name: str, args: list | None = None, kwargs: dict | None = None) -> Callable:
     """
     Get the initialized module based on the given name, arguments, and keyword arguments.
 
@@ -46,8 +47,8 @@ def get_init_module(name: str, args: Optional[List] = None, kwargs: Optional[Dic
 
 
 def load_yaml(
-    config: Dict[str, Any] | List | str, hyperpar: Optional[Dict[str, Any] | str] = None
-) -> Dict[str, Any] | List:
+    config: dict[str, Any] | list | str, hyperpar: dict[str, Any] | str | None = None
+) -> dict[str, Any] | list:
     """
     Load a YAML configuration file and apply optional hyperparameters.
 
@@ -88,8 +89,8 @@ def load_yaml(
 
 
 def _iter_rec_bottomup(
-    d: Dict[str, Any] | List,
-) -> Iterator[Tuple[Dict[str, Any] | List, str | int, Any]]:
+    d: dict[str, Any] | list,
+) -> Iterator[tuple[dict[str, Any] | list, str | int, Any]]:
     if isinstance(d, list):
         it = enumerate(d)
     elif isinstance(d, dict):
@@ -102,9 +103,7 @@ def _iter_rec_bottomup(
         yield d, k, v
 
 
-def build_module(
-    config: Union[str, Dict, List], hyperpar: Union[str, Dict, None] = None
-) -> Union[List, Dict, Callable]:
+def build_module(config: str | dict | list, hyperpar: str | dict | None = None) -> list | dict | Callable:
     """
     Build a module based on the provided configuration.
     Every (possibly nested) dictionary with a 'class' key will be replaced by an instance initialized with

aimnet/data/sgdataset.py

@@ -1,6 +1,7 @@
 import os
+from collections.abc import Callable, Sequence
 from glob import glob
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing import Any
 
 import h5py
 import numpy as np
@@ -23,12 +24,12 @@ class DataGroup:
 
     def __init__(
         self,
-        data: Optional[str | Dict[str, np.ndarray] | h5py.Group] = None,
-        keys: Optional[List[str]] = None,
-        shard: Optional[Tuple[int, int]] = None,
+        data: str | dict[str, np.ndarray] | h5py.Group | None = None,
+        keys: list[str] | None = None,
+        shard: tuple[int, int] | None = None,
     ):
         # main container for data
-        self._data: Dict[str, np.ndarray] = {}
+        self._data: dict[str, np.ndarray] = {}
 
         if data is None:
             data = {}
@@ -184,13 +185,13 @@ class DataGroup:
 class SizeGroupedDataset:
     def __init__(
         self,
-        data: Optional[str | List[str] | Dict[int, Dict[str, np.ndarray]] | Dict[int, DataGroup]] = None,
-        keys: Optional[List[str]] = None,
-        shard: Optional[Tuple[int, int]] = None,
+        data: str | list[str] | dict[int, dict[str, np.ndarray]] | dict[int, DataGroup] | None = None,
+        keys: list[str] | None = None,
+        shard: tuple[int, int] | None = None,
     ):
         # main containers
-        self._data: Dict[int, DataGroup] = {}
-        self._meta: Dict[str, str] = {}
+        self._data: dict[int, DataGroup] = {}
+        self._meta: dict[str, str] = {}
 
         # load data
         if isinstance(data, str):
@@ -203,17 +204,17 @@ class SizeGroupedDataset:
         elif isinstance(data, dict):
             self.load_dict(data)
         self.loader_mode = False
-        self.x: List[str] = []
-        self.y: List[str] = []
+        self.x: list[str] = []
+        self.y: list[str] = []
 
-    def load_datadir(self, path, keys=None, shard: Optional[Tuple[int, int]] = None):
+    def load_datadir(self, path, keys=None, shard: tuple[int, int] | None = None):
         if not os.path.isdir(path):
             raise FileNotFoundError(f"{path} does not exist or not a directory.")
         for f in glob(os.path.join(path, "???.npz")):
             k = int(os.path.basename(f)[:3])
             self[k] = DataGroup(f, keys=keys, shard=shard)
 
-    def load_files(self, files, keys=None, shard: Optional[Tuple[int, int]] = None):
+    def load_files(self, files, keys=None, shard: tuple[int, int] | None = None):
         for fil in files:
             if not os.path.isfile(fil):
                 raise FileNotFoundError(f"{fil} does not exist or not a file.")
@@ -224,27 +225,27 @@ class SizeGroupedDataset:
         for k, v in data.items():
             self[k] = DataGroup(v, keys=keys)
 
-    def load_h5(self, data, keys=None, shard: Optional[Tuple[int, int]] = None):
+    def load_h5(self, data, keys=None, shard: tuple[int, int] | None = None):
         with h5py.File(data, "r") as f:
             for k, g in f.items():
                 k = int(k)
                 self[k] = DataGroup(g, keys=keys, shard=shard)
             self._meta = dict(f.attrs)  # type: ignore[attr-defined]
 
-    def keys(self) -> List[int]:
+    def keys(self) -> list[int]:
         return sorted(self._data.keys())
 
-    def values(self) -> List:
+    def values(self) -> list:
         return [self[k] for k in self.keys()]
 
-    def items(self) -> List[Tuple[int, Any]]:
+    def items(self) -> list[tuple[int, Any]]:
         return [(k, self[k]) for k in self.keys()]
 
     def datakeys(self):
         return next(iter(self._data.values())).keys() if self._data else set()
 
     @property
-    def groups(self) -> List[DataGroup]:
+    def groups(self) -> list[DataGroup]:
         return self.values()
 
     def __len__(self):
@@ -259,7 +260,7 @@ class SizeGroupedDataset:
             raise ValueError("Wrong set of data keys.")
         self._data[key] = value
 
-    def __getitem__(self, item: int | Tuple[int, Sequence]) -> Dict | Tuple[Dict, Dict]:
+    def __getitem__(self, item: int | tuple[int, Sequence]) -> dict | tuple[dict, dict]:
         if isinstance(item, int):
             ret = self._data[item]
         else:
@@ -332,7 +333,7 @@ class SizeGroupedDataset:
         for v in self.values():
             v.shuffle(seed)
 
-    def save(self, dirname, namemap_fn: Optional[Callable] = None, compress: bool = False):
+    def save(self, dirname, namemap_fn: Callable | None = None, compress: bool = False):
         os.makedirs(dirname, exist_ok=True)
         if namemap_fn is None:
             namemap_fn = lambda x: f"{x:03d}.npz"
@@ -441,7 +442,7 @@ class SizeGroupedDataset:
         for g in self.values():
             yield from g.iter_batched(batch_size, keys)
 
-    def get_loader(self, sampler, x: List[str], y: Optional[List[str]] = None, **loader_kwargs):
+    def get_loader(self, sampler, x: list[str], y: list[str] | None = None, **loader_kwargs):
         self.loader_mode = True
         self.x = x
         self.y = y or []

aimnet/kernels/__init__.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+"""AIMNet Kernels Package - GPU-accelerated custom operations."""
+
+import torch
+
+from .conv_sv_2d_sp_wp import conv_sv_2d_sp
+
+
+def load_ops():
+    """
+    Load and register all custom ops.
+
+    This function ensures that all custom operations are properly
+    registered with PyTorch's operator registry.
+
+    Should be called before using any of the custom kernels to ensure
+    proper registration with the PyTorch dispatcher.
+
+    Returns:
+        list: Available ops that were registered.
+    """
+    available_ops = []
+
+    # Import warp kernels to trigger registration
+    # Import custom ops module to trigger registration
+    from aimnet.modules import ops as _ops  # noqa: F401
+
+    from . import conv_sv_2d_sp_wp  # noqa: F401
+
+    # Verify ops are available
+    if hasattr(torch.ops, "aimnet"):
+        if hasattr(torch.ops.aimnet, "conv_sv_2d_sp_fwd"):
+            available_ops.append("aimnet::conv_sv_2d_sp_fwd")
+        if hasattr(torch.ops.aimnet, "conv_sv_2d_sp_bwd"):
+            available_ops.append("aimnet::conv_sv_2d_sp_bwd")
+        if hasattr(torch.ops.aimnet, "conv_sv_2d_sp_bwd_bwd"):
+            available_ops.append("aimnet::conv_sv_2d_sp_bwd_bwd")
+        if hasattr(torch.ops.aimnet, "dftd3_fwd"):
+            available_ops.append("aimnet::dftd3_fwd")
+
+    return available_ops
+
+
+__all__ = [
+    "conv_sv_2d_sp",
+    "load_ops",
+]