Load MACE-OFF pretrained models into DeePMD-GNN for MD simulations

README.md

@@ -188,7 +188,31 @@ Type maps that starts with `m` (such as `mH`) or `OW` or `HW` will be recognized
 Two MM atoms will not build edges with each other.
 Such GNN+DPRc model can be directly used in AmberTools24.
 
+## MACE-OFF Pretrained Models
+
+DeePMD-GNN supports loading MACE-OFF foundation models for use with DeePMD-kit and MD packages:
+
+```python
+from deepmd_gnn import load_mace_off_model, convert_mace_off_to_deepmd
+
+# Load MACE-OFF model into DeePMD-GNN wrapper
+model = load_mace_off_model("small")  # or "medium", "large"
+
+# Convert to frozen format for MD simulations
+convert_mace_off_to_deepmd("small", "frozen_model.pth")
+# Now use with LAMMPS, AMBER/sander through DeePMD-kit
+```
+
+The frozen model can be used with:
+
+- **LAMMPS**: Set `DP_PLUGIN_PATH` and use `pair_style deepmd`
+- **AMBER/sander**: Through DeePMD-kit's AMBER interface (supports QM/MM with DPRc)
+- **Other MD packages**: Through DeePMD-kit's C++ interface
+
+For QM/MM simulations, use the DPRc mechanism: QM atoms use standard symbols (H, C, O), MM atoms use 'm' prefix (mH, mC) or HW/OW.
+
 ## Examples
 
-- [examples/water](examples/water)
-- [examples/dprc](examples/dprc)
+- [examples/water](examples/water) - Basic MACE and NequIP training
+- [examples/dprc](examples/dprc) - DPRc (QM/MM) with mixed atom types
+- [examples/mace_off](examples/mace_off) - Using MACE-OFF pretrained models with DeePMD-kit

deepmd_gnn/__init__.py

@@ -4,11 +4,19 @@
 
 from ._version import __version__
 from .argcheck import mace_model_args
+from .mace_off import (
+    convert_mace_off_to_deepmd,
+    download_mace_off_model,
+    load_mace_off_model,
+)
 
 __email__ = "jinzhe.zeng@ustc.edu.cn"
 
 __all__ = [
     "__version__",
+    "convert_mace_off_to_deepmd",
+    "download_mace_off_model",
+    "load_mace_off_model",
     "mace_model_args",
 ]
 

deepmd_gnn/__main__.py

@@ -1,5 +1,99 @@
 """Main entry point for the command line interface."""
 
+import argparse
+import logging
+import sys
+from pathlib import Path
+
+from deepmd_gnn.mace_off import convert_mace_off_to_deepmd, download_mace_off_model
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+def main() -> int:
+    """Run the main CLI."""
+    parser = argparse.ArgumentParser(
+        description="DeePMD-GNN utilities",
+        prog="deepmd-gnn",
+    )
+    subparsers = parser.add_subparsers(dest="command", help="Available commands")
+
+    # MACE-OFF download command
+    download_parser = subparsers.add_parser(
+        "download-mace-off",
+        help="Download MACE-OFF pretrained models",
+    )
+    download_parser.add_argument(
+        "model",
+        choices=["small", "medium", "large"],
+        help="MACE-OFF model size to download",
+    )
+    download_parser.add_argument(
+        "--cache-dir",
+        type=str,
+        default=None,
+        help="Directory to cache the downloaded model (default: ~/.cache/deepmd-gnn/mace-off/)",
+    )
+    download_parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Force re-download even if file exists",
+    )
+
+    # MACE-OFF convert command
+    convert_parser = subparsers.add_parser(
+        "convert-mace-off",
+        help="Convert MACE-OFF model to DeePMD format",
+    )
+    convert_parser.add_argument(
+        "model",
+        choices=["small", "medium", "large"],
+        help="MACE-OFF model size to convert",
+    )
+    convert_parser.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        default="frozen_model.pth",
+        help="Output file name for the frozen model (default: frozen_model.pth)",
+    )
+    convert_parser.add_argument(
+        "--cache-dir",
+        type=str,
+        default=None,
+        help="Directory where MACE-OFF models are cached",
+    )
+
+    args = parser.parse_args()
+
+    if args.command == "download-mace-off":
+        cache_dir = Path(args.cache_dir) if args.cache_dir else None
+        model_path = download_mace_off_model(
+            model_name=args.model,
+            cache_dir=cache_dir,
+            force_download=args.force,
+        )
+        logger.info("Model downloaded to: %s", model_path)
+        return 0
+
+    if args.command == "convert-mace-off":
+        cache_dir = Path(args.cache_dir) if args.cache_dir else None
+        output_path = convert_mace_off_to_deepmd(
+            model_name=args.model,
+            output_file=args.output,
+            cache_dir=cache_dir,
+        )
+        logger.info("Converted model saved to: %s", output_path)
+        return 0
+
+    parser.print_help()
+    return 1
+
+
 if __name__ == "__main__":
-    msg = "This module is not meant to be executed directly."
-    raise NotImplementedError(msg)
+    sys.exit(main())