{"manifests":[{"schema_version":5,"cluster":"aurora","root":"/lus/flare/projects/Rootstock/rootstock","maintainer":{"name":"Owen Price Skelly","email":"OwenPriceSkelly@uchicago.edu"},"rootstock_version":"1.6.4","python_version":"unknown","last_updated":"2026-09-01T20:25:23.757506+00:00","environments":{"allscaip":{"built_at":"2026-09-01T20:22:15.697733+00:00","source_hash":"sha256:4f471ec85da7b66f1b72c24d3c0d8f52cdf23e6b0ea3ccb0ee3ed7dda9856fe0","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core\",\n# \"ase>=3.26\",\n# # Intel XPU (Aurora PVC) torch build. >=2.13: older XPU wheels have far\n# # slower FP64 kernels (see uma.py).\n# \"torch>=2.13\",\n# # torch's XPU wheels depend on this; it lives only on the XPU index, so it\n# # must be a direct dep for [tool.uv.sources] to route it (the index is\n# # explicit, so transitive-only deps aren't fetched from it).\n# \"triton-xpu\",\n# ]\n#\n# [tool.uv.sources]\n# # Experimental: fairchem-core from a fork whose xpu-support branch adds\n# # native Intel-GPU device handling (device=\"xpu\", torch.xpu seeding/cache\n# # management, XCCL collectives) -- PyPI fairchem-core accepts only\n# # \"cpu\"/\"cuda\" device strings. The branch tracks upstream main; switch back\n# # to a PyPI fairchem-core once XPU support merges upstream.\n# fairchem-core = { git = \"https://github.com/abagusetty/fairchem.git\", branch = \"xpu-support\", subdirectory = \"packages/fairchem-core\" }\n# torch = { index = \"pytorch-xpu\" }\n# triton-xpu = { index = \"pytorch-xpu\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-xpu\"\n# url = \"https://download.pytorch.org/whl/xpu\"\n# explicit = true\n# ///\n\"\"\"AllScAIP env (Intel XPU / Aurora) - FAIRChem scalable attention MLIP.\n\nSame as nvidia_configs/allscaip.py except (1) torch resolves from the Intel\nXPU wheel index and (2) fairchem-core installs from a fork with native XPU\nsupport.\n\nUnlike uma.py/esen.py, this env does NOT force\nInferenceSettings(base_precision_dtype=float64): AllScAIP does not support\nFP64 inference. Its radius-graph construction creates tensors at the torch\ndefault dtype that crash against a double batch (torch.mm and index_put\ndtype mismatches), and even past those, the backbone hard-casts its node\nrepresentations to float32 before the output heads\n(fairchem models/allscaip/AllScAIP.py), which then mismatches the doubled\nhead weights. So this env runs the fairchem default float32 -- the same\nprecision the NVIDIA deployments verify at. If XPU float32 kernels prove\nnumerically inadequate here (the reason uma.py forces FP64), that will\nsurface as a verification failure, and FP64 support has to land in\nfairchem first.\n\nPin one PVC tile with ZE_AFFINITY_MASK in the job (the worker inherits it).\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"allscaip-md-conserving-all-omol\": \"allscaip-md-conserving-all-omol\",\n \"allscaip-md-direct-all-omol\": \"allscaip-md-direct-all-omol\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"allscaip:custom\": None,\n}\n\n\ndef _fairchem_device(device: str) -> str:\n \"\"\"Translate an indexed device (\"xpu:2\") into what fairchem accepts.\n\n MLIPPredictUnit normalizes the requested device to a bare type and resolves\n the actual GPU itself from torch's current-device state\n (torch.xpu.current_device()), so an index has to travel through\n torch.xpu.set_device, not the argument -- the same constraint the CUDA\n configs shim around for multi-GPU verifies.\n \"\"\"\n if device.startswith(\"xpu:\"):\n import torch\n\n torch.xpu.set_device(int(device.split(\":\", 1)[1]))\n return \"xpu\"\n return device\n\n\ndef setup(checkpoint: str, device: str = \"xpu\", **kwargs):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(\n CHECKPOINTS[checkpoint],\n device=_fairchem_device(device),\n )\n return FAIRChemCalculator(predictor, **kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"xpu\", **kwargs):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads\n # through load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=_fairchem_device(device))\n return FAIRChemCalculator(predictor, **kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.22.1.dev17+gc54d9d20","rootstock":"1.6.4","torch":"2.13.0+xpu","triton-xpu":"3.7.2"},"lock_hash":"sha256:9058bfd4e353be7da7ee5f8e41bd9c22ec719a7ea3815946f628883c8ece6060","checkpoints":{"allscaip-md-direct-all-omol":{"fetched_at":"2026-08-28T22:12:52.602322+00:00","verified_at":"2026-09-01T20:25:23.538760+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/5517360cf5a23ca141f1b51b708d131289e35080bfa79f27b990695e3e78d18e","size":694955368},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-09-01T20:25:23.538783+00:00"},"allscaip-md-conserving-all-omol":{"fetched_at":"2026-08-28T22:12:58.225467+00:00","verified_at":"2026-09-01T20:24:35.245217+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/14fe8626e9e75337769ecf98a765cdc02667763ec4dc1103d958e1c9f56b26e9","size":688432239},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-09-01T20:24:35.245236+00:00"}}},"esen":{"built_at":"2026-08-28T22:07:58.741952+00:00","source_hash":"sha256:c5d979385eca6ca408d18b167ff6ef0426f302eb224443a0ddc8b0d6800ba3c5","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core\",\n# \"ase>=3.26\",\n# # Intel XPU (Aurora PVC) torch build. >=2.13: older XPU wheels have far\n# # slower FP64 kernels (see uma.py).\n# \"torch>=2.13\",\n# # torch's XPU wheels depend on this; it lives only on the XPU index, so it\n# # must be a direct dep for [tool.uv.sources] to route it (the index is\n# # explicit, so transitive-only deps aren't fetched from it).\n# \"triton-xpu\",\n# ]\n#\n# [tool.uv.sources]\n# # Experimental: fairchem-core from a fork whose xpu-support branch adds\n# # native Intel-GPU device handling (device=\"xpu\", torch.xpu seeding/cache\n# # management, XCCL collectives) -- PyPI fairchem-core accepts only\n# # \"cpu\"/\"cuda\" device strings. The branch tracks upstream main; switch back\n# # to a PyPI fairchem-core once XPU support merges upstream.\n# fairchem-core = { git = \"https://github.com/abagusetty/fairchem.git\", branch = \"xpu-support\", subdirectory = \"packages/fairchem-core\" }\n# torch = { index = \"pytorch-xpu\" }\n# triton-xpu = { index = \"pytorch-xpu\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-xpu\"\n# url = \"https://download.pytorch.org/whl/xpu\"\n# explicit = true\n# ///\n\"\"\"eSEN env (Intel XPU / Aurora) - FAIRChem eSEN single-task checkpoints.\n\nSame as nvidia_configs/esen.py except (1) torch resolves from the Intel XPU\nwheel index, (2) fairchem-core installs from a fork with native XPU support,\nand (3) InferenceSettings defaults to float32, so we set\nbase_precision_dtype=float64 to match the FP64 reference.\n\nPin one PVC tile with ZE_AFFINITY_MASK in the job (the worker inherits it).\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"esen-md-direct-all-omol\": \"esen-md-direct-all-omol\",\n \"esen-sm-conserving-all-omol\": \"esen-sm-conserving-all-omol\",\n \"esen-sm-direct-all-omol\": \"esen-sm-direct-all-omol\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"esen:custom\": None,\n}\n\n\ndef _fairchem_device(device: str) -> str:\n \"\"\"Translate an indexed device (\"xpu:2\") into what fairchem accepts.\n\n MLIPPredictUnit normalizes the requested device to a bare type and resolves\n the actual GPU itself from torch's current-device state\n (torch.xpu.current_device()), so an index has to travel through\n torch.xpu.set_device, not the argument -- the same constraint the CUDA\n configs shim around for multi-GPU verifies.\n \"\"\"\n if device.startswith(\"xpu:\"):\n import torch\n\n torch.xpu.set_device(int(device.split(\":\", 1)[1]))\n return \"xpu\"\n return device\n\n\ndef _fp64_settings():\n import torch\n from fairchem.core.units.mlip_unit.api.inference import InferenceSettings\n\n return InferenceSettings(base_precision_dtype=torch.float64, tf32=False)\n\n\ndef setup(checkpoint: str, device: str = \"xpu\", **kwargs):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(\n CHECKPOINTS[checkpoint],\n device=_fairchem_device(device),\n inference_settings=_fp64_settings(),\n )\n return FAIRChemCalculator(predictor, **kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"xpu\", **kwargs):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads\n # through load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(\n path, device=_fairchem_device(device), inference_settings=_fp64_settings()\n )\n return FAIRChemCalculator(predictor, **kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.22.1.dev17+gc54d9d20","rootstock":"1.6.3","torch":"2.13.0+xpu","triton-xpu":"3.7.2"},"lock_hash":"sha256:9058bfd4e353be7da7ee5f8e41bd9c22ec719a7ea3815946f628883c8ece6060","checkpoints":{"esen-sm-conserving-all-omol":{"fetched_at":"2026-08-28T22:12:45.087678+00:00","verified_at":"2026-08-28T22:28:23.467643+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/01f63da2d071e39fc46a5f22f8369d0fc9de317ab2ef361a76603f5661238025","size":50958067},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-08-28T22:28:23.467662+00:00"},"esen-md-direct-all-omol":{"fetched_at":"2026-08-28T22:12:46.073360+00:00","verified_at":"2026-08-28T22:25:18.924392+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/e0cd2ac37d2992b385a39778469c443ccbc679c7c4fb3b40b38aa2ee23c29b1d","size":405789591},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-08-28T22:25:18.924415+00:00"},"esen-sm-direct-all-omol":{"fetched_at":"2026-08-28T22:15:03.785140+00:00","verified_at":"2026-08-28T22:30:16.130354+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/a1b7109bcca6e1288c5f17c68776d3c6a1ab0766f6d7fdd45f901605355d2831","size":50960291},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-08-28T22:30:16.130371+00:00"}}},"mace":{"built_at":"2026-08-28T22:07:07.290997+00:00","source_hash":"sha256:3e5283739d4adc2c8bea71b299aa429c6ffe253e8f6f9ccb5f7bc0fb8ecee82a","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # 0.3.15+ needed for the mh-1 registry entry (matpes needs 0.3.13,\n# # omol needs 0.3.14, mpa-0 needs 0.3.10).\n# \"mace-torch>=0.3.15\",\n# \"ase>=3.22\",\n# # Intel XPU (Aurora PVC) torch build. >=2.13: older XPU wheels have much\n# # slower FP64 kernels on PVC.\n# \"torch>=2.13\",\n# # torch's XPU wheels depend on this; it lives only on the XPU index, so it\n# # must be a direct dep for [tool.uv.sources] to route it.\n# \"triton-xpu\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-xpu\" }\n# triton-xpu = { index = \"pytorch-xpu\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-xpu\"\n# url = \"https://download.pytorch.org/whl/xpu\"\n# explicit = true\n# ///\n\"\"\"MACE env (Intel XPU / Aurora) - MACE checkpoints on Intel Data Center GPU Max.\n\nIdentical to nvidia_configs/mace.py except torch resolves from the Intel XPU\nwheel index and setup() defaults to device=\"xpu\". MACE's plain e3nn/torch path\nruns on XPU as-is; cuEquivariance acceleration is CUDA-only and is not used.\nPin one PVC tile with ZE_AFFINITY_MASK in the job (the worker inherits it).\n\nAll checkpoints ship in the same `mace-torch` package, so they share an\nenvironment. Upstream-string routing in CHECKPOINTS: an `off:` prefix routes to\nmace_off() and an `omol:` prefix to mace_omol() (float64, molecules only); an\n`mh:` prefix marks a multi-head model (float64, per the MACE-MH-1 model card).\n\nMulti-head checkpoints select a head via the `head` kwarg on setup()\n(setup_kwargs={\"head\": ...} / --kwarg head=...), named by upstream's training\ncorpus - see MH1_HEADS. Selection is required: the multi-head checkpoint\nhas no default head, and setup() errors when none is given.\n\nThe OMOL checkpoint expects `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-mp-0-small\": \"small\",\n \"mace-mp-0-medium\": \"medium\",\n \"mace-mp-0-large\": \"large\",\n \"mace-off23-small\": \"off:small\",\n \"mace-off23-medium\": \"off:medium\",\n \"mace-off23-large\": \"off:large\",\n # Only a medium MPA-0 has been released, but upstream names the weights\n # file mace-mpa-0-medium.model - keep the size explicit like mace-mp-0.\n \"mace-mpa-0-medium\": \"medium-mpa-0\",\n \"mace-matpes-r2scan-0\": \"mace-matpes-r2scan-0\",\n # One entry per weights file: MH-1's heads are selected by setup(head=...).\n \"mace-mh-1\": \"mh:mh-1\",\n # Only the extra-large OMOL model has been released.\n \"mace-omol-0-extra-large\": \"omol:extra_large\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace:custom\": None,\n}\n\n# MH-1's heads, named by training corpus. The released weights file is the\n# authority (mace_select_head --list_heads; the model card also lists a\n# rgd1_b3lyp head, but that shipped only in mh-0 - ACEsuit/mace#1462).\n# Validated here because upstream only warns on an unknown head and silently\n# falls back to the last one.\nMH1_HEADS = (\n \"omat_pbe\",\n \"omol\",\n \"spice_wB97M\",\n \"oc20_usemppbe\",\n \"mp_pbe_refit_add\",\n \"matpes_r2scan\",\n)\n\n# Verification-only head selection for the multi-head checkpoint:\n# smoke-test and a bare `rootstock add` verify with this\n# (setup() itself has no default head).\nVERIFY_KWARGS = {\n \"mace-mh-1\": {\"head\": \"omat_pbe\"},\n}\n\n\ndef setup(checkpoint: str, device: str = \"xpu\", head: str | None = None, **kwargs):\n arg = CHECKPOINTS[checkpoint]\n if arg.startswith(\"mh:\"):\n if head is None:\n raise ValueError(\n f\"{checkpoint} is multi-head and has no default - select one \"\n f\"with setup_kwargs={{'head': ...}} (or --kwarg head=...): \"\n f\"one of {', '.join(MH1_HEADS)}\"\n )\n if head not in MH1_HEADS:\n raise ValueError(f\"unknown head {head!r}; expected one of {', '.join(MH1_HEADS)}\")\n from mace.calculators import mace_mp\n\n kwargs.setdefault(\"default_dtype\", \"float64\")\n return mace_mp(model=arg[3:], device=device, head=head, **kwargs)\n if head is not None:\n raise ValueError(f\"'head' selects a head of a multi-head model; {checkpoint} has one head\")\n if arg.startswith(\"off:\"):\n from mace.calculators import mace_off\n\n kwargs.setdefault(\"default_dtype\", \"float32\")\n return mace_off(model=arg[4:], device=device, **kwargs)\n if arg.startswith(\"omol:\"):\n from mace.calculators import mace_omol\n\n kwargs.setdefault(\"default_dtype\", \"float64\")\n return mace_omol(model=arg[5:], device=device, **kwargs)\n from mace.calculators import mace_mp\n\n kwargs.setdefault(\"default_dtype\", \"float32\")\n return mace_mp(model=arg, device=device, **kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"xpu\", head: str | None = None, **kwargs):\n # Custom checkpoints (`:custom` ids with user weights): fine-tunes load through\n # MACECalculator directly - the mp/off dispatch in setup() only exists\n # to pick which pretrained file to download. `head` is for fine-tunes that\n # keep multiple heads; single-head weights load without it.\n from mace.calculators import MACECalculator\n\n kwargs.setdefault(\"default_dtype\", \"float32\")\n return MACECalculator(model_paths=path, device=device, head=head, **kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"1.6.3","torch":"2.13.0+xpu","triton-xpu":"3.7.2"},"lock_hash":"sha256:55c014f0bf87d1e5a5ffdb31e64fb8cd9e5511ae7ad709fef95bc781b548c960","checkpoints":{"mace-mh-1":{"fetched_at":"2026-09-01T19:35:20.465011+00:00","verified_at":"2026-09-01T19:44:03.646365+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/mace/macemh1model","size":59208139}],"weights_recorded_at":"2026-09-01T19:44:03.646384+00:00"},"mace-matpes-r2scan-0":{"fetched_at":"2026-08-28T22:15:57.315178+00:00","verified_at":"2026-08-28T22:32:05.700578+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/mace/MACEmatpesr2scanomatftmodel","size":79470738}],"weights_recorded_at":"2026-08-28T22:32:05.700599+00:00"},"mace-mp-0-medium":{"fetched_at":"2026-08-28T22:15:57.428766+00:00","verified_at":"2026-08-28T22:35:51.429862+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/mace/20231203mace128L1_epoch199model","size":44422970}],"weights_recorded_at":"2026-08-28T22:35:51.429883+00:00"},"mace-mp-0-large":{"fetched_at":"2026-08-28T22:16:06.143682+00:00","verified_at":"2026-08-28T22:34:09.756516+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/mace/MACE_MPtrj_20229model","size":133803220}],"weights_recorded_at":"2026-08-28T22:34:09.756532+00:00"},"mace-mp-0-small":{"fetched_at":"2026-08-28T22:16:20.569819+00:00","verified_at":"2026-08-28T22:37:28.524287+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/mace/20231210mace128L0_energy_epoch249model","size":32581838}],"weights_recorded_at":"2026-08-28T22:37:28.524303+00:00"},"mace-off23-medium":{"fetched_at":"2026-08-28T22:17:42.056246+00:00","verified_at":"2026-08-28T22:43:21.354581+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_medium.model","size":18350596}],"weights_recorded_at":"2026-08-28T22:43:21.354596+00:00"},"mace-off23-small":{"fetched_at":"2026-08-28T22:17:42.570166+00:00","verified_at":"2026-08-28T22:45:42.048460+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_small.model","size":7347350}],"weights_recorded_at":"2026-08-28T22:45:42.048476+00:00"},"mace-mpa-0-medium":{"fetched_at":"2026-08-28T22:17:42.605598+00:00","verified_at":"2026-08-28T22:39:36.589438+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/mace/macempa0mediummodel","size":79462305}],"weights_recorded_at":"2026-08-28T22:39:36.589457+00:00"},"mace-off23-large":{"fetched_at":"2026-08-28T22:17:42.835955+00:00","verified_at":"2026-08-28T22:41:58.157425+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_large.model","size":55492786}],"weights_recorded_at":"2026-08-28T22:41:58.157444+00:00"},"mace-omol-0-extra-large":{"fetched_at":"2026-08-28T22:18:53.071216+00:00","verified_at":"2026-08-28T22:47:48.146899+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/mace/MACE-omol-0-extra-large-1024.model","size":422242640}],"weights_recorded_at":"2026-08-28T22:47:48.146915+00:00"}}},"uma":{"built_at":"2026-08-28T22:07:42.346127+00:00","source_hash":"sha256:80b80f520e7ee9e22b2d59307f3e72a953ea539b6a45b672ad54e249152ae899","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core\",\n# \"ase>=3.26\",\n# # Intel XPU (Aurora PVC) torch build. >=2.13: older XPU wheels (e.g. 2.8)\n# # have far slower FP64 kernels -- UMA's first forward took >40 min on\n# # 2.8.0+xpu vs ~2 min on 2.13.0+xpu (PVC tile).\n# \"torch>=2.13\",\n# # torch's XPU wheels depend on this; it lives only on the XPU index, so it\n# # must be a direct dep for [tool.uv.sources] to route it (the index is\n# # explicit, so transitive-only deps aren't fetched from it).\n# \"triton-xpu\",\n# ]\n#\n# [tool.uv.sources]\n# # Experimental: fairchem-core from a fork whose xpu-support branch adds\n# # native Intel-GPU device handling (device=\"xpu\", torch.xpu seeding/cache\n# # management, XCCL collectives) -- PyPI fairchem-core accepts only\n# # \"cpu\"/\"cuda\" device strings. The branch tracks upstream main; switch back\n# # to a PyPI fairchem-core once XPU support merges upstream.\n# fairchem-core = { git = \"https://github.com/abagusetty/fairchem.git\", branch = \"xpu-support\", subdirectory = \"packages/fairchem-core\" }\n# torch = { index = \"pytorch-xpu\" }\n# triton-xpu = { index = \"pytorch-xpu\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-xpu\"\n# url = \"https://download.pytorch.org/whl/xpu\"\n# explicit = true\n# ///\n\"\"\"UMA env (Intel XPU / Aurora) - Meta's UMA foundation model via FAIRChem.\n\nSame as nvidia_configs/uma.py except (1) torch resolves from the Intel XPU wheel\nindex, (2) fairchem-core installs from a fork with native XPU support (replacing\nthe monkeypatch this env previously carried), and (3) InferenceSettings defaults\nto float32, so we set base_precision_dtype=float64 to match the FP64 reference\n(fp32 is wrong at ~1e-7).\n\nPin one PVC tile with ZE_AFFINITY_MASK in the job (the worker inherits it).\nRequires HF_TOKEN for the gated facebook/UMA checkpoints (download on a login\nnode; `rootstock add ... --no-verify`, then verify on a compute node).\n\nUMA is multi-task: setup() requires an explicit `task`\n(setup_kwargs={\"task\": ...} / --kwarg task=...) and errors without one.\nVerification picks its own head via VERIFY_KWARGS below.\n\"\"\"\n\nCHECKPOINTS = {\n \"uma-s-1p1\": \"uma-s-1p1\",\n # uma-s-1p2 had a known major bug; uma-s-1p2p1 is the fixed,\n # upstream-recommended small model and replaces it here.\n \"uma-s-1p2p1\": \"uma-s-1p2p1\",\n \"uma-m-1p1\": \"uma-m-1p1\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"uma:custom\": None,\n}\n\nUMA_TASKS = (\"omat\", \"omol\", \"oc20\", \"odac\", \"omc\")\n\n# Verification-only head selection: smoke-test and a bare `rootstock add`\n# verify with these (setup() itself has no default task).\nVERIFY_KWARGS = {\n \"uma-s-1p1\": {\"task\": \"omat\"},\n \"uma-s-1p2p1\": {\"task\": \"omat\"},\n \"uma-m-1p1\": {\"task\": \"omat\"},\n \"uma:custom\": {\"task\": \"omat\"},\n}\n\n\ndef _require_task(task, checkpoint):\n if task is None:\n raise ValueError(\n f\"{checkpoint} is multi-task and has no default head - select one \"\n f'with setup_kwargs={{\"task\": ...}} (or --kwarg task=...): '\n f\"one of {', '.join(UMA_TASKS)}\"\n )\n if task not in UMA_TASKS:\n raise ValueError(f\"unknown task {task!r}; expected one of {', '.join(UMA_TASKS)}\")\n return task\n\n\ndef _fairchem_device(device: str) -> str:\n \"\"\"Translate an indexed device (\"xpu:2\") into what fairchem accepts.\n\n MLIPPredictUnit normalizes the requested device to a bare type and resolves\n the actual GPU itself from torch's current-device state\n (torch.xpu.current_device()), so an index has to travel through\n torch.xpu.set_device, not the argument -- the same constraint the CUDA\n configs shim around for multi-GPU verifies.\n \"\"\"\n if device.startswith(\"xpu:\"):\n import torch\n\n torch.xpu.set_device(int(device.split(\":\", 1)[1]))\n return \"xpu\"\n return device\n\n\ndef _fp64_settings():\n import torch\n from fairchem.core.units.mlip_unit.api.inference import InferenceSettings\n\n return InferenceSettings(base_precision_dtype=torch.float64, tf32=False)\n\n\ndef setup(checkpoint: str, device: str = \"xpu\", task: str | None = None, **kwargs):\n task = _require_task(task, checkpoint)\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(\n CHECKPOINTS[checkpoint],\n device=_fairchem_device(device),\n inference_settings=_fp64_settings(),\n )\n return FAIRChemCalculator(predictor, task_name=task, **kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"xpu\", task: str | None = None, **kwargs):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads\n # through load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(\n path, device=_fairchem_device(device), inference_settings=_fp64_settings()\n )\n return FAIRChemCalculator(predictor, task_name=task, **kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.22.1.dev17+gc54d9d20","rootstock":"1.6.3","torch":"2.13.0+xpu","triton-xpu":"3.7.2"},"lock_hash":"sha256:9058bfd4e353be7da7ee5f8e41bd9c22ec719a7ea3815946f628883c8ece6060","checkpoints":{"uma-s-1p2p1":{"fetched_at":"2026-09-01T19:38:05.042220+00:00","verified_at":"2026-09-01T19:49:20.900494+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/b2673b85037b075674c25f55c34ffe1ff1e15db924be977b10a184765df0d5ce","size":2333393615},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-09-01T19:49:20.900517+00:00"},"uma-s-1p1":{"fetched_at":"2026-09-01T19:38:04.379997+00:00","verified_at":"2026-09-01T19:48:09.312149+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/07068e9c76702ca173d13155095f2117c1b327ec228557e64cd2709c777b824a","size":1173759575},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-09-01T19:48:09.312164+00:00"},"uma-m-1p1":{"fetched_at":"2026-09-01T19:39:13.671213+00:00","verified_at":"2026-09-01T19:46:49.584507+00:00","verified_device":"xpu","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/c30034edbf2e127f703f814cacb632661767da99b3e71c5b2ee5290510a52d68","size":11174706771},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-09-01T19:46:49.584535+00:00"}}}}},{"schema_version":5,"cluster":"della","root":"/scratch/gpfs/ROSENGROUP/common/rootstock","maintainer":{"name":"Will Engler","email":"willengler@uchicago.edu"},"rootstock_version":"1.6.0","python_version":"3.12.12","last_updated":"2026-08-09T21:31:43.124307+00:00","environments":{"mace":{"built_at":"2026-08-09T21:28:40.327274+00:00","source_hash":"sha256:b6cf58f573787e39bbc099deb93aa77918c4679f7bdd9a974df2966a58bafc2e","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # 0.3.15+ needed for the mh-1 registry entry (matpes needs 0.3.13,\n# # omol needs 0.3.14, mpa-0 needs 0.3.10).\n# \"mace-torch>=0.3.15\",\n# \"ase>=3.22\",\n# # 2.4.1 is explicitly unsupported by mace-torch.\n# \"torch>=2.4.0,!=2.4.1,<2.10\",\n# ]\n# ///\n\"\"\"MACE env — hosts MACE-MP-0, MACE-OFF23, MPA-0, MATPES, MH-1, and OMOL checkpoints.\n\nAll ship in the same `mace-torch` package, so they share an environment.\nUpstream-string routing in CHECKPOINTS: an `off:` prefix routes to mace_off()\nand an `omol:` prefix to mace_omol() (float64, molecules only); an `mh:`\nprefix marks a multi-head model (float64, per the MACE-MH-1 model card).\n\nMulti-head checkpoints select a head via the `head` kwarg on setup()\n(setup_kwargs={\"head\": ...} / --kwarg head=...), named by upstream's training\ncorpus — see MH1_HEADS; omat_pbe is the default.\n\nThe OMOL checkpoint expects `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-mp-0-small\": \"small\",\n \"mace-mp-0-medium\": \"medium\",\n \"mace-mp-0-large\": \"large\",\n \"mace-off23-small\": \"off:small\",\n \"mace-off23-medium\": \"off:medium\",\n \"mace-off23-large\": \"off:large\",\n # Only a medium MPA-0 has been released, but upstream names the weights\n # file mace-mpa-0-medium.model — keep the size explicit like mace-mp-0.\n \"mace-mpa-0-medium\": \"medium-mpa-0\",\n \"mace-matpes-r2scan-0\": \"mace-matpes-r2scan-0\",\n # One entry per weights file: MH-1's heads are selected by setup(head=...).\n \"mace-mh-1\": \"mh:mh-1\",\n # Only the extra-large OMOL model has been released.\n \"mace-omol-0-extra-large\": \"omol:extra_large\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace:custom\": None,\n}\n\n# MH-1's heads, named by training corpus. The released weights file is the\n# authority (mace_select_head --list_heads; the model card also lists a\n# rgd1_b3lyp head, but that shipped only in mh-0 — ACEsuit/mace#1462).\n# Validated here because upstream only warns on an unknown head and silently\n# falls back to the last one.\nMH1_HEADS = (\n \"omat_pbe\",\n \"omol\",\n \"spice_wB97M\",\n \"oc20_usemppbe\",\n \"mp_pbe_refit_add\",\n \"matpes_r2scan\",\n)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", head: str | None = None):\n arg = CHECKPOINTS[checkpoint]\n if arg.startswith(\"mh:\"):\n head = head or \"omat_pbe\"\n if head not in MH1_HEADS:\n raise ValueError(f\"unknown head {head!r}; expected one of {', '.join(MH1_HEADS)}\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg[3:], device=device, default_dtype=\"float64\", head=head)\n if head is not None:\n raise ValueError(f\"'head' selects a head of a multi-head model; {checkpoint} has one head\")\n if arg.startswith(\"off:\"):\n from mace.calculators import mace_off\n\n return mace_off(model=arg[4:], device=device, default_dtype=\"float32\")\n if arg.startswith(\"omol:\"):\n from mace.calculators import mace_omol\n\n return mace_omol(model=arg[5:], device=device, default_dtype=\"float64\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg, device=device, default_dtype=\"float32\")\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", head: str | None = None):\n # Custom checkpoints (`:custom` ids with user weights): fine-tunes load through\n # MACECalculator directly — the mp/off dispatch in setup() only exists\n # to pick which pretrained file to download. `head` is for fine-tunes that\n # keep multiple heads; single-head weights load without it.\n from mace.calculators import MACECalculator\n\n return MACECalculator(model_paths=path, device=device, default_dtype=\"float32\", head=head)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"1.6.0","torch":"2.9.1"},"lock_hash":"sha256:2ef43c5e21b8c2dc7b3830833e9ae38555d87ade585a064e9fba3d63ca792931","checkpoints":{"mace-mp-0-medium":{"fetched_at":"2026-07-15T00:14:09.994980+00:00","verified_at":"2026-08-02T21:06:31.311259+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-mp-0-small":{"fetched_at":"2026-07-15T00:48:09.548852+00:00","verified_at":"2026-08-02T21:06:43.067393+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-off23-medium":{"fetched_at":"2026-07-15T00:50:11.161724+00:00","verified_at":"2026-08-02T21:07:17.054646+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-mp-0-large":{"fetched_at":"2026-07-31T21:09:33.369252+00:00","verified_at":"2026-08-02T21:06:20.009869+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-mh-1-matpes-r2scan":{"fetched_at":"2026-07-31T21:09:33.000240+00:00","verified_at":"2026-07-31T21:31:01.481791+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-matpes-r2scan-0":{"fetched_at":"2026-07-31T21:09:33.068800+00:00","verified_at":"2026-08-02T21:05:56.550536+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-omol-0-extra-large":{"fetched_at":"2026-07-31T21:09:56.093195+00:00","verified_at":"2026-08-02T21:07:40.196129+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-off23-small":{"fetched_at":"2026-07-31T21:09:54.640390+00:00","verified_at":"2026-08-02T21:07:28.279614+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-off23-large":{"fetched_at":"2026-07-31T21:09:55.121989+00:00","verified_at":"2026-08-02T21:07:05.735741+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-mpa-0-medium":{"fetched_at":"2026-07-31T21:09:32.996896+00:00","verified_at":"2026-08-02T21:06:54.459145+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-mh-1":{"fetched_at":"2026-08-02T21:04:03.555433+00:00","verified_at":"2026-08-02T21:06:08.459835+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"uma":{"built_at":"2026-08-09T21:31:42.891716+00:00","source_hash":"sha256:90fa56797dbc112f801242475202c036de053f8cc9b27d7d6aeb2410a37b8149","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"UMA env — hosts Meta's UMA foundation model via FAIRChem.\n\nfairchem-core v2 dropped the torch-geometric / pyg-find-links install dance, so\nthis env is a plain PyPI install. The original uma-s-1 had an extensivity bug\nand was removed from the fairchem 2.20 registry — use uma-s-1p1 or uma-s-1p2.\n\"\"\"\n\nCHECKPOINTS = {\n \"uma-s-1p1\": \"uma-s-1p1\",\n \"uma-s-1p2\": \"uma-s-1p2\",\n \"uma-m-1p1\": \"uma-m-1p1\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", task: str = \"omat\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(CHECKPOINTS[checkpoint], device=device)\n return FAIRChemCalculator(predictor, task_name=task)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.6.0","torch":"2.8.0"},"lock_hash":"sha256:c99d413acdd9b5b2cdd0241ba29966574fe6800af2b63376abacd0ad28fed8b0","checkpoints":{"uma-s-1p1":{"fetched_at":"2026-07-31T22:57:09.889223+00:00","verified_at":"2026-07-31T22:59:02.723456+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"uma-s-1p2":{"fetched_at":"2026-07-31T21:15:16.614742+00:00","verified_at":"2026-07-31T21:41:53.804658+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"uma-m-1p1":{"fetched_at":"2026-07-31T21:16:00.209873+00:00","verified_at":"2026-07-31T22:26:47.733347+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"chgnet":{"built_at":"2026-07-31T21:00:03.489865+00:00","source_hash":"sha256:09ca4efcb414150c6fc3f3f55b5c987f7ea279a740e11bd6e5a948cd8866adfa","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"chgnet>=0.3.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"CHGNet env — hosts pretrained charge-informed universal potentials.\"\"\"\n\nCHECKPOINTS = {\n \"chgnet-default\": \"chgnet-default\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a CHGNet calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\")\n\n Returns:\n ASE-compatible calculator\n \"\"\"\n from chgnet.model import CHGNet, CHGNetCalculator\n\n model_name = CHECKPOINTS[checkpoint]\n model = CHGNet.load() if model_name == \"chgnet-default\" else CHGNet.load(model_name)\n return CHGNetCalculator(model=model, use_device=device)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","chgnet":"0.4.2","rootstock":"1.4.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:b5bc1f87f4c596d669c56f6a5053d21b488b1e141dc8e38e7a7a964b887fcc73","checkpoints":{"chgnet-default":{"fetched_at":"2026-07-15T00:48:55.366181+00:00","verified_at":"2026-07-31T21:28:52.426760+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"tensornet":{"built_at":"2026-07-15T00:17:28.973758+00:00","source_hash":"sha256:ab0259c5151c890fa587448f3a36455046aa6c5e64a0aa9127e64c56e1c505dd","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"torch>=2.4.0,<2.5\",\n# \"ase>=3.22\",\n# \"huggingface_hub\",\n# \"matgl\",\n# # 0.4+ needs torch>=2.8 at runtime (custom-op registration uses string\n# # annotations infer_schema can't parse on older torch) but only declares\n# # the constraint on its extras, so the resolver won't catch it.\n# \"nvalchemi-toolkit-ops<0.4\",\n# \"pymatgen\",\n# \"monty\",\n# \"ruamel.yaml\",\n# \"scipy\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# \"torch-spline-conv\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n#\n# [tool.uv.sources]\n# matgl = { git = \"https://github.com/materialsvirtuallab/matgl.git\" }\n# ///\n\"\"\"TensorNet env — hosts MatPES TensorNet checkpoints via MatGL.\"\"\"\n\nCHECKPOINTS = {\n \"tensornet-matpes-pbe-2025-2\": \"materialyze/TensorNet-PES-MatPES-PBE-2025.2\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n import torch\n\n torch.set_default_device(device)\n\n # matgl 1.0.0 imports ExpCellFilter from ase.constraints, but it moved to\n # ase.filters in ASE 3.23. Patch it in before matgl imports.\n import ase.constraints\n\n if not hasattr(ase.constraints, \"ExpCellFilter\"):\n from ase.filters import ExpCellFilter\n\n ase.constraints.ExpCellFilter = ExpCellFilter\n\n # DGL 2.x graphbolt imports torchdata submodules removed in torchdata>=0.7.\n # Stub the entire graphbolt subpackage before `import dgl` runs; DGL's\n # __init__ will use our empty stub and skip the real graphbolt initialisation.\n # matgl only uses DGL for graph construction — graphbolt is never called.\n import sys, types\n\n for _name in [\n \"dgl.graphbolt\",\n \"dgl.graphbolt.base\",\n \"dgl.graphbolt.dataloader\",\n \"dgl.graphbolt.feature_fetcher\",\n \"dgl.graphbolt.minibatch_transformer\",\n ]:\n if _name not in sys.modules:\n sys.modules[_name] = types.ModuleType(_name)\n\n from huggingface_hub import snapshot_download\n\n import matgl\n from matgl.ext.ase import PESCalculator\n\n # matgl 1.0.0 load_model only checks the GitHub manifest; HF models must\n # be downloaded explicitly and passed as a local path.\n local_path = snapshot_download(repo_id=CHECKPOINTS[checkpoint])\n pot = matgl.load_model(local_path)\n return PESCalculator(potential=pot)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","huggingface-hub":"1.23.0","matgl":"4.0.3","monty":"2026.5.18","nvalchemi-toolkit-ops":"0.3.1","pymatgen":"2026.5.4","rootstock":"0.9.7.post13.dev0+58005d7","ruamel-yaml":"0.19.1","scipy":"1.17.1","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121","torch-spline-conv":"1.2.2+pt24cu121"},"lock_hash":null,"checkpoints":{"tensornet-matpes-pbe-2025-2":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"download: ValueError: infer_schema(func): Parameter positions has unsupported type torch.Tensor. The valid types are: dict_keys([, typing.Optional[torch.Tensor], typing.Sequence[torch.Tensor], typing.List[torch.Tensor], typing.Sequence[typing.Optional[torch.Tensor]], typing.List[typing.Optional[torch.Tensor]], , typing.Optional[int], typing.Sequence[int], typing.List[int], typing.Optional[typing.Sequence[int]], typing.Optional[typing.List[int]], , typing.Optional[float], typing.Sequence[float], typing.List[float], typing.Optional[typing.Sequence[float]], typing.Optional[typing.List[float]], , typing.Optional[bool], typing.Sequence[bool], typing.List[bool], typing.Optional[typing.Sequence[bool]], typing.Optional[typing.List[bool]], , typing.Optional[str], typing.Union[int, float, bool], typing.Union[int, float, bool, NoneType], typing.Sequence[typing.Union[int, float, bool]], typing.List[typing.Union[int, float, bool]], , typing.Optional[torch.dtype], , typing.Optional[torch.device]]). Got func with signature (positions: 'torch.Tensor', cutoff: 'float', cell: 'torch.Tensor', pbc: 'torch.Tensor', batch_idx: 'torch.Tensor', cells_per_dimension: 'torch.Tensor', atom_periodic_shifts: 'torch.Tensor', atom_to_cell_mapping: 'torch.Tensor', atoms_per_cell_count: 'torch.Tensor', cell_atom_start_indices: 'torch.Tensor', cell_atom_list: 'torch.Tensor') -> 'None')","weight_files":null,"weights_recorded_at":null}}},"orb_v3":{"built_at":"2026-07-31T21:01:30.729114+00:00","source_hash":"sha256:305f3d2e8183d7fb0912e6d4bf3a737c0f42e383cf12e4fec4289a67c00be837","source":"# /// script\n# requires-python = \">=3.12\"\n# dependencies = [\n# \"orb-models>=0.6.2\",\n# \"ase>=3.25\",\n# \"torch>=2.8\",\n# # Not imported here — constrains orb-models' transitive dep. setup()'s\n# # no-lock serve path relies on cached_path returning local files without\n# # locking or writing, verified against exactly this version (#67).\n# \"cached_path==1.8.10\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"Orb v3 env — Orbital Materials' Orb v3 universal potentials.\n\nSeparate from orb.py because orb-models>=0.5 changed the loader API\n(returns a tuple, requires `atoms_adapter` on ORBCalculator, moved calculator\nimport path) and 0.6.x bumped the Python floor to 3.12 and torch to 2.8.\n\"\"\"\n\nimport os\nimport shutil\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"orb-v3-conservative-inf-omat\": \"orb-v3-conservative-inf-omat\",\n \"orb-v3-conservative-20-omat\": \"orb-v3-conservative-20-omat\",\n \"orb-v3-direct-inf-omat\": \"orb-v3-direct-inf-omat\",\n \"orb-v3-direct-20-omat\": \"orb-v3-direct-20-omat\",\n \"orb-v3-conservative-inf-mpa\": \"orb-v3-conservative-inf-mpa\",\n \"orb-v3-conservative-20-mpa\": \"orb-v3-conservative-20-mpa\",\n \"orb-v3-direct-inf-mpa\": \"orb-v3-direct-inf-mpa\",\n \"orb-v3-direct-20-mpa\": \"orb-v3-direct-20-mpa\",\n \"orb-v3-conservative-omol\": \"orb-v3-conservative-omol\",\n \"orb-v3-direct-omol\": \"orb-v3-direct-omol\",\n}\n\n\ndef _default_weights_url(load_fn) -> str:\n \"\"\"The upstream URL baked into the loader's ``weights_path`` default.\"\"\"\n import inspect\n\n default = inspect.signature(load_fn).parameters[\"weights_path\"].default\n if not isinstance(default, str) or not default.startswith((\"http://\", \"https://\")):\n raise RuntimeError(\n f\"{load_fn.__name__} has no URL default for weights_path \"\n f\"(got {default!r}); update this env file for the installed orb-models\"\n )\n return default\n\n\ndef _local_weights_path(url: str) -> Path:\n \"\"\"Where the checkpoint lives in the shared model cache.\"\"\"\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"orb\" / os.path.basename(url)\n\n\ndef _fetch(url: str, dest: Path) -> None:\n \"\"\"Download ``url`` to ``dest`` atomically (tmp file + rename).\"\"\"\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp = dest.with_name(f\"{dest.name}.tmp.{os.getpid()}\")\n try:\n with urllib.request.urlopen(url) as resp, open(tmp, \"wb\") as out:\n shutil.copyfileobj(resp, out)\n os.replace(tmp, dest)\n finally:\n tmp.unlink(missing_ok=True)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", precision: str = \"float32-high\"):\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.inference.calculator import ORBCalculator\n\n fn_name = CHECKPOINTS[checkpoint].replace(\"-\", \"_\")\n load_fn = getattr(pretrained, fn_name)\n\n # orb-models resolves its default weights URL through `cached_path`, which\n # write-locks its cache dir even on warm hits — EACCES for anyone who can't\n # write the shared install (Garden-AI/rootstock#67). Handed a *local* path\n # instead, cached_path returns it without locking. So the weights are\n # pre-fetched into the shared model cache at `rootstock add` time\n # (maintainer, cache writable) and every later serve loads that file.\n url = _default_weights_url(load_fn)\n weights = _local_weights_path(url)\n if not weights.exists():\n _fetch(url, weights)\n\n orbff, atoms_adapter = load_fn(\n weights_path=str(weights), device=torch.device(device), precision=precision\n )\n return ORBCalculator(orbff, atoms_adapter=atoms_adapter, device=torch.device(device))\n","python_requires":">=3.12","dependencies":{"ase":"3.29.0","cached-path":"1.8.10","orb-models":"0.7.0","rootstock":"1.4.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:0d382198e1fa29caa752ed08f843d75272d961be187cdc97a27b3386cf340847","checkpoints":{"orb-v3-conservative-inf-omat":{"fetched_at":"2026-07-15T00:48:42.107577+00:00","verified_at":"2026-07-31T21:36:59.595150+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-conservative-20-mpa":{"fetched_at":"2026-07-31T21:11:58.857405+00:00","verified_at":"2026-07-31T21:35:45.539672+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-conservative-inf-mpa":{"fetched_at":"2026-07-31T21:11:59.427563+00:00","verified_at":"2026-07-31T21:36:34.966333+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-conservative-20-omat":{"fetched_at":"2026-07-31T21:11:59.829956+00:00","verified_at":"2026-07-31T21:36:10.271438+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-conservative-omol":{"fetched_at":"2026-07-31T21:12:00.265816+00:00","verified_at":"2026-07-31T22:01:55.008333+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-direct-inf-omat":{"fetched_at":"2026-07-31T21:12:41.396063+00:00","verified_at":"2026-07-31T21:38:10.891224+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-direct-20-omat":{"fetched_at":"2026-07-31T21:12:41.462024+00:00","verified_at":"2026-07-31T21:37:41.631359+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-direct-inf-mpa":{"fetched_at":"2026-07-31T21:12:41.898243+00:00","verified_at":"2026-07-31T21:37:56.723053+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-direct-20-mpa":{"fetched_at":"2026-07-31T21:12:41.902478+00:00","verified_at":"2026-07-31T21:37:27.365314+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-direct-omol":{"fetched_at":"2026-07-31T21:13:04.896673+00:00","verified_at":"2026-07-31T22:02:11.263619+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"mattersim":{"built_at":"2026-07-31T21:00:58.617938+00:00","source_hash":"sha256:b94a0a3fe92ec7591278986b3085c4d9d6f3ecc359545752a544a6bf8f99e715","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"mattersim>=1.1.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# \"torchvision\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n# torchvision = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"\nMatterSim environment for Rootstock.\n\nProvides access to Microsoft's MatterSim universal potential. MatterSim-v1\ncovers ~100 elements and supports periodic and non-periodic systems.\n\nModels:\n - \"MatterSim-v1.0.0-5M\": 5M parameter model (faster, default)\n - \"MatterSim-v1.0.0-1M\": 1M parameter model (smallest)\n\"\"\"\n\nCHECKPOINTS = {\n \"mattersim-v1-0-0-5m\": \"MatterSim-v1.0.0-5M\",\n \"mattersim-v1-0-0-1m\": \"MatterSim-v1.0.0-1M\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a MatterSim calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n from mattersim.forcefield import MatterSimCalculator\n\n return MatterSimCalculator(load_path=CHECKPOINTS[checkpoint], device=device)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mattersim":"1.2.3","rootstock":"1.4.0","torch":"2.11.0+cu128","torchvision":"0.26.0+cu128"},"lock_hash":"sha256:0c96011b274c96b203b4ff488cb18c5e7c09c2061833a6ae5116043c86961176","checkpoints":{"mattersim-v1-0-0-5m":{"fetched_at":"2026-07-16T21:19:48.160641+00:00","verified_at":"2026-07-31T21:34:08.665284+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mattersim-v1-0-0-1m":{"fetched_at":"2026-07-31T21:10:48.553228+00:00","verified_at":"2026-07-31T21:33:53.253923+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"allscaip":{"built_at":"2026-08-09T21:29:39.726688+00:00","source_hash":"sha256:774cd05d08292476738237cf379892e77753c5ab3d35da294c240e7f3a041b53","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"AllScAIP env — FAIRChem scalable attention MLIP trained on OMol25.\n\nallscaip-md-conserving-all-omol is an energy-conserving, all-to-all node\nattention model served through fairchem-core's get_predict_unit — the same\nAPI as eSEN. fairchem v2 carries the architecture in-package, so no\nflash-attention or custom CUDA kernels are needed.\n\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"allscaip-md-conserving-all-omol\": \"allscaip-md-conserving-all-omol\",\n \"allscaip-md-direct-all-omol\": \"allscaip-md-direct-all-omol\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"allscaip:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(CHECKPOINTS[checkpoint], device=device)\n return FAIRChemCalculator(predictor)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=device)\n return FAIRChemCalculator(predictor)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.6.0","torch":"2.8.0"},"lock_hash":"sha256:c99d413acdd9b5b2cdd0241ba29966574fe6800af2b63376abacd0ad28fed8b0","checkpoints":{"allscaip-md-direct-all-omol":{"fetched_at":"2026-07-31T21:06:04.006305+00:00","verified_at":"2026-07-31T21:28:18.759757+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"allscaip-md-conserving-all-omol":{"fetched_at":"2026-07-31T21:06:08.856088+00:00","verified_at":"2026-07-31T21:28:04.556356+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"grace":{"built_at":"2026-07-31T21:00:41.219372+00:00","source_hash":"sha256:f46207cbbd702910ea889c09cae3f9933de6c48baf22d75c4039caf02467b635","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"tensorpotential>=0.6.0\",\n# \"ase>=3.22\",\n# ]\n# ///\n\"\"\"GRACE env — hosts GRACE foundation checkpoints via tensorpotential.\n\nGRACE runs on TensorFlow, not torch. TPCalculator has no device argument:\nTF grabs whatever GPU it sees, so device selection happens via\nCUDA_VISIBLE_DEVICES, and both it and TF_USE_LEGACY_KERAS must be set\nbefore the first TF import. The first calculation triggers an XLA compile —\na slow first step is expected.\n\"\"\"\n\nCHECKPOINTS = {\n \"grace-2l-smax-omat-large\": \"GRACE-2L-SMAX-OMAT-large\",\n \"grace-3l-omat-large-ft-am\": \"GRACE-3L-OMAT-large-ft-AM\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n import os\n\n os.environ[\"TF_USE_LEGACY_KERAS\"] = \"1\"\n if device == \"cpu\":\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n elif device.startswith(\"cuda:\"):\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = device.split(\":\", 1)[1]\n\n from tensorpotential.calculator import grace_fm\n\n return grace_fm(CHECKPOINTS[checkpoint])\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.4.0","tensorpotential":"0.6.0"},"lock_hash":"sha256:1ea87e2d3e8620111788e4f3804932ac53f399ab75df08ee6b2b38195a82d50e","checkpoints":{"grace-2l-smax-omat-large":{"fetched_at":"2026-07-31T21:06:12.581904+00:00","verified_at":"2026-07-31T21:29:43.701929+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"grace-3l-omat-large-ft-am":{"fetched_at":"2026-07-31T21:06:14.537247+00:00","verified_at":"2026-07-31T21:30:24.207905+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"mace_polar":{"built_at":"2026-07-31T21:00:27.713198+00:00","source_hash":"sha256:8db4cfd29a59c415e3762b04195e30b9e08d5f044ff6aa21ebcf01bcb78bc1ef","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"ase>=3.22\",\n# \"torch>=2.4.0,<2.10\",\n# # 0.3.16 is the first PyPI release with mace_polar().\n# \"mace-torch>=0.3.16\",\n# # PolarMACE imports graph_longrange at runtime; the distribution is named\n# # graph-longrange and exists only as this git repo (no PyPI release).\n# \"graph-longrange @ git+https://github.com/WillBaldwin0/graph_electrostatics.git\",\n# ]\n# ///\n\"\"\"MACE-POLAR env — electrostatic/polarizable MACE foundation models (OMol25).\n\nKept separate from the stable `mace` env because of the extra git-only\ngraph-longrange dependency.\n\nPOLAR checkpoints expect `charge`, `spin`, and `external_field` in atoms.info.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-polar-1-s\": \"polar-1-s\",\n \"mace-polar-1-m\": \"polar-1-m\",\n \"mace-polar-1-l\": \"polar-1-l\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace-polar:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from mace.calculators import mace_polar\n\n return mace_polar(model=CHECKPOINTS[checkpoint], device=device, default_dtype=\"float32\")\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): mace_polar() accepts a\n # weights file directly, keeping the PolarMACE model-type wiring.\n from mace.calculators import mace_polar\n\n return mace_polar(model=path, device=device, default_dtype=\"float32\")\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"1.4.0","torch":"2.9.1"},"lock_hash":"sha256:bf60a2e3dd4771b8eb5b17a99053c1bf90448615005aab598328165f8da4f5aa","checkpoints":{"mace-polar-1-s":{"fetched_at":"2026-07-31T21:11:21.262548+00:00","verified_at":"2026-07-31T21:33:20.108166+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-polar-1-m":{"fetched_at":"2026-07-31T21:11:21.629524+00:00","verified_at":"2026-07-31T21:33:08.407911+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-polar-1-l":{"fetched_at":"2026-07-31T21:11:22.256531+00:00","verified_at":"2026-07-31T21:32:56.539784+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"pet":{"built_at":"2026-07-31T23:02:00.228670+00:00","source_hash":"sha256:91f3da0a95e220bbfa5e8870c3236ea61322ef70f7912a2ad1d323b7c3a75544","source":"# /// script\n# requires-python = \">=3.11,<3.15\"\n# dependencies = [\n# \"upet>=0.2.6\",\n# \"ase>=3.22\",\n# # upet pulls nvalchemi-toolkit-ops unpinned; 0.4+ needs torch>=2.8 at\n# # runtime but only declares the constraint on its extras, so the\n# # resolver won't catch it (same trap as the tensornet env).\n# \"torch>=2.8,<2.14\",\n# ]\n# ///\n\"\"\"PET env — hosts lab-cosmo's UPET foundation checkpoints (PET-MAD successor).\n\nThe upstream string encodes model@version; versions are pinned rather than\n\"latest\" so rebuilds serve the same weights. pet-omatpes-l is trained at the\nr2SCAN level of theory — its energies are not comparable to the PBE-level\npet-oam models.\n\"\"\"\n\nCHECKPOINTS = {\n \"pet-oam-xl\": \"pet-oam-xl@1.0.0\",\n \"pet-omatpes-l\": \"pet-omatpes-l@0.1.0\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from huggingface_hub import hf_hub_download\n from upet.calculator import UPETCalculator\n\n # Passing model=/version= makes UPETCalculator resolve the name by listing\n # the hub repo — an uncached API call that fails on workers, which run\n # with HF_HUB_OFFLINE=1 (and on any node without internet). Fetch the\n # pinned file ourselves — a cache hit needs no network even offline — and\n # hand it over as checkpoint_path, which skips the resolve entirely.\n model, version = CHECKPOINTS[checkpoint].split(\"@\", 1)\n path = hf_hub_download(\n repo_id=\"lab-cosmo/upet\",\n filename=f\"{model}-v{version}.ckpt\",\n subfolder=\"models\",\n )\n return UPETCalculator(checkpoint_path=path, device=device)\n","python_requires":">=3.11,<3.15","dependencies":{"ase":"3.29.0","rootstock":"1.4.3","torch":"2.13.0","upet":"0.2.6"},"lock_hash":"sha256:9809737b1bd5fea93de01e7f7340e8140d88dbda6fa66c09ef1d7144435a6f5c","checkpoints":{"pet-omatpes-l":{"fetched_at":"2026-07-31T21:14:06.700916+00:00","verified_at":"2026-07-31T23:03:45.522243+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"pet-oam-xl":{"fetched_at":"2026-07-31T21:14:40.106175+00:00","verified_at":null,"verified_device":null,"last_error":"verify: RuntimeError: Worker calculation failed:\nTraceback (most recent call last):\n File \"/scratch/gpfs/ROSENGROUP/common/rootstock/envs/pet/lib/python3.11/site-packages/rootstock/worker.py\", line 275, in run\n energy, forces, virial = self._calculate(positions, cell)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/scratch/gpfs/ROSENGROUP/common/rootstock/envs/pet/lib/python3.11/site-packages/rootstock/worker.py\", line 182, in _calculate\n energy = atoms.get_potential_energy()\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/scratch/gpfs/ROSENGROUP/common/rootstock/envs/pet/lib/python3.11/site-packages/ase/atoms.py\", line 1971, in get_potential_energy\n energy = self._calc.get_potential_energy(self)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/scratch/gpfs/ROSENGROUP/common/rootstock/envs/pet/lib/python3.11/site-packages/ase/calculators/abc.py\", line 25, in get_potential_energy\n return self.get_property(name, atoms)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/scratch/gpfs/ROSENGROUP/common/rootstock/envs/pet/lib/python3.11/site-packages/ase/calculators/calculator.py\", line 517, in get_property\n self.calculate(atoms, [name], system_changes)\n File \"/scratch/gpfs/ROSENGROUP/common/rootstock/envs/pet/lib/python3.11/site-packages/upet/calculator.py\", line 221, in calculate\n self.calculator.calculate(atoms, properties, system_changes)\n File \"/scratch/gpfs/ROSENGROUP/common/rootstock/envs/pet/lib/python3.11/site-packages/metatomic_ase/_calculator.py\", line 698, in calculate\n input_system = self._nl_calculators.compute(systems=[system])[0]\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/scratch/gpfs/ROSENGROUP/common/rootstock/envs/pet/lib/python3.11/site-packages/metatomic_ase/_neighbors.py\", line 76, in compute\n systems = _compute_requested_neighbors_nvalchemi(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/scratch/gpfs/ROSENGROUP/common/rootstock/envs/pet/lib/python3.11/site-packages/metatomic_ase/_neighbors.py\", line 137, in _compute_requested_neighbors_nvalchemi\n edge_index, _, S = nvalchemi_neighbor_list(\n ^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/scratch/gpfs/ROSENGROUP/common/rootstock/envs/pet/lib/python3.11/site-packages/nvalchemiops/torch/neighbors/__init__.py\", line 431, in neighbor_list\n return naive_neighbor_list(\n ^^^^^^^^^^^^^^^^^^^^\n File \"/scratch/gpfs/ROSENGROUP/common/rootstock/envs/pet/lib/python3.11/site-packages/nvalchemiops/torch/neighbors/naive.py\", line 1539, in naive_neighbor_list\n neighbor_matrix = torch.full(\n ^^^^^^^^^^^\nTypeError: full(): argument 'size' failed to unpack the object at pos 2 with error \"type must be tuple of ints,but got float\"\n","weight_files":null,"weights_recorded_at":null}}},"tace":{"built_at":"2026-07-31T21:01:19.068632+00:00","source_hash":"sha256:1942bf8d57e869cefac328435f093792c94591ce3a24da6e38050e420385189c","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"tace\",\n# \"ase>=3.22\",\n# \"torch>=2.4,<2.14\",\n# ]\n#\n# # The TECE architecture and its foundation registry ship in tace 0.2.0,\n# # which is not on PyPI (latest release there is 0.1.0). This is the commit\n# # the Matbench Discovery submission pinned.\n# [tool.uv.sources.tace]\n# git = \"https://github.com/xvzemin/tace\"\n# rev = \"81f65a4c188bd09cec8d1419388f7afdcc1b6fd0\"\n# ///\n\"\"\"TACE env — hosts TECE/TACE foundation checkpoints (Xu, Xie & Hu).\n\nopenequivariance CUDA-kernel acceleration is optional upstream and not\ninstalled here — TACE runs on the e3nn path.\n\"\"\"\n\nCHECKPOINTS = {\n \"tece-oam-rra-1.0\": \"TECE-OAM-RRA-1.0\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from tace.foundations import tace_foundations\n from tace.interface.ase import TACEAseCalc\n\n model_path = tace_foundations[CHECKPOINTS[checkpoint]]\n return TACEAseCalc(model_path, device=device, dtype=\"float32\")\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.4.0","tace":"0.2.0","torch":"2.13.0"},"lock_hash":"sha256:8d8042b567b0cffa695062cf4306219ba0990530d21b7d1be7d795b315d9c1b9","checkpoints":{"tece-oam-rra-1.0":{"fetched_at":"2026-07-31T21:14:44.622276+00:00","verified_at":"2026-07-31T21:40:42.185062+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}}}},{"schema_version":5,"cluster":"delta","root":"/work/hdd/data/rootstock","maintainer":{"name":"Owen Price-Skelly","email":"OwenPriceSkelly@uchicago.edu"},"rootstock_version":"1.6.4","python_version":"unknown","last_updated":"2026-10-24T03:29:46.925055+00:00","environments":{"ani":{"built_at":"2026-07-17T22:28:36.697190+00:00","source_hash":"sha256:16ebe0eda0074fa275e0ce2c2d68c84dc10d2069ef0290d19802bd46bcf96d0e","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"torchani>=2.2\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n# ///\n\"\"\"\nANI-2x environment for Rootstock.\n\nANI-2x is a neural network potential for organic molecules containing\nH, C, N, O, F, S, Cl. It is not a universal potential — do not use it\nfor inorganic or periodic systems.\n\nModels:\n - \"ANI2x\": ANI-2x ensemble (default, 8 networks)\n - \"ANI1ccx\": ANI-1ccx, trained on CCSD(T)/CBS data (H, C, N, O only)\n - \"ANI1x\": ANI-1x (H, C, N, O only)\n\"\"\"\n\nCHECKPOINTS = {\n \"ani-2x\": \"ANI2x\",\n \"ani-1ccx\": \"ANI1ccx\",\n \"ani-1x\": \"ANI1x\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an ANI calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n import torchani\n\n model_map = {\n \"ANI2x\": torchani.models.ANI2x,\n \"ANI1ccx\": torchani.models.ANI1ccx,\n \"ANI1x\": torchani.models.ANI1x,\n }\n model = CHECKPOINTS[checkpoint]\n\n return model_map[model](periodic_table_index=True).to(device).ase()\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.0.0","torch":"2.8.0","torchani":"2.8.2"},"lock_hash":"sha256:475e7562a113e093673026a125aeb2c935d91fc941beb50d8bcc183566cb8386","checkpoints":{"ani-2x":{"fetched_at":"2026-07-18T01:07:01.392973+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 34998 files, 7.6 GB in 1753.4s; weights: 55 MB (manifest)\n/work/hdd/data/rootstock/envs/ani/lib/python3.11/site-packages/torchani/csrc/__init__.py:56: UserWarning: The extensions: ['cuaev', 'mnp', 'cell_list'] are not installed and will not be available. To install the extensions first install the CUDA Toolkit, and afterwards run `ani build-extensions` To suppress warn set the env var TORCHANI_NO_WARN_EXTENSIONS=1 For example, if using bash, you may add `export TORCHANI_NO_WARN_EXTENSIONS=1` to your .bashrc\n warnings.warn(","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.local/share/Torchani/StateDicts/.cache/huggingface/download/ani2x_state_dict.pt.metadata","size":125},{"path":"home/.local/share/Torchani/StateDicts/ani2x_state_dict.pt","size":54982957}],"weights_recorded_at":"2026-09-26T05:40:47.735716+00:00"},"ani-1ccx":{"fetched_at":"2026-07-18T01:07:52.739750+00:00","verified_at":"2026-10-24T03:28:45.589177+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.local/share/Torchani/StateDicts/.cache/huggingface/download/ani1ccx_state_dict.pt.metadata","size":125},{"path":"home/.local/share/Torchani/StateDicts/ani1ccx_state_dict.pt","size":10546029}],"weights_recorded_at":"2026-10-24T03:28:45.589196+00:00"},"ani-1x":{"fetched_at":"2026-07-18T01:08:01.222324+00:00","verified_at":"2026-10-24T03:28:45.589206+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.local/share/Torchani/StateDicts/.cache/huggingface/download/ani1x_state_dict.pt.metadata","size":124},{"path":"home/.local/share/Torchani/StateDicts/ani1x_state_dict.pt","size":10546029}],"weights_recorded_at":"2026-10-24T03:28:45.589214+00:00"}}},"chgnet":{"built_at":"2026-07-30T23:26:54.008571+00:00","source_hash":"sha256:b6ea6ac08a17cb3dc17d5455571dc5c21d513cd802916015de7f272e83eefe37","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"chgnet>=0.3.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"CHGNet env — hosts pretrained charge-informed universal potentials.\"\"\"\n\nCHECKPOINTS = {\n \"chgnet-default\": \"chgnet-default\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"chgnet:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a CHGNet calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\")\n\n Returns:\n ASE-compatible calculator\n \"\"\"\n from chgnet.model import CHGNet, CHGNetCalculator\n\n model_name = CHECKPOINTS[checkpoint]\n model = CHGNet.load() if model_name == \"chgnet-default\" else CHGNet.load(model_name)\n return CHGNetCalculator(model=model, use_device=device)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # CHGNet.from_file, not the named-model CHGNet.load() setup() uses.\n from chgnet.model import CHGNet, CHGNetCalculator\n\n model = CHGNet.from_file(path)\n return CHGNetCalculator(model=model, use_device=device)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","chgnet":"0.4.2","rootstock":"1.3.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:5e76831d9d411c768ccbf9542377f9eabec12498c024336b085270b24c65082f","checkpoints":{"chgnet-default":{"fetched_at":"2026-07-18T01:14:26.156946+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null}}},"dimenet":{"built_at":"2026-07-20T16:25:08.850502+00:00","source_hash":"sha256:84b1aa2101b0cdfa5f74270cfc5d7a53854a84d4b8f68945e11692563c24b322","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\nDimeNet++ environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 DimeNet++ checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"DimeNet++-S2EF-OC20-All\": default\n - \"DimeNet++-S2EF-OC20-20M\"\n - \"DimeNet++-S2EF-OC20-2M\"\n - \"DimeNet++-S2EF-OC20-200k\"\n\"\"\"\n\nCHECKPOINTS = {\n \"dimenet-plus-plus-s2ef-oc20-all\": \"DimeNet++-S2EF-OC20-All\",\n \"dimenet-plus-plus-s2ef-oc20-20m\": \"DimeNet++-S2EF-OC20-20M\",\n \"dimenet-plus-plus-s2ef-oc20-2m\": \"DimeNet++-S2EF-OC20-2M\",\n \"dimenet-plus-plus-s2ef-oc20-200k\": \"DimeNet++-S2EF-OC20-200k\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a DimeNet++ OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.0.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"dimenet-plus-plus-s2ef-oc20-all":{"fetched_at":"2026-07-20T20:04:28.057409+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/dimenetpp_all.pt","size":21841849}],"weights_recorded_at":"2026-09-18T20:25:36.277728+00:00"},"dimenet-plus-plus-s2ef-oc20-20m":{"fetched_at":"2026-07-20T20:08:11.040759+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/dimenetpp_20M.pt","size":21841849}],"weights_recorded_at":"2026-09-26T05:40:47.735759+00:00"},"dimenet-plus-plus-s2ef-oc20-2m":{"fetched_at":"2026-07-20T20:08:53.776120+00:00","verified_at":"2026-10-24T03:28:45.589226+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/dimenetpp_2M.pt","size":21841593}],"weights_recorded_at":"2026-10-24T03:28:45.589230+00:00"},"dimenet-plus-plus-s2ef-oc20-200k":{"fetched_at":"2026-07-20T20:09:37.960754+00:00","verified_at":"2026-10-24T03:28:45.589235+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/dimenetpp_200k.pt","size":21841593}],"weights_recorded_at":"2026-10-24T03:28:45.589239+00:00"}}},"equiformer":{"built_at":"2026-07-30T23:26:40.858583+00:00","source_hash":"sha256:ac01aec3be646cbd2c63a2e9f6a0cc8308aa03f5d8528150c5f664b2aa06426f","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\nEquiformerV2 environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 EquiformerV2 checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"EquiformerV2-153M-S2EF-OC20-All+MD\": default\n - \"EquiformerV2-31M-S2EF-OC20-All+MD\"\n - \"EquiformerV2-83M-S2EF-OC20-2M\"\n\"\"\"\n\nCHECKPOINTS = {\n \"equiformer-v2-153m-s2ef-oc20-all-md\": \"EquiformerV2-153M-S2EF-OC20-All+MD\",\n \"equiformer-v2-31m-s2ef-oc20-all-md\": \"EquiformerV2-31M-S2EF-OC20-All+MD\",\n \"equiformer-v2-83m-s2ef-oc20-2m\": \"EquiformerV2-83M-S2EF-OC20-2M\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"equiformer:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an EquiformerV2 OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): OCPCalculator loads a\n # checkpoint file natively — this is setup() minus the registry download.\n from fairchem.core import OCPCalculator\n\n return OCPCalculator(checkpoint_path=path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.3.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"equiformer-v2-153m-s2ef-oc20-all-md":{"fetched_at":"2026-07-20T20:19:46.718895+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/eq2_153M_ec4_allmd.pt","size":619855697}],"weights_recorded_at":"2026-09-18T20:25:36.277757+00:00"},"equiformer-v2-31m-s2ef-oc20-all-md":{"fetched_at":"2026-07-20T20:21:06.384332+00:00","verified_at":"2026-10-24T03:28:45.589244+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_31M_ec4_allmd.pt","size":125685325}],"weights_recorded_at":"2026-10-24T03:28:45.589248+00:00"},"equiformer-v2-83m-s2ef-oc20-2m":{"fetched_at":"2026-07-20T20:21:33.323971+00:00","verified_at":"2026-10-24T03:28:45.589252+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_83M_2M.pt","size":336843381}],"weights_recorded_at":"2026-10-24T03:28:45.589256+00:00"},"equiformer:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null}}},"escn":{"built_at":"2026-07-30T23:26:43.867330+00:00","source_hash":"sha256:e1301f3e20ac38e355e0b36e2ab9fb0f0a78f58059b5e2c57ab280aa9dad7ff5","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\neSCN environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 eSCN checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"eSCN-L6-M2-Lay12-S2EF-OC20-All+MD\": default\n - \"eSCN-L6-M3-Lay20-S2EF-OC20-All+MD\"\n - \"eSCN-L6-M2-Lay12-S2EF-OC20-2M\"\n - \"eSCN-L4-M2-Lay12-S2EF-OC20-2M\"\n\"\"\"\n\nCHECKPOINTS = {\n \"escn-l6-m2-lay12-s2ef-oc20-all-md\": \"eSCN-L6-M2-Lay12-S2EF-OC20-All+MD\",\n \"escn-l6-m3-lay20-s2ef-oc20-all-md\": \"eSCN-L6-M3-Lay20-S2EF-OC20-All+MD\",\n \"escn-l6-m2-lay12-s2ef-oc20-2m\": \"eSCN-L6-M2-Lay12-S2EF-OC20-2M\",\n \"escn-l4-m2-lay12-s2ef-oc20-2m\": \"eSCN-L4-M2-Lay12-S2EF-OC20-2M\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"escn:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an eSCN OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): OCPCalculator loads a\n # checkpoint file natively — this is setup() minus the registry download.\n from fairchem.core import OCPCalculator\n\n return OCPCalculator(checkpoint_path=path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.3.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"escn-l6-m2-lay12-s2ef-oc20-all-md":{"fetched_at":"2026-07-20T20:31:34.523398+00:00","verified_at":"2026-10-24T03:28:45.589260+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l6_m2_lay12_all_md_s2ef.pt","size":207735109}],"weights_recorded_at":"2026-10-24T03:28:45.589264+00:00"},"escn-l6-m3-lay20-s2ef-oc20-all-md":{"fetched_at":"2026-07-20T20:34:01.649114+00:00","verified_at":"2026-10-24T03:28:45.589267+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l6_m3_lay20_all_md_s2ef.pt","size":801642787}],"weights_recorded_at":"2026-10-24T03:28:45.589270+00:00"},"escn-l6-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-07-20T20:34:48.217991+00:00","verified_at":"2026-10-24T03:28:45.589274+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l6_m2_lay12_2M_s2ef.pt","size":207735045}],"weights_recorded_at":"2026-10-24T03:28:45.589278+00:00"},"escn-l4-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-07-20T20:35:22.280265+00:00","verified_at":"2026-10-24T03:28:45.589292+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l4_m2_lay12_2M_s2ef.pt","size":144761989}],"weights_recorded_at":"2026-10-24T03:28:45.589295+00:00"},"escn:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null}}},"esen":{"built_at":"2026-07-30T23:28:11.546183+00:00","source_hash":"sha256:0abd4bac1f4ee2550adb62c935eafc15d76bf294c67f1a03c6c40c7f3809cfe8","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"eSEN env — hosts FAIRChem eSEN single-task checkpoints.\n\nfairchem-core v2 dropped the torch-geometric / pyg-find-links install dance\n(simplified here 2026-07-30 to match uma.py/allscaip.py — the find-links were\nalso blocking universal lockfile resolution).\n\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"esen-md-direct-all-omol\": \"esen-md-direct-all-omol\",\n \"esen-sm-conserving-all-omol\": \"esen-sm-conserving-all-omol\",\n \"esen-sm-direct-all-omol\": \"esen-sm-direct-all-omol\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"esen:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(CHECKPOINTS[checkpoint], device=device)\n return FAIRChemCalculator(predictor)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=device)\n return FAIRChemCalculator(predictor)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.3.0","torch":"2.8.0"},"lock_hash":"sha256:82d65965c5b3d3bf22c7a444a3ed1f22382c3c2b20c35e80f30f8d5df775e272","checkpoints":{"esen-md-direct-all-omol":{"fetched_at":"2026-07-18T01:53:43.151954+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/e0cd2ac37d2992b385a39778469c443ccbc679c7c4fb3b40b38aa2ee23c29b1d","size":405789591},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-09-18T20:25:36.277804+00:00"},"esen-sm-conserving-all-omol":{"fetched_at":"2026-07-18T01:54:34.092323+00:00","verified_at":"2026-10-24T03:28:45.589301+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/01f63da2d071e39fc46a5f22f8369d0fc9de317ab2ef361a76603f5661238025","size":50958067},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-24T03:28:45.589306+00:00"},"esen-sm-direct-all-omol":{"fetched_at":"2026-07-18T01:55:21.266324+00:00","verified_at":"2026-10-24T03:28:45.589310+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/a1b7109bcca6e1288c5f17c68776d3c6a1ab0766f6d7fdd45f901605355d2831","size":50960291},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-24T03:28:45.589315+00:00"},"esen:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null}}},"m3gnet":{"built_at":"2026-07-17T23:36:02.432534+00:00","source_hash":"sha256:3d22f362208d5a4297208395d7f188c943aee8866ed9f111a4b7c8d3801bbc14","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"chgnet>=0.4.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"\nM3GNet environment for Rootstock — redirected to CHGNet.\n\nM3GNet-MP-2021.2.8-PES (the universal inorganic PES) is no longer accessible\nvia any modern Python package:\n - matgl 2.x (PyG backend) only has TensorNet on HuggingFace (materialyze)\n - matgl 1.x (DGL backend) pointed to a GitHub URL that was removed\n - The original m3gnet package is archived TensorFlow code\n - materialyze HuggingFace has only M3GNet-Eform (formation energy, not PES)\n\nFor universal inorganic PES, use:\n - tensornet.py: TensorNet-MatPES (same authors, newer, better)\n - chgnet.py: CHGNet (charge-informed, strong on magnetic materials)\n - orb.py: Orb v3 (universal, supports periodic systems)\n\nThis file loads CHGNet as the practical substitute for M3GNet-PES.\n\"\"\"\n\nCHECKPOINTS = {\n \"m3gnet-mp-2021-2-8-pes\": \"chgnet-default\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load CHGNet as a substitute for M3GNet-PES.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible CHGNetCalculator.\n \"\"\"\n from chgnet.model import CHGNetCalculator\n\n return CHGNetCalculator(use_device=device)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","chgnet":"0.4.2","rootstock":"1.0.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:bae5a0b584d53536db75471c513af80d87d410b12c8928b00ef388a053a92600","checkpoints":{"m3gnet-mp-2021-2-8-pes":{"fetched_at":"2026-07-18T02:00:48.833246+00:00","verified_at":"2026-10-24T03:28:45.589319+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"mace":{"built_at":"2026-08-03T17:57:32.486289+00:00","source_hash":"sha256:b6cf58f573787e39bbc099deb93aa77918c4679f7bdd9a974df2966a58bafc2e","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # 0.3.15+ needed for the mh-1 registry entry (matpes needs 0.3.13,\n# # omol needs 0.3.14, mpa-0 needs 0.3.10).\n# \"mace-torch>=0.3.15\",\n# \"ase>=3.22\",\n# # 2.4.1 is explicitly unsupported by mace-torch.\n# \"torch>=2.4.0,!=2.4.1,<2.10\",\n# ]\n# ///\n\"\"\"MACE env — hosts MACE-MP-0, MACE-OFF23, MPA-0, MATPES, MH-1, and OMOL checkpoints.\n\nAll ship in the same `mace-torch` package, so they share an environment.\nUpstream-string routing in CHECKPOINTS: an `off:` prefix routes to mace_off()\nand an `omol:` prefix to mace_omol() (float64, molecules only); an `mh:`\nprefix marks a multi-head model (float64, per the MACE-MH-1 model card).\n\nMulti-head checkpoints select a head via the `head` kwarg on setup()\n(setup_kwargs={\"head\": ...} / --kwarg head=...), named by upstream's training\ncorpus — see MH1_HEADS; omat_pbe is the default.\n\nThe OMOL checkpoint expects `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-mp-0-small\": \"small\",\n \"mace-mp-0-medium\": \"medium\",\n \"mace-mp-0-large\": \"large\",\n \"mace-off23-small\": \"off:small\",\n \"mace-off23-medium\": \"off:medium\",\n \"mace-off23-large\": \"off:large\",\n # Only a medium MPA-0 has been released, but upstream names the weights\n # file mace-mpa-0-medium.model — keep the size explicit like mace-mp-0.\n \"mace-mpa-0-medium\": \"medium-mpa-0\",\n \"mace-matpes-r2scan-0\": \"mace-matpes-r2scan-0\",\n # One entry per weights file: MH-1's heads are selected by setup(head=...).\n \"mace-mh-1\": \"mh:mh-1\",\n # Only the extra-large OMOL model has been released.\n \"mace-omol-0-extra-large\": \"omol:extra_large\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace:custom\": None,\n}\n\n# MH-1's heads, named by training corpus. The released weights file is the\n# authority (mace_select_head --list_heads; the model card also lists a\n# rgd1_b3lyp head, but that shipped only in mh-0 — ACEsuit/mace#1462).\n# Validated here because upstream only warns on an unknown head and silently\n# falls back to the last one.\nMH1_HEADS = (\n \"omat_pbe\",\n \"omol\",\n \"spice_wB97M\",\n \"oc20_usemppbe\",\n \"mp_pbe_refit_add\",\n \"matpes_r2scan\",\n)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", head: str | None = None):\n arg = CHECKPOINTS[checkpoint]\n if arg.startswith(\"mh:\"):\n head = head or \"omat_pbe\"\n if head not in MH1_HEADS:\n raise ValueError(f\"unknown head {head!r}; expected one of {', '.join(MH1_HEADS)}\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg[3:], device=device, default_dtype=\"float64\", head=head)\n if head is not None:\n raise ValueError(f\"'head' selects a head of a multi-head model; {checkpoint} has one head\")\n if arg.startswith(\"off:\"):\n from mace.calculators import mace_off\n\n return mace_off(model=arg[4:], device=device, default_dtype=\"float32\")\n if arg.startswith(\"omol:\"):\n from mace.calculators import mace_omol\n\n return mace_omol(model=arg[5:], device=device, default_dtype=\"float64\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg, device=device, default_dtype=\"float32\")\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", head: str | None = None):\n # Custom checkpoints (`:custom` ids with user weights): fine-tunes load through\n # MACECalculator directly — the mp/off dispatch in setup() only exists\n # to pick which pretrained file to download. `head` is for fine-tunes that\n # keep multiple heads; single-head weights load without it.\n from mace.calculators import MACECalculator\n\n return MACECalculator(model_paths=path, device=device, default_dtype=\"float32\", head=head)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"1.4.3","torch":"2.9.1"},"lock_hash":"sha256:8e4fadbc59d2f9c8a12c8e0b32ff42e43b08c3415206c5306ecdb7f08449f1e9","checkpoints":{"mace-mp-0-small":{"fetched_at":"2026-07-18T02:09:46.072270+00:00","verified_at":"2026-10-24T03:28:45.589323+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231210mace128L0_energy_epoch249model","size":32581838}],"weights_recorded_at":"2026-10-24T03:28:45.589327+00:00"},"mace-mp-0-medium":{"fetched_at":"2026-07-18T02:10:38.115353+00:00","verified_at":"2026-10-24T03:28:45.589331+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231203mace128L1_epoch199model","size":44422970}],"weights_recorded_at":"2026-10-24T03:28:45.589334+00:00"},"mace-mp-0-large":{"fetched_at":"2026-07-18T02:10:53.149921+00:00","verified_at":"2026-10-24T03:28:45.589338+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE_MPtrj_20229model","size":133803220}],"weights_recorded_at":"2026-10-24T03:28:45.589341+00:00"},"mace-off23-small":{"fetched_at":"2026-07-18T02:11:07.975353+00:00","verified_at":"2026-10-24T03:28:45.589345+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_small.model","size":7347350}],"weights_recorded_at":"2026-10-24T03:28:45.589348+00:00"},"mace-off23-medium":{"fetched_at":"2026-07-18T02:11:22.830181+00:00","verified_at":"2026-10-24T03:28:45.589352+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_medium.model","size":18350596}],"weights_recorded_at":"2026-10-24T03:28:45.589356+00:00"},"mace-off23-large":{"fetched_at":"2026-07-18T02:11:37.728101+00:00","verified_at":"2026-10-24T03:28:45.589362+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_large.model","size":55492786}],"weights_recorded_at":"2026-10-24T03:28:45.589365+00:00"},"mace-mh-0":{"fetched_at":"2026-07-31T17:53:34.196632+00:00","verified_at":"2026-07-31T18:40:24.897683+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-matpes-pbe-0-medium":{"fetched_at":"2026-07-30T23:57:06.064509+00:00","verified_at":null,"verified_device":null,"last_error":"verify: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null},"mace-matpes-r2scan-0-medium":{"fetched_at":"2026-07-30T23:57:06.541659+00:00","verified_at":"2026-07-31T18:40:09.943335+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-mh-1":{"fetched_at":"2026-07-31T17:53:34.618127+00:00","verified_at":"2026-10-24T03:28:45.589383+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macemh1model","size":59208139}],"weights_recorded_at":"2026-10-24T03:28:45.589386+00:00"},"mace-mpa-0-medium":{"fetched_at":"2026-07-30T23:57:26.498782+00:00","verified_at":"2026-10-24T03:28:45.589368+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macempa0mediummodel","size":79462305}],"weights_recorded_at":"2026-10-24T03:28:45.589372+00:00"},"mace-mp-0b3-medium":{"fetched_at":"2026-07-30T23:57:26.652156+00:00","verified_at":"2026-07-31T18:41:39.407408+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-omat-0-medium":{"fetched_at":"2026-07-30T23:57:42.588121+00:00","verified_at":"2026-07-31T18:42:53.216953+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-omat-0-small":{"fetched_at":"2026-07-30T23:57:42.698847+00:00","verified_at":"2026-07-31T18:43:08.068964+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-matpes-r2scan-0":{"fetched_at":"2026-08-03T18:02:58.175952+00:00","verified_at":"2026-10-24T03:28:45.589375+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEmatpesr2scanomatftmodel","size":79470738}],"weights_recorded_at":"2026-10-24T03:28:45.589379+00:00"},"mace-omol-0-extra-large":{"fetched_at":"2026-08-03T18:03:27.581411+00:00","verified_at":"2026-10-24T03:28:45.589390+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-omol-0-extra-large-1024.model","size":422242640}],"weights_recorded_at":"2026-10-24T03:28:45.589393+00:00"},"mace:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null}}},"orb":{"built_at":"2026-07-31T17:52:13.028131+00:00","source_hash":"sha256:4ae1362a989170b29d5082a23f9586c50d4f04b700bd29b8c65609efcbe52bef","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # >=0.5,<0.6: 0.5.5 is what the verified Delta env resolved — the v2\n# # loaders keep their single-return API through 0.5.x. 0.4.x is broken\n# # for us: it imports pynanoflann, which is git-only and undeclared, so\n# # a fresh build dies at import (Delta, 2026-07-31). 0.6 raises the\n# # Python floor to 3.12 — that line lives in orb_v3.py.\n# \"orb-models>=0.5,<0.6\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# # Not imported here — constrains orb-models' transitive dep. setup()'s\n# # no-lock serve path relies on cached_path returning local files without\n# # locking or writing, verified against exactly this version (#67).\n# \"cached_path==1.8.10\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"Orb v2 env — kept only for the built-in-D3 dispersion variant.\n\nOrb v3 (orb_v3.py) is the primary orb env; it supersedes the v2 checkpoints\nexcept orb-d3-v2, which has no v3 equivalent (v3 ships no dispersion-corrected\nmodel). Catalog trimmed to that one id 2026-07-30. The two lines can't share\nan env: the v3 loaders need orb-models>=0.6, which raises the Python floor\nto 3.12 and torch to 2.8 (the v2 loaders here are fine through 0.5.x).\n\"\"\"\n\nimport os\nimport shutil\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"orb-d3-v2\": \"orb-d3-v2\",\n # Your own fine-tuned v2-architecture weights: pair with weights=\n # (loaded via setup_from_path). v3 fine-tunes go to orb-v3:custom.\n \"orb-v2:custom\": None,\n}\n\n\ndef _default_weights_url(load_fn) -> str:\n \"\"\"The upstream URL baked into the loader's ``weights_path`` default.\"\"\"\n import inspect\n\n default = inspect.signature(load_fn).parameters[\"weights_path\"].default\n if not isinstance(default, str) or not default.startswith((\"http://\", \"https://\")):\n raise RuntimeError(\n f\"{load_fn.__name__} has no URL default for weights_path \"\n f\"(got {default!r}); update this env file for the installed orb-models\"\n )\n return default\n\n\ndef _local_weights_path(url: str) -> Path:\n \"\"\"Where the checkpoint lives in the shared model cache.\"\"\"\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"orb\" / os.path.basename(url)\n\n\ndef _fetch(url: str, dest: Path) -> None:\n \"\"\"Download ``url`` to ``dest`` atomically (tmp file + rename).\"\"\"\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp = dest.with_name(f\"{dest.name}.tmp.{os.getpid()}\")\n try:\n with urllib.request.urlopen(url) as resp, open(tmp, \"wb\") as out:\n shutil.copyfileobj(resp, out)\n os.replace(tmp, dest)\n finally:\n tmp.unlink(missing_ok=True)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.calculator import ORBCalculator\n\n # orb-models exposes one function per checkpoint, e.g. pretrained.orb_v2().\n fn_name = CHECKPOINTS[checkpoint].replace(\"-\", \"_\")\n load_fn = getattr(pretrained, fn_name)\n\n # orb-models resolves its default weights URL through `cached_path`, which\n # write-locks its cache dir even on warm hits — EACCES for anyone who can't\n # write the shared install (Garden-AI/rootstock#67). Handed a *local* path\n # instead, cached_path returns it without locking. So the weights are\n # pre-fetched into the shared model cache at `rootstock add` time\n # (maintainer, cache writable) and every later serve loads that file.\n url = _default_weights_url(load_fn)\n weights = _local_weights_path(url)\n if not weights.exists():\n _fetch(url, weights)\n\n orbff = load_fn(weights_path=str(weights), device=torch.device(device))\n return ORBCalculator(orbff, device=torch.device(device))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", arch: str = \"orb-v2\"):\n # Custom checkpoints (`:custom` ids with user weights). A weights file doesn't say\n # which orb architecture produced it, so `arch` names the pretrained\n # loader to instantiate — pass the right one at call time\n # (setup_kwargs={\"arch\": ...} / --kwarg arch=...). Handing the loader a local path also means no\n # network and no cached_path locking (see setup()).\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.calculator import ORBCalculator\n\n fn_name = arch.replace(\"-\", \"_\")\n try:\n load_fn = getattr(pretrained, fn_name)\n except AttributeError:\n raise ValueError(\n f\"unknown orb architecture {arch!r}; expected a loader name from \"\n f\"orb_models.forcefield.pretrained, e.g. orb-v2, orb-d3-v2\"\n ) from None\n\n orbff = load_fn(weights_path=path, device=torch.device(device))\n return ORBCalculator(orbff, device=torch.device(device))\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","cached-path":"1.8.10","orb-models":"0.5.5","rootstock":"1.3.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:9be9e6ab5d689b7be4211640bc378cf4821a9217f6924afa99fbab81c7c818cb","checkpoints":{"orb-v2":{"fetched_at":"2026-07-18T02:18:34.284102+00:00","verified_at":"2026-07-18T03:53:42.849085+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-d3-v2":{"fetched_at":"2026-07-18T02:20:54.558117+00:00","verified_at":"2026-10-24T03:28:45.589397+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-d3-v2-20241011.ckpt","size":100967743}],"weights_recorded_at":"2026-10-24T03:28:45.589400+00:00"},"orb-mptraj-only-v2":{"fetched_at":"2026-07-18T02:21:23.960431+00:00","verified_at":"2026-07-18T03:53:42.849089+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"orb_v3":{"built_at":"2026-07-31T17:52:07.485384+00:00","source_hash":"sha256:1f63f36fe8c791313c04bbb5bf534a1a7168c327c1a35f98ecd3d80e3ff74122","source":"# /// script\n# # <3.13: orb-models pins dm-tree==0.1.8, which has no cp313 wheel and whose\n# # sdist doesn't compile against modern GCC (vendored abseil).\n# requires-python = \">=3.12,<3.13\"\n# dependencies = [\n# \"orb-models>=0.6.2\",\n# \"ase>=3.25\",\n# \"torch>=2.8\",\n# # Not imported here — constrains orb-models' transitive dep. setup()'s\n# # no-lock serve path relies on cached_path returning local files without\n# # locking or writing, verified against exactly this version (#67).\n# \"cached_path==1.8.10\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"Orb v3 env — the primary orb env, Orbital Materials' Orb v3 potentials.\n\nSeparate from orb.py because the v3 loaders need orb-models>=0.6.2, which\nbumped the Python floor to 3.12 and torch to 2.8; the v3 loader API also\ndiffers (returns a tuple, requires `atoms_adapter` on ORBCalculator, imports\nthe calculator from forcefield.inference). orb.py (v2) survives only for\norb-d3-v2, the dispersion-corrected variant with no v3 equivalent.\n\"\"\"\n\nimport os\nimport shutil\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"orb-v3-conservative-inf-omat\": \"orb-v3-conservative-inf-omat\",\n \"orb-v3-conservative-20-omat\": \"orb-v3-conservative-20-omat\",\n \"orb-v3-direct-inf-omat\": \"orb-v3-direct-inf-omat\",\n \"orb-v3-direct-20-omat\": \"orb-v3-direct-20-omat\",\n \"orb-v3-conservative-inf-mpa\": \"orb-v3-conservative-inf-mpa\",\n \"orb-v3-conservative-20-mpa\": \"orb-v3-conservative-20-mpa\",\n \"orb-v3-direct-inf-mpa\": \"orb-v3-direct-inf-mpa\",\n \"orb-v3-direct-20-mpa\": \"orb-v3-direct-20-mpa\",\n # The omol ids (orb-v3-{conservative,direct}-omol) are dropped from the\n # catalog 2026-07-30: they had been failing verify on every cluster since\n # 2026-05. Re-add once the failure is understood.\n # Your own fine-tuned v3 weights: pair with weights= (loaded via\n # setup_from_path).\n \"orb-v3:custom\": None,\n}\n\n\ndef _default_weights_url(load_fn) -> str:\n \"\"\"The upstream URL baked into the loader's ``weights_path`` default.\"\"\"\n import inspect\n\n default = inspect.signature(load_fn).parameters[\"weights_path\"].default\n if not isinstance(default, str) or not default.startswith((\"http://\", \"https://\")):\n raise RuntimeError(\n f\"{load_fn.__name__} has no URL default for weights_path \"\n f\"(got {default!r}); update this env file for the installed orb-models\"\n )\n return default\n\n\ndef _local_weights_path(url: str) -> Path:\n \"\"\"Where the checkpoint lives in the shared model cache.\"\"\"\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"orb\" / os.path.basename(url)\n\n\ndef _fetch(url: str, dest: Path) -> None:\n \"\"\"Download ``url`` to ``dest`` atomically (tmp file + rename).\"\"\"\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp = dest.with_name(f\"{dest.name}.tmp.{os.getpid()}\")\n try:\n with urllib.request.urlopen(url) as resp, open(tmp, \"wb\") as out:\n shutil.copyfileobj(resp, out)\n os.replace(tmp, dest)\n finally:\n tmp.unlink(missing_ok=True)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", precision: str = \"float32-high\"):\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.inference.calculator import ORBCalculator\n\n fn_name = CHECKPOINTS[checkpoint].replace(\"-\", \"_\")\n load_fn = getattr(pretrained, fn_name)\n\n # orb-models resolves its default weights URL through `cached_path`, which\n # write-locks its cache dir even on warm hits — EACCES for anyone who can't\n # write the shared install (Garden-AI/rootstock#67). Handed a *local* path\n # instead, cached_path returns it without locking. So the weights are\n # pre-fetched into the shared model cache at `rootstock add` time\n # (maintainer, cache writable) and every later serve loads that file.\n url = _default_weights_url(load_fn)\n weights = _local_weights_path(url)\n if not weights.exists():\n _fetch(url, weights)\n\n orbff, atoms_adapter = load_fn(\n weights_path=str(weights), device=torch.device(device), precision=precision\n )\n return ORBCalculator(orbff, atoms_adapter=atoms_adapter, device=torch.device(device))\n\n\ndef setup_from_path(\n path: str,\n device: str = \"cuda\",\n arch: str = \"orb-v3-conservative-inf-omat\",\n precision: str = \"float32-high\",\n):\n # Custom checkpoints (`:custom` ids with user weights). A weights file doesn't say\n # which orb architecture produced it, so `arch` names the pretrained\n # loader to instantiate — pass the right one at call time\n # (setup_kwargs={\"arch\": ...} / --kwarg arch=...). Handing the loader a\n # local path also means no network and no cached_path locking (see setup()).\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.inference.calculator import ORBCalculator\n\n fn_name = arch.replace(\"-\", \"_\")\n try:\n load_fn = getattr(pretrained, fn_name)\n except AttributeError:\n raise ValueError(\n f\"unknown orb architecture {arch!r}; expected a loader name from \"\n f\"orb_models.forcefield.pretrained, e.g. orb-v3-conservative-inf-omat\"\n ) from None\n\n orbff, atoms_adapter = load_fn(\n weights_path=path, device=torch.device(device), precision=precision\n )\n return ORBCalculator(orbff, atoms_adapter=atoms_adapter, device=torch.device(device))\n","python_requires":">=3.12,<3.13","dependencies":{"ase":"3.29.0","cached-path":"1.8.10","orb-models":"0.7.0","rootstock":"1.3.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:91b7a9b8df00a500d5e6147d7378259fe4e665c36954470bb5f97207c8c96279","checkpoints":{"orb-v3-conservative-inf-omat":{"fetched_at":"2026-07-18T02:30:03.511974+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/orb/orb-v3-conservative-inf-omat-20250404.ckpt","size":102097517}],"weights_recorded_at":"2026-09-26T05:40:47.735908+00:00"},"orb-v3-conservative-20-omat":{"fetched_at":"2026-07-18T02:33:01.402092+00:00","verified_at":"2026-10-24T03:28:45.589405+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-20-omat-20250404.ckpt","size":102097517}],"weights_recorded_at":"2026-10-24T03:28:45.589409+00:00"},"orb-v3-direct-inf-omat":{"fetched_at":"2026-07-18T02:33:48.238058+00:00","verified_at":"2026-10-24T03:28:45.589413+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-inf-omat-20250404.ckpt","size":102636658}],"weights_recorded_at":"2026-10-24T03:28:45.589416+00:00"},"orb-v3-direct-20-omat":{"fetched_at":"2026-07-18T02:34:23.077709+00:00","verified_at":"2026-10-24T03:28:45.589419+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-20-omat-20250404.ckpt","size":102636510}],"weights_recorded_at":"2026-10-24T03:28:45.589423+00:00"},"orb-v3-conservative-inf-mpa":{"fetched_at":"2026-07-18T02:34:57.953047+00:00","verified_at":"2026-10-24T03:28:45.589426+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-inf-mpa-20250404.ckpt","size":102097517}],"weights_recorded_at":"2026-10-24T03:28:45.589430+00:00"},"orb-v3-conservative-20-mpa":{"fetched_at":"2026-07-18T02:35:37.786464+00:00","verified_at":"2026-10-24T03:28:45.589433+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-20-mpa-20250404.ckpt","size":102097380}],"weights_recorded_at":"2026-10-24T03:28:45.589437+00:00"},"orb-v3-direct-inf-mpa":{"fetched_at":"2026-07-18T02:36:11.635141+00:00","verified_at":"2026-10-24T03:28:45.589440+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-inf-mpa-20250404.ckpt","size":102636510}],"weights_recorded_at":"2026-10-24T03:28:45.589444+00:00"},"orb-v3-direct-20-mpa":{"fetched_at":"2026-07-18T02:36:34.279807+00:00","verified_at":"2026-10-24T03:28:45.589447+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-20-mpa-20250404.ckpt","size":102636362}],"weights_recorded_at":"2026-10-24T03:28:45.589450+00:00"},"orb-v3:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null}}},"painn":{"built_at":"2026-07-20T18:33:36.548881+00:00","source_hash":"sha256:1cb2a0a04abfb53c98beb463cd40138791d5d0febe4d5cefb83730decf46edd1","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\nPaiNN environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 PaiNN checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"PaiNN-S2EF-OC20-All\": default\n\"\"\"\n\nCHECKPOINTS = {\n \"painn-s2ef-oc20-all\": \"PaiNN-S2EF-OC20-All\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a PaiNN OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.0.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"painn-s2ef-oc20-all":{"fetched_at":"2026-07-20T20:43:48.015191+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/painn_h512_s2ef_all.pt","size":80340525}],"weights_recorded_at":"2026-09-18T20:25:36.277968+00:00"}}},"scn":{"built_at":"2026-07-20T19:12:52.391743+00:00","source_hash":"sha256:cba233c494f7a5a6c9b675e024176c9755a7394bfff2a8d29dc5233d6a83c2f8","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\nSCN environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 SCN checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"SCN-S2EF-OC20-All+MD\": default\n - \"SCN-t4-b2-S2EF-OC20-2M\"\n - \"SCN-S2EF-OC20-2M\"\n\"\"\"\n\nCHECKPOINTS = {\n \"scn-s2ef-oc20-all-md\": \"SCN-S2EF-OC20-All+MD\",\n \"scn-t4-b2-s2ef-oc20-2m\": \"SCN-t4-b2-S2EF-OC20-2M\",\n \"scn-s2ef-oc20-2m\": \"SCN-S2EF-OC20-2M\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an SCN OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.0.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"scn-s2ef-oc20-all-md":{"fetched_at":"2026-07-20T20:51:58.830409+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null},"scn-t4-b2-s2ef-oc20-2m":{"fetched_at":"2026-07-20T20:53:26.667532+00:00","verified_at":"2026-10-24T03:28:45.589456+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/scn_t4_b2_s2ef_2M.pt","size":506994109}],"weights_recorded_at":"2026-10-24T03:28:45.589459+00:00"},"scn-s2ef-oc20-2m":{"fetched_at":"2026-07-20T20:53:50.713253+00:00","verified_at":"2026-10-24T03:28:45.589462+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/scn_t1_b1_s2ef_2M.pt","size":494398909}],"weights_recorded_at":"2026-10-24T03:28:45.589466+00:00"}}},"sevennet":{"built_at":"2026-07-30T23:31:06.759675+00:00","source_hash":"sha256:560f24594f3048803e4527d1a3761784e11c9a1d198c14781697f323a1384057","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"sevenn>=0.10.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"SevenNet env — hosts pretrained SevenNet universal potentials.\n\nSevenNet (SCalable EquiVariance-Enabled Neural Network) ships several\npretrained models, loaded by keyword through ``SevenNetCalculator``.\n\nMulti-fidelity models (``7net-omni``, ``7net-mf-ompa``) take a ``modal``\nargument selecting the training fidelity (e.g. ``\"mpa\"`` or ``\"omat24\"``),\nand sevenn releases after mid-2026 make it mandatory (older ones defaulted\nsilently). When not given, setup() defaults them to ``\"mpa\"`` — the\nMPtrj+sAlex fidelity, consistent with the lineage of the other checkpoints\nhere — so no-kwarg paths (smoke-test, plain ``rootstock add``) keep working.\nOverride at runtime via ``setup_kwargs={\"modal\": ...}`` on\nRootstockCalculator (or ``--kwarg modal=...`` for ``rootstock add``);\nsingle-fidelity models ignore it.\n\"\"\"\n\nCHECKPOINTS = {\n \"sevennet-0\": \"7net-0\",\n \"sevennet-l3i5\": \"7net-l3i5\",\n \"sevennet-omat\": \"7net-omat\",\n \"sevennet-mf-ompa\": \"7net-mf-ompa\",\n \"sevennet-omni\": \"7net-omni\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"sevennet:custom\": None,\n}\n\n\n# Multi-fidelity models and the fidelity used when modal isn't specified.\nMULTI_FIDELITY_DEFAULT_MODAL = {\n \"sevennet-mf-ompa\": \"mpa\",\n \"sevennet-omni\": \"mpa\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", modal: str | None = None):\n \"\"\"\n Load a SevenNet calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\").\n modal: Fidelity selector for multi-fidelity models (e.g. \"mpa\",\n \"omat24\"). Defaults to \"mpa\" for 7net-omni / 7net-mf-ompa;\n ignored by the single-fidelity models.\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n from sevenn.calculator import SevenNetCalculator\n\n if modal is None:\n modal = MULTI_FIDELITY_DEFAULT_MODAL.get(checkpoint)\n kwargs = {\"modal\": modal} if modal is not None else {}\n return SevenNetCalculator(model=CHECKPOINTS[checkpoint], device=device, **kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", modal: str | None = None):\n # Custom checkpoints (`:custom` ids with user weights): SevenNetCalculator loads a\n # checkpoint file directly. A multi-fidelity fine-tune must pass its\n # fidelity (setup_kwargs={\"modal\": \"mpa\"} / --kwarg modal=mpa);\n # single-fidelity ones need no kwargs.\n from sevenn.calculator import SevenNetCalculator\n\n kwargs = {\"modal\": modal} if modal is not None else {}\n return SevenNetCalculator(model=path, device=device, **kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.3.0","sevenn":"0.13.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:41daf60ec5bc217e733c33e027e6b81682dcc11e45354105449f167b7197551c","checkpoints":{"sevennet-0":{"fetched_at":"2026-07-18T02:57:16.494630+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null},"sevennet-l3i5":{"fetched_at":"2026-07-18T02:57:57.379501+00:00","verified_at":"2026-10-24T03:28:45.589471+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-omat":{"fetched_at":"2026-07-18T02:58:15.617503+00:00","verified_at":"2026-10-24T03:28:45.589475+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-mf-ompa":{"fetched_at":"2026-07-20T21:01:40.608149+00:00","verified_at":"2026-10-24T03:28:45.589478+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-omni":{"fetched_at":"2026-07-20T21:02:30.262147+00:00","verified_at":"2026-10-24T03:28:45.589482+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"uma":{"built_at":"2026-07-31T17:52:41.555061+00:00","source_hash":"sha256:bdc59e56e59ea9ce61a5c73af4bfed857582ee956d334aaf4ac3c9880069be8f","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"UMA env — hosts Meta's UMA foundation model via FAIRChem.\n\nfairchem-core v2 dropped the torch-geometric / pyg-find-links install dance, so\nthis env is a plain PyPI install. The original uma-s-1 had an extensivity bug\nand was removed from the fairchem 2.20 registry — use uma-s-1p1 or uma-s-1p2p1.\n\"\"\"\n\nCHECKPOINTS = {\n \"uma-s-1p1\": \"uma-s-1p1\",\n # uma-s-1p2 has a known major bug; uma-s-1p2p1 fixes it and is the\n # upstream-recommended small model. 1p2 stays listed for reproducibility\n # of existing runs.\n \"uma-s-1p2\": \"uma-s-1p2\",\n # uma-s-1p2p1 is in fairchem's registry on git main but NOT in any\n # release yet (latest fairchem-core 2.21.0, 2026-06-08, lacks it — the\n # 2026-07-30 sync failed on exactly this). Re-add when the next\n # fairchem-core ships, and bump the dependency floor to that version.\n # \"uma-s-1p2p1\": \"uma-s-1p2p1\",\n \"uma-m-1p1\": \"uma-m-1p1\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"uma:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", task: str = \"omat\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(CHECKPOINTS[checkpoint], device=device)\n return FAIRChemCalculator(predictor, task_name=task)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", task: str = \"omat\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=device)\n return FAIRChemCalculator(predictor, task_name=task)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.3.0","torch":"2.8.0"},"lock_hash":"sha256:e016d95f2bf9abb6f6827d56e8ba3db9b39747307b05404bc4f7e0258c082848","checkpoints":{"uma-s-1p1":{"fetched_at":"2026-07-18T03:06:29.158095+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/07068e9c76702ca173d13155095f2117c1b327ec228557e64cd2709c777b824a","size":1173759575},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-09-18T20:25:36.278008+00:00"},"uma-s-1p2":{"fetched_at":"2026-07-18T03:07:41.520380+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/ba5c0d912efa22dc238e5fb1b5b7f66ee2e68c48c1b95b7cfd7fe1da5938398b","size":2333393167},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-09-18T20:25:36.278020+00:00"},"uma-m-1p1":{"fetched_at":"2026-07-18T03:08:53.351374+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/c30034edbf2e127f703f814cacb632661767da99b3e71c5b2ee5290510a52d68","size":11174706771},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-09-18T20:25:36.278028+00:00"},"uma-s-1p2p1":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"download: KeyError: \"Model 'uma-s-1p2p1' not found. Available models: ('uma-s-1p2', 'uma-s-1p1', 'uma-m-1p1', 'esen-md-direct-all-omol', 'esen-sm-conserving-all-omol', 'esen-sm-direct-all-omol', 'allscaip-md-conserving-all-omol', 'allscaip-md-direct-all-omol', 'esen-sm-conserving-all-oc25', 'esen-md-direct-all-oc25', 'esen-sm-filtered-odac25', 'esen-sm-full-odac25')\"","weight_files":null,"weights_recorded_at":null},"uma:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null}}},"allscaip":{"built_at":"2026-08-03T17:58:06.805068+00:00","source_hash":"sha256:774cd05d08292476738237cf379892e77753c5ab3d35da294c240e7f3a041b53","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"AllScAIP env — FAIRChem scalable attention MLIP trained on OMol25.\n\nallscaip-md-conserving-all-omol is an energy-conserving, all-to-all node\nattention model served through fairchem-core's get_predict_unit — the same\nAPI as eSEN. fairchem v2 carries the architecture in-package, so no\nflash-attention or custom CUDA kernels are needed.\n\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"allscaip-md-conserving-all-omol\": \"allscaip-md-conserving-all-omol\",\n \"allscaip-md-direct-all-omol\": \"allscaip-md-direct-all-omol\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"allscaip:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(CHECKPOINTS[checkpoint], device=device)\n return FAIRChemCalculator(predictor)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=device)\n return FAIRChemCalculator(predictor)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.4.3","torch":"2.8.0"},"lock_hash":"sha256:8e7948fa1775faff6db4a1214f6e014923d5f5836a262f0e61d2e7bae6cb2b68","checkpoints":{"allscaip-md-conserving-all-omol":{"fetched_at":"2026-07-31T18:03:32.290693+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/14fe8626e9e75337769ecf98a765cdc02667763ec4dc1103d958e1c9f56b26e9","size":688432239},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-09-18T20:25:36.278035+00:00"},"allscaip-md-direct-all-omol":{"fetched_at":"2026-08-03T18:03:18.239815+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/5517360cf5a23ca141f1b51b708d131289e35080bfa79f27b990695e3e78d18e","size":694955368},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-09-18T20:25:36.278043+00:00"},"allscaip:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null}}},"mace_polar":{"built_at":"2026-08-03T17:57:34.313503+00:00","source_hash":"sha256:8db4cfd29a59c415e3762b04195e30b9e08d5f044ff6aa21ebcf01bcb78bc1ef","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"ase>=3.22\",\n# \"torch>=2.4.0,<2.10\",\n# # 0.3.16 is the first PyPI release with mace_polar().\n# \"mace-torch>=0.3.16\",\n# # PolarMACE imports graph_longrange at runtime; the distribution is named\n# # graph-longrange and exists only as this git repo (no PyPI release).\n# \"graph-longrange @ git+https://github.com/WillBaldwin0/graph_electrostatics.git\",\n# ]\n# ///\n\"\"\"MACE-POLAR env — electrostatic/polarizable MACE foundation models (OMol25).\n\nKept separate from the stable `mace` env because of the extra git-only\ngraph-longrange dependency.\n\nPOLAR checkpoints expect `charge`, `spin`, and `external_field` in atoms.info.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-polar-1-s\": \"polar-1-s\",\n \"mace-polar-1-m\": \"polar-1-m\",\n \"mace-polar-1-l\": \"polar-1-l\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace-polar:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from mace.calculators import mace_polar\n\n return mace_polar(model=CHECKPOINTS[checkpoint], device=device, default_dtype=\"float32\")\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): mace_polar() accepts a\n # weights file directly, keeping the PolarMACE model-type wiring.\n from mace.calculators import mace_polar\n\n return mace_polar(model=path, device=device, default_dtype=\"float32\")\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"1.4.3","torch":"2.9.1"},"lock_hash":"sha256:69dbb533819b90456ff3c1ee14c6dcbefdc8567bca144abd5b9c41ffa833fe63","checkpoints":{"mace-polar-1-s":{"fetched_at":"2026-07-31T00:02:47.536619+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/mace/MACEPOLAR1Smodel","size":33375439}],"weights_recorded_at":"2026-09-18T20:25:36.278050+00:00"},"mace-polar-1-m":{"fetched_at":"2026-07-31T00:02:47.844007+00:00","verified_at":"2026-10-24T03:28:45.589493+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Mmodel","size":68133235}],"weights_recorded_at":"2026-10-24T03:28:45.589496+00:00"},"mace-polar-1-l":{"fetched_at":"2026-07-31T00:02:48.613868+00:00","verified_at":"2026-10-24T03:28:45.589499+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Lmodel","size":129808715}],"weights_recorded_at":"2026-10-24T03:28:45.589503+00:00"},"mace-polar:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null}}},"mattersim":{"built_at":"2026-07-30T23:28:56.962927+00:00","source_hash":"sha256:c1706582d1cb4c355cc8d2e3381e399215e7e60886e4b3c09fbc39a7fb8c5611","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"mattersim>=1.1.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# # Not imported here. mattersim -> torchmetrics -> torchvision pulls it in\n# # transitively, and torchvision ships compiled ops ABI-locked to one exact\n# # torch build. An `explicit` index is only consulted for packages listed\n# # in `dependencies`, so without this line torchvision silently resolves\n# # from PyPI against a different torch and the mismatch surfaces only at\n# # `import torchvision` as:\n# # RuntimeError: operator torchvision::nms does not exist\n# # Leave it unpinned: the index hosts one build per torch release and uv\n# # enforces the torch<->torchvision pairing.\n# \"torchvision\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n# torchvision = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"\nMatterSim environment for Rootstock.\n\nProvides access to Microsoft's MatterSim universal potential. MatterSim-v1\ncovers ~100 elements and supports periodic and non-periodic systems.\n\nModels:\n - \"MatterSim-v1.0.0-5M\": 5M parameter model (faster, default)\n - \"MatterSim-v1.0.0-1M\": 1M parameter model (smallest)\n\"\"\"\n\nCHECKPOINTS = {\n \"mattersim-v1-0-0-5m\": \"MatterSim-v1.0.0-5M\",\n \"mattersim-v1-0-0-1m\": \"MatterSim-v1.0.0-1M\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mattersim:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a MatterSim calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n from mattersim.forcefield import MatterSimCalculator\n\n return MatterSimCalculator(load_path=CHECKPOINTS[checkpoint], device=device)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): load_path accepts a filesystem\n # path as well as a model name, so this is setup() minus the name mapping.\n from mattersim.forcefield import MatterSimCalculator\n\n return MatterSimCalculator(load_path=path, device=device)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mattersim":"1.2.3","rootstock":"1.3.0","torch":"2.11.0+cu128","torchvision":"0.26.0+cu128"},"lock_hash":"sha256:50d8510a961a3daf57e21dec70ffba0863e7517a55e9dca4bec5c0e243b7fc36","checkpoints":{"mattersim-v1-0-0-1m":{"fetched_at":"2026-07-31T00:15:48.223712+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null},"mattersim-v1-0-0-5m":{"fetched_at":"2026-07-31T00:15:49.594071+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null}}},"tensornet":{"built_at":"2026-07-30T23:32:01.045317+00:00","source_hash":"sha256:cddb5a364bae14466f61a30d7525f2f8c42db66c830977c95f547feceb8ce235","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"torch>=2.4.0,<2.5\",\n# \"ase>=3.22\",\n# \"huggingface_hub\",\n# \"matgl\",\n# # nvalchemi-toolkit-ops is deliberately absent. It's only an optional\n# # matgl extra (accelerated neighbor lists), and no version can work on\n# # this env's torch pin: 0.3.x registers torch custom ops in modules that\n# # use `from __future__ import annotations`, which torch 2.4's\n# # infer_schema can't parse — ValueError at import, and matgl's optional-\n# # import guard in matgl/ext/ase.py only catches ImportError, so merely\n# # importing PESCalculator crashes — while 0.4+ requires torch>=2.8.\n# # Without it, matgl falls back to its own neighbor list.\n# \"pymatgen\",\n# \"monty\",\n# \"ruamel.yaml\",\n# \"scipy\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# \"torch-spline-conv\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n#\n# [tool.uv.sources]\n# # Pinned: this recipe was originally written against matgl 1.0.0, and the\n# # unpinned git HEAD silently started building 4.x. setup() below is written\n# # for 4.0.3 — bump the tag deliberately, not by rebuild accident.\n# matgl = { git = \"https://github.com/materialsvirtuallab/matgl.git\", tag = \"v4.0.3\" }\n# ///\n\"\"\"TensorNet env — hosts MatPES TensorNet checkpoints via MatGL.\"\"\"\n\nCHECKPOINTS = {\n \"tensornet-matpes-pbe-2025-2\": \"materialyze/TensorNet-PES-MatPES-PBE-2025.2\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from huggingface_hub import snapshot_download\n\n import matgl\n from matgl.ext.ase import PESCalculator\n\n # load_model only resolves names against matgl's own manifest; HF models\n # must be downloaded explicitly and passed as a local path.\n local_path = snapshot_download(repo_id=CHECKPOINTS[checkpoint])\n\n # Move with .to(device), never torch.set_default_device: under matgl 4.x\n # the default-device hack splits the model across devices at load —\n # Potential.__init__ registers data_mean from a constructor-kwarg tensor\n # torch.load restored to cpu, while _eye3 (a persistent=False buffer, not\n # in the state dict) is created fresh on the default device — and\n # forward() then crashes with a cuda/cpu mismatch at\n # `lat @ (self._eye3 + st)` in matgl/apps/pes.py. Module.to() moves\n # params and all buffers coherently, and forward() migrates inputs to the\n # model's device itself.\n pot = matgl.load_model(local_path).to(device)\n return PESCalculator(potential=pot)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","huggingface-hub":"1.26.0","matgl":"4.0.3","monty":"2026.7.16","pymatgen":"2026.5.4","rootstock":"1.5.0","ruamel-yaml":"0.19.1","scipy":"1.17.1","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121","torch-spline-conv":"1.2.2+pt24cu121"},"lock_hash":null,"checkpoints":{"tensornet-matpes-pbe-2025-2":{"fetched_at":"2026-08-03T23:57:34.185051+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/blobs/19a0a506a58ea454e407040d55c393a174a5165834f22ea95c4c923ad9a406cc","size":4495},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/blobs/544966013206da7faa1966ffa360a82fb2d31c3e","size":5298},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/blobs/b9efa8c708014dd6c37e348f65dc1d8ae5c7bd5c7290de4d9f210e16c7fc3766","size":3402824},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/refs/main","size":40},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/trees/7ac2b90400a87eca7f44c1e7956dca6b1c85e418.json","size":906}],"weights_recorded_at":"2026-08-04T03:27:02.056431+00:00"}}},"grace":{"built_at":"2026-08-03T17:57:24.922622+00:00","source_hash":"sha256:f46207cbbd702910ea889c09cae3f9933de6c48baf22d75c4039caf02467b635","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"tensorpotential>=0.6.0\",\n# \"ase>=3.22\",\n# ]\n# ///\n\"\"\"GRACE env — hosts GRACE foundation checkpoints via tensorpotential.\n\nGRACE runs on TensorFlow, not torch. TPCalculator has no device argument:\nTF grabs whatever GPU it sees, so device selection happens via\nCUDA_VISIBLE_DEVICES, and both it and TF_USE_LEGACY_KERAS must be set\nbefore the first TF import. The first calculation triggers an XLA compile —\na slow first step is expected.\n\"\"\"\n\nCHECKPOINTS = {\n \"grace-2l-smax-omat-large\": \"GRACE-2L-SMAX-OMAT-large\",\n \"grace-3l-omat-large-ft-am\": \"GRACE-3L-OMAT-large-ft-AM\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n import os\n\n os.environ[\"TF_USE_LEGACY_KERAS\"] = \"1\"\n if device == \"cpu\":\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n elif device.startswith(\"cuda:\"):\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = device.split(\":\", 1)[1]\n\n from tensorpotential.calculator import grace_fm\n\n return grace_fm(CHECKPOINTS[checkpoint])\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.4.3","tensorpotential":"0.6.0"},"lock_hash":"sha256:1ea87e2d3e8620111788e4f3804932ac53f399ab75df08ee6b2b38195a82d50e","checkpoints":{"grace-2l-smax-omat-large":{"fetched_at":"2026-08-03T20:38:07.699352+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null},"grace-3l-omat-large-ft-am":{"fetched_at":"2026-08-03T18:03:21.272720+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":null,"weights_recorded_at":null}}},"pet":{"built_at":"2026-08-04T21:38:35.114654+00:00","source_hash":"sha256:8e3b96c758fe0ab8effbfc77bace1fd6f67c03daba959fd5a73d10420db7c0f6","source":"# /// script\n# requires-python = \">=3.11,<3.15\"\n# dependencies = [\n# \"upet>=0.2.6\",\n# \"ase>=3.22\",\n# # upet pulls nvalchemi-toolkit-ops unpinned; 0.4+ needs torch>=2.8 at\n# # runtime but only declares the constraint on its extras, so the\n# # resolver won't catch it (same trap as the tensornet env).\n# \"torch>=2.8,<2.14\",\n# ]\n#\n# # Delta's driver caps at CUDA 12.8 — route torch to the cu128 index, or the\n# # default PyPI wheel (cu130 for torch >=2.9) fails CUDA init on every checkpoint.\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"PET env — hosts lab-cosmo's UPET foundation checkpoints (PET-MAD successor).\n\nThe upstream string encodes model@version; versions are pinned rather than\n\"latest\" so rebuilds serve the same weights. pet-omatpes-l is trained at the\nr2SCAN level of theory — its energies are not comparable to the PBE-level\npet-oam models.\n\"\"\"\n\nCHECKPOINTS = {\n \"pet-oam-xl\": \"pet-oam-xl@1.0.0\",\n \"pet-omatpes-l\": \"pet-omatpes-l@0.1.0\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from huggingface_hub import hf_hub_download\n from upet.calculator import UPETCalculator\n\n # metatomic-ase 0.1.2's nvalchemi fast path passes a float max_neighbors\n # (len * max(128, cutoff**3)) into torch.full — TypeError for cutoffs > ~5 Å.\n # Force the vesin fallback; the flag is read at call time.\n import metatomic_ase._neighbors as _mta_n\n _mta_n.HAS_NVALCHEMIOPS = False\n\n # Passing model=/version= makes UPETCalculator resolve the name by listing\n # the hub repo — an uncached API call that fails on workers, which run\n # with HF_HUB_OFFLINE=1 (and on any node without internet). Fetch the\n # pinned file ourselves — a cache hit needs no network even offline — and\n # hand it over as checkpoint_path, which skips the resolve entirely.\n model, version = CHECKPOINTS[checkpoint].split(\"@\", 1)\n path = hf_hub_download(\n repo_id=\"lab-cosmo/upet\",\n filename=f\"{model}-v{version}.ckpt\",\n subfolder=\"models\",\n )\n return UPETCalculator(checkpoint_path=path, device=device)\n","python_requires":">=3.11,<3.15","dependencies":{"ase":"3.29.0","rootstock":"1.6.0","torch":"2.11.0+cu128","upet":"0.2.6"},"lock_hash":"sha256:c2c9d5767b95babe1be5d7d227f74fb98073b749444d88a8fe75c1c800c065f5","checkpoints":{"pet-omatpes-l":{"fetched_at":"2026-08-03T18:04:22.511196+00:00","verified_at":"2026-10-24T03:28:45.589514+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/blobs/879b1045391d88869522605a8b8b3cedeed74668e7062fdd7487548ab7b08004","size":771692879},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/refs/main","size":40}],"weights_recorded_at":"2026-10-24T03:28:45.589519+00:00"},"pet-oam-xl":{"fetched_at":"2026-08-03T18:04:29.291262+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/blobs/c3a67cd019969dfd4dcabe9574682fe035f861d3f1c10190989b36c983699409","size":2920687712},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/refs/main","size":40}],"weights_recorded_at":"2026-10-10T05:36:58.615609+00:00"}}},"tace":{"built_at":"2026-08-04T21:38:38.449879+00:00","source_hash":"sha256:638dbb140dd069d50c7c26a4797c413585355b8b75656cd35e4db8db5f1ddcba","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"tace\",\n# \"ase>=3.22\",\n# \"torch>=2.4,<2.14\",\n# ]\n#\n# # Delta's driver caps at CUDA 12.8 — route torch to the cu128 index, or the\n# # default PyPI wheel (cu130 for torch >=2.9) fails CUDA init on every checkpoint.\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# # The TECE architecture and its foundation registry ship in tace 0.2.0,\n# # which is not on PyPI (latest release there is 0.1.0). This is the commit\n# # the Matbench Discovery submission pinned.\n# [tool.uv.sources.tace]\n# git = \"https://github.com/xvzemin/tace\"\n# rev = \"81f65a4c188bd09cec8d1419388f7afdcc1b6fd0\"\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"TACE env — hosts TECE/TACE foundation checkpoints (Xu, Xie & Hu).\n\nopenequivariance CUDA-kernel acceleration is optional upstream and not\ninstalled here — TACE runs on the e3nn path.\n\"\"\"\n\nCHECKPOINTS = {\n \"tece-oam-rra-1.0\": \"TECE-OAM-RRA-1.0\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from tace.foundations import tace_foundations\n from tace.interface.ase import TACEAseCalc\n\n model_path = tace_foundations[CHECKPOINTS[checkpoint]]\n return TACEAseCalc(model_path, device=device, dtype=\"float32\")\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.6.0","tace":"0.2.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:427fc12575e2e2db79ba4c51df4aeb999bb88733785dd978e093de3f4516166d","checkpoints":{"tece-oam-rra-1.0":{"fetched_at":"2026-08-03T18:05:10.526573+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"home/.cache/tace/TECE-OAM-RRA-1.0.pt","size":890523407}],"weights_recorded_at":"2026-10-10T05:36:58.615628+00:00"}}},"deepmd":{"built_at":"2026-09-04T20:11:51.045203+00:00","source_hash":"sha256:8b68512349a498392bfffb8e6b13b21746c81e8f7fe57c3e1a2c6d645de58745","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"deepmd-kit[torch]>=3.2.0\",\n# \"ase>=3.23\",\n# # deepmd-kit's torch extra pins torch==2.11.0.*. Listed here too so the\n# # explicit cu128 index below routes it: an `explicit` index is only\n# # consulted for packages named in `dependencies`, and PyPI's torch 2.11\n# # wheels bundle CUDA 13, which needs a newer driver than most clusters run.\n# \"torch>=2.11,<2.12\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"DeePMD-kit env — hosts the DPA foundation models built into deepmd-kit.\n\nThe upstream strings are the names deepmd-kit's own pretrained-model registry\nknows (the same ones ``dp pretrained download`` accepts). Weights are fetched\nfrom Hugging Face with sha256 verification into the shared model cache and\nloaded from that local file, so serving needs no network.\n\nMultitask checkpoints (the OpenLAM DPA-2.4 / DPA-3.x models) carry one\nfitting head per training dataset and require ``head=`` selecting one —\n``setup_kwargs={\"head\": \"OMat24\"}`` on RootstockCalculator, or\n``--kwarg head=OMat24`` for ``rootstock add``. Valid heads per checkpoint\nare listed in HEADS; pick the dataset closest to your system and check it\ncovers your elements. Single-task checkpoints take no head.\n\nCharge and spin: checkpoints trained with frame-level charge/spin inputs\n(DPA3-Omol-Large, DPA-3.2-5M, DPA-3.3-1M) read them from\n``atoms.info[\"charge\"]`` and ``atoms.info[\"spin\"]`` (spin = multiplicity\n2S+1), the convention the other OMol-trained envs use. Without them the\nmodel's trained-in default (neutral singlet) applies. deepmd's native\n``fparam`` / ``charge_spin`` info keys still work and take precedence.\n\nDevice: deepmd's PyTorch backend fixes its device at first import from the\nenvironment (``DEVICE=cpu`` for CPU, otherwise ``cuda:{LOCAL_RANK}``), so\nsetup() sets those variables before importing deepmd.\n\nMPI: the deepmd-kit wheel preloads ``libmpi.so.12`` from the ``mpich`` wheel\nat import (its custom-op library links MPI, used only by the LAMMPS\nplugin). That libmpi needs the libfabric bundled next to it under\n``lib/mpich/``; on Cray systems the module environment puts the system\nlibfabric on ``LD_LIBRARY_PATH``, which outranks the wheel's RUNPATH and\nlacks the ``FABRIC_1.9`` symbol version libmpi wants. setup() therefore\nloads the bundled libfabric by absolute path first, so the dynamic linker\nreuses it when libmpi asks for ``libfabric.so.1``.\n\nLicenses: the DPA-2 / DPA-3 checkpoints are CC-BY-4.0; the DPA4 checkpoints\nare CC-BY-NC-4.0 (non-commercial).\n\"\"\"\n\nimport ctypes\nimport os\nimport sys\nfrom importlib import metadata\nfrom pathlib import Path\n\nCHECKPOINTS = {\n # Single-task checkpoints: no head selection.\n \"dpa3-omol-large\": \"DPA3-Omol-Large\",\n \"dpa4-nano-omat24\": \"DPA4-Nano-OMat24-v20260805\",\n \"dpa4-mini-omat24\": \"DPA4-Mini-OMat24-v20260805\",\n \"dpa4-neo-omat24\": \"DPA4-Neo-OMat24-v20260805\",\n \"dpa4-air-omat24\": \"DPA4-Air-OMat24-v20260805\",\n \"dpa4-plus-omat24\": \"DPA4-Plus-OMat24-v20260805\",\n # Multitask OpenLAM checkpoints: head= required (see HEADS).\n \"dpa-3.3-1m\": \"DPA-3.3-1M\",\n \"dpa-3.2-5m\": \"DPA-3.2-5M\",\n \"dpa-3.1-3m\": \"DPA-3.1-3M\",\n \"dpa-2.4-7m\": \"DPA-2.4-7M\",\n # Your own DeePMD-kit model — a training checkpoint (.pt) or frozen model\n # (.pth): pair with weights= (loaded via setup_from_path).\n \"dpa:custom\": None,\n}\n\n# Fitting heads of the multitask checkpoints (the upstream `model-branch`\n# names). Aliases such as \"Default\" or \"materials\" are accepted too, and\n# matching is case-insensitive.\n_OPENLAM_V2_HEADS = (\n \"OMat24\",\n \"OMol25\",\n \"MPTrj\",\n \"OC20M\",\n \"OC22\",\n \"ODAC23\",\n \"Alex2D\",\n \"MPGen_OpenCSP\",\n \"Domains_Alloy\",\n \"Domains_Anode\",\n \"Domains_Cluster\",\n \"Domains_FerroEle\",\n \"Domains_SSE_PBE\",\n \"Domains_SemiCond\",\n \"H2O_H2O_PD\",\n \"Metals_AlMgCu\",\n \"Metals_AgAu_PBED3\",\n \"Others_In2Se3\",\n \"Alloy_APEX\",\n \"SSE_ABACUS\",\n \"Hybrid_Perovskite\",\n \"Electrolyte\",\n \"Organic_Reactions\",\n)\n_OPENLAM_V1_HEADS = (\n \"Omat24\",\n \"MP_traj_v024_alldata_mixu\",\n \"OC20M\",\n \"OC22\",\n \"ODAC23\",\n \"Alex2D\",\n \"SPICE2\",\n \"Domains_Alloy\",\n \"Domains_Anode\",\n \"Domains_Cluster\",\n \"Domains_Drug\",\n \"Domains_FerroEle\",\n \"Domains_SSE_PBE\",\n \"Domains_SSE_PBESol\",\n \"Domains_SemiCond\",\n \"Domains_Transition1x\",\n \"H2O_H2O_PD\",\n \"Metals_AlMgCu\",\n \"Metals_Sn\",\n \"Metals_Ti\",\n \"Metals_V\",\n \"Metals_W\",\n \"Metals_AgAu_PBED3\",\n \"Others_HfO2\",\n \"Others_In2Se3\",\n \"Alloy_tongqi\",\n \"SSE_ABACUS\",\n \"Hybrid_Perovskite\",\n \"solvated_protein_fragments\",\n \"Electrolyte\",\n \"Organic_Reactions\",\n)\nHEADS = {\n \"dpa-3.3-1m\": _OPENLAM_V2_HEADS,\n \"dpa-3.2-5m\": _OPENLAM_V2_HEADS,\n \"dpa-3.1-3m\": _OPENLAM_V1_HEADS,\n \"dpa-2.4-7m\": _OPENLAM_V1_HEADS,\n}\n\n# Head used when verification runs with no user selection. Users still select\n# explicitly. The custom entry's value is forwarded to whichever shipped\n# checkpoint the weights= leg borrows: single-task checkpoints ignore it.\nVERIFY_KWARGS = {\n \"dpa-3.3-1m\": {\"head\": \"OMat24\"},\n \"dpa-3.2-5m\": {\"head\": \"OMat24\"},\n \"dpa-3.1-3m\": {\"head\": \"Omat24\"},\n \"dpa-2.4-7m\": {\"head\": \"Omat24\"},\n \"dpa:custom\": {\"head\": \"OMat24\"},\n}\n\n\ndef _cache_dir() -> Path:\n \"\"\"Where the registry downloads land: deepmd's default layout, relocated\n into the shared model cache.\"\"\"\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"deepmd\" / \"pretrained\" / \"models\"\n\n\ndef _select_device(device: str) -> None:\n \"\"\"Point deepmd's PyTorch backend at ``device``.\n\n The backend reads its device once, when ``deepmd.pt.utils.env`` is first\n imported: ``DEVICE=cpu`` forces CPU, otherwise it uses\n ``cuda:{LOCAL_RANK}`` (LOCAL_RANK defaulting to 0). There is no per-call\n device argument, so this must run before the first deepmd import.\n \"\"\"\n if device == \"cpu\":\n os.environ[\"DEVICE\"] = \"cpu\"\n elif device.startswith(\"cuda:\"):\n os.environ[\"LOCAL_RANK\"] = device.split(\":\", 1)[1]\n\n\ndef _bundled_libfabric() -> Path | None:\n \"\"\"The libfabric shipped inside the ``mpich`` wheel, if installed.\"\"\"\n try:\n files = metadata.files(\"mpich\") or []\n except metadata.PackageNotFoundError:\n files = []\n for entry in files:\n if entry.match(\"mpich/libfabric.so.1\"):\n return Path(entry.locate()).resolve()\n fallback = Path(sys.prefix) / \"lib\" / \"mpich\" / \"libfabric.so.1\"\n return fallback if fallback.is_file() else None\n\n\ndef _preload_bundled_libfabric() -> None:\n \"\"\"Load the wheel's own libfabric before deepmd pulls in libmpi.\n\n Must run before the first deepmd import: once ``libmpi.so.12`` has\n resolved ``libfabric.so.1`` against whatever LD_LIBRARY_PATH offered\n (the Cray system copy, on Delta), the choice is fixed for the process.\n A library already loaded under that soname is reused instead.\n \"\"\"\n lib = _bundled_libfabric()\n if lib is not None:\n ctypes.CDLL(str(lib), mode=ctypes.RTLD_GLOBAL)\n\n\ndef _calculator_class():\n \"\"\"deepmd's ASE calculator, reading charge/spin the way the other envs do.\"\"\"\n from ase.calculators.calculator import all_changes\n from deepmd.calculator import DP\n\n class RootstockDP(DP):\n def __init__(self, *args, **kwargs):\n super().__init__(*args, **kwargs)\n # Two ways a DPA model can take charge/spin: as the two frame\n # parameters (DPA3-Omol-Large, DPA-3.2-5M) or through a dedicated\n # charge/spin embedding (DPA-3.3-1M, DPA4-OMol).\n self._charge_spin_as_fparam = self.dp.get_dim_fparam() == 2\n self._charge_spin_embedded = self.dp.has_chg_spin_ebd()\n\n def calculate(\n self,\n atoms=None,\n properties=(\"energy\", \"forces\", \"virial\"),\n system_changes=all_changes,\n ):\n if atoms is not None and (\"charge\" in atoms.info or \"spin\" in atoms.info):\n atoms = atoms.copy()\n charge_spin = [\n float(atoms.info.get(\"charge\", 0)),\n float(atoms.info.get(\"spin\", 1)),\n ]\n if self._charge_spin_as_fparam:\n atoms.info.setdefault(\"fparam\", charge_spin)\n if self._charge_spin_embedded:\n atoms.info.setdefault(\"charge_spin\", charge_spin)\n super().calculate(atoms, list(properties), system_changes)\n\n return RootstockDP\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", head: str | None = None, **kwargs):\n \"\"\"\n Load a DeePMD-kit calculator for a built-in pretrained model.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\").\n head: Fitting head of a multitask checkpoint (see HEADS). Required\n for those; ignored by single-task checkpoints.\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n if checkpoint in HEADS and head is None:\n raise ValueError(\n f\"{checkpoint} is multitask and has no default head - select one \"\n f'with setup_kwargs={{\"head\": ...}}: one of {\", \".join(HEADS[checkpoint])}'\n )\n _select_device(device)\n _preload_bundled_libfabric()\n\n from deepmd.pretrained.download import resolve_model_path\n\n # Resolves the registry name to a sha256-verified local file, downloading\n # only when it is missing from the cache (at `rootstock add` time).\n weights = resolve_model_path(CHECKPOINTS[checkpoint], cache_dir=_cache_dir())\n return _calculator_class()(model=str(weights), head=head, **kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", head: str | None = None, **kwargs):\n # Custom checkpoints (`:custom` ids with user weights): DP loads a training\n # checkpoint (.pt) or a frozen model (.pth) straight from a path. A\n # multitask fine-tune needs its head (setup_kwargs={\"head\": ...} /\n # --kwarg head=...) unless it declares a default; single-task ones\n # take none.\n _select_device(device)\n _preload_bundled_libfabric()\n return _calculator_class()(model=path, head=head, **kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","deepmd-kit":"3.2.0","rootstock":"1.6.4","torch":"2.11.0+cu128"},"lock_hash":"sha256:3b2adea6c46292acddbb21dd50059f218c16b8b950099984fd152869e18e6f55","checkpoints":{"dpa-3.3-1m":{"fetched_at":"2026-09-04T21:10:29.970025+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/deepmd/pretrained/models/DPA-3.3-1M.pt","size":41612295}],"weights_recorded_at":"2026-10-10T05:36:58.615674+00:00"},"dpa-3.1-3m":{"fetched_at":"2026-09-04T21:10:30.947721+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/deepmd/pretrained/models/DPA-3.1-3M.pt","size":47176032}],"weights_recorded_at":"2026-10-10T05:36:58.615688+00:00"},"dpa-3.2-5m":{"fetched_at":"2026-09-04T21:10:32.190658+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/deepmd/pretrained/models/DPA-3.2-5M.pt","size":64800782}],"weights_recorded_at":"2026-10-10T05:36:58.615681+00:00"},"dpa4-air-omat24":{"fetched_at":"2026-09-04T21:11:03.916056+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/deepmd/pretrained/models/DPA4-Air-OMat24-v20260805.pt","size":20713035}],"weights_recorded_at":"2026-10-10T05:36:58.615662+00:00"},"dpa3-omol-large":{"fetched_at":"2026-09-04T21:11:02.892992+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/deepmd/pretrained/models/DPA3-Omol-Large.pt","size":224057430}],"weights_recorded_at":"2026-10-10T05:36:58.615634+00:00"},"dpa4-mini-omat24":{"fetched_at":"2026-09-04T21:11:00.583238+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/deepmd/pretrained/models/DPA4-Mini-OMat24-v20260805.pt","size":2701347}],"weights_recorded_at":"2026-10-10T05:36:58.615648+00:00"},"dpa4-nano-omat24":{"fetched_at":"2026-09-04T21:11:00.030849+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/deepmd/pretrained/models/DPA4-Nano-OMat24-v20260805.pt","size":1997147}],"weights_recorded_at":"2026-10-10T05:36:58.615641+00:00"},"dpa4-neo-omat24":{"fetched_at":"2026-09-04T21:11:19.407213+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/deepmd/pretrained/models/DPA4-Neo-OMat24-v20260805.pt","size":4583819}],"weights_recorded_at":"2026-10-10T05:36:58.615655+00:00"},"dpa4-plus-omat24":{"fetched_at":"2026-09-04T21:11:20.501886+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/deepmd/pretrained/models/DPA4-Plus-OMat24-v20260805.pt","size":35550605}],"weights_recorded_at":"2026-10-10T05:36:58.615669+00:00"},"dpa-2.4-7m":{"fetched_at":"2026-09-04T21:10:30.286150+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (1800s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/deepmd/pretrained/models/dpa-2.4-7M.pt","size":169224765}],"weights_recorded_at":"2026-10-10T05:36:58.615695+00:00"}}}}},{"schema_version":5,"cluster":"frontier","root":"/sw/frontier/ums/ums047/rootstock","maintainer":{"name":"Owen Price Skelly","email":"owenpriceskelly@uchicago.edu"},"rootstock_version":"1.6.4","python_version":"3.10","last_updated":"2026-10-12T06:41:35.441612+00:00","environments":{"mace":{"built_at":"2026-08-03T17:39:02.047579+00:00","source_hash":"sha256:72c1ec14e2a04a0e5afe6f25aaf472dcaf5afe1d4858cb29454fa3783dcfee84","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # 0.3.15+ needed for the mh-1 registry entry (matpes needs 0.3.13,\n# # omol needs 0.3.14, mpa-0 needs 0.3.10).\n# \"mace-torch>=0.3.15\",\n# \"ase>=3.22\",\n# # 2.4.1 is explicitly unsupported by mace-torch.\n# \"torch>=2.4.0,!=2.4.1,<2.10\",\n# # torch's ROCm wheels depend on this; it lives only on the ROCm\n# # index, so it must be a direct dep for [tool.uv.sources] to route it.\n# \"pytorch-triton-rocm\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-rocm\" }\n# pytorch-triton-rocm = { index = \"pytorch-rocm\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-rocm\"\n# url = \"https://download.pytorch.org/whl/rocm6.4\"\n# explicit = true\n# ///\n\"\"\"MACE env (ROCm) - MACE checkpoints on AMD GPUs.\n\nIdentical to nvidia_configs/mace.py except torch resolves from the ROCm wheel\nindex. PyTorch ROCm builds expose AMD GPUs as device=\"cuda\", so setup() is\nunchanged. cuEquivariance acceleration is CUDA-only and is not used here -\nMACE falls back to the pure e3nn/torch path.\n\nUpstream-string routing in CHECKPOINTS: an `off:` prefix routes to mace_off()\nand an `omol:` prefix to mace_omol() (float64, molecules only); an `mh:`\nprefix marks a multi-head model (float64, per the MACE-MH-1 model card).\n\nMulti-head checkpoints select a head via the `head` kwarg on setup()\n(setup_kwargs={\"head\": ...} / --kwarg head=...), named by upstream's training\ncorpus — see MH1_HEADS; omat_pbe is the default.\n\nThe OMOL checkpoint expects `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-mp-0-small\": \"small\",\n \"mace-mp-0-medium\": \"medium\",\n \"mace-mp-0-large\": \"large\",\n \"mace-off23-small\": \"off:small\",\n \"mace-off23-medium\": \"off:medium\",\n \"mace-off23-large\": \"off:large\",\n # Only a medium MPA-0 has been released, but upstream names the weights\n # file mace-mpa-0-medium.model — keep the size explicit like mace-mp-0.\n \"mace-mpa-0-medium\": \"medium-mpa-0\",\n \"mace-matpes-r2scan-0\": \"mace-matpes-r2scan-0\",\n # One entry per weights file: MH-1's heads are selected by setup(head=...).\n \"mace-mh-1\": \"mh:mh-1\",\n # Only the extra-large OMOL model has been released.\n \"mace-omol-0-extra-large\": \"omol:extra_large\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace:custom\": None,\n}\n\n# MH-1's heads, named by training corpus. The released weights file is the\n# authority (mace_select_head --list_heads; the model card also lists a\n# rgd1_b3lyp head, but that shipped only in mh-0 — ACEsuit/mace#1462).\n# Validated here because upstream only warns on an unknown head and silently\n# falls back to the last one.\nMH1_HEADS = (\n \"omat_pbe\",\n \"omol\",\n \"spice_wB97M\",\n \"oc20_usemppbe\",\n \"mp_pbe_refit_add\",\n \"matpes_r2scan\",\n)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", head: str | None = None):\n arg = CHECKPOINTS[checkpoint]\n if arg.startswith(\"mh:\"):\n head = head or \"omat_pbe\"\n if head not in MH1_HEADS:\n raise ValueError(f\"unknown head {head!r}; expected one of {', '.join(MH1_HEADS)}\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg[3:], device=device, default_dtype=\"float64\", head=head)\n if head is not None:\n raise ValueError(f\"'head' selects a head of a multi-head model; {checkpoint} has one head\")\n if arg.startswith(\"off:\"):\n from mace.calculators import mace_off\n\n return mace_off(model=arg[4:], device=device, default_dtype=\"float32\")\n if arg.startswith(\"omol:\"):\n from mace.calculators import mace_omol\n\n return mace_omol(model=arg[5:], device=device, default_dtype=\"float64\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg, device=device, default_dtype=\"float32\")\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", head: str | None = None):\n # Custom checkpoints (`:custom` ids with user weights): fine-tunes load through\n # MACECalculator directly — the mp/off dispatch in setup() only exists\n # to pick which pretrained file to download. `head` is for fine-tunes that\n # keep multiple heads; single-head weights load without it.\n from mace.calculators import MACECalculator\n\n return MACECalculator(model_paths=path, device=device, default_dtype=\"float32\", head=head)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","pytorch-triton-rocm":"3.5.1","rootstock":"1.4.3","torch":"2.9.1+rocm6.4"},"lock_hash":"sha256:3c4a11c77df8da2739f51e91984fe029c50d6734cf417327282c1a8501ad3de8","checkpoints":{"mace-mp-0-medium":{"fetched_at":"2026-07-20T19:36:12.439451+00:00","verified_at":"2026-10-12T06:41:33.314934+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231203mace128L1_epoch199model","size":44422970}],"weights_recorded_at":"2026-10-12T06:41:33.314942+00:00"},"mace-mp-0-small":{"fetched_at":"2026-07-20T19:56:00.318387+00:00","verified_at":"2026-10-12T06:41:33.314902+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231210mace128L0_energy_epoch249model","size":32581838}],"weights_recorded_at":"2026-10-12T06:41:33.314917+00:00"},"mace-mp-0-large":{"fetched_at":"2026-07-20T19:59:39.486721+00:00","verified_at":"2026-10-12T06:41:33.314948+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE_MPtrj_20229model","size":133803220}],"weights_recorded_at":"2026-10-12T06:41:33.314953+00:00"},"mace-off23-small":{"fetched_at":"2026-07-20T20:01:57.226028+00:00","verified_at":"2026-10-12T06:41:33.314957+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_small.model","size":7347350}],"weights_recorded_at":"2026-10-12T06:41:33.314961+00:00"},"mace-off23-medium":{"fetched_at":"2026-07-20T20:04:28.688189+00:00","verified_at":"2026-10-12T06:41:33.314965+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_medium.model","size":18350596}],"weights_recorded_at":"2026-10-12T06:41:33.314969+00:00"},"mace-off23-large":{"fetched_at":"2026-07-20T20:07:12.519950+00:00","verified_at":"2026-10-12T06:41:33.314972+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_large.model","size":55492786}],"weights_recorded_at":"2026-10-12T06:41:33.314976+00:00"},"mace-matpes-r2scan-0":{"fetched_at":"2026-08-03T17:54:44.773719+00:00","verified_at":"2026-10-12T06:41:33.314987+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEmatpesr2scanomatftmodel","size":79470738}],"weights_recorded_at":"2026-10-12T06:41:33.314992+00:00"},"mace-mh-1":{"fetched_at":"2026-08-03T17:54:45.010039+00:00","verified_at":"2026-10-12T06:41:33.314996+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macemh1model","size":59208139}],"weights_recorded_at":"2026-10-12T06:41:33.315000+00:00"},"mace-mpa-0-medium":{"fetched_at":"2026-08-03T17:55:21.444729+00:00","verified_at":"2026-10-12T06:41:33.314980+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macempa0mediummodel","size":79462305}],"weights_recorded_at":"2026-10-12T06:41:33.314984+00:00"},"mace-omol-0-extra-large":{"fetched_at":"2026-08-03T17:55:23.154847+00:00","verified_at":"2026-10-12T06:41:33.315004+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-omol-0-extra-large-1024.model","size":422242640}],"weights_recorded_at":"2026-10-12T06:41:33.315008+00:00"},"mace:custom":{"fetched_at":null,"verified_at":"2026-10-12T06:41:33.315121+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"chgnet":{"built_at":"2026-08-03T17:35:11.370248+00:00","source_hash":"sha256:bfd563395ac0baca67ae203cf487cb6a1481641c145aae92775673ae8bd36428","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"chgnet>=0.3.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# # torch's ROCm wheels depend on this; it lives only on the ROCm\n# # index, so it must be a direct dep for [tool.uv.sources] to route it.\n# \"pytorch-triton-rocm\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-rocm\" }\n# pytorch-triton-rocm = { index = \"pytorch-rocm\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-rocm\"\n# url = \"https://download.pytorch.org/whl/rocm6.4\"\n# explicit = true\n# ///\n\"\"\"CHGNet env (ROCm) - pretrained charge-informed universal potentials on AMD GPUs.\"\"\"\n\nCHECKPOINTS = {\n \"chgnet-default\": \"chgnet-default\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"chgnet:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from chgnet.model import CHGNet, CHGNetCalculator\n\n model_name = CHECKPOINTS[checkpoint]\n model = CHGNet.load() if model_name == \"chgnet-default\" else CHGNet.load(model_name)\n return CHGNetCalculator(model=model, use_device=device)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # CHGNet.from_file, not the named-model CHGNet.load() setup() uses.\n from chgnet.model import CHGNet, CHGNetCalculator\n\n model = CHGNet.from_file(path)\n return CHGNetCalculator(model=model, use_device=device)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","chgnet":"0.4.2","pytorch-triton-rocm":"3.5.1","rootstock":"1.4.3","torch":"2.9.1+rocm6.4"},"lock_hash":"sha256:e9c1453559916d27477776b51f5e1856e3b582ad70d880bbc8dd8772aa8f049c","checkpoints":{"chgnet-default":{"fetched_at":"2026-07-20T19:52:58.472650+00:00","verified_at":"2026-10-12T06:41:33.315011+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"mattersim":{"built_at":"2026-08-03T17:42:32.732907+00:00","source_hash":"sha256:88635158c005031b00c0cb9e206d17dc087f2847d0d5336d83325bccc5c3a571","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"mattersim>=1.1.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# # torch's ROCm wheels depend on this; it lives only on the ROCm\n# # index, so it must be a direct dep for [tool.uv.sources] to route it.\n# \"pytorch-triton-rocm\",\n# # Not imported here. mattersim -> torchmetrics -> torchvision pulls it in\n# # transitively, and torchvision ships compiled ops ABI-locked to one exact\n# # torch build. An `explicit` index is only consulted for packages listed\n# # in `dependencies`, so without this line torchvision silently resolves\n# # from PyPI (a CUDA build!) against a different torch, and the mismatch\n# # surfaces only at `import torchvision` as:\n# # RuntimeError: operator torchvision::nms does not exist\n# # Leave it unpinned: the index hosts one build per torch release and uv\n# # enforces the torch<->torchvision pairing.\n# \"torchvision\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-rocm\" }\n# pytorch-triton-rocm = { index = \"pytorch-rocm\" }\n# torchvision = { index = \"pytorch-rocm\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-rocm\"\n# url = \"https://download.pytorch.org/whl/rocm6.4\"\n# explicit = true\n# ///\n\"\"\"MatterSim env (ROCm) - Microsoft's MatterSim universal potential on AMD GPUs.\"\"\"\n\nCHECKPOINTS = {\n \"mattersim-v1-0-0-5m\": \"MatterSim-v1.0.0-5M\",\n \"mattersim-v1-0-0-1m\": \"MatterSim-v1.0.0-1M\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mattersim:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from mattersim.forcefield import MatterSimCalculator\n\n return MatterSimCalculator(load_path=CHECKPOINTS[checkpoint], device=device)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): load_path accepts a filesystem\n # path as well as a model name, so this is setup() minus the name mapping.\n from mattersim.forcefield import MatterSimCalculator\n\n return MatterSimCalculator(load_path=path, device=device)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mattersim":"1.2.3","pytorch-triton-rocm":"3.5.1","rootstock":"1.4.3","torch":"2.9.1+rocm6.4","torchvision":"0.24.1+rocm6.4"},"lock_hash":"sha256:2a80807521fec162c671dbfea8f41bfda7bf786ed3504f383eb57869859df7f7","checkpoints":{"mattersim-v1-0-0-5m":{"fetched_at":"2026-07-20T20:11:10.579876+00:00","verified_at":"2026-10-12T06:41:33.315015+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"home/.local/mattersim/pretrained_models/mattersim-v1.0.0-5M.pth","size":91176875}],"weights_recorded_at":"2026-10-12T06:41:33.315019+00:00"},"mattersim-v1-0-0-1m":{"fetched_at":"2026-07-20T20:13:50.936758+00:00","verified_at":"2026-10-12T06:41:33.315023+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"home/.local/mattersim/pretrained_models/mattersim-v1.0.0-1M.pth","size":17932943}],"weights_recorded_at":"2026-10-12T06:41:33.315026+00:00"},"mattersim:custom":{"fetched_at":null,"verified_at":"2026-10-12T06:41:33.315125+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"orb":{"built_at":"2026-08-03T17:42:20.125187+00:00","source_hash":"sha256:8f41e48388ca69a6931510f5470ecbb12bc4a3046559dfd07a5015a784b5c7dc","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # >=0.5,<0.6: 0.5.5 is what the verified Delta env resolved — the v2\n# # loaders keep their single-return API through 0.5.x. 0.4.x is broken\n# # for us: it imports pynanoflann, which is git-only and undeclared, so\n# # a fresh build dies at import (Delta, 2026-07-31). 0.6 raises the\n# # Python floor to 3.12 — see nvidia_configs/orb_v3.py (no ROCm v3\n# # config yet).\n# \"orb-models>=0.5,<0.6\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# # Not imported here - constrains orb-models' transitive dep (see\n# # nvidia_configs/orb.py and Garden-AI/rootstock#67).\n# \"cached_path==1.8.10\",\n# # torch's ROCm wheels depend on this; it lives only on the ROCm\n# # index, so it must be a direct dep for [tool.uv.sources] to route it.\n# \"pytorch-triton-rocm\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-rocm\" }\n# pytorch-triton-rocm = { index = \"pytorch-rocm\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-rocm\"\n# url = \"https://download.pytorch.org/whl/rocm6.4\"\n# explicit = true\n# ///\n\"\"\"Orb v2 env (ROCm) — kept only for the built-in-D3 dispersion variant.\n\nMirrors nvidia_configs/orb.py: Orb v3 supersedes the other v2 checkpoints,\nso the catalog was trimmed to orb-d3-v2 (2026-07-30). There is no ROCm orb\nv3 config yet (v3 needs torch>=2.8 ROCm wheels — untested on Frontier).\n\"\"\"\n\nimport os\nimport shutil\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"orb-d3-v2\": \"orb-d3-v2\",\n # Your own fine-tuned v2-architecture weights: pair with weights=\n # (loaded via setup_from_path). v3 fine-tunes go to orb-v3:custom.\n \"orb-v2:custom\": None,\n}\n\n\ndef _default_weights_url(load_fn) -> str:\n import inspect\n\n default = inspect.signature(load_fn).parameters[\"weights_path\"].default\n if not isinstance(default, str) or not default.startswith((\"http://\", \"https://\")):\n raise RuntimeError(\n f\"{load_fn.__name__} has no URL default for weights_path \"\n f\"(got {default!r}); update this env file for the installed orb-models\"\n )\n return default\n\n\ndef _local_weights_path(url: str) -> Path:\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"orb\" / os.path.basename(url)\n\n\ndef _fetch(url: str, dest: Path) -> None:\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp = dest.with_name(f\"{dest.name}.tmp.{os.getpid()}\")\n try:\n with urllib.request.urlopen(url) as resp, open(tmp, \"wb\") as out:\n shutil.copyfileobj(resp, out)\n os.replace(tmp, dest)\n finally:\n tmp.unlink(missing_ok=True)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.calculator import ORBCalculator\n\n fn_name = CHECKPOINTS[checkpoint].replace(\"-\", \"_\")\n load_fn = getattr(pretrained, fn_name)\n\n url = _default_weights_url(load_fn)\n weights = _local_weights_path(url)\n if not weights.exists():\n _fetch(url, weights)\n\n orbff = load_fn(weights_path=str(weights), device=torch.device(device))\n return ORBCalculator(orbff, device=torch.device(device))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", arch: str = \"orb-v2\"):\n # Custom checkpoints (`:custom` ids with user weights). A weights file doesn't say\n # which orb architecture produced it, so `arch` names the pretrained\n # loader to instantiate — pass the right one at call time\n # (setup_kwargs={\"arch\": ...} / --kwarg arch=...). Handing the loader a local path also means no\n # network and no cached_path locking (see setup()).\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.calculator import ORBCalculator\n\n fn_name = arch.replace(\"-\", \"_\")\n try:\n load_fn = getattr(pretrained, fn_name)\n except AttributeError:\n raise ValueError(\n f\"unknown orb architecture {arch!r}; expected a loader name from \"\n f\"orb_models.forcefield.pretrained, e.g. orb-v2, orb-d3-v2\"\n ) from None\n\n orbff = load_fn(weights_path=path, device=torch.device(device))\n return ORBCalculator(orbff, device=torch.device(device))\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","cached-path":"1.8.10","orb-models":"0.5.5","pytorch-triton-rocm":"3.5.1","rootstock":"1.4.3","torch":"2.9.1+rocm6.4"},"lock_hash":"sha256:e4ad47ac657fea1458740ed50122e530f4de8f1155c3ba5cff3ab84c0c12e238","checkpoints":{"orb-d3-v2":{"fetched_at":"2026-07-20T20:21:13.941169+00:00","verified_at":"2026-10-12T06:41:33.315029+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-d3-v2-20241011.ckpt","size":100967743}],"weights_recorded_at":"2026-10-12T06:41:33.315032+00:00"}}},"uma":{"built_at":"2026-08-03T17:47:35.543245+00:00","source_hash":"sha256:060e943dbf3b2e43a57f6684da36c6db7bb697b70fb340c8f31a3e76ce2a3ca5","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# # torch's ROCm wheels depend on this; it lives only on the ROCm\n# # index, so it must be a direct dep for [tool.uv.sources] to route it.\n# \"pytorch-triton-rocm\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-rocm\" }\n# pytorch-triton-rocm = { index = \"pytorch-rocm\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-rocm\"\n# url = \"https://download.pytorch.org/whl/rocm6.4\"\n# explicit = true\n# ///\n\"\"\"UMA env (ROCm) - Meta's UMA foundation model via FAIRChem on AMD GPUs.\n\nfairchem-core v2 is a plain PyPI install (no torch-geometric/pyg-find-links),\nso the only ROCm change is the torch wheel index. Requires HF_TOKEN for the\ngated facebook/UMA checkpoints.\n\"\"\"\n\nCHECKPOINTS = {\n \"uma-s-1p1\": \"uma-s-1p1\",\n \"uma-s-1p2\": \"uma-s-1p2\",\n \"uma-m-1p1\": \"uma-m-1p1\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"uma:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", task: str = \"omat\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(CHECKPOINTS[checkpoint], device=device)\n return FAIRChemCalculator(predictor, task_name=task)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", task: str = \"omat\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=device)\n return FAIRChemCalculator(predictor, task_name=task)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","pytorch-triton-rocm":"3.4.0","rootstock":"1.4.3","torch":"2.8.0+rocm6.4"},"lock_hash":"sha256:ea64c4ba64b390d7ea6fec4fa89a85f14ff47b5352954bc7868cd6327a4717a0","checkpoints":{"uma-s-1p1":{"fetched_at":"2026-07-20T20:27:50.207764+00:00","verified_at":"2026-10-12T06:41:33.315036+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/07068e9c76702ca173d13155095f2117c1b327ec228557e64cd2709c777b824a","size":1173759575},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-10-12T06:41:33.315042+00:00"},"uma-s-1p2":{"fetched_at":"2026-07-20T20:30:39.154652+00:00","verified_at":"2026-10-12T06:41:33.315048+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/ba5c0d912efa22dc238e5fb1b5b7f66ee2e68c48c1b95b7cfd7fe1da5938398b","size":2333393167},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-10-12T06:41:33.315053+00:00"},"uma-m-1p1":{"fetched_at":"2026-07-20T20:34:07.816107+00:00","verified_at":"2026-10-12T06:41:33.315056+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/c30034edbf2e127f703f814cacb632661767da99b3e71c5b2ee5290510a52d68","size":11174706771},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-10-12T06:41:33.315062+00:00"},"uma:custom":{"fetched_at":null,"verified_at":"2026-10-12T06:41:33.315127+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"allscaip":{"built_at":"2026-08-03T17:34:45.442663+00:00","source_hash":"sha256:b4918b317d335431a92e4b0ea71ded5cbeaf5aa7ad1d445df57c00e30227dc22","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# # torch's ROCm wheels depend on this; it lives only on the ROCm\n# # index, so it must be a direct dep for [tool.uv.sources] to route it.\n# \"pytorch-triton-rocm\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-rocm\" }\n# pytorch-triton-rocm = { index = \"pytorch-rocm\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-rocm\"\n# url = \"https://download.pytorch.org/whl/rocm6.4\"\n# explicit = true\n# ///\n\"\"\"AllScAIP env (ROCm) — FAIRChem scalable attention MLIP trained on OMol25.\n\nIdentical to nvidia_configs/allscaip.py except torch resolves from the ROCm\nwheel index; AllScAIP has no flash-attention or custom CUDA kernels, so\nnothing else changes. On ROCm fairchem falls back to math attention for\ngradient-based forces: same numbers, slower attention path.\n\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"allscaip-md-conserving-all-omol\": \"allscaip-md-conserving-all-omol\",\n \"allscaip-md-direct-all-omol\": \"allscaip-md-direct-all-omol\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"allscaip:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(CHECKPOINTS[checkpoint], device=device)\n return FAIRChemCalculator(predictor)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=device)\n return FAIRChemCalculator(predictor)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","pytorch-triton-rocm":"3.4.0","rootstock":"1.4.3","torch":"2.8.0+rocm6.4"},"lock_hash":"sha256:5bfd2fc9891e838f114e174ceea24cd0ba971c79711abcf5de9f3e333fd418fa","checkpoints":{"allscaip-md-direct-all-omol":{"fetched_at":"2026-08-03T17:51:42.913179+00:00","verified_at":"2026-10-12T06:41:33.315073+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/5517360cf5a23ca141f1b51b708d131289e35080bfa79f27b990695e3e78d18e","size":694955368},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-12T06:41:33.315077+00:00"},"allscaip-md-conserving-all-omol":{"fetched_at":"2026-08-03T17:51:44.672552+00:00","verified_at":"2026-10-12T06:41:33.315065+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/14fe8626e9e75337769ecf98a765cdc02667763ec4dc1103d958e1c9f56b26e9","size":688432239},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-12T06:41:33.315070+00:00"},"allscaip:custom":{"fetched_at":null,"verified_at":"2026-10-12T06:41:33.315130+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"mace_polar":{"built_at":"2026-08-03T17:38:25.860597+00:00","source_hash":"sha256:965ffabb3457ef81e1c9d47173072ecfcbb0bf3d539d9340afb54a839574a70a","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"ase>=3.22\",\n# \"torch>=2.4.0,<2.10\",\n# # 0.3.16 is the first PyPI release with mace_polar().\n# \"mace-torch>=0.3.16\",\n# # PolarMACE imports graph_longrange at runtime; the distribution is named\n# # graph-longrange and exists only as this git repo (no PyPI release).\n# \"graph-longrange @ git+https://github.com/WillBaldwin0/graph_electrostatics.git\",\n# # torch's ROCm wheels depend on this; it lives only on the ROCm\n# # index, so it must be a direct dep for [tool.uv.sources] to route it.\n# \"pytorch-triton-rocm\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-rocm\" }\n# pytorch-triton-rocm = { index = \"pytorch-rocm\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-rocm\"\n# url = \"https://download.pytorch.org/whl/rocm6.4\"\n# explicit = true\n# ///\n\"\"\"MACE-POLAR env (ROCm) — electrostatic/polarizable MACE foundation models.\n\nIdentical to nvidia_configs/mace_polar.py except torch resolves from the ROCm\nwheel index. Kept separate from the stable `mace` env because of the extra\ngit-only graph-longrange dependency.\n\nPOLAR checkpoints expect `charge`, `spin`, and `external_field` in atoms.info.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-polar-1-s\": \"polar-1-s\",\n \"mace-polar-1-m\": \"polar-1-m\",\n \"mace-polar-1-l\": \"polar-1-l\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace-polar:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from mace.calculators import mace_polar\n\n return mace_polar(model=CHECKPOINTS[checkpoint], device=device, default_dtype=\"float32\")\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): mace_polar() accepts a\n # weights file directly, keeping the PolarMACE model-type wiring.\n from mace.calculators import mace_polar\n\n return mace_polar(model=path, device=device, default_dtype=\"float32\")\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","pytorch-triton-rocm":"3.5.1","rootstock":"1.4.3","torch":"2.9.1+rocm6.4"},"lock_hash":"sha256:418448181f95203fb6b46d07bfacc06eb4be206d7aff5db9bbbb87dc37c21fc1","checkpoints":{"mace-polar-1-m":{"fetched_at":"2026-08-03T17:58:22.377431+00:00","verified_at":"2026-10-12T06:41:33.315093+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Mmodel","size":68133235}],"weights_recorded_at":"2026-10-12T06:41:33.315099+00:00"},"mace-polar-1-l":{"fetched_at":"2026-08-03T17:58:22.706652+00:00","verified_at":"2026-10-12T06:41:33.315102+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Lmodel","size":129808715}],"weights_recorded_at":"2026-10-12T06:41:33.315105+00:00"},"mace-polar-1-s":{"fetched_at":"2026-08-03T17:58:58.732967+00:00","verified_at":"2026-10-12T06:41:33.315080+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Smodel","size":33375439}],"weights_recorded_at":"2026-10-12T06:41:33.315088+00:00"},"mace-polar:custom":{"fetched_at":null,"verified_at":"2026-10-12T06:41:33.315133+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"orbmol":{"built_at":"2026-08-12T21:31:11.658887+00:00","source_hash":"sha256:435b33b70b0c798f43b93f421d5310063df3a530313f5154d274c466aec564cc","source":"# /// script\n# requires-python = \">=3.12,<3.13\"\n# dependencies = [\n# # 0.7.0 (2026-05-26) introduces orbmol_v2 and the orbmol-v1-* aliases\n# # (orbmol-v1-conservative == orb-v3-conservative-omol). Upstream still\n# # pins dm-tree==0.1.8 (no cp313 wheel), hence the <3.13 cap.\n# \"orb-models>=0.7,<0.8\",\n# \"ase>=3.25\",\n# \"torch>=2.8,<3.0\",\n# \"cached_path==1.8.10\",\n# # torch's ROCm wheels depend on this; it lives only on the ROCm\n# # index, so it must be a direct dep for [tool.uv.sources] to route it.\n# \"pytorch-triton-rocm\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-rocm\" }\n# pytorch-triton-rocm = { index = \"pytorch-rocm\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-rocm\"\n# url = \"https://download.pytorch.org/whl/rocm6.4\"\n# explicit = true\n# ///\n\"\"\"OrbMol env (ROCm) — Orbital Materials' molecular potentials (OMol25/OPoly26).\n\norbmol-v1-conservative is upstream's alias for orb-v3-conservative-omol;\norbmol-v2 (2026-05) adds learnable long-range electrostatics (CoulombModule).\n\nROCm caveats (orb-models 0.7 hard-depends on nvalchemi-toolkit-ops, NVIDIA\nWarp/CUDA kernels — pure-python wheel, so it *installs* fine here):\n\n- Neighbor lists: the 0.7 default edge_method is knn_alchemi (Warp/CUDA);\n torch ROCm reports device.type == \"cuda\", so the default would try to\n launch CUDA kernels on an MI250X. knn_scipy is rejected upstream off-CPU,\n so we pin edge_method=\"knn_brute_force\" (pure torch cdist/topk, runs on\n the device; O(N^2) but these are molecular systems) — override via setup\n kwarg if a ROCm-capable path appears. Upstream deprecates it in favor of\n knn_alchemi, which is not an option here.\n- torch.compile defaults off (compile=False): the triton-ROCm inductor\n path is unvalidated on MI250X; pass compile=True to opt in.\n- orbmol-v2 electrostatics: non-periodic systems use a pure-torch direct\n Coulomb sum (fine on ROCm). Periodic systems go through nvalchemiops\n Particle Mesh Ewald — expect failure on Frontier. Molecules only.\n\nCharge/spin: these are OMol-style conditioned models — ORBCalculator raises if\natoms.info lacks \"charge\"/\"spin\" (spin = multiplicity). Set them per structure\nvia atoms.info; absent that, we default to neutral singlet (charge=0, spin=1).\n\"\"\"\n\nimport os\nimport shutil\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"orbmol-v1-conservative\": \"orbmol-v1-conservative\",\n \"orbmol-v2\": \"orbmol-v2\",\n # Your own fine-tuned OrbMol weights: pair with weights= (loaded via\n # setup_from_path); pass arch= to pick the base architecture.\n \"orbmol:custom\": None,\n}\n\n\ndef _default_weights_url(load_fn) -> str:\n \"\"\"The upstream URL baked into the loader's ``weights_path`` default.\"\"\"\n import inspect\n\n default = inspect.signature(load_fn).parameters[\"weights_path\"].default\n if not isinstance(default, str) or not default.startswith((\"http://\", \"https://\")):\n raise RuntimeError(\n f\"{load_fn.__name__} has no URL default for weights_path \"\n f\"(got {default!r}); update this env file for the installed orb-models\"\n )\n return default\n\n\ndef _local_weights_path(url: str) -> Path:\n \"\"\"Where the checkpoint lives in the shared model cache.\"\"\"\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"orb\" / os.path.basename(url)\n\n\ndef _fetch(url: str, dest: Path) -> None:\n \"\"\"Download ``url`` to ``dest`` atomically (tmp file + rename).\"\"\"\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp = dest.with_name(f\"{dest.name}.tmp.{os.getpid()}\")\n try:\n with urllib.request.urlopen(url) as resp, open(tmp, \"wb\") as out:\n shutil.copyfileobj(resp, out)\n os.replace(tmp, dest)\n finally:\n tmp.unlink(missing_ok=True)\n\n\ndef _make_calculator(orbff, atoms_adapter, device, edge_method, **kwargs):\n from ase.calculators.calculator import all_changes\n from orb_models.forcefield.inference.calculator import ORBCalculator\n\n class OrbMolCalculator(ORBCalculator):\n \"\"\"ORBCalculator that defaults missing charge/spin to neutral singlet.\"\"\"\n\n def calculate(self, atoms=None, properties=None, system_changes=all_changes):\n if atoms is not None:\n atoms.info.setdefault(\"charge\", 0)\n atoms.info.setdefault(\"spin\", 1)\n super().calculate(atoms, properties, system_changes)\n\n return OrbMolCalculator(\n orbff, atoms_adapter, device=device, edge_method=edge_method, **kwargs\n )\n\n\ndef setup(\n checkpoint: str,\n device: str = \"cuda\",\n precision: str = \"float32-high\",\n compile: bool = False,\n edge_method: str = \"knn_brute_force\",\n **kwargs,\n):\n # Extra **kwargs go to ORBCalculator (e.g. max_num_neighbors=,\n # half_supercell=); precision/compile go to the checkpoint loader.\n import torch\n from orb_models.forcefield import pretrained\n\n fn_name = CHECKPOINTS[checkpoint].replace(\"-\", \"_\")\n load_fn = getattr(pretrained, fn_name)\n\n # orb-models resolves its default weights URL through `cached_path`, which\n # write-locks its cache dir even on warm hits — EACCES for anyone who can't\n # write the shared install. Handed a *local* path instead, cached_path\n # returns it without locking. So the weights are pre-fetched into the\n # shared model cache at `rootstock add` time (maintainer, cache writable)\n # and every later serve loads that file.\n url = _default_weights_url(load_fn)\n weights = _local_weights_path(url)\n if not weights.exists():\n _fetch(url, weights)\n\n orbff, atoms_adapter = load_fn(\n weights_path=str(weights),\n device=torch.device(device),\n precision=precision,\n compile=compile,\n )\n return _make_calculator(\n orbff, atoms_adapter, torch.device(device), edge_method, **kwargs\n )\n\n\ndef setup_from_path(\n path: str,\n device: str = \"cuda\",\n arch: str = \"orbmol-v2\",\n precision: str = \"float32-high\",\n compile: bool = False,\n edge_method: str = \"knn_brute_force\",\n **kwargs,\n):\n # Custom checkpoints (`:custom` ids with user weights). A weights file\n # doesn't say which architecture produced it, so `arch` names the\n # pretrained loader to instantiate — pass the right one at call time\n # (setup_kwargs={\"arch\": ...} / --kwarg arch=...). Handing the loader a\n # local path also means no network and no cached_path locking (see setup()).\n import torch\n from orb_models.forcefield import pretrained\n\n fn_name = arch.replace(\"-\", \"_\").replace(\":\", \"_\")\n try:\n load_fn = getattr(pretrained, fn_name)\n except AttributeError:\n raise ValueError(\n f\"unknown OrbMol architecture {arch!r}; expected a loader name from \"\n f\"orb_models.forcefield.pretrained, e.g. orbmol-v2, orbmol-v1-conservative\"\n ) from None\n\n orbff, atoms_adapter = load_fn(\n weights_path=path,\n device=torch.device(device),\n precision=precision,\n compile=compile,\n )\n return _make_calculator(\n orbff, atoms_adapter, torch.device(device), edge_method, **kwargs\n )\n\n# The smoke-test's weights= leg borrows the first fetched family checkpoint\n# (CHECKPOINTS order -> orbmol-v1-conservative) and reloads its weights through\n# setup_from_path, whose default arch is orbmol-v2. The v2 energy head has 257\n# input columns (CoulombModule adds one) against v1's 256, so the default arch\n# fails with a state_dict size mismatch (Frontier, 2026-08-24). Name the\n# matching architecture for the leg; your own weights still take arch= per call.\nVERIFY_KWARGS = {\n \"orbmol:custom\": {\"arch\": \"orbmol-v1-conservative\"},\n}\n","python_requires":">=3.12,<3.13","dependencies":{"ase":"3.29.0","cached-path":"1.8.10","orb-models":"0.7.0","pytorch-triton-rocm":"3.5.1","rootstock":"1.6.2","torch":"2.9.1+rocm6.4"},"lock_hash":"sha256:f352076fde94ccbdceef25206ed4951664b1f555f5b74248d23e8206efc3efca","checkpoints":{"orbmol-v2":{"fetched_at":"2026-08-12T20:24:41.330460+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: RuntimeError: Worker calculation failed:\nTraceback (most recent call last):\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/rootstock/worker.py\", line 275, in run\n energy, forces, virial = self._calculate(positions, cell)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/rootstock/worker.py\", line 182, in _calculate\n energy = atoms.get_potential_energy()\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/ase/atoms.py\", line 1971, in get_potential_energy\n energy = self._calc.get_potential_energy(self)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/ase/calculators/abc.py\", line 25, in get_potential_energy\n return self.get_property(name, atoms)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/ase/calculators/calculator.py\", line 517, in get_property\n self.calculate(atoms, [name], system_changes)\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/env_source.py\", line 107, in calculate\n super().calculate(atoms, properties, system_changes)\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/orb_models/forcefield/inference/calculator.py\", line 108, in calculate\n out = self.model.predict(batch) # type: ignore\n ^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/orb_models/forcefield/models/conservative_regressor.py\", line 271, in predict\n preds = self(batch)\n ^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/torch/nn/modules/module.py\", line 1775, in _wrapped_call_impl\n return self._call_impl(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/torch/nn/modules/module.py\", line 1786, in _call_impl\n return forward_call(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/orb_models/forcefield/models/conservative_regressor.py\", line 211, in forward\n coulomb_energy, coulomb_explicit_forces, coulomb_explicit_virial = self.coulomb_module(\n ^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/torch/nn/modules/module.py\", line 1775, in _wrapped_call_impl\n return self._call_impl(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/torch/nn/modules/module.py\", line 1786, in _call_impl\n return forward_call(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/orb_models/forcefield/models/coulomb_module.py\", line 129, in forward\n p_energy, p_forces, p_virial = self._particle_mesh_ewald(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/orb_models/forcefield/models/coulomb_module.py\", line 203, in _particle_mesh_ewald\n _estimate_pme_params_and_neighbors(\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py\", line 1044, in _fn\n return fn(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/orb_models/forcefield/models/coulomb_module.py\", line 285, in _estimate_pme_params_and_neighbors\n _compute_neighbor_list_with_fallback(\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/orb_models/common/atoms/graph_featurization.py\", line 742, in _compute_neighbor_list_with_fallback\n neighbor_matrix, num_neighbors, neighbor_shift_matrix = nva_neighbor_list(\n ^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/nvalchemiops/torch/neighbors/__init__.py\", line 324, in neighbor_list\n return batch_naive_neighbor_list(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/nvalchemiops/torch/neighbors/batch_naive.py\", line 430, in batch_naive_neighbor_list\n compute_naive_num_shifts(cell, cutoff, pbc)\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/nvalchemiops/torch/neighbors/neighbor_utils.py\", line 90, in compute_naive_num_shifts\n wp_device = wp.device_from_torch(device)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/warp/_src/torch.py\", line 39, in device_from_torch\n return warp._src.context.runtime.cuda_devices[torch_device.index]\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^\nIndexError: list index out of range\n","weight_files":[{"path":"cache/orb/orbmol-v2-teqabfhg-20260523.ckpt","size":103552643}],"weights_recorded_at":"2026-10-12T06:41:33.315117+00:00"},"orbmol-v1-conservative":{"fetched_at":"2026-08-12T20:24:41.375905+00:00","verified_at":"2026-10-12T06:41:33.315108+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-omol-20250820.ckpt","size":103417970}],"weights_recorded_at":"2026-10-12T06:41:33.315111+00:00"},"orbmol:custom":{"fetched_at":null,"verified_at":"2026-10-12T06:41:33.315135+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}}}},{"schema_version":5,"cluster":"perlmutter","root":"/global/cfs/cdirs/m5268/rootstock","maintainer":{"name":"Owen Price Skelly","email":"owenpriceskelly@uchicago.edu"},"rootstock_version":"1.6.4","python_version":"unknown","last_updated":"2026-10-25T19:14:16.891189+00:00","environments":{"mace":{"built_at":"2026-08-07T18:08:38.875863+00:00","source_hash":"sha256:eb8e9790141ad43e2735656a942a3b3af3a551ae6e89aef0b99c72ff89f80360","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # 0.3.15+ needed for the mh-1 registry entry (matpes needs 0.3.13,\n# # omol needs 0.3.14, mpa-0 needs 0.3.10).\n# \"mace-torch>=0.3.15\",\n# \"ase>=3.22\",\n# # 2.4.1 is explicitly unsupported by mace-torch.\n# \"torch>=2.4.0,!=2.4.1,<2.10\",\n# ]\n# ///\n\"\"\"MACE env — hosts MACE-MP-0, MACE-OFF23, MPA-0, MATPES, MH-1, and OMOL checkpoints.\n\nAll ship in the same `mace-torch` package, so they share an environment.\nUpstream-string routing in CHECKPOINTS: an `off:` prefix routes to mace_off()\nand an `omol:` prefix to mace_omol() (float64, molecules only); an `mh:`\nprefix marks a multi-head model (float64, per the MACE-MH-1 model card).\n\nMulti-head checkpoints select a head via the `head` kwarg on setup()\n(setup_kwargs={\"head\": ...} / --kwarg head=...), named by upstream's training\ncorpus — see MH1_HEADS; omat_pbe is the default.\n\nThe OMOL checkpoint expects `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-mp-0-small\": \"small\",\n \"mace-mp-0-medium\": \"medium\",\n \"mace-mp-0-large\": \"large\",\n \"mace-off23-small\": \"off:small\",\n \"mace-off23-medium\": \"off:medium\",\n \"mace-off23-large\": \"off:large\",\n # Only a medium MPA-0 has been released, but upstream names the weights\n # file mace-mpa-0-medium.model — keep the size explicit like mace-mp-0.\n \"mace-mpa-0-medium\": \"medium-mpa-0\",\n \"mace-matpes-r2scan-0\": \"mace-matpes-r2scan-0\",\n # One entry per weights file: MH-1's heads are selected by setup(head=...).\n \"mace-mh-1\": \"mh:mh-1\",\n # Only the extra-large OMOL model has been released.\n \"mace-omol-0-extra-large\": \"omol:extra_large\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace:custom\": None,\n}\n\n# MH-1's heads, named by training corpus. The released weights file is the\n# authority (mace_select_head --list_heads; the model card also lists a\n# rgd1_b3lyp head, but that shipped only in mh-0 — ACEsuit/mace#1462).\n# Validated here because upstream only warns on an unknown head and silently\n# falls back to the last one.\nMH1_HEADS = (\n \"omat_pbe\",\n \"omol\",\n \"spice_wB97M\",\n \"oc20_usemppbe\",\n \"mp_pbe_refit_add\",\n \"matpes_r2scan\",\n)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", head: str | None = None):\n arg = CHECKPOINTS[checkpoint]\n if arg.startswith(\"mh:\"):\n head = head or \"omat_pbe\"\n if head not in MH1_HEADS:\n raise ValueError(f\"unknown head {head!r}; expected one of {', '.join(MH1_HEADS)}\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg[3:], device=device, default_dtype=\"float64\", head=head)\n if head is not None:\n raise ValueError(f\"'head' selects a head of a multi-head model; {checkpoint} has one head\")\n if arg.startswith(\"off:\"):\n from mace.calculators import mace_off\n\n return mace_off(model=arg[4:], device=device, default_dtype=\"float32\")\n if arg.startswith(\"omol:\"):\n from mace.calculators import mace_omol\n\n return mace_omol(model=arg[5:], device=device, default_dtype=\"float64\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg, device=device, default_dtype=\"float32\")\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", default_dtype: str = \"float32\", head: str | None = None):\n # Custom checkpoints (`:custom` ids with user weights): fine-tunes load through\n # MACECalculator directly — the mp/off dispatch in setup() only exists\n # to pick which pretrained file to download. `head` is for fine-tunes that\n # keep multiple heads; single-head weights load without it.\n from mace.calculators import MACECalculator\n\n return MACECalculator(model_paths=path, device=device, default_dtype=default_dtype, head=head)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"1.6.0","torch":"2.9.1"},"lock_hash":"sha256:c570c03b9916ce0e672c3eb8408b9aa605abc4fb2cf1aae1c9e846e42471cc9c","checkpoints":{"mace-mh-1":{"fetched_at":"2026-08-07T18:38:54.862793+00:00","verified_at":"2026-10-25T19:13:04.576764+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macemh1model","size":59208139}],"weights_recorded_at":"2026-10-25T19:13:04.576768+00:00"},"mace-matpes-r2scan-0":{"fetched_at":"2026-08-07T18:38:55.030933+00:00","verified_at":"2026-10-25T19:13:04.576751+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEmatpesr2scanomatftmodel","size":79470738}],"weights_recorded_at":"2026-10-25T19:13:04.576757+00:00"},"mace-mp-0-medium":{"fetched_at":"2026-08-07T18:39:11.227407+00:00","verified_at":"2026-10-25T19:13:04.576698+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231203mace128L1_epoch199model","size":44422970}],"weights_recorded_at":"2026-10-25T19:13:04.576704+00:00"},"mace-mp-0-large":{"fetched_at":"2026-08-07T18:39:12.628241+00:00","verified_at":"2026-10-25T19:13:04.576709+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE_MPtrj_20229model","size":133803220}],"weights_recorded_at":"2026-10-25T19:13:04.576714+00:00"},"mace-mp-0-small":{"fetched_at":"2026-08-07T18:39:27.560268+00:00","verified_at":"2026-10-25T19:13:04.576666+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231210mace128L0_energy_epoch249model","size":32581838}],"weights_recorded_at":"2026-10-25T19:13:04.576687+00:00"},"mace-mpa-0-medium":{"fetched_at":"2026-08-07T18:39:27.578362+00:00","verified_at":"2026-10-25T19:13:04.576744+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macempa0mediummodel","size":79462305}],"weights_recorded_at":"2026-10-25T19:13:04.576747+00:00"},"mace-off23-large":{"fetched_at":"2026-08-07T18:39:34.646763+00:00","verified_at":"2026-10-25T19:13:04.576734+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_large.model","size":55492786}],"weights_recorded_at":"2026-10-25T19:13:04.576737+00:00"},"mace-off23-medium":{"fetched_at":"2026-08-07T18:39:39.457787+00:00","verified_at":"2026-10-25T19:13:04.576726+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_medium.model","size":18350596}],"weights_recorded_at":"2026-10-25T19:13:04.576730+00:00"},"mace-off23-small":{"fetched_at":"2026-08-07T18:39:41.383497+00:00","verified_at":"2026-10-25T19:13:04.576718+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_small.model","size":7347350}],"weights_recorded_at":"2026-10-25T19:13:04.576722+00:00"},"mace-omol-0-extra-large":{"fetched_at":"2026-08-07T18:39:43.466419+00:00","verified_at":"2026-10-25T19:13:04.576773+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-omol-0-extra-large-1024.model","size":422242640}],"weights_recorded_at":"2026-10-25T19:13:04.576776+00:00"},"mace:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:13:04.577072+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"allscaip":{"built_at":"2026-08-07T18:11:29.802163+00:00","source_hash":"sha256:bde4548ac275b9af7d17c3103e3da8008038f306ab5cb03e898c95ce6ad994d9","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"AllScAIP env — FAIRChem scalable attention MLIP trained on OMol25.\n\nallscaip-md-conserving-all-omol is an energy-conserving, all-to-all node\nattention model served through fairchem-core's get_predict_unit — the same\nAPI as eSEN. fairchem v2 carries the architecture in-package, so no\nflash-attention or custom CUDA kernels are needed.\n\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"allscaip-md-conserving-all-omol\": \"allscaip-md-conserving-all-omol\",\n \"allscaip-md-direct-all-omol\": \"allscaip-md-direct-all-omol\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"allscaip:custom\": None,\n}\n\n\ndef _fairchem_device(device: str) -> str:\n \"\"\"Translate an indexed device (\"cuda:2\") into what fairchem v2 accepts.\n\n MLIPPredictUnit._setup_device asserts `device in [\"cpu\", \"cuda\"]` and then\n resolves the real GPU itself via get_device_for_local_rank(), which returns\n f\"cuda:{torch.cuda.current_device()}\". So an index has to travel through\n torch's current-device state, not the argument. Verifying several\n checkpoints at once on a multi-GPU node hands each worker \"cuda:N\" — that\n killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync\n (4x A100, VERIFY_JOBS=4), while single-GPU Sophia never hit it.\n \"\"\"\n if device.startswith(\"cuda:\"):\n import torch\n\n torch.cuda.set_device(int(device.split(\":\", 1)[1]))\n return \"cuda\"\n return device\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(\n CHECKPOINTS[checkpoint], device=_fairchem_device(device)\n )\n return FAIRChemCalculator(predictor)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=_fairchem_device(device))\n return FAIRChemCalculator(predictor)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.6.0","torch":"2.8.0"},"lock_hash":"sha256:37e7cd799c8736978027d1cb9918aa06410e3bac86a82dfba59cbf31c5fd7eab","checkpoints":{"allscaip-md-conserving-all-omol":{"fetched_at":"2026-08-07T19:00:42.320790+00:00","verified_at":"2026-10-25T19:13:04.576781+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/14fe8626e9e75337769ecf98a765cdc02667763ec4dc1103d958e1c9f56b26e9","size":688432239},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:13:04.576786+00:00"},"allscaip-md-direct-all-omol":{"fetched_at":"2026-08-07T19:00:39.080181+00:00","verified_at":"2026-10-25T19:13:04.576790+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/5517360cf5a23ca141f1b51b708d131289e35080bfa79f27b990695e3e78d18e","size":694955368},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:13:04.576794+00:00"},"allscaip:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:13:04.577075+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"ani":{"built_at":"2026-08-07T18:11:31.469453+00:00","source_hash":"sha256:16ebe0eda0074fa275e0ce2c2d68c84dc10d2069ef0290d19802bd46bcf96d0e","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"torchani>=2.2\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n# ///\n\"\"\"\nANI-2x environment for Rootstock.\n\nANI-2x is a neural network potential for organic molecules containing\nH, C, N, O, F, S, Cl. It is not a universal potential — do not use it\nfor inorganic or periodic systems.\n\nModels:\n - \"ANI2x\": ANI-2x ensemble (default, 8 networks)\n - \"ANI1ccx\": ANI-1ccx, trained on CCSD(T)/CBS data (H, C, N, O only)\n - \"ANI1x\": ANI-1x (H, C, N, O only)\n\"\"\"\n\nCHECKPOINTS = {\n \"ani-2x\": \"ANI2x\",\n \"ani-1ccx\": \"ANI1ccx\",\n \"ani-1x\": \"ANI1x\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an ANI calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n import torchani\n\n model_map = {\n \"ANI2x\": torchani.models.ANI2x,\n \"ANI1ccx\": torchani.models.ANI1ccx,\n \"ANI1x\": torchani.models.ANI1x,\n }\n model = CHECKPOINTS[checkpoint]\n\n return model_map[model](periodic_table_index=True).to(device).ase()\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.6.0","torch":"2.13.0","torchani":"2.8.4"},"lock_hash":"sha256:387dd9038769210056cbc5601717edbb8d2f9b55a8f0b184d513fb9bd10b0bd5","checkpoints":{"ani-1x":{"fetched_at":"2026-08-07T18:16:55.906080+00:00","verified_at":"2026-10-25T19:13:04.576812+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.local/share/Torchani/StateDicts/.cache/huggingface/download/ani1x_state_dict.pt.metadata","size":124},{"path":"home/.local/share/Torchani/StateDicts/ani1x_state_dict.pt","size":10546029}],"weights_recorded_at":"2026-10-25T19:13:04.576816+00:00"},"ani-1ccx":{"fetched_at":"2026-08-07T18:16:56.177074+00:00","verified_at":"2026-10-25T19:13:04.576805+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.local/share/Torchani/StateDicts/.cache/huggingface/download/ani1ccx_state_dict.pt.metadata","size":124},{"path":"home/.local/share/Torchani/StateDicts/ani1ccx_state_dict.pt","size":10546029}],"weights_recorded_at":"2026-10-25T19:13:04.576809+00:00"},"ani-2x":{"fetched_at":"2026-08-07T18:17:09.535866+00:00","verified_at":"2026-10-25T19:13:04.576798+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.local/share/Torchani/StateDicts/.cache/huggingface/download/ani2x_state_dict.pt.metadata","size":125},{"path":"home/.local/share/Torchani/StateDicts/ani2x_state_dict.pt","size":54982957}],"weights_recorded_at":"2026-10-25T19:13:04.576802+00:00"}}},"chgnet":{"built_at":"2026-08-07T18:11:42.922865+00:00","source_hash":"sha256:b6ea6ac08a17cb3dc17d5455571dc5c21d513cd802916015de7f272e83eefe37","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"chgnet>=0.3.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"CHGNet env — hosts pretrained charge-informed universal potentials.\"\"\"\n\nCHECKPOINTS = {\n \"chgnet-default\": \"chgnet-default\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"chgnet:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a CHGNet calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\")\n\n Returns:\n ASE-compatible calculator\n \"\"\"\n from chgnet.model import CHGNet, CHGNetCalculator\n\n model_name = CHECKPOINTS[checkpoint]\n model = CHGNet.load() if model_name == \"chgnet-default\" else CHGNet.load(model_name)\n return CHGNetCalculator(model=model, use_device=device)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # CHGNet.from_file, not the named-model CHGNet.load() setup() uses.\n from chgnet.model import CHGNet, CHGNetCalculator\n\n model = CHGNet.from_file(path)\n return CHGNetCalculator(model=model, use_device=device)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","chgnet":"0.4.2","rootstock":"1.6.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:aa0b72f3991bfa8c4d8465982e75df9635fa1b626ece8e541900f0ed8637abb6","checkpoints":{"chgnet-default":{"fetched_at":"2026-08-07T18:18:20.972193+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 38686 files, 7.6 GB in 577.8s; weights: none recorded","weight_files":null,"weights_recorded_at":null}}},"dimenet":{"built_at":"2026-08-07T18:12:13.887953+00:00","source_hash":"sha256:84b1aa2101b0cdfa5f74270cfc5d7a53854a84d4b8f68945e11692563c24b322","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\nDimeNet++ environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 DimeNet++ checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"DimeNet++-S2EF-OC20-All\": default\n - \"DimeNet++-S2EF-OC20-20M\"\n - \"DimeNet++-S2EF-OC20-2M\"\n - \"DimeNet++-S2EF-OC20-200k\"\n\"\"\"\n\nCHECKPOINTS = {\n \"dimenet-plus-plus-s2ef-oc20-all\": \"DimeNet++-S2EF-OC20-All\",\n \"dimenet-plus-plus-s2ef-oc20-20m\": \"DimeNet++-S2EF-OC20-20M\",\n \"dimenet-plus-plus-s2ef-oc20-2m\": \"DimeNet++-S2EF-OC20-2M\",\n \"dimenet-plus-plus-s2ef-oc20-200k\": \"DimeNet++-S2EF-OC20-200k\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a DimeNet++ OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.6.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"dimenet-plus-plus-s2ef-oc20-20m":{"fetched_at":"2026-08-07T18:19:32.658646+00:00","verified_at":"2026-10-25T19:13:04.576824+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/dimenetpp_20M.pt","size":21841849}],"weights_recorded_at":"2026-10-25T19:13:04.576827+00:00"},"dimenet-plus-plus-s2ef-oc20-2m":{"fetched_at":"2026-08-07T18:33:36.719332+00:00","verified_at":"2026-10-25T19:13:04.576831+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/dimenetpp_2M.pt","size":21841593}],"weights_recorded_at":"2026-10-25T19:13:04.576834+00:00"},"dimenet-plus-plus-s2ef-oc20-all":{"fetched_at":"2026-08-07T18:33:37.904281+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 44478 files, 6.4 GB in 544.0s; weights: 22 MB (manifest)","weight_files":[{"path":"cache/dimenetpp_all.pt","size":21841849}],"weights_recorded_at":"2026-10-18T16:47:56.675519+00:00"},"dimenet-plus-plus-s2ef-oc20-200k":{"fetched_at":"2026-08-07T18:33:38.187393+00:00","verified_at":"2026-10-25T19:13:04.576838+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/dimenetpp_200k.pt","size":21841593}],"weights_recorded_at":"2026-10-25T19:13:04.576841+00:00"}}},"equiformer":{"built_at":"2026-08-07T18:12:16.198733+00:00","source_hash":"sha256:ac01aec3be646cbd2c63a2e9f6a0cc8308aa03f5d8528150c5f664b2aa06426f","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\nEquiformerV2 environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 EquiformerV2 checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"EquiformerV2-153M-S2EF-OC20-All+MD\": default\n - \"EquiformerV2-31M-S2EF-OC20-All+MD\"\n - \"EquiformerV2-83M-S2EF-OC20-2M\"\n\"\"\"\n\nCHECKPOINTS = {\n \"equiformer-v2-153m-s2ef-oc20-all-md\": \"EquiformerV2-153M-S2EF-OC20-All+MD\",\n \"equiformer-v2-31m-s2ef-oc20-all-md\": \"EquiformerV2-31M-S2EF-OC20-All+MD\",\n \"equiformer-v2-83m-s2ef-oc20-2m\": \"EquiformerV2-83M-S2EF-OC20-2M\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"equiformer:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an EquiformerV2 OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): OCPCalculator loads a\n # checkpoint file natively — this is setup() minus the registry download.\n from fairchem.core import OCPCalculator\n\n return OCPCalculator(checkpoint_path=path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.6.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"equiformer-v2-31m-s2ef-oc20-all-md":{"fetched_at":"2026-08-07T18:34:48.975897+00:00","verified_at":"2026-10-25T19:13:04.576849+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_31M_ec4_allmd.pt","size":125685325}],"weights_recorded_at":"2026-10-25T19:13:04.576852+00:00"},"equiformer-v2-83m-s2ef-oc20-2m":{"fetched_at":"2026-08-07T18:34:56.875435+00:00","verified_at":"2026-10-25T19:13:04.576856+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_83M_2M.pt","size":336843381}],"weights_recorded_at":"2026-10-25T19:13:04.576859+00:00"},"equiformer-v2-153m-s2ef-oc20-all-md":{"fetched_at":"2026-08-07T18:35:08.566741+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 44478 files, 7.0 GB in 529.9s; weights: 620 MB (manifest)\n/global/cfs/cdirs/m5268/rootstock/envs/equiformer/lib/python3.11/site-packages/fairchem/core/models/scn/spherical_harmonics.py:23: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n _Jd = torch.load(os.path.join(os.path.dirname(__file__), \"Jd.pt\"))\n/global/cfs/cdirs/m5268/rootstock/envs/equiformer/lib/python3.11/site-packages/fairchem/core/models/equiformer_v2/wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n _Jd = torch.load(os.path.join(os.path.dirname(__file__), \"Jd.pt\"))\n/global/cfs/cdirs/m5268/rootstock/envs/equiformer/lib/python3.11/site-packages/fairchem/core/models/escn/so3.py:23: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n _Jd = torch.load(os.path.join(os.path.dirname(__file__), \"Jd.pt\"))","weight_files":[{"path":"cache/eq2_153M_ec4_allmd.pt","size":619855697}],"weights_recorded_at":"2026-10-18T16:47:56.675548+00:00"},"equiformer:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: weights= run diverges from equiformer-v2-153m-s2ef-oc20-all-md: energy differs by 3.455e-04 eV (tolerance 1e-04)","weight_files":null,"weights_recorded_at":null}}},"escn":{"built_at":"2026-08-07T18:12:17.215004+00:00","source_hash":"sha256:e1301f3e20ac38e355e0b36e2ab9fb0f0a78f58059b5e2c57ab280aa9dad7ff5","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\neSCN environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 eSCN checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"eSCN-L6-M2-Lay12-S2EF-OC20-All+MD\": default\n - \"eSCN-L6-M3-Lay20-S2EF-OC20-All+MD\"\n - \"eSCN-L6-M2-Lay12-S2EF-OC20-2M\"\n - \"eSCN-L4-M2-Lay12-S2EF-OC20-2M\"\n\"\"\"\n\nCHECKPOINTS = {\n \"escn-l6-m2-lay12-s2ef-oc20-all-md\": \"eSCN-L6-M2-Lay12-S2EF-OC20-All+MD\",\n \"escn-l6-m3-lay20-s2ef-oc20-all-md\": \"eSCN-L6-M3-Lay20-S2EF-OC20-All+MD\",\n \"escn-l6-m2-lay12-s2ef-oc20-2m\": \"eSCN-L6-M2-Lay12-S2EF-OC20-2M\",\n \"escn-l4-m2-lay12-s2ef-oc20-2m\": \"eSCN-L4-M2-Lay12-S2EF-OC20-2M\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"escn:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an eSCN OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): OCPCalculator loads a\n # checkpoint file natively — this is setup() minus the registry download.\n from fairchem.core import OCPCalculator\n\n return OCPCalculator(checkpoint_path=path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.6.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"escn-l6-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-08-07T18:35:27.546087+00:00","verified_at":"2026-10-25T19:13:04.576877+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l6_m2_lay12_2M_s2ef.pt","size":207735045}],"weights_recorded_at":"2026-10-25T19:13:04.576880+00:00"},"escn-l6-m2-lay12-s2ef-oc20-all-md":{"fetched_at":"2026-08-07T18:35:29.989827+00:00","verified_at":"2026-10-25T19:13:04.576863+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l6_m2_lay12_all_md_s2ef.pt","size":207735109}],"weights_recorded_at":"2026-10-25T19:13:04.576866+00:00"},"escn-l6-m3-lay20-s2ef-oc20-all-md":{"fetched_at":"2026-08-07T18:35:33.656833+00:00","verified_at":"2026-10-25T19:13:04.576869+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l6_m3_lay20_all_md_s2ef.pt","size":801642787}],"weights_recorded_at":"2026-10-25T19:13:04.576873+00:00"},"escn-l4-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-08-07T18:35:53.977072+00:00","verified_at":"2026-10-25T19:13:04.576884+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l4_m2_lay12_2M_s2ef.pt","size":144761989}],"weights_recorded_at":"2026-10-25T19:13:04.576887+00:00"},"escn:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: weights= run diverges from escn-l6-m2-lay12-s2ef-oc20-all-md: energy differs by 1.327e-03 eV (tolerance 1e-04)","weight_files":null,"weights_recorded_at":null}}},"esen":{"built_at":"2026-08-07T18:12:33.037324+00:00","source_hash":"sha256:d07f4d08c7ce3ed5aa36b795281d2be931b17cd6d342e2b569ef0fd1f53e59db","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"eSEN env — hosts FAIRChem eSEN single-task checkpoints.\n\nfairchem-core v2 dropped the torch-geometric / pyg-find-links install dance\n(simplified here 2026-07-30 to match uma.py/allscaip.py — the find-links were\nalso blocking universal lockfile resolution).\n\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"esen-md-direct-all-omol\": \"esen-md-direct-all-omol\",\n \"esen-sm-conserving-all-omol\": \"esen-sm-conserving-all-omol\",\n \"esen-sm-direct-all-omol\": \"esen-sm-direct-all-omol\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"esen:custom\": None,\n}\n\n\ndef _fairchem_device(device: str) -> str:\n \"\"\"Translate an indexed device (\"cuda:2\") into what fairchem v2 accepts.\n\n MLIPPredictUnit._setup_device asserts `device in [\"cpu\", \"cuda\"]` and then\n resolves the real GPU itself via get_device_for_local_rank(), which returns\n f\"cuda:{torch.cuda.current_device()}\". So an index has to travel through\n torch's current-device state, not the argument. Verifying several\n checkpoints at once on a multi-GPU node hands each worker \"cuda:N\" — that\n killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync\n (4x A100, VERIFY_JOBS=4), while single-GPU Sophia never hit it.\n \"\"\"\n if device.startswith(\"cuda:\"):\n import torch\n\n torch.cuda.set_device(int(device.split(\":\", 1)[1]))\n return \"cuda\"\n return device\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(\n CHECKPOINTS[checkpoint], device=_fairchem_device(device)\n )\n return FAIRChemCalculator(predictor)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=_fairchem_device(device))\n return FAIRChemCalculator(predictor)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.6.0","torch":"2.8.0"},"lock_hash":"sha256:37e7cd799c8736978027d1cb9918aa06410e3bac86a82dfba59cbf31c5fd7eab","checkpoints":{"esen-sm-conserving-all-omol":{"fetched_at":"2026-08-07T19:00:30.424489+00:00","verified_at":"2026-10-25T19:13:04.576892+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/01f63da2d071e39fc46a5f22f8369d0fc9de317ab2ef361a76603f5661238025","size":50958067},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:13:04.576897+00:00"},"esen-md-direct-all-omol":{"fetched_at":"2026-08-07T19:00:36.064518+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 44626 files, 8.3 GB in 574.7s; weights: 406 MB (manifest)","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/e0cd2ac37d2992b385a39778469c443ccbc679c7c4fb3b40b38aa2ee23c29b1d","size":405789591},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-18T16:47:56.675602+00:00"},"esen-sm-direct-all-omol":{"fetched_at":"2026-08-07T19:00:56.382873+00:00","verified_at":"2026-10-25T19:13:04.576900+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/a1b7109bcca6e1288c5f17c68776d3c6a1ab0766f6d7fdd45f901605355d2831","size":50960291},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:13:04.576904+00:00"},"esen:custom":{"fetched_at":null,"verified_at":"2026-10-18T16:47:56.675848+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"m3gnet":{"built_at":"2026-08-07T18:12:50.297885+00:00","source_hash":"sha256:3d22f362208d5a4297208395d7f188c943aee8866ed9f111a4b7c8d3801bbc14","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"chgnet>=0.4.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"\nM3GNet environment for Rootstock — redirected to CHGNet.\n\nM3GNet-MP-2021.2.8-PES (the universal inorganic PES) is no longer accessible\nvia any modern Python package:\n - matgl 2.x (PyG backend) only has TensorNet on HuggingFace (materialyze)\n - matgl 1.x (DGL backend) pointed to a GitHub URL that was removed\n - The original m3gnet package is archived TensorFlow code\n - materialyze HuggingFace has only M3GNet-Eform (formation energy, not PES)\n\nFor universal inorganic PES, use:\n - tensornet.py: TensorNet-MatPES (same authors, newer, better)\n - chgnet.py: CHGNet (charge-informed, strong on magnetic materials)\n - orb.py: Orb v3 (universal, supports periodic systems)\n\nThis file loads CHGNet as the practical substitute for M3GNet-PES.\n\"\"\"\n\nCHECKPOINTS = {\n \"m3gnet-mp-2021-2-8-pes\": \"chgnet-default\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load CHGNet as a substitute for M3GNet-PES.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible CHGNetCalculator.\n \"\"\"\n from chgnet.model import CHGNetCalculator\n\n return CHGNetCalculator(use_device=device)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","chgnet":"0.4.2","rootstock":"1.6.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:1e7fb7c8dc56e7a942e52d5e1d50ff69579f7e581ce3e64c9cc62266f2300a20","checkpoints":{"m3gnet-mp-2021-2-8-pes":{"fetched_at":"2026-08-07T18:39:18.387057+00:00","verified_at":"2026-10-25T19:13:04.576909+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"mace_polar":{"built_at":"2026-08-07T18:12:57.135495+00:00","source_hash":"sha256:8db4cfd29a59c415e3762b04195e30b9e08d5f044ff6aa21ebcf01bcb78bc1ef","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"ase>=3.22\",\n# \"torch>=2.4.0,<2.10\",\n# # 0.3.16 is the first PyPI release with mace_polar().\n# \"mace-torch>=0.3.16\",\n# # PolarMACE imports graph_longrange at runtime; the distribution is named\n# # graph-longrange and exists only as this git repo (no PyPI release).\n# \"graph-longrange @ git+https://github.com/WillBaldwin0/graph_electrostatics.git\",\n# ]\n# ///\n\"\"\"MACE-POLAR env — electrostatic/polarizable MACE foundation models (OMol25).\n\nKept separate from the stable `mace` env because of the extra git-only\ngraph-longrange dependency.\n\nPOLAR checkpoints expect `charge`, `spin`, and `external_field` in atoms.info.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-polar-1-s\": \"polar-1-s\",\n \"mace-polar-1-m\": \"polar-1-m\",\n \"mace-polar-1-l\": \"polar-1-l\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace-polar:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from mace.calculators import mace_polar\n\n return mace_polar(model=CHECKPOINTS[checkpoint], device=device, default_dtype=\"float32\")\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): mace_polar() accepts a\n # weights file directly, keeping the PolarMACE model-type wiring.\n from mace.calculators import mace_polar\n\n return mace_polar(model=path, device=device, default_dtype=\"float32\")\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"1.6.0","torch":"2.9.1"},"lock_hash":"sha256:82da6f7c8c18c84cf9c0bf34830b86029a965670658e380aeeb38d50d38511fd","checkpoints":{"mace-polar-1-s":{"fetched_at":"2026-08-07T18:40:24.843571+00:00","verified_at":"2026-10-25T19:13:04.576913+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Smodel","size":33375439}],"weights_recorded_at":"2026-10-25T19:13:04.576917+00:00"},"mace-polar-1-m":{"fetched_at":"2026-08-07T18:40:25.467111+00:00","verified_at":"2026-10-25T19:13:04.576920+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Mmodel","size":68133235}],"weights_recorded_at":"2026-10-25T19:13:04.576924+00:00"},"mace-polar-1-l":{"fetched_at":"2026-08-07T18:40:26.911489+00:00","verified_at":"2026-10-25T19:13:04.576927+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Lmodel","size":129808715}],"weights_recorded_at":"2026-10-25T19:13:04.576931+00:00"},"mace-polar:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:13:04.577079+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"mattersim":{"built_at":"2026-08-07T18:13:22.432426+00:00","source_hash":"sha256:c1706582d1cb4c355cc8d2e3381e399215e7e60886e4b3c09fbc39a7fb8c5611","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"mattersim>=1.1.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# # Not imported here. mattersim -> torchmetrics -> torchvision pulls it in\n# # transitively, and torchvision ships compiled ops ABI-locked to one exact\n# # torch build. An `explicit` index is only consulted for packages listed\n# # in `dependencies`, so without this line torchvision silently resolves\n# # from PyPI against a different torch and the mismatch surfaces only at\n# # `import torchvision` as:\n# # RuntimeError: operator torchvision::nms does not exist\n# # Leave it unpinned: the index hosts one build per torch release and uv\n# # enforces the torch<->torchvision pairing.\n# \"torchvision\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n# torchvision = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"\nMatterSim environment for Rootstock.\n\nProvides access to Microsoft's MatterSim universal potential. MatterSim-v1\ncovers ~100 elements and supports periodic and non-periodic systems.\n\nModels:\n - \"MatterSim-v1.0.0-5M\": 5M parameter model (faster, default)\n - \"MatterSim-v1.0.0-1M\": 1M parameter model (smallest)\n\"\"\"\n\nCHECKPOINTS = {\n \"mattersim-v1-0-0-5m\": \"MatterSim-v1.0.0-5M\",\n \"mattersim-v1-0-0-1m\": \"MatterSim-v1.0.0-1M\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mattersim:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a MatterSim calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n from mattersim.forcefield import MatterSimCalculator\n\n return MatterSimCalculator(load_path=CHECKPOINTS[checkpoint], device=device)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): load_path accepts a filesystem\n # path as well as a model name, so this is setup() minus the name mapping.\n from mattersim.forcefield import MatterSimCalculator\n\n return MatterSimCalculator(load_path=path, device=device)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mattersim":"1.2.3","rootstock":"1.6.0","torch":"2.11.0+cu128","torchvision":"0.26.0+cu128"},"lock_hash":"sha256:ee44da043e144d6d63a0a78476c14c799690487796a26396f2b95feb807d4228","checkpoints":{"mattersim-v1-0-0-1m":{"fetched_at":"2026-08-07T18:41:34.575484+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"home/.local/mattersim/pretrained_models/mattersim-v1.0.0-1M.pth","size":17932943}],"weights_recorded_at":"2026-10-18T16:47:56.675656+00:00"},"mattersim-v1-0-0-5m":{"fetched_at":"2026-08-07T18:41:35.132251+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"home/.local/mattersim/pretrained_models/mattersim-v1.0.0-5M.pth","size":91176875}],"weights_recorded_at":"2026-09-26T19:24:12.820464+00:00"},"mattersim:custom":{"fetched_at":null,"verified_at":"2026-09-26T19:24:12.820662+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"orb":{"built_at":"2026-08-07T18:13:20.783480+00:00","source_hash":"sha256:4ae1362a989170b29d5082a23f9586c50d4f04b700bd29b8c65609efcbe52bef","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # >=0.5,<0.6: 0.5.5 is what the verified Delta env resolved — the v2\n# # loaders keep their single-return API through 0.5.x. 0.4.x is broken\n# # for us: it imports pynanoflann, which is git-only and undeclared, so\n# # a fresh build dies at import (Delta, 2026-07-31). 0.6 raises the\n# # Python floor to 3.12 — that line lives in orb_v3.py.\n# \"orb-models>=0.5,<0.6\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# # Not imported here — constrains orb-models' transitive dep. setup()'s\n# # no-lock serve path relies on cached_path returning local files without\n# # locking or writing, verified against exactly this version (#67).\n# \"cached_path==1.8.10\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"Orb v2 env — kept only for the built-in-D3 dispersion variant.\n\nOrb v3 (orb_v3.py) is the primary orb env; it supersedes the v2 checkpoints\nexcept orb-d3-v2, which has no v3 equivalent (v3 ships no dispersion-corrected\nmodel). Catalog trimmed to that one id 2026-07-30. The two lines can't share\nan env: the v3 loaders need orb-models>=0.6, which raises the Python floor\nto 3.12 and torch to 2.8 (the v2 loaders here are fine through 0.5.x).\n\"\"\"\n\nimport os\nimport shutil\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"orb-d3-v2\": \"orb-d3-v2\",\n # Your own fine-tuned v2-architecture weights: pair with weights=\n # (loaded via setup_from_path). v3 fine-tunes go to orb-v3:custom.\n \"orb-v2:custom\": None,\n}\n\n\ndef _default_weights_url(load_fn) -> str:\n \"\"\"The upstream URL baked into the loader's ``weights_path`` default.\"\"\"\n import inspect\n\n default = inspect.signature(load_fn).parameters[\"weights_path\"].default\n if not isinstance(default, str) or not default.startswith((\"http://\", \"https://\")):\n raise RuntimeError(\n f\"{load_fn.__name__} has no URL default for weights_path \"\n f\"(got {default!r}); update this env file for the installed orb-models\"\n )\n return default\n\n\ndef _local_weights_path(url: str) -> Path:\n \"\"\"Where the checkpoint lives in the shared model cache.\"\"\"\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"orb\" / os.path.basename(url)\n\n\ndef _fetch(url: str, dest: Path) -> None:\n \"\"\"Download ``url`` to ``dest`` atomically (tmp file + rename).\"\"\"\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp = dest.with_name(f\"{dest.name}.tmp.{os.getpid()}\")\n try:\n with urllib.request.urlopen(url) as resp, open(tmp, \"wb\") as out:\n shutil.copyfileobj(resp, out)\n os.replace(tmp, dest)\n finally:\n tmp.unlink(missing_ok=True)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.calculator import ORBCalculator\n\n # orb-models exposes one function per checkpoint, e.g. pretrained.orb_v2().\n fn_name = CHECKPOINTS[checkpoint].replace(\"-\", \"_\")\n load_fn = getattr(pretrained, fn_name)\n\n # orb-models resolves its default weights URL through `cached_path`, which\n # write-locks its cache dir even on warm hits — EACCES for anyone who can't\n # write the shared install (Garden-AI/rootstock#67). Handed a *local* path\n # instead, cached_path returns it without locking. So the weights are\n # pre-fetched into the shared model cache at `rootstock add` time\n # (maintainer, cache writable) and every later serve loads that file.\n url = _default_weights_url(load_fn)\n weights = _local_weights_path(url)\n if not weights.exists():\n _fetch(url, weights)\n\n orbff = load_fn(weights_path=str(weights), device=torch.device(device))\n return ORBCalculator(orbff, device=torch.device(device))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", arch: str = \"orb-v2\"):\n # Custom checkpoints (`:custom` ids with user weights). A weights file doesn't say\n # which orb architecture produced it, so `arch` names the pretrained\n # loader to instantiate — pass the right one at call time\n # (setup_kwargs={\"arch\": ...} / --kwarg arch=...). Handing the loader a local path also means no\n # network and no cached_path locking (see setup()).\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.calculator import ORBCalculator\n\n fn_name = arch.replace(\"-\", \"_\")\n try:\n load_fn = getattr(pretrained, fn_name)\n except AttributeError:\n raise ValueError(\n f\"unknown orb architecture {arch!r}; expected a loader name from \"\n f\"orb_models.forcefield.pretrained, e.g. orb-v2, orb-d3-v2\"\n ) from None\n\n orbff = load_fn(weights_path=path, device=torch.device(device))\n return ORBCalculator(orbff, device=torch.device(device))\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","cached-path":"1.8.10","orb-models":"0.5.5","rootstock":"1.6.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:8282748753e7201545484a51c53d2e7c9e1d5c38c66511efd2ee579ec62b71c6","checkpoints":{"orb-d3-v2":{"fetched_at":"2026-08-07T18:41:18.788701+00:00","verified_at":"2026-10-25T19:13:04.576938+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-d3-v2-20241011.ckpt","size":100967743}],"weights_recorded_at":"2026-10-25T19:13:04.576941+00:00"}}},"grace":{"built_at":"2026-08-07T18:14:52.989048+00:00","source_hash":"sha256:f46207cbbd702910ea889c09cae3f9933de6c48baf22d75c4039caf02467b635","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"tensorpotential>=0.6.0\",\n# \"ase>=3.22\",\n# ]\n# ///\n\"\"\"GRACE env — hosts GRACE foundation checkpoints via tensorpotential.\n\nGRACE runs on TensorFlow, not torch. TPCalculator has no device argument:\nTF grabs whatever GPU it sees, so device selection happens via\nCUDA_VISIBLE_DEVICES, and both it and TF_USE_LEGACY_KERAS must be set\nbefore the first TF import. The first calculation triggers an XLA compile —\na slow first step is expected.\n\"\"\"\n\nCHECKPOINTS = {\n \"grace-2l-smax-omat-large\": \"GRACE-2L-SMAX-OMAT-large\",\n \"grace-3l-omat-large-ft-am\": \"GRACE-3L-OMAT-large-ft-AM\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n import os\n\n os.environ[\"TF_USE_LEGACY_KERAS\"] = \"1\"\n if device == \"cpu\":\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n elif device.startswith(\"cuda:\"):\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = device.split(\":\", 1)[1]\n\n from tensorpotential.calculator import grace_fm\n\n return grace_fm(CHECKPOINTS[checkpoint])\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.6.0","tensorpotential":"0.6.0"},"lock_hash":"sha256:b192eb5027b93f6fc506427581731d2ce617095364e767bec8ee98cfffc3f3bc","checkpoints":{"grace-2l-smax-omat-large":{"fetched_at":"2026-08-07T18:53:53.338619+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"home/.cache/grace/GRACE-2L-SMAX-OMAT-large/fingerprint.pb","size":76},{"path":"home/.cache/grace/GRACE-2L-SMAX-OMAT-large/metadata.json","size":175792},{"path":"home/.cache/grace/GRACE-2L-SMAX-OMAT-large/metadata.yaml","size":93326},{"path":"home/.cache/grace/GRACE-2L-SMAX-OMAT-large/saved_model.pb","size":11604577},{"path":"home/.cache/grace/GRACE-2L-SMAX-OMAT-large/variables/variables.data-00000-of-00001","size":137443668},{"path":"home/.cache/grace/GRACE-2L-SMAX-OMAT-large/variables/variables.index","size":3280},{"path":"home/.keras/keras.json","size":120}],"weights_recorded_at":"2026-08-07T18:53:53.338649+00:00"},"grace-3l-omat-large-ft-am":{"fetched_at":"2026-08-07T18:39:24.759378+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 46662 files, 7.5 GB in 511.2s; weights: 274 MB (manifest)\n2026-10-25 09:03:34.143729: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\nTo enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\nWARNING:tensorflow:UserWarning: enabling the new type promotion must happen at the beginning of the program. Please ensure no TF APIs have been used yet.","weight_files":[{"path":"home/.cache/grace/GRACE-3L-OMAT-large-ft-AM/fingerprint.pb","size":98},{"path":"home/.cache/grace/GRACE-3L-OMAT-large-ft-AM/metadata.json","size":1305},{"path":"home/.cache/grace/GRACE-3L-OMAT-large-ft-AM/metadata.yaml","size":474},{"path":"home/.cache/grace/GRACE-3L-OMAT-large-ft-AM/saved_model.pb","size":11406881},{"path":"home/.cache/grace/GRACE-3L-OMAT-large-ft-AM/variables/variables.data-00000-of-00001","size":262276271},{"path":"home/.cache/grace/GRACE-3L-OMAT-large-ft-AM/variables/variables.index","size":5877},{"path":"home/.keras/keras.json","size":120}],"weights_recorded_at":"2026-08-07T18:39:24.759396+00:00"}}},"orb_v3":{"built_at":"2026-08-07T18:14:21.747240+00:00","source_hash":"sha256:1f63f36fe8c791313c04bbb5bf534a1a7168c327c1a35f98ecd3d80e3ff74122","source":"# /// script\n# # <3.13: orb-models pins dm-tree==0.1.8, which has no cp313 wheel and whose\n# # sdist doesn't compile against modern GCC (vendored abseil).\n# requires-python = \">=3.12,<3.13\"\n# dependencies = [\n# \"orb-models>=0.6.2\",\n# \"ase>=3.25\",\n# \"torch>=2.8\",\n# # Not imported here — constrains orb-models' transitive dep. setup()'s\n# # no-lock serve path relies on cached_path returning local files without\n# # locking or writing, verified against exactly this version (#67).\n# \"cached_path==1.8.10\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"Orb v3 env — the primary orb env, Orbital Materials' Orb v3 potentials.\n\nSeparate from orb.py because the v3 loaders need orb-models>=0.6.2, which\nbumped the Python floor to 3.12 and torch to 2.8; the v3 loader API also\ndiffers (returns a tuple, requires `atoms_adapter` on ORBCalculator, imports\nthe calculator from forcefield.inference). orb.py (v2) survives only for\norb-d3-v2, the dispersion-corrected variant with no v3 equivalent.\n\"\"\"\n\nimport os\nimport shutil\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"orb-v3-conservative-inf-omat\": \"orb-v3-conservative-inf-omat\",\n \"orb-v3-conservative-20-omat\": \"orb-v3-conservative-20-omat\",\n \"orb-v3-direct-inf-omat\": \"orb-v3-direct-inf-omat\",\n \"orb-v3-direct-20-omat\": \"orb-v3-direct-20-omat\",\n \"orb-v3-conservative-inf-mpa\": \"orb-v3-conservative-inf-mpa\",\n \"orb-v3-conservative-20-mpa\": \"orb-v3-conservative-20-mpa\",\n \"orb-v3-direct-inf-mpa\": \"orb-v3-direct-inf-mpa\",\n \"orb-v3-direct-20-mpa\": \"orb-v3-direct-20-mpa\",\n # The omol ids (orb-v3-{conservative,direct}-omol) are dropped from the\n # catalog 2026-07-30: they had been failing verify on every cluster since\n # 2026-05. Re-add once the failure is understood.\n # Your own fine-tuned v3 weights: pair with weights= (loaded via\n # setup_from_path).\n \"orb-v3:custom\": None,\n}\n\n\ndef _default_weights_url(load_fn) -> str:\n \"\"\"The upstream URL baked into the loader's ``weights_path`` default.\"\"\"\n import inspect\n\n default = inspect.signature(load_fn).parameters[\"weights_path\"].default\n if not isinstance(default, str) or not default.startswith((\"http://\", \"https://\")):\n raise RuntimeError(\n f\"{load_fn.__name__} has no URL default for weights_path \"\n f\"(got {default!r}); update this env file for the installed orb-models\"\n )\n return default\n\n\ndef _local_weights_path(url: str) -> Path:\n \"\"\"Where the checkpoint lives in the shared model cache.\"\"\"\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"orb\" / os.path.basename(url)\n\n\ndef _fetch(url: str, dest: Path) -> None:\n \"\"\"Download ``url`` to ``dest`` atomically (tmp file + rename).\"\"\"\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp = dest.with_name(f\"{dest.name}.tmp.{os.getpid()}\")\n try:\n with urllib.request.urlopen(url) as resp, open(tmp, \"wb\") as out:\n shutil.copyfileobj(resp, out)\n os.replace(tmp, dest)\n finally:\n tmp.unlink(missing_ok=True)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", precision: str = \"float32-high\"):\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.inference.calculator import ORBCalculator\n\n fn_name = CHECKPOINTS[checkpoint].replace(\"-\", \"_\")\n load_fn = getattr(pretrained, fn_name)\n\n # orb-models resolves its default weights URL through `cached_path`, which\n # write-locks its cache dir even on warm hits — EACCES for anyone who can't\n # write the shared install (Garden-AI/rootstock#67). Handed a *local* path\n # instead, cached_path returns it without locking. So the weights are\n # pre-fetched into the shared model cache at `rootstock add` time\n # (maintainer, cache writable) and every later serve loads that file.\n url = _default_weights_url(load_fn)\n weights = _local_weights_path(url)\n if not weights.exists():\n _fetch(url, weights)\n\n orbff, atoms_adapter = load_fn(\n weights_path=str(weights), device=torch.device(device), precision=precision\n )\n return ORBCalculator(orbff, atoms_adapter=atoms_adapter, device=torch.device(device))\n\n\ndef setup_from_path(\n path: str,\n device: str = \"cuda\",\n arch: str = \"orb-v3-conservative-inf-omat\",\n precision: str = \"float32-high\",\n):\n # Custom checkpoints (`:custom` ids with user weights). A weights file doesn't say\n # which orb architecture produced it, so `arch` names the pretrained\n # loader to instantiate — pass the right one at call time\n # (setup_kwargs={\"arch\": ...} / --kwarg arch=...). Handing the loader a\n # local path also means no network and no cached_path locking (see setup()).\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.inference.calculator import ORBCalculator\n\n fn_name = arch.replace(\"-\", \"_\")\n try:\n load_fn = getattr(pretrained, fn_name)\n except AttributeError:\n raise ValueError(\n f\"unknown orb architecture {arch!r}; expected a loader name from \"\n f\"orb_models.forcefield.pretrained, e.g. orb-v3-conservative-inf-omat\"\n ) from None\n\n orbff, atoms_adapter = load_fn(\n weights_path=path, device=torch.device(device), precision=precision\n )\n return ORBCalculator(orbff, atoms_adapter=atoms_adapter, device=torch.device(device))\n","python_requires":">=3.12,<3.13","dependencies":{"ase":"3.29.0","cached-path":"1.8.10","orb-models":"0.7.0","rootstock":"1.6.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:1b14ab95f8183b2845ff04c3d145511c15b378f2a2a6f660d83545fa10594eeb","checkpoints":{"orb-v3-conservative-inf-omat":{"fetched_at":"2026-08-07T18:42:05.251153+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 36559 files, 8.1 GB in 503.4s; weights: 102 MB (manifest)","weight_files":[{"path":"cache/orb/orb-v3-conservative-inf-omat-20250404.ckpt","size":102097517}],"weights_recorded_at":"2026-10-18T16:47:56.675696+00:00"},"orb-v3-conservative-inf-mpa":{"fetched_at":"2026-08-07T18:42:05.487867+00:00","verified_at":"2026-10-25T19:13:04.576970+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-inf-mpa-20250404.ckpt","size":102097517}],"weights_recorded_at":"2026-10-25T19:13:04.576973+00:00"},"orb-v3-conservative-20-mpa":{"fetched_at":"2026-08-07T18:42:05.833271+00:00","verified_at":"2026-10-25T19:13:04.576977+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-20-mpa-20250404.ckpt","size":102097380}],"weights_recorded_at":"2026-10-25T19:13:04.576980+00:00"},"orb-v3-conservative-20-omat":{"fetched_at":"2026-08-07T18:42:06.048676+00:00","verified_at":"2026-10-25T19:13:04.576949+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-20-omat-20250404.ckpt","size":102097517}],"weights_recorded_at":"2026-10-25T19:13:04.576953+00:00"},"orb-v3-direct-20-mpa":{"fetched_at":"2026-08-07T18:42:26.108950+00:00","verified_at":"2026-10-25T19:13:04.576990+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-20-mpa-20250404.ckpt","size":102636362}],"weights_recorded_at":"2026-10-25T19:13:04.576993+00:00"},"orb-v3-direct-20-omat":{"fetched_at":"2026-08-07T18:42:26.198770+00:00","verified_at":"2026-10-25T19:13:04.576963+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-20-omat-20250404.ckpt","size":102636510}],"weights_recorded_at":"2026-10-25T19:13:04.576966+00:00"},"orb-v3-direct-inf-mpa":{"fetched_at":"2026-08-07T18:42:26.341487+00:00","verified_at":"2026-10-25T19:13:04.576983+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-inf-mpa-20250404.ckpt","size":102636510}],"weights_recorded_at":"2026-10-25T19:13:04.576986+00:00"},"orb-v3-direct-inf-omat":{"fetched_at":"2026-08-07T18:42:26.487865+00:00","verified_at":"2026-10-25T19:13:04.576956+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-inf-omat-20250404.ckpt","size":102636658}],"weights_recorded_at":"2026-10-25T19:13:04.576960+00:00"},"orb-v3:custom":{"fetched_at":null,"verified_at":"2026-10-18T16:47:56.675854+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"painn":{"built_at":"2026-08-07T18:14:08.548797+00:00","source_hash":"sha256:1cb2a0a04abfb53c98beb463cd40138791d5d0febe4d5cefb83730decf46edd1","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\nPaiNN environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 PaiNN checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"PaiNN-S2EF-OC20-All\": default\n\"\"\"\n\nCHECKPOINTS = {\n \"painn-s2ef-oc20-all\": \"PaiNN-S2EF-OC20-All\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a PaiNN OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.6.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"painn-s2ef-oc20-all":{"fetched_at":"2026-08-07T18:44:12.251376+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/painn_h512_s2ef_all.pt","size":80340525}],"weights_recorded_at":"2026-10-18T16:47:56.675754+00:00"}}},"pet":{"built_at":"2026-08-07T18:14:11.117504+00:00","source_hash":"sha256:2ed7f37a981d7b622c8f25c58a80b4269a919554a0fd61e36ddee4a7678654fd","source":"# /// script\n# requires-python = \">=3.11,<3.15\"\n# dependencies = [\n# \"upet>=0.2.6\",\n# \"ase>=3.22\",\n# # upet pulls nvalchemi-toolkit-ops unpinned; 0.4+ needs torch>=2.8 at\n# # runtime but only declares the constraint on its extras, so the\n# # resolver won't catch it (same trap as the tensornet env).\n# \"torch>=2.8,<2.14\",\n# ]\n# ///\n\"\"\"PET env — hosts lab-cosmo's UPET foundation checkpoints (PET-MAD successor).\n\nThe upstream string encodes model@version; versions are pinned rather than\n\"latest\" so rebuilds serve the same weights. pet-omatpes-l is trained at the\nr2SCAN level of theory — its energies are not comparable to the PBE-level\npet-oam models.\n\"\"\"\n\nCHECKPOINTS = {\n \"pet-oam-xl\": \"pet-oam-xl@1.0.0\",\n \"pet-omatpes-l\": \"pet-omatpes-l@0.1.0\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from huggingface_hub import hf_hub_download\n from upet.calculator import UPETCalculator\n\n # metatomic-ase 0.1.2's nvalchemi fast path passes a float max_neighbors\n # (len * max(128, cutoff**3)) into torch.full — TypeError for cutoffs > ~5 Å.\n # Force the vesin fallback; the flag is read at call time.\n import metatomic_ase._neighbors as _mta_n\n _mta_n.HAS_NVALCHEMIOPS = False\n\n # Passing model=/version= makes UPETCalculator resolve the name by listing\n # the hub repo — an uncached API call that fails on workers, which run\n # with HF_HUB_OFFLINE=1 (and on any node without internet). Fetch the\n # pinned file ourselves — a cache hit needs no network even offline — and\n # hand it over as checkpoint_path, which skips the resolve entirely.\n model, version = CHECKPOINTS[checkpoint].split(\"@\", 1)\n path = hf_hub_download(\n repo_id=\"lab-cosmo/upet\",\n filename=f\"{model}-v{version}.ckpt\",\n subfolder=\"models\",\n )\n return UPETCalculator(checkpoint_path=path, device=device)\n","python_requires":">=3.11,<3.15","dependencies":{"ase":"3.29.0","rootstock":"1.6.0","torch":"2.13.0","upet":"0.2.6"},"lock_hash":"sha256:a70651437572b0f9850932bae01a42e6cd89feef1661ccbcc687386448ea65e9","checkpoints":{"pet-omatpes-l":{"fetched_at":"2026-08-07T18:44:13.271159+00:00","verified_at":"2026-10-25T19:13:04.577006+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/blobs/879b1045391d88869522605a8b8b3cedeed74668e7062fdd7487548ab7b08004","size":771692879},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:13:04.577011+00:00"},"pet-oam-xl":{"fetched_at":"2026-08-07T18:44:27.538890+00:00","verified_at":"2026-10-25T19:13:04.576997+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/blobs/c3a67cd019969dfd4dcabe9574682fe035f861d3f1c10190989b36c983699409","size":2920687712},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:13:04.577001+00:00"}}},"scn":{"built_at":"2026-08-07T18:14:49.117957+00:00","source_hash":"sha256:cba233c494f7a5a6c9b675e024176c9755a7394bfff2a8d29dc5233d6a83c2f8","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\nSCN environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 SCN checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"SCN-S2EF-OC20-All+MD\": default\n - \"SCN-t4-b2-S2EF-OC20-2M\"\n - \"SCN-S2EF-OC20-2M\"\n\"\"\"\n\nCHECKPOINTS = {\n \"scn-s2ef-oc20-all-md\": \"SCN-S2EF-OC20-All+MD\",\n \"scn-t4-b2-s2ef-oc20-2m\": \"SCN-t4-b2-S2EF-OC20-2M\",\n \"scn-s2ef-oc20-2m\": \"SCN-S2EF-OC20-2M\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an SCN OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.6.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"scn-s2ef-oc20-2m":{"fetched_at":"2026-08-07T18:44:24.117756+00:00","verified_at":"2026-10-25T19:13:04.577028+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/scn_t1_b1_s2ef_2M.pt","size":494398909}],"weights_recorded_at":"2026-10-25T19:13:04.577031+00:00"},"scn-t4-b2-s2ef-oc20-2m":{"fetched_at":"2026-08-07T18:44:48.572433+00:00","verified_at":"2026-10-25T19:13:04.577021+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/scn_t4_b2_s2ef_2M.pt","size":506994109}],"weights_recorded_at":"2026-10-25T19:13:04.577024+00:00"},"scn-s2ef-oc20-all-md":{"fetched_at":"2026-08-07T18:44:53.943614+00:00","verified_at":"2026-10-25T19:13:04.577014+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/scn_all_md_s2ef.pt","size":675885597}],"weights_recorded_at":"2026-10-25T19:13:04.577018+00:00"}}},"sevennet":{"built_at":"2026-08-07T18:14:51.903219+00:00","source_hash":"sha256:560f24594f3048803e4527d1a3761784e11c9a1d198c14781697f323a1384057","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"sevenn>=0.10.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"SevenNet env — hosts pretrained SevenNet universal potentials.\n\nSevenNet (SCalable EquiVariance-Enabled Neural Network) ships several\npretrained models, loaded by keyword through ``SevenNetCalculator``.\n\nMulti-fidelity models (``7net-omni``, ``7net-mf-ompa``) take a ``modal``\nargument selecting the training fidelity (e.g. ``\"mpa\"`` or ``\"omat24\"``),\nand sevenn releases after mid-2026 make it mandatory (older ones defaulted\nsilently). When not given, setup() defaults them to ``\"mpa\"`` — the\nMPtrj+sAlex fidelity, consistent with the lineage of the other checkpoints\nhere — so no-kwarg paths (smoke-test, plain ``rootstock add``) keep working.\nOverride at runtime via ``setup_kwargs={\"modal\": ...}`` on\nRootstockCalculator (or ``--kwarg modal=...`` for ``rootstock add``);\nsingle-fidelity models ignore it.\n\"\"\"\n\nCHECKPOINTS = {\n \"sevennet-0\": \"7net-0\",\n \"sevennet-l3i5\": \"7net-l3i5\",\n \"sevennet-omat\": \"7net-omat\",\n \"sevennet-mf-ompa\": \"7net-mf-ompa\",\n \"sevennet-omni\": \"7net-omni\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"sevennet:custom\": None,\n}\n\n\n# Multi-fidelity models and the fidelity used when modal isn't specified.\nMULTI_FIDELITY_DEFAULT_MODAL = {\n \"sevennet-mf-ompa\": \"mpa\",\n \"sevennet-omni\": \"mpa\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", modal: str | None = None):\n \"\"\"\n Load a SevenNet calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\").\n modal: Fidelity selector for multi-fidelity models (e.g. \"mpa\",\n \"omat24\"). Defaults to \"mpa\" for 7net-omni / 7net-mf-ompa;\n ignored by the single-fidelity models.\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n from sevenn.calculator import SevenNetCalculator\n\n if modal is None:\n modal = MULTI_FIDELITY_DEFAULT_MODAL.get(checkpoint)\n kwargs = {\"modal\": modal} if modal is not None else {}\n return SevenNetCalculator(model=CHECKPOINTS[checkpoint], device=device, **kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", modal: str | None = None):\n # Custom checkpoints (`:custom` ids with user weights): SevenNetCalculator loads a\n # checkpoint file directly. A multi-fidelity fine-tune must pass its\n # fidelity (setup_kwargs={\"modal\": \"mpa\"} / --kwarg modal=mpa);\n # single-fidelity ones need no kwargs.\n from sevenn.calculator import SevenNetCalculator\n\n kwargs = {\"modal\": modal} if modal is not None else {}\n return SevenNetCalculator(model=path, device=device, **kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.6.0","sevenn":"0.13.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:f0907ac3ba945602b4ecef938a6f6b838750486f828f87e55248231beda8c649","checkpoints":{"sevennet-0":{"fetched_at":"2026-08-07T18:45:55.176461+00:00","verified_at":"2026-10-25T19:13:04.577034+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-l3i5":{"fetched_at":"2026-08-07T18:45:56.221601+00:00","verified_at":"2026-10-25T19:13:04.577038+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-omat":{"fetched_at":"2026-08-07T18:45:59.080797+00:00","verified_at":"2026-10-25T19:13:04.577041+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-mf-ompa":{"fetched_at":"2026-08-07T18:46:00.558499+00:00","verified_at":"2026-10-25T19:13:04.577044+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-omni":{"fetched_at":"2026-08-07T18:46:16.657397+00:00","verified_at":"2026-10-25T19:13:04.577048+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"tace":{"built_at":"2026-08-07T18:15:02.652567+00:00","source_hash":"sha256:1942bf8d57e869cefac328435f093792c94591ce3a24da6e38050e420385189c","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"tace\",\n# \"ase>=3.22\",\n# \"torch>=2.4,<2.14\",\n# ]\n#\n# # The TECE architecture and its foundation registry ship in tace 0.2.0,\n# # which is not on PyPI (latest release there is 0.1.0). This is the commit\n# # the Matbench Discovery submission pinned.\n# [tool.uv.sources.tace]\n# git = \"https://github.com/xvzemin/tace\"\n# rev = \"81f65a4c188bd09cec8d1419388f7afdcc1b6fd0\"\n# ///\n\"\"\"TACE env — hosts TECE/TACE foundation checkpoints (Xu, Xie & Hu).\n\nopenequivariance CUDA-kernel acceleration is optional upstream and not\ninstalled here — TACE runs on the e3nn path.\n\"\"\"\n\nCHECKPOINTS = {\n \"tece-oam-rra-1.0\": \"TECE-OAM-RRA-1.0\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from tace.foundations import tace_foundations\n from tace.interface.ase import TACEAseCalc\n\n model_path = tace_foundations[CHECKPOINTS[checkpoint]]\n return TACEAseCalc(model_path, device=device, dtype=\"float32\")\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.6.0","tace":"0.2.0","torch":"2.13.0"},"lock_hash":"sha256:8181d8b7b1f5a6290854e5c32c163264540e28911815088cee590964b035c880","checkpoints":{"tece-oam-rra-1.0":{"fetched_at":"2026-08-07T18:47:47.347403+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 38504 files, 6.3 GB in 521.4s; weights: 891 MB (manifest)\nWARNING:root:AOTI is not enabled for energy, forces, stress, or virials outputs. For acceleration options, see https://tace.readthedocs.io/en/latest/guide/acceleration.html","weight_files":[{"path":"home/.cache/tace/TECE-OAM-RRA-1.0.pt","size":890523407}],"weights_recorded_at":"2026-10-18T16:47:56.675803+00:00"}}},"tensornet":{"built_at":"2026-08-07T18:15:39.402219+00:00","source_hash":"sha256:cddb5a364bae14466f61a30d7525f2f8c42db66c830977c95f547feceb8ce235","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"torch>=2.4.0,<2.5\",\n# \"ase>=3.22\",\n# \"huggingface_hub\",\n# \"matgl\",\n# # nvalchemi-toolkit-ops is deliberately absent. It's only an optional\n# # matgl extra (accelerated neighbor lists), and no version can work on\n# # this env's torch pin: 0.3.x registers torch custom ops in modules that\n# # use `from __future__ import annotations`, which torch 2.4's\n# # infer_schema can't parse — ValueError at import, and matgl's optional-\n# # import guard in matgl/ext/ase.py only catches ImportError, so merely\n# # importing PESCalculator crashes — while 0.4+ requires torch>=2.8.\n# # Without it, matgl falls back to its own neighbor list.\n# \"pymatgen\",\n# \"monty\",\n# \"ruamel.yaml\",\n# \"scipy\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# \"torch-spline-conv\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n#\n# [tool.uv.sources]\n# # Pinned: this recipe was originally written against matgl 1.0.0, and the\n# # unpinned git HEAD silently started building 4.x. setup() below is written\n# # for 4.0.3 — bump the tag deliberately, not by rebuild accident.\n# matgl = { git = \"https://github.com/materialsvirtuallab/matgl.git\", tag = \"v4.0.3\" }\n# ///\n\"\"\"TensorNet env — hosts MatPES TensorNet checkpoints via MatGL.\"\"\"\n\nCHECKPOINTS = {\n \"tensornet-matpes-pbe-2025-2\": \"materialyze/TensorNet-PES-MatPES-PBE-2025.2\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from huggingface_hub import snapshot_download\n\n import matgl\n from matgl.ext.ase import PESCalculator\n\n # load_model only resolves names against matgl's own manifest; HF models\n # must be downloaded explicitly and passed as a local path.\n local_path = snapshot_download(repo_id=CHECKPOINTS[checkpoint])\n\n # Move with .to(device), never torch.set_default_device: under matgl 4.x\n # the default-device hack splits the model across devices at load —\n # Potential.__init__ registers data_mean from a constructor-kwarg tensor\n # torch.load restored to cpu, while _eye3 (a persistent=False buffer, not\n # in the state dict) is created fresh on the default device — and\n # forward() then crashes with a cuda/cpu mismatch at\n # `lat @ (self._eye3 + st)` in matgl/apps/pes.py. Module.to() moves\n # params and all buffers coherently, and forward() migrates inputs to the\n # model's device itself.\n pot = matgl.load_model(local_path).to(device)\n return PESCalculator(potential=pot)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","huggingface-hub":"1.27.0","matgl":"4.0.3","monty":"2026.7.16","pymatgen":"2026.5.4","rootstock":"1.6.0","ruamel-yaml":"0.19.1","scipy":"1.17.1","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121","torch-spline-conv":"1.2.2+pt24cu121"},"lock_hash":null,"checkpoints":{"tensornet-matpes-pbe-2025-2":{"fetched_at":"2026-08-07T18:47:11.391749+00:00","verified_at":"2026-10-25T19:13:04.577052+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/blobs/19a0a506a58ea454e407040d55c393a174a5165834f22ea95c4c923ad9a406cc","size":4495},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/blobs/544966013206da7faa1966ffa360a82fb2d31c3e","size":5298},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/blobs/b9efa8c708014dd6c37e348f65dc1d8ae5c7bd5c7290de4d9f210e16c7fc3766","size":3402824},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/refs/main","size":40},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/trees/7ac2b90400a87eca7f44c1e7956dca6b1c85e418.json","size":906}],"weights_recorded_at":"2026-10-25T19:13:04.577057+00:00"}}},"uma":{"built_at":"2026-08-07T18:15:30.164548+00:00","source_hash":"sha256:a21a78ddc6bc95aaabc4e17428d768b8d713671906246b4baf30a4e965b7d1f5","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"UMA env — hosts Meta's UMA foundation model via FAIRChem.\n\nfairchem-core v2 dropped the torch-geometric / pyg-find-links install dance, so\nthis env is a plain PyPI install. The original uma-s-1 had an extensivity bug\nand was removed from the fairchem 2.20 registry — use uma-s-1p1 or uma-s-1p2p1.\n\"\"\"\n\nCHECKPOINTS = {\n \"uma-s-1p1\": \"uma-s-1p1\",\n # uma-s-1p2 has a known major bug; uma-s-1p2p1 fixes it and is the\n # upstream-recommended small model. 1p2 stays listed for reproducibility\n # of existing runs.\n \"uma-s-1p2\": \"uma-s-1p2\",\n # uma-s-1p2p1 is in fairchem's registry on git main but NOT in any\n # release yet (latest fairchem-core 2.21.0, 2026-06-08, lacks it — the\n # 2026-07-30 sync failed on exactly this). Re-add when the next\n # fairchem-core ships, and bump the dependency floor to that version.\n # \"uma-s-1p2p1\": \"uma-s-1p2p1\",\n \"uma-m-1p1\": \"uma-m-1p1\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"uma:custom\": None,\n}\n\n\ndef _fairchem_device(device: str) -> str:\n \"\"\"Translate an indexed device (\"cuda:2\") into what fairchem v2 accepts.\n\n MLIPPredictUnit._setup_device asserts `device in [\"cpu\", \"cuda\"]` and then\n resolves the real GPU itself via get_device_for_local_rank(), which returns\n f\"cuda:{torch.cuda.current_device()}\". So an index has to travel through\n torch's current-device state, not the argument. Verifying several\n checkpoints at once on a multi-GPU node hands each worker \"cuda:N\" — that\n killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync\n (4x A100, VERIFY_JOBS=4), while single-GPU Sophia never hit it.\n \"\"\"\n if device.startswith(\"cuda:\"):\n import torch\n\n torch.cuda.set_device(int(device.split(\":\", 1)[1]))\n return \"cuda\"\n return device\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", task: str = \"omat\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(\n CHECKPOINTS[checkpoint], device=_fairchem_device(device)\n )\n return FAIRChemCalculator(predictor, task_name=task)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", task: str = \"omat\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=_fairchem_device(device))\n return FAIRChemCalculator(predictor, task_name=task)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.6.0","torch":"2.8.0"},"lock_hash":"sha256:37e7cd799c8736978027d1cb9918aa06410e3bac86a82dfba59cbf31c5fd7eab","checkpoints":{"uma-m-1p1":{"fetched_at":"2026-08-07T19:02:45.891498+00:00","verified_at":"2026-10-25T19:13:04.577063+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/c30034edbf2e127f703f814cacb632661767da99b3e71c5b2ee5290510a52d68","size":11174706771},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:13:04.577067+00:00"},"uma-s-1p1":{"fetched_at":"2026-08-07T19:01:33.536471+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/07068e9c76702ca173d13155095f2117c1b327ec228557e64cd2709c777b824a","size":1173759575},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-10-18T16:47:56.675820+00:00"},"uma-s-1p2":{"fetched_at":"2026-08-07T19:01:40.358988+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker did not connect within timeout (600s). The worker never finished setup() — a slow or stalled model load (cold cache, degraded filesystem). If the load is genuinely slow, raise `timeout`: worker process is still running (hung, or blocked on the device?).\n(worker produced no output)","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/ba5c0d912efa22dc238e5fb1b5b7f66ee2e68c48c1b95b7cfd7fe1da5938398b","size":2333393167},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-10-18T16:47:56.675828+00:00"},"uma:custom":{"fetched_at":null,"verified_at":"2026-10-18T16:47:56.675857+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}}}},{"schema_version":5,"cluster":"polaris","root":"/eagle/projects/Rootstock/rootstock","maintainer":{"name":"Owen Price Skelly","email":"OwenPriceSkelly@uchicago.edu"},"rootstock_version":"1.6.4","python_version":"unknown","last_updated":"2026-10-25T20:00:19.052187+00:00","environments":{"allscaip":{"built_at":"2026-08-06T19:38:53.967540+00:00","source_hash":"sha256:bde4548ac275b9af7d17c3103e3da8008038f306ab5cb03e898c95ce6ad994d9","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"AllScAIP env — FAIRChem scalable attention MLIP trained on OMol25.\n\nallscaip-md-conserving-all-omol is an energy-conserving, all-to-all node\nattention model served through fairchem-core's get_predict_unit — the same\nAPI as eSEN. fairchem v2 carries the architecture in-package, so no\nflash-attention or custom CUDA kernels are needed.\n\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"allscaip-md-conserving-all-omol\": \"allscaip-md-conserving-all-omol\",\n \"allscaip-md-direct-all-omol\": \"allscaip-md-direct-all-omol\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"allscaip:custom\": None,\n}\n\n\ndef _fairchem_device(device: str) -> str:\n \"\"\"Translate an indexed device (\"cuda:2\") into what fairchem v2 accepts.\n\n MLIPPredictUnit._setup_device asserts `device in [\"cpu\", \"cuda\"]` and then\n resolves the real GPU itself via get_device_for_local_rank(), which returns\n f\"cuda:{torch.cuda.current_device()}\". So an index has to travel through\n torch's current-device state, not the argument. Verifying several\n checkpoints at once on a multi-GPU node hands each worker \"cuda:N\" — that\n killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync\n (4x A100, VERIFY_JOBS=4), while single-GPU Sophia never hit it.\n \"\"\"\n if device.startswith(\"cuda:\"):\n import torch\n\n torch.cuda.set_device(int(device.split(\":\", 1)[1]))\n return \"cuda\"\n return device\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(\n CHECKPOINTS[checkpoint], device=_fairchem_device(device)\n )\n return FAIRChemCalculator(predictor)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=_fairchem_device(device))\n return FAIRChemCalculator(predictor)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.6.0","torch":"2.8.0"},"lock_hash":"sha256:c99d413acdd9b5b2cdd0241ba29966574fe6800af2b63376abacd0ad28fed8b0","checkpoints":{"allscaip-md-direct-all-omol":{"fetched_at":"2026-08-03T17:03:56.967701+00:00","verified_at":"2026-10-25T19:59:55.994642+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/5517360cf5a23ca141f1b51b708d131289e35080bfa79f27b990695e3e78d18e","size":694955368},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.994651+00:00"},"allscaip-md-conserving-all-omol":{"fetched_at":"2026-08-03T17:03:57.183635+00:00","verified_at":"2026-10-25T19:59:55.993395+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/14fe8626e9e75337769ecf98a765cdc02667763ec4dc1103d958e1c9f56b26e9","size":688432239},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.993953+00:00"},"allscaip:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:59:55.999565+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"ani":{"built_at":"2026-07-31T18:31:40.101251+00:00","source_hash":"sha256:16ebe0eda0074fa275e0ce2c2d68c84dc10d2069ef0290d19802bd46bcf96d0e","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"torchani>=2.2\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n# ///\n\"\"\"\nANI-2x environment for Rootstock.\n\nANI-2x is a neural network potential for organic molecules containing\nH, C, N, O, F, S, Cl. It is not a universal potential — do not use it\nfor inorganic or periodic systems.\n\nModels:\n - \"ANI2x\": ANI-2x ensemble (default, 8 networks)\n - \"ANI1ccx\": ANI-1ccx, trained on CCSD(T)/CBS data (H, C, N, O only)\n - \"ANI1x\": ANI-1x (H, C, N, O only)\n\"\"\"\n\nCHECKPOINTS = {\n \"ani-2x\": \"ANI2x\",\n \"ani-1ccx\": \"ANI1ccx\",\n \"ani-1x\": \"ANI1x\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an ANI calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n import torchani\n\n model_map = {\n \"ANI2x\": torchani.models.ANI2x,\n \"ANI1ccx\": torchani.models.ANI1ccx,\n \"ANI1x\": torchani.models.ANI1x,\n }\n model = CHECKPOINTS[checkpoint]\n\n return model_map[model](periodic_table_index=True).to(device).ase()\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.3.0","torch":"2.13.0","torchani":"2.8.4"},"lock_hash":"sha256:8864bbf63ae2a9aa45ee0fc265b0c30e3c4d7b67fb9631b3efd90b5b7f774934","checkpoints":{}},"chgnet":{"built_at":"2026-07-31T18:31:25+00:00","source_hash":"sha256:b6ea6ac08a17cb3dc17d5455571dc5c21d513cd802916015de7f272e83eefe37","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"chgnet>=0.3.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"CHGNet env — hosts pretrained charge-informed universal potentials.\"\"\"\n\nCHECKPOINTS = {\n \"chgnet-default\": \"chgnet-default\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"chgnet:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a CHGNet calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\")\n\n Returns:\n ASE-compatible calculator\n \"\"\"\n from chgnet.model import CHGNet, CHGNetCalculator\n\n model_name = CHECKPOINTS[checkpoint]\n model = CHGNet.load() if model_name == \"chgnet-default\" else CHGNet.load(model_name)\n return CHGNetCalculator(model=model, use_device=device)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # CHGNet.from_file, not the named-model CHGNet.load() setup() uses.\n from chgnet.model import CHGNet, CHGNetCalculator\n\n model = CHGNet.from_file(path)\n return CHGNetCalculator(model=model, use_device=device)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","chgnet":"0.4.2","rootstock":"1.3.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:b5bc1f87f4c596d669c56f6a5053d21b488b1e141dc8e38e7a7a964b887fcc73","checkpoints":{"chgnet-default":{"fetched_at":"2026-08-03T17:04:29.877040+00:00","verified_at":"2026-10-25T19:59:55.994677+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"equiformer":{"built_at":"2026-07-31T18:31:42.640544+00:00","source_hash":"sha256:ac01aec3be646cbd2c63a2e9f6a0cc8308aa03f5d8528150c5f664b2aa06426f","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\nEquiformerV2 environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 EquiformerV2 checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"EquiformerV2-153M-S2EF-OC20-All+MD\": default\n - \"EquiformerV2-31M-S2EF-OC20-All+MD\"\n - \"EquiformerV2-83M-S2EF-OC20-2M\"\n\"\"\"\n\nCHECKPOINTS = {\n \"equiformer-v2-153m-s2ef-oc20-all-md\": \"EquiformerV2-153M-S2EF-OC20-All+MD\",\n \"equiformer-v2-31m-s2ef-oc20-all-md\": \"EquiformerV2-31M-S2EF-OC20-All+MD\",\n \"equiformer-v2-83m-s2ef-oc20-2m\": \"EquiformerV2-83M-S2EF-OC20-2M\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"equiformer:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an EquiformerV2 OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): OCPCalculator loads a\n # checkpoint file natively — this is setup() minus the registry download.\n from fairchem.core import OCPCalculator\n\n return OCPCalculator(checkpoint_path=path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.3.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"equiformer-v2-31m-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:07:10.587007+00:00","verified_at":"2026-10-25T19:59:55.994686+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_31M_ec4_allmd.pt","size":125685325}],"weights_recorded_at":"2026-10-25T19:59:55.994689+00:00"},"equiformer-v2-83m-s2ef-oc20-2m":{"fetched_at":"2026-08-03T17:07:10.622992+00:00","verified_at":"2026-10-25T19:59:55.994692+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_83M_2M.pt","size":336843381}],"weights_recorded_at":"2026-10-25T19:59:55.994695+00:00"},"equiformer-v2-153m-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:07:12.126551+00:00","verified_at":"2026-10-25T19:59:55.994680+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_153M_ec4_allmd.pt","size":619855697}],"weights_recorded_at":"2026-10-25T19:59:55.994683+00:00"},"equiformer:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: weights= run diverges from equiformer-v2-153m-s2ef-oc20-all-md: energy differs by 1.459e-04 eV (tolerance 1e-04)","weight_files":null,"weights_recorded_at":null}}},"escn":{"built_at":"2026-07-31T18:32:58.772603+00:00","source_hash":"sha256:e1301f3e20ac38e355e0b36e2ab9fb0f0a78f58059b5e2c57ab280aa9dad7ff5","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\neSCN environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 eSCN checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"eSCN-L6-M2-Lay12-S2EF-OC20-All+MD\": default\n - \"eSCN-L6-M3-Lay20-S2EF-OC20-All+MD\"\n - \"eSCN-L6-M2-Lay12-S2EF-OC20-2M\"\n - \"eSCN-L4-M2-Lay12-S2EF-OC20-2M\"\n\"\"\"\n\nCHECKPOINTS = {\n \"escn-l6-m2-lay12-s2ef-oc20-all-md\": \"eSCN-L6-M2-Lay12-S2EF-OC20-All+MD\",\n \"escn-l6-m3-lay20-s2ef-oc20-all-md\": \"eSCN-L6-M3-Lay20-S2EF-OC20-All+MD\",\n \"escn-l6-m2-lay12-s2ef-oc20-2m\": \"eSCN-L6-M2-Lay12-S2EF-OC20-2M\",\n \"escn-l4-m2-lay12-s2ef-oc20-2m\": \"eSCN-L4-M2-Lay12-S2EF-OC20-2M\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"escn:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an eSCN OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): OCPCalculator loads a\n # checkpoint file natively — this is setup() minus the registry download.\n from fairchem.core import OCPCalculator\n\n return OCPCalculator(checkpoint_path=path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.3.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"escn-l4-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-08-03T17:07:19.768803+00:00","verified_at":"2026-10-25T19:59:55.995303+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l4_m2_lay12_2M_s2ef.pt","size":144761989}],"weights_recorded_at":"2026-10-25T19:59:55.995306+00:00"},"escn-l6-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-08-03T17:08:03.373277+00:00","verified_at":"2026-10-25T19:59:55.995297+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l6_m2_lay12_2M_s2ef.pt","size":207735045}],"weights_recorded_at":"2026-10-25T19:59:55.995300+00:00"},"escn-l6-m2-lay12-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:08:03.392125+00:00","verified_at":"2026-10-25T19:59:55.994697+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l6_m2_lay12_all_md_s2ef.pt","size":207735109}],"weights_recorded_at":"2026-10-25T19:59:55.994700+00:00"},"escn-l6-m3-lay20-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:08:14.400434+00:00","verified_at":"2026-10-25T19:59:55.995286+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l6_m3_lay20_all_md_s2ef.pt","size":801642787}],"weights_recorded_at":"2026-10-25T19:59:55.995293+00:00"},"escn:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: weights= run diverges from escn-l6-m2-lay12-s2ef-oc20-all-md: energy differs by 4.449e-03 eV (tolerance 1e-04)","weight_files":null,"weights_recorded_at":null}}},"esen":{"built_at":"2026-08-06T19:39:02.621368+00:00","source_hash":"sha256:d07f4d08c7ce3ed5aa36b795281d2be931b17cd6d342e2b569ef0fd1f53e59db","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"eSEN env — hosts FAIRChem eSEN single-task checkpoints.\n\nfairchem-core v2 dropped the torch-geometric / pyg-find-links install dance\n(simplified here 2026-07-30 to match uma.py/allscaip.py — the find-links were\nalso blocking universal lockfile resolution).\n\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"esen-md-direct-all-omol\": \"esen-md-direct-all-omol\",\n \"esen-sm-conserving-all-omol\": \"esen-sm-conserving-all-omol\",\n \"esen-sm-direct-all-omol\": \"esen-sm-direct-all-omol\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"esen:custom\": None,\n}\n\n\ndef _fairchem_device(device: str) -> str:\n \"\"\"Translate an indexed device (\"cuda:2\") into what fairchem v2 accepts.\n\n MLIPPredictUnit._setup_device asserts `device in [\"cpu\", \"cuda\"]` and then\n resolves the real GPU itself via get_device_for_local_rank(), which returns\n f\"cuda:{torch.cuda.current_device()}\". So an index has to travel through\n torch's current-device state, not the argument. Verifying several\n checkpoints at once on a multi-GPU node hands each worker \"cuda:N\" — that\n killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync\n (4x A100, VERIFY_JOBS=4), while single-GPU Sophia never hit it.\n \"\"\"\n if device.startswith(\"cuda:\"):\n import torch\n\n torch.cuda.set_device(int(device.split(\":\", 1)[1]))\n return \"cuda\"\n return device\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(\n CHECKPOINTS[checkpoint], device=_fairchem_device(device)\n )\n return FAIRChemCalculator(predictor)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=_fairchem_device(device))\n return FAIRChemCalculator(predictor)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.6.0","torch":"2.8.0"},"lock_hash":"sha256:c99d413acdd9b5b2cdd0241ba29966574fe6800af2b63376abacd0ad28fed8b0","checkpoints":{"esen-sm-direct-all-omol":{"fetched_at":"2026-08-03T17:11:00.904446+00:00","verified_at":"2026-10-25T19:59:55.995322+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/a1b7109bcca6e1288c5f17c68776d3c6a1ab0766f6d7fdd45f901605355d2831","size":50960291},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.995326+00:00"},"esen-sm-conserving-all-omol":{"fetched_at":"2026-08-03T17:11:00.922460+00:00","verified_at":"2026-10-25T19:59:55.995316+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/01f63da2d071e39fc46a5f22f8369d0fc9de317ab2ef361a76603f5661238025","size":50958067},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.995320+00:00"},"esen-md-direct-all-omol":{"fetched_at":"2026-08-03T17:11:00.992499+00:00","verified_at":"2026-10-25T19:59:55.995309+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/e0cd2ac37d2992b385a39778469c443ccbc679c7c4fb3b40b38aa2ee23c29b1d","size":405789591},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.995313+00:00"},"esen:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:59:55.999572+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"grace":{"built_at":"2026-07-31T18:34:21.902859+00:00","source_hash":"sha256:f46207cbbd702910ea889c09cae3f9933de6c48baf22d75c4039caf02467b635","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"tensorpotential>=0.6.0\",\n# \"ase>=3.22\",\n# ]\n# ///\n\"\"\"GRACE env — hosts GRACE foundation checkpoints via tensorpotential.\n\nGRACE runs on TensorFlow, not torch. TPCalculator has no device argument:\nTF grabs whatever GPU it sees, so device selection happens via\nCUDA_VISIBLE_DEVICES, and both it and TF_USE_LEGACY_KERAS must be set\nbefore the first TF import. The first calculation triggers an XLA compile —\na slow first step is expected.\n\"\"\"\n\nCHECKPOINTS = {\n \"grace-2l-smax-omat-large\": \"GRACE-2L-SMAX-OMAT-large\",\n \"grace-3l-omat-large-ft-am\": \"GRACE-3L-OMAT-large-ft-AM\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n import os\n\n os.environ[\"TF_USE_LEGACY_KERAS\"] = \"1\"\n if device == \"cpu\":\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n elif device.startswith(\"cuda:\"):\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = device.split(\":\", 1)[1]\n\n from tensorpotential.calculator import grace_fm\n\n return grace_fm(CHECKPOINTS[checkpoint])\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.3.0","tensorpotential":"0.6.0"},"lock_hash":"sha256:1ea87e2d3e8620111788e4f3804932ac53f399ab75df08ee6b2b38195a82d50e","checkpoints":{"grace-2l-smax-omat-large":{"fetched_at":"2026-08-03T17:10:49.308918+00:00","verified_at":"2026-10-25T19:59:55.995328+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"grace-3l-omat-large-ft-am":{"fetched_at":"2026-08-03T17:11:39.418528+00:00","verified_at":"2026-10-25T19:59:55.995332+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"mace":{"built_at":"2026-08-03T17:00:55.121628+00:00","source_hash":"sha256:eb8e9790141ad43e2735656a942a3b3af3a551ae6e89aef0b99c72ff89f80360","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # 0.3.15+ needed for the mh-1 registry entry (matpes needs 0.3.13,\n# # omol needs 0.3.14, mpa-0 needs 0.3.10).\n# \"mace-torch>=0.3.15\",\n# \"ase>=3.22\",\n# # 2.4.1 is explicitly unsupported by mace-torch.\n# \"torch>=2.4.0,!=2.4.1,<2.10\",\n# ]\n# ///\n\"\"\"MACE env — hosts MACE-MP-0, MACE-OFF23, MPA-0, MATPES, MH-1, and OMOL checkpoints.\n\nAll ship in the same `mace-torch` package, so they share an environment.\nUpstream-string routing in CHECKPOINTS: an `off:` prefix routes to mace_off()\nand an `omol:` prefix to mace_omol() (float64, molecules only); an `mh:`\nprefix marks a multi-head model (float64, per the MACE-MH-1 model card).\n\nMulti-head checkpoints select a head via the `head` kwarg on setup()\n(setup_kwargs={\"head\": ...} / --kwarg head=...), named by upstream's training\ncorpus — see MH1_HEADS; omat_pbe is the default.\n\nThe OMOL checkpoint expects `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-mp-0-small\": \"small\",\n \"mace-mp-0-medium\": \"medium\",\n \"mace-mp-0-large\": \"large\",\n \"mace-off23-small\": \"off:small\",\n \"mace-off23-medium\": \"off:medium\",\n \"mace-off23-large\": \"off:large\",\n # Only a medium MPA-0 has been released, but upstream names the weights\n # file mace-mpa-0-medium.model — keep the size explicit like mace-mp-0.\n \"mace-mpa-0-medium\": \"medium-mpa-0\",\n \"mace-matpes-r2scan-0\": \"mace-matpes-r2scan-0\",\n # One entry per weights file: MH-1's heads are selected by setup(head=...).\n \"mace-mh-1\": \"mh:mh-1\",\n # Only the extra-large OMOL model has been released.\n \"mace-omol-0-extra-large\": \"omol:extra_large\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace:custom\": None,\n}\n\n# MH-1's heads, named by training corpus. The released weights file is the\n# authority (mace_select_head --list_heads; the model card also lists a\n# rgd1_b3lyp head, but that shipped only in mh-0 — ACEsuit/mace#1462).\n# Validated here because upstream only warns on an unknown head and silently\n# falls back to the last one.\nMH1_HEADS = (\n \"omat_pbe\",\n \"omol\",\n \"spice_wB97M\",\n \"oc20_usemppbe\",\n \"mp_pbe_refit_add\",\n \"matpes_r2scan\",\n)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", head: str | None = None):\n arg = CHECKPOINTS[checkpoint]\n if arg.startswith(\"mh:\"):\n head = head or \"omat_pbe\"\n if head not in MH1_HEADS:\n raise ValueError(f\"unknown head {head!r}; expected one of {', '.join(MH1_HEADS)}\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg[3:], device=device, default_dtype=\"float64\", head=head)\n if head is not None:\n raise ValueError(f\"'head' selects a head of a multi-head model; {checkpoint} has one head\")\n if arg.startswith(\"off:\"):\n from mace.calculators import mace_off\n\n return mace_off(model=arg[4:], device=device, default_dtype=\"float32\")\n if arg.startswith(\"omol:\"):\n from mace.calculators import mace_omol\n\n return mace_omol(model=arg[5:], device=device, default_dtype=\"float64\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg, device=device, default_dtype=\"float32\")\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", default_dtype: str = \"float32\", head: str | None = None):\n # Custom checkpoints (`:custom` ids with user weights): fine-tunes load through\n # MACECalculator directly — the mp/off dispatch in setup() only exists\n # to pick which pretrained file to download. `head` is for fine-tunes that\n # keep multiple heads; single-head weights load without it.\n from mace.calculators import MACECalculator\n\n return MACECalculator(model_paths=path, device=device, default_dtype=default_dtype, head=head)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"1.4.3","torch":"2.9.1"},"lock_hash":"sha256:2ef43c5e21b8c2dc7b3830833e9ae38555d87ade585a064e9fba3d63ca792931","checkpoints":{"mace-mp-0-medium":{"fetched_at":"2026-08-03T17:12:25.735382+00:00","verified_at":"2026-10-25T19:59:55.995341+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231203mace128L1_epoch199model","size":44422970}],"weights_recorded_at":"2026-10-25T19:59:55.995344+00:00"},"mace-mh-1":{"fetched_at":"2026-08-03T17:12:25.771339+00:00","verified_at":"2026-10-25T19:59:55.995379+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macemh1model","size":59208139}],"weights_recorded_at":"2026-10-25T19:59:55.995382+00:00"},"mace-matpes-r2scan-0":{"fetched_at":"2026-08-03T17:12:26.162740+00:00","verified_at":"2026-10-25T19:59:55.995374+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEmatpesr2scanomatftmodel","size":79470738}],"weights_recorded_at":"2026-10-25T19:59:55.995377+00:00"},"mace-mp-0-large":{"fetched_at":"2026-08-03T17:12:26.906848+00:00","verified_at":"2026-10-25T19:59:55.995346+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE_MPtrj_20229model","size":133803220}],"weights_recorded_at":"2026-10-25T19:59:55.995349+00:00"},"mace-off23-medium":{"fetched_at":"2026-08-03T17:12:44.272336+00:00","verified_at":"2026-10-25T19:59:55.995358+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_medium.model","size":18350596}],"weights_recorded_at":"2026-10-25T19:59:55.995361+00:00"},"mace-mp-0-small":{"fetched_at":"2026-08-03T17:12:44.380401+00:00","verified_at":"2026-10-25T19:59:55.995335+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231210mace128L0_energy_epoch249model","size":32581838}],"weights_recorded_at":"2026-10-25T19:59:55.995338+00:00"},"mace-mpa-0-medium":{"fetched_at":"2026-08-03T17:12:44.712260+00:00","verified_at":"2026-10-25T19:59:55.995368+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macempa0mediummodel","size":79462305}],"weights_recorded_at":"2026-10-25T19:59:55.995371+00:00"},"mace-off23-large":{"fetched_at":"2026-08-03T17:12:44.790639+00:00","verified_at":"2026-10-25T19:59:55.995363+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_large.model","size":55492786}],"weights_recorded_at":"2026-10-25T19:59:55.995366+00:00"},"mace-off23-small":{"fetched_at":"2026-08-03T17:13:00.903240+00:00","verified_at":"2026-10-25T19:59:55.995351+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_small.model","size":7347350}],"weights_recorded_at":"2026-10-25T19:59:55.995354+00:00"},"mace-omol-0-extra-large":{"fetched_at":"2026-08-03T17:13:02.181806+00:00","verified_at":"2026-10-25T19:59:55.995385+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-omol-0-extra-large-1024.model","size":422242640}],"weights_recorded_at":"2026-10-25T19:59:55.995387+00:00"},"mace:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:59:55.999575+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"mace_polar":{"built_at":"2026-07-31T18:34:29.831867+00:00","source_hash":"sha256:8db4cfd29a59c415e3762b04195e30b9e08d5f044ff6aa21ebcf01bcb78bc1ef","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"ase>=3.22\",\n# \"torch>=2.4.0,<2.10\",\n# # 0.3.16 is the first PyPI release with mace_polar().\n# \"mace-torch>=0.3.16\",\n# # PolarMACE imports graph_longrange at runtime; the distribution is named\n# # graph-longrange and exists only as this git repo (no PyPI release).\n# \"graph-longrange @ git+https://github.com/WillBaldwin0/graph_electrostatics.git\",\n# ]\n# ///\n\"\"\"MACE-POLAR env — electrostatic/polarizable MACE foundation models (OMol25).\n\nKept separate from the stable `mace` env because of the extra git-only\ngraph-longrange dependency.\n\nPOLAR checkpoints expect `charge`, `spin`, and `external_field` in atoms.info.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-polar-1-s\": \"polar-1-s\",\n \"mace-polar-1-m\": \"polar-1-m\",\n \"mace-polar-1-l\": \"polar-1-l\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace-polar:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from mace.calculators import mace_polar\n\n return mace_polar(model=CHECKPOINTS[checkpoint], device=device, default_dtype=\"float32\")\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): mace_polar() accepts a\n # weights file directly, keeping the PolarMACE model-type wiring.\n from mace.calculators import mace_polar\n\n return mace_polar(model=path, device=device, default_dtype=\"float32\")\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"1.3.0","torch":"2.9.1"},"lock_hash":"sha256:bf60a2e3dd4771b8eb5b17a99053c1bf90448615005aab598328165f8da4f5aa","checkpoints":{"mace-polar-1-s":{"fetched_at":"2026-08-03T17:15:38.414289+00:00","verified_at":"2026-10-25T19:59:55.995390+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Smodel","size":33375439}],"weights_recorded_at":"2026-10-25T19:59:55.995393+00:00"},"mace-polar-1-m":{"fetched_at":"2026-08-03T17:15:38.434404+00:00","verified_at":"2026-10-25T19:59:55.995395+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Mmodel","size":68133235}],"weights_recorded_at":"2026-10-25T19:59:55.995398+00:00"},"mace-polar-1-l":{"fetched_at":"2026-08-03T17:15:38.454464+00:00","verified_at":"2026-10-25T19:59:55.995401+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Lmodel","size":129808715}],"weights_recorded_at":"2026-10-25T19:59:55.995404+00:00"},"mace-polar:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:59:55.999578+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"orb_v3":{"built_at":"2026-07-31T18:36:13.727862+00:00","source_hash":"sha256:1f63f36fe8c791313c04bbb5bf534a1a7168c327c1a35f98ecd3d80e3ff74122","source":"# /// script\n# # <3.13: orb-models pins dm-tree==0.1.8, which has no cp313 wheel and whose\n# # sdist doesn't compile against modern GCC (vendored abseil).\n# requires-python = \">=3.12,<3.13\"\n# dependencies = [\n# \"orb-models>=0.6.2\",\n# \"ase>=3.25\",\n# \"torch>=2.8\",\n# # Not imported here — constrains orb-models' transitive dep. setup()'s\n# # no-lock serve path relies on cached_path returning local files without\n# # locking or writing, verified against exactly this version (#67).\n# \"cached_path==1.8.10\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"Orb v3 env — the primary orb env, Orbital Materials' Orb v3 potentials.\n\nSeparate from orb.py because the v3 loaders need orb-models>=0.6.2, which\nbumped the Python floor to 3.12 and torch to 2.8; the v3 loader API also\ndiffers (returns a tuple, requires `atoms_adapter` on ORBCalculator, imports\nthe calculator from forcefield.inference). orb.py (v2) survives only for\norb-d3-v2, the dispersion-corrected variant with no v3 equivalent.\n\"\"\"\n\nimport os\nimport shutil\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"orb-v3-conservative-inf-omat\": \"orb-v3-conservative-inf-omat\",\n \"orb-v3-conservative-20-omat\": \"orb-v3-conservative-20-omat\",\n \"orb-v3-direct-inf-omat\": \"orb-v3-direct-inf-omat\",\n \"orb-v3-direct-20-omat\": \"orb-v3-direct-20-omat\",\n \"orb-v3-conservative-inf-mpa\": \"orb-v3-conservative-inf-mpa\",\n \"orb-v3-conservative-20-mpa\": \"orb-v3-conservative-20-mpa\",\n \"orb-v3-direct-inf-mpa\": \"orb-v3-direct-inf-mpa\",\n \"orb-v3-direct-20-mpa\": \"orb-v3-direct-20-mpa\",\n # The omol ids (orb-v3-{conservative,direct}-omol) are dropped from the\n # catalog 2026-07-30: they had been failing verify on every cluster since\n # 2026-05. Re-add once the failure is understood.\n # Your own fine-tuned v3 weights: pair with weights= (loaded via\n # setup_from_path).\n \"orb-v3:custom\": None,\n}\n\n\ndef _default_weights_url(load_fn) -> str:\n \"\"\"The upstream URL baked into the loader's ``weights_path`` default.\"\"\"\n import inspect\n\n default = inspect.signature(load_fn).parameters[\"weights_path\"].default\n if not isinstance(default, str) or not default.startswith((\"http://\", \"https://\")):\n raise RuntimeError(\n f\"{load_fn.__name__} has no URL default for weights_path \"\n f\"(got {default!r}); update this env file for the installed orb-models\"\n )\n return default\n\n\ndef _local_weights_path(url: str) -> Path:\n \"\"\"Where the checkpoint lives in the shared model cache.\"\"\"\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"orb\" / os.path.basename(url)\n\n\ndef _fetch(url: str, dest: Path) -> None:\n \"\"\"Download ``url`` to ``dest`` atomically (tmp file + rename).\"\"\"\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp = dest.with_name(f\"{dest.name}.tmp.{os.getpid()}\")\n try:\n with urllib.request.urlopen(url) as resp, open(tmp, \"wb\") as out:\n shutil.copyfileobj(resp, out)\n os.replace(tmp, dest)\n finally:\n tmp.unlink(missing_ok=True)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", precision: str = \"float32-high\"):\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.inference.calculator import ORBCalculator\n\n fn_name = CHECKPOINTS[checkpoint].replace(\"-\", \"_\")\n load_fn = getattr(pretrained, fn_name)\n\n # orb-models resolves its default weights URL through `cached_path`, which\n # write-locks its cache dir even on warm hits — EACCES for anyone who can't\n # write the shared install (Garden-AI/rootstock#67). Handed a *local* path\n # instead, cached_path returns it without locking. So the weights are\n # pre-fetched into the shared model cache at `rootstock add` time\n # (maintainer, cache writable) and every later serve loads that file.\n url = _default_weights_url(load_fn)\n weights = _local_weights_path(url)\n if not weights.exists():\n _fetch(url, weights)\n\n orbff, atoms_adapter = load_fn(\n weights_path=str(weights), device=torch.device(device), precision=precision\n )\n return ORBCalculator(orbff, atoms_adapter=atoms_adapter, device=torch.device(device))\n\n\ndef setup_from_path(\n path: str,\n device: str = \"cuda\",\n arch: str = \"orb-v3-conservative-inf-omat\",\n precision: str = \"float32-high\",\n):\n # Custom checkpoints (`:custom` ids with user weights). A weights file doesn't say\n # which orb architecture produced it, so `arch` names the pretrained\n # loader to instantiate — pass the right one at call time\n # (setup_kwargs={\"arch\": ...} / --kwarg arch=...). Handing the loader a\n # local path also means no network and no cached_path locking (see setup()).\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.inference.calculator import ORBCalculator\n\n fn_name = arch.replace(\"-\", \"_\")\n try:\n load_fn = getattr(pretrained, fn_name)\n except AttributeError:\n raise ValueError(\n f\"unknown orb architecture {arch!r}; expected a loader name from \"\n f\"orb_models.forcefield.pretrained, e.g. orb-v3-conservative-inf-omat\"\n ) from None\n\n orbff, atoms_adapter = load_fn(\n weights_path=path, device=torch.device(device), precision=precision\n )\n return ORBCalculator(orbff, atoms_adapter=atoms_adapter, device=torch.device(device))\n","python_requires":">=3.12,<3.13","dependencies":{"ase":"3.29.0","cached-path":"1.8.10","orb-models":"0.7.0","rootstock":"1.3.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:cf6466e7518601cbbb8b296939d1804006146ec5a541c64179818df29b2ecda4","checkpoints":{"orb-v3-conservative-20-mpa":{"fetched_at":"2026-08-03T17:18:18.499544+00:00","verified_at":"2026-10-25T19:59:55.995433+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-20-mpa-20250404.ckpt","size":102097380}],"weights_recorded_at":"2026-10-25T19:59:55.995435+00:00"},"orb-v3-conservative-20-omat":{"fetched_at":"2026-08-03T17:18:18.543840+00:00","verified_at":"2026-10-25T19:59:55.995411+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-20-omat-20250404.ckpt","size":102097517}],"weights_recorded_at":"2026-10-25T19:59:55.995414+00:00"},"orb-v3-direct-20-mpa":{"fetched_at":"2026-08-03T17:18:51.755916+00:00","verified_at":"2026-10-25T19:59:55.995443+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-20-mpa-20250404.ckpt","size":102636362}],"weights_recorded_at":"2026-10-25T19:59:55.995447+00:00"},"orb-v3-conservative-inf-omat":{"fetched_at":"2026-08-03T17:18:52.017384+00:00","verified_at":"2026-10-25T19:59:55.995406+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-inf-omat-20250404.ckpt","size":102097517}],"weights_recorded_at":"2026-10-25T19:59:55.995409+00:00"},"orb-v3-conservative-inf-mpa":{"fetched_at":"2026-08-03T17:18:52.075825+00:00","verified_at":"2026-10-25T19:59:55.995428+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-inf-mpa-20250404.ckpt","size":102097517}],"weights_recorded_at":"2026-10-25T19:59:55.995430+00:00"},"orb-v3-direct-20-omat":{"fetched_at":"2026-08-03T17:18:53.566555+00:00","verified_at":"2026-10-25T19:59:55.995422+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-20-omat-20250404.ckpt","size":102636510}],"weights_recorded_at":"2026-10-25T19:59:55.995425+00:00"},"orb-v3-direct-inf-omat":{"fetched_at":"2026-08-03T17:19:14.412381+00:00","verified_at":"2026-10-25T19:59:55.995417+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-inf-omat-20250404.ckpt","size":102636658}],"weights_recorded_at":"2026-10-25T19:59:55.995420+00:00"},"orb-v3-direct-inf-mpa":{"fetched_at":"2026-08-03T17:19:16.914781+00:00","verified_at":"2026-10-25T19:59:55.995438+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-inf-mpa-20250404.ckpt","size":102636510}],"weights_recorded_at":"2026-10-25T19:59:55.995441+00:00"},"orb-v3:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:59:55.999580+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"orb":{"built_at":"2026-07-31T18:36:40.771968+00:00","source_hash":"sha256:4ae1362a989170b29d5082a23f9586c50d4f04b700bd29b8c65609efcbe52bef","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # >=0.5,<0.6: 0.5.5 is what the verified Delta env resolved — the v2\n# # loaders keep their single-return API through 0.5.x. 0.4.x is broken\n# # for us: it imports pynanoflann, which is git-only and undeclared, so\n# # a fresh build dies at import (Delta, 2026-07-31). 0.6 raises the\n# # Python floor to 3.12 — that line lives in orb_v3.py.\n# \"orb-models>=0.5,<0.6\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# # Not imported here — constrains orb-models' transitive dep. setup()'s\n# # no-lock serve path relies on cached_path returning local files without\n# # locking or writing, verified against exactly this version (#67).\n# \"cached_path==1.8.10\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"Orb v2 env — kept only for the built-in-D3 dispersion variant.\n\nOrb v3 (orb_v3.py) is the primary orb env; it supersedes the v2 checkpoints\nexcept orb-d3-v2, which has no v3 equivalent (v3 ships no dispersion-corrected\nmodel). Catalog trimmed to that one id 2026-07-30. The two lines can't share\nan env: the v3 loaders need orb-models>=0.6, which raises the Python floor\nto 3.12 and torch to 2.8 (the v2 loaders here are fine through 0.5.x).\n\"\"\"\n\nimport os\nimport shutil\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"orb-d3-v2\": \"orb-d3-v2\",\n # Your own fine-tuned v2-architecture weights: pair with weights=\n # (loaded via setup_from_path). v3 fine-tunes go to orb-v3:custom.\n \"orb-v2:custom\": None,\n}\n\n\ndef _default_weights_url(load_fn) -> str:\n \"\"\"The upstream URL baked into the loader's ``weights_path`` default.\"\"\"\n import inspect\n\n default = inspect.signature(load_fn).parameters[\"weights_path\"].default\n if not isinstance(default, str) or not default.startswith((\"http://\", \"https://\")):\n raise RuntimeError(\n f\"{load_fn.__name__} has no URL default for weights_path \"\n f\"(got {default!r}); update this env file for the installed orb-models\"\n )\n return default\n\n\ndef _local_weights_path(url: str) -> Path:\n \"\"\"Where the checkpoint lives in the shared model cache.\"\"\"\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"orb\" / os.path.basename(url)\n\n\ndef _fetch(url: str, dest: Path) -> None:\n \"\"\"Download ``url`` to ``dest`` atomically (tmp file + rename).\"\"\"\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp = dest.with_name(f\"{dest.name}.tmp.{os.getpid()}\")\n try:\n with urllib.request.urlopen(url) as resp, open(tmp, \"wb\") as out:\n shutil.copyfileobj(resp, out)\n os.replace(tmp, dest)\n finally:\n tmp.unlink(missing_ok=True)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.calculator import ORBCalculator\n\n # orb-models exposes one function per checkpoint, e.g. pretrained.orb_v2().\n fn_name = CHECKPOINTS[checkpoint].replace(\"-\", \"_\")\n load_fn = getattr(pretrained, fn_name)\n\n # orb-models resolves its default weights URL through `cached_path`, which\n # write-locks its cache dir even on warm hits — EACCES for anyone who can't\n # write the shared install (Garden-AI/rootstock#67). Handed a *local* path\n # instead, cached_path returns it without locking. So the weights are\n # pre-fetched into the shared model cache at `rootstock add` time\n # (maintainer, cache writable) and every later serve loads that file.\n url = _default_weights_url(load_fn)\n weights = _local_weights_path(url)\n if not weights.exists():\n _fetch(url, weights)\n\n orbff = load_fn(weights_path=str(weights), device=torch.device(device))\n return ORBCalculator(orbff, device=torch.device(device))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", arch: str = \"orb-v2\"):\n # Custom checkpoints (`:custom` ids with user weights). A weights file doesn't say\n # which orb architecture produced it, so `arch` names the pretrained\n # loader to instantiate — pass the right one at call time\n # (setup_kwargs={\"arch\": ...} / --kwarg arch=...). Handing the loader a local path also means no\n # network and no cached_path locking (see setup()).\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.calculator import ORBCalculator\n\n fn_name = arch.replace(\"-\", \"_\")\n try:\n load_fn = getattr(pretrained, fn_name)\n except AttributeError:\n raise ValueError(\n f\"unknown orb architecture {arch!r}; expected a loader name from \"\n f\"orb_models.forcefield.pretrained, e.g. orb-v2, orb-d3-v2\"\n ) from None\n\n orbff = load_fn(weights_path=path, device=torch.device(device))\n return ORBCalculator(orbff, device=torch.device(device))\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","cached-path":"1.8.10","orb-models":"0.5.5","rootstock":"1.3.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:58ce179e368a874745398cf83be4cf4a5a96dc6e28669213de8e201b67be49d3","checkpoints":{"orb-d3-v2":{"fetched_at":"2026-08-03T17:18:30.895591+00:00","verified_at":"2026-10-25T19:59:55.995449+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-d3-v2-20241011.ckpt","size":100967743}],"weights_recorded_at":"2026-10-25T19:59:55.995452+00:00"}}},"pet":{"built_at":"2026-08-03T20:28:04.390522+00:00","source_hash":"sha256:2ed7f37a981d7b622c8f25c58a80b4269a919554a0fd61e36ddee4a7678654fd","source":"# /// script\n# requires-python = \">=3.11,<3.15\"\n# dependencies = [\n# \"upet>=0.2.6\",\n# \"ase>=3.22\",\n# # upet pulls nvalchemi-toolkit-ops unpinned; 0.4+ needs torch>=2.8 at\n# # runtime but only declares the constraint on its extras, so the\n# # resolver won't catch it (same trap as the tensornet env).\n# \"torch>=2.8,<2.14\",\n# ]\n# ///\n\"\"\"PET env — hosts lab-cosmo's UPET foundation checkpoints (PET-MAD successor).\n\nThe upstream string encodes model@version; versions are pinned rather than\n\"latest\" so rebuilds serve the same weights. pet-omatpes-l is trained at the\nr2SCAN level of theory — its energies are not comparable to the PBE-level\npet-oam models.\n\"\"\"\n\nCHECKPOINTS = {\n \"pet-oam-xl\": \"pet-oam-xl@1.0.0\",\n \"pet-omatpes-l\": \"pet-omatpes-l@0.1.0\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from huggingface_hub import hf_hub_download\n from upet.calculator import UPETCalculator\n\n # metatomic-ase 0.1.2's nvalchemi fast path passes a float max_neighbors\n # (len * max(128, cutoff**3)) into torch.full — TypeError for cutoffs > ~5 Å.\n # Force the vesin fallback; the flag is read at call time.\n import metatomic_ase._neighbors as _mta_n\n _mta_n.HAS_NVALCHEMIOPS = False\n\n # Passing model=/version= makes UPETCalculator resolve the name by listing\n # the hub repo — an uncached API call that fails on workers, which run\n # with HF_HUB_OFFLINE=1 (and on any node without internet). Fetch the\n # pinned file ourselves — a cache hit needs no network even offline — and\n # hand it over as checkpoint_path, which skips the resolve entirely.\n model, version = CHECKPOINTS[checkpoint].split(\"@\", 1)\n path = hf_hub_download(\n repo_id=\"lab-cosmo/upet\",\n filename=f\"{model}-v{version}.ckpt\",\n subfolder=\"models\",\n )\n return UPETCalculator(checkpoint_path=path, device=device)\n","python_requires":">=3.11,<3.15","dependencies":{"ase":"3.29.0","rootstock":"1.4.3","torch":"2.13.0","upet":"0.2.6"},"lock_hash":"sha256:9809737b1bd5fea93de01e7f7340e8140d88dbda6fa66c09ef1d7144435a6f5c","checkpoints":{}},"sevennet":{"built_at":"2026-07-31T18:37:57.316160+00:00","source_hash":"sha256:560f24594f3048803e4527d1a3761784e11c9a1d198c14781697f323a1384057","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"sevenn>=0.10.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"SevenNet env — hosts pretrained SevenNet universal potentials.\n\nSevenNet (SCalable EquiVariance-Enabled Neural Network) ships several\npretrained models, loaded by keyword through ``SevenNetCalculator``.\n\nMulti-fidelity models (``7net-omni``, ``7net-mf-ompa``) take a ``modal``\nargument selecting the training fidelity (e.g. ``\"mpa\"`` or ``\"omat24\"``),\nand sevenn releases after mid-2026 make it mandatory (older ones defaulted\nsilently). When not given, setup() defaults them to ``\"mpa\"`` — the\nMPtrj+sAlex fidelity, consistent with the lineage of the other checkpoints\nhere — so no-kwarg paths (smoke-test, plain ``rootstock add``) keep working.\nOverride at runtime via ``setup_kwargs={\"modal\": ...}`` on\nRootstockCalculator (or ``--kwarg modal=...`` for ``rootstock add``);\nsingle-fidelity models ignore it.\n\"\"\"\n\nCHECKPOINTS = {\n \"sevennet-0\": \"7net-0\",\n \"sevennet-l3i5\": \"7net-l3i5\",\n \"sevennet-omat\": \"7net-omat\",\n \"sevennet-mf-ompa\": \"7net-mf-ompa\",\n \"sevennet-omni\": \"7net-omni\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"sevennet:custom\": None,\n}\n\n\n# Multi-fidelity models and the fidelity used when modal isn't specified.\nMULTI_FIDELITY_DEFAULT_MODAL = {\n \"sevennet-mf-ompa\": \"mpa\",\n \"sevennet-omni\": \"mpa\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", modal: str | None = None):\n \"\"\"\n Load a SevenNet calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\").\n modal: Fidelity selector for multi-fidelity models (e.g. \"mpa\",\n \"omat24\"). Defaults to \"mpa\" for 7net-omni / 7net-mf-ompa;\n ignored by the single-fidelity models.\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n from sevenn.calculator import SevenNetCalculator\n\n if modal is None:\n modal = MULTI_FIDELITY_DEFAULT_MODAL.get(checkpoint)\n kwargs = {\"modal\": modal} if modal is not None else {}\n return SevenNetCalculator(model=CHECKPOINTS[checkpoint], device=device, **kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", modal: str | None = None):\n # Custom checkpoints (`:custom` ids with user weights): SevenNetCalculator loads a\n # checkpoint file directly. A multi-fidelity fine-tune must pass its\n # fidelity (setup_kwargs={\"modal\": \"mpa\"} / --kwarg modal=mpa);\n # single-fidelity ones need no kwargs.\n from sevenn.calculator import SevenNetCalculator\n\n kwargs = {\"modal\": modal} if modal is not None else {}\n return SevenNetCalculator(model=path, device=device, **kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.3.0","sevenn":"0.13.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:5e341191eeda23990bfa2217b40a922d8c0e31eb5361ebb1233fc10c4edced4d","checkpoints":{"sevennet-0":{"fetched_at":"2026-08-03T17:21:35.263018+00:00","verified_at":"2026-10-25T19:59:55.999483+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-l3i5":{"fetched_at":"2026-08-03T17:21:35.332510+00:00","verified_at":"2026-10-25T19:59:55.999486+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-omat":{"fetched_at":"2026-08-03T17:22:03.544304+00:00","verified_at":"2026-10-25T19:59:55.999489+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-mf-ompa":{"fetched_at":"2026-08-03T17:22:04.480485+00:00","verified_at":"2026-10-25T19:59:55.999492+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-omni":{"fetched_at":"2026-08-03T17:22:06.723092+00:00","verified_at":"2026-10-25T19:59:55.999495+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"uma":{"built_at":"2026-08-06T19:38:53.367289+00:00","source_hash":"sha256:a21a78ddc6bc95aaabc4e17428d768b8d713671906246b4baf30a4e965b7d1f5","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"UMA env — hosts Meta's UMA foundation model via FAIRChem.\n\nfairchem-core v2 dropped the torch-geometric / pyg-find-links install dance, so\nthis env is a plain PyPI install. The original uma-s-1 had an extensivity bug\nand was removed from the fairchem 2.20 registry — use uma-s-1p1 or uma-s-1p2p1.\n\"\"\"\n\nCHECKPOINTS = {\n \"uma-s-1p1\": \"uma-s-1p1\",\n # uma-s-1p2 has a known major bug; uma-s-1p2p1 fixes it and is the\n # upstream-recommended small model. 1p2 stays listed for reproducibility\n # of existing runs.\n \"uma-s-1p2\": \"uma-s-1p2\",\n # uma-s-1p2p1 is in fairchem's registry on git main but NOT in any\n # release yet (latest fairchem-core 2.21.0, 2026-06-08, lacks it — the\n # 2026-07-30 sync failed on exactly this). Re-add when the next\n # fairchem-core ships, and bump the dependency floor to that version.\n # \"uma-s-1p2p1\": \"uma-s-1p2p1\",\n \"uma-m-1p1\": \"uma-m-1p1\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"uma:custom\": None,\n}\n\n\ndef _fairchem_device(device: str) -> str:\n \"\"\"Translate an indexed device (\"cuda:2\") into what fairchem v2 accepts.\n\n MLIPPredictUnit._setup_device asserts `device in [\"cpu\", \"cuda\"]` and then\n resolves the real GPU itself via get_device_for_local_rank(), which returns\n f\"cuda:{torch.cuda.current_device()}\". So an index has to travel through\n torch's current-device state, not the argument. Verifying several\n checkpoints at once on a multi-GPU node hands each worker \"cuda:N\" — that\n killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync\n (4x A100, VERIFY_JOBS=4), while single-GPU Sophia never hit it.\n \"\"\"\n if device.startswith(\"cuda:\"):\n import torch\n\n torch.cuda.set_device(int(device.split(\":\", 1)[1]))\n return \"cuda\"\n return device\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", task: str = \"omat\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(\n CHECKPOINTS[checkpoint], device=_fairchem_device(device)\n )\n return FAIRChemCalculator(predictor, task_name=task)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", task: str = \"omat\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=_fairchem_device(device))\n return FAIRChemCalculator(predictor, task_name=task)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.6.0","torch":"2.8.0"},"lock_hash":"sha256:c99d413acdd9b5b2cdd0241ba29966574fe6800af2b63376abacd0ad28fed8b0","checkpoints":{"uma-s-1p1":{"fetched_at":"2026-08-03T17:25:15.360399+00:00","verified_at":"2026-10-25T19:59:55.999498+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/07068e9c76702ca173d13155095f2117c1b327ec228557e64cd2709c777b824a","size":1173759575},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.999502+00:00"},"uma-m-1p1":{"fetched_at":"2026-08-03T17:25:21.596525+00:00","verified_at":"2026-10-25T19:59:55.999512+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/c30034edbf2e127f703f814cacb632661767da99b3e71c5b2ee5290510a52d68","size":11174706771},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.999516+00:00"},"uma-s-1p2":{"fetched_at":"2026-08-03T17:26:22.685148+00:00","verified_at":"2026-10-25T19:59:55.999505+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/ba5c0d912efa22dc238e5fb1b5b7f66ee2e68c48c1b95b7cfd7fe1da5938398b","size":2333393167},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.999509+00:00"},"uma:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:59:55.999583+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"tensornet":{"built_at":"2026-08-03T20:28:00.083517+00:00","source_hash":"sha256:659d5d93612741c0a62befa641f1f5d2259ab760837e261d8c84ecf5d4564014","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"torch>=2.4.0,<2.5\",\n# \"ase>=3.22\",\n# \"huggingface_hub\",\n# \"matgl\",\n# # 0.4+ needs torch>=2.8 at runtime (custom-op registration uses string\n# # annotations infer_schema can't parse on older torch) but only declares\n# # the constraint on its extras, so the resolver won't catch it.\n# # \"nvalchemi-toolkit-ops<0.4\",\n# \"pymatgen\",\n# \"monty\",\n# \"ruamel.yaml\",\n# \"scipy\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# \"torch-spline-conv\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n#\n# [tool.uv.sources]\n# matgl = { git = \"https://github.com/materialsvirtuallab/matgl.git\", rev = \"35f3ea2d2028febfeb99bb6707f4c6f5223a0851\" }\n# ///\n\"\"\"TensorNet env — hosts MatPES TensorNet checkpoints via MatGL.\"\"\"\n\nCHECKPOINTS = {\n \"tensornet-matpes-pbe-2025-2\": \"materialyze/TensorNet-PES-MatPES-PBE-2025.2\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n # matgl 1.0.0 imports ExpCellFilter from ase.constraints, but it moved to\n # ase.filters in ASE 3.23. Patch it in before matgl imports.\n import ase.constraints\n\n if not hasattr(ase.constraints, \"ExpCellFilter\"):\n from ase.filters import ExpCellFilter\n\n ase.constraints.ExpCellFilter = ExpCellFilter\n\n # DGL 2.x graphbolt imports torchdata submodules removed in torchdata>=0.7.\n # Stub the entire graphbolt subpackage before `import dgl` runs; DGL's\n # __init__ will use our empty stub and skip the real graphbolt initialisation.\n # matgl only uses DGL for graph construction — graphbolt is never called.\n import sys, types\n\n for _name in [\n \"dgl.graphbolt\",\n \"dgl.graphbolt.base\",\n \"dgl.graphbolt.dataloader\",\n \"dgl.graphbolt.feature_fetcher\",\n \"dgl.graphbolt.minibatch_transformer\",\n ]:\n if _name not in sys.modules:\n sys.modules[_name] = types.ModuleType(_name)\n\n from huggingface_hub import snapshot_download\n\n import matgl\n from matgl.ext.ase import PESCalculator\n\n # matgl 1.0.0 load_model only checks the GitHub manifest; HF models must\n # be downloaded explicitly and passed as a local path.\n local_path = snapshot_download(repo_id=CHECKPOINTS[checkpoint])\n pot = matgl.load_model(local_path).to(device)\n return PESCalculator(potential=pot)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","huggingface-hub":"1.26.0","matgl":"4.0.3","monty":"2026.7.16","pymatgen":"2026.5.4","rootstock":"1.4.3","ruamel-yaml":"0.19.1","scipy":"1.17.1","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121","torch-spline-conv":"1.2.2+pt24cu121"},"lock_hash":null,"checkpoints":{"tensornet-matpes-pbe-2025-2":{"fetched_at":"2026-08-03T19:34:11.158461+00:00","verified_at":"2026-10-25T19:59:55.999519+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/blobs/19a0a506a58ea454e407040d55c393a174a5165834f22ea95c4c923ad9a406cc","size":4495},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/blobs/544966013206da7faa1966ffa360a82fb2d31c3e","size":5298},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/blobs/b9efa8c708014dd6c37e348f65dc1d8ae5c7bd5c7290de4d9f210e16c7fc3766","size":3402824},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/refs/main","size":40},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/trees/7ac2b90400a87eca7f44c1e7956dca6b1c85e418.json","size":906}],"weights_recorded_at":"2026-10-25T19:59:55.999524+00:00"}}},"polaris-ani":{"built_at":"2026-08-05T20:22:42.062311+00:00","source_hash":"sha256:e758e872f6f077df2062463eb9bf78d00baacdf535008e8b563fb34579edc35f","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"torchani>=2.2\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"\nANI-2x environment for Rootstock — Polaris variant.\n\nIdentical to ani.py except torch is pinned to the cu128 index: Polaris's\ndriver stack tops out at CUDA 12.8, and the default PyPI torch wheel is\nbuilt against a newer CUDA, so cuda init fails at setup() (smoke-test\n2026-08-04). Sophia keeps the unrestricted env.\n\nANI-2x is a neural network potential for organic molecules containing\nH, C, N, O, F, S, Cl. It is not a universal potential — do not use it\nfor inorganic or periodic systems.\n\nModels:\n - \"ANI2x\": ANI-2x ensemble (default, 8 networks)\n - \"ANI1ccx\": ANI-1ccx, trained on CCSD(T)/CBS data (H, C, N, O only)\n - \"ANI1x\": ANI-1x (H, C, N, O only)\n\"\"\"\n\nCHECKPOINTS = {\n \"ani-2x\": \"ANI2x\",\n \"ani-1ccx\": \"ANI1ccx\",\n \"ani-1x\": \"ANI1x\",\n}\n\nCLUSTERS = [\"polaris\"]\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an ANI calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n import torchani\n\n model_map = {\n \"ANI2x\": torchani.models.ANI2x,\n \"ANI1ccx\": torchani.models.ANI1ccx,\n \"ANI1x\": torchani.models.ANI1x,\n }\n model = CHECKPOINTS[checkpoint]\n\n return model_map[model](periodic_table_index=True).to(device).ase()\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.6.0","torch":"2.11.0+cu128","torchani":"2.8.4"},"lock_hash":"sha256:1eee8e0a85fade92fc8f8c30e314b8e91bc8a3fc6b3c4d9375cc599fff02cc46","checkpoints":{"ani-1x":{"fetched_at":"2026-08-06T18:03:01.286673+00:00","verified_at":"2026-10-25T19:59:55.994670+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.local/share/Torchani/StateDicts/.cache/huggingface/download/ani1x_state_dict.pt.metadata","size":124},{"path":"home/.local/share/Torchani/StateDicts/ani1x_state_dict.pt","size":10546029}],"weights_recorded_at":"2026-10-25T19:59:55.994674+00:00"},"ani-1ccx":{"fetched_at":"2026-08-06T18:03:01.300850+00:00","verified_at":"2026-10-25T19:59:55.994664+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.local/share/Torchani/StateDicts/.cache/huggingface/download/ani1ccx_state_dict.pt.metadata","size":125},{"path":"home/.local/share/Torchani/StateDicts/ani1ccx_state_dict.pt","size":10546029}],"weights_recorded_at":"2026-10-25T19:59:55.994668+00:00"},"ani-2x":{"fetched_at":"2026-08-06T18:03:01.427309+00:00","verified_at":"2026-10-25T19:59:55.994656+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.local/share/Torchani/StateDicts/.cache/huggingface/download/ani2x_state_dict.pt.metadata","size":125},{"path":"home/.local/share/Torchani/StateDicts/ani2x_state_dict.pt","size":54982957}],"weights_recorded_at":"2026-10-25T19:59:55.994660+00:00"}}},"polaris-pet":{"built_at":"2026-08-05T20:22:41.090075+00:00","source_hash":"sha256:095563066044cf6c727de4a2cbf78c1d0f61142f16ce47d66fa3f4fca01eb312","source":"# /// script\n# requires-python = \">=3.11,<3.15\"\n# dependencies = [\n# \"upet>=0.2.6\",\n# \"ase>=3.22\",\n# # upet pulls nvalchemi-toolkit-ops unpinned; 0.4+ needs torch>=2.8 at\n# # runtime but only declares the constraint on its extras, so the\n# # resolver won't catch it (same trap as the tensornet env).\n# \"torch>=2.8,<2.14\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"PET env — Polaris variant of pet.py (UPET foundation checkpoints).\n\nIdentical to pet.py except torch is pinned to the cu128 index: Polaris's\ndriver stack tops out at CUDA 12.8, and the default PyPI torch wheel is\nbuilt against a newer CUDA, so metatomic's pick_device finds no usable\ncuda at setup() (smoke-test 2026-08-04). Sophia keeps the unrestricted env.\n\nThe upstream string encodes model@version; versions are pinned rather than\n\"latest\" so rebuilds serve the same weights. pet-omatpes-l is trained at the\nr2SCAN level of theory — its energies are not comparable to the PBE-level\npet-oam models.\n\"\"\"\n\nCHECKPOINTS = {\n \"pet-oam-xl\": \"pet-oam-xl@1.0.0\",\n \"pet-omatpes-l\": \"pet-omatpes-l@0.1.0\",\n}\n\nCLUSTERS = [\"polaris\"]\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from huggingface_hub import hf_hub_download\n from upet.calculator import UPETCalculator\n\n # metatomic-ase 0.1.2's nvalchemi fast path passes a float max_neighbors\n # (len * max(128, cutoff**3)) into torch.full — TypeError for cutoffs > ~5 Å.\n # Force the vesin fallback; the flag is read at call time.\n import metatomic_ase._neighbors as _mta_n\n _mta_n.HAS_NVALCHEMIOPS = False\n\n # Passing model=/version= makes UPETCalculator resolve the name by listing\n # the hub repo — an uncached API call that fails on workers, which run\n # with HF_HUB_OFFLINE=1 (and on any node without internet). Fetch the\n # pinned file ourselves — a cache hit needs no network even offline — and\n # hand it over as checkpoint_path, which skips the resolve entirely.\n model, version = CHECKPOINTS[checkpoint].split(\"@\", 1)\n path = hf_hub_download(\n repo_id=\"lab-cosmo/upet\",\n filename=f\"{model}-v{version}.ckpt\",\n subfolder=\"models\",\n )\n return UPETCalculator(checkpoint_path=path, device=device)\n","python_requires":">=3.11,<3.15","dependencies":{"ase":"3.29.0","rootstock":"1.6.0","torch":"2.11.0+cu128","upet":"0.2.6"},"lock_hash":"sha256:db5d5de3a74b8df105d3b146acc6fe5a09f3a6e9be7caa4c179969dd7b8e3c27","checkpoints":{"pet-omatpes-l":{"fetched_at":"2026-08-06T18:03:21.426504+00:00","verified_at":"2026-10-25T19:59:55.999474+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/blobs/879b1045391d88869522605a8b8b3cedeed74668e7062fdd7487548ab7b08004","size":771692879},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.999480+00:00"},"pet-oam-xl":{"fetched_at":"2026-08-06T18:03:24.011446+00:00","verified_at":"2026-10-25T19:59:55.995455+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/blobs/c3a67cd019969dfd4dcabe9574682fe035f861d3f1c10190989b36c983699409","size":2920687712},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.996527+00:00"}}},"allegro":{"built_at":"2026-10-16T21:43:39.464385+00:00","source_hash":"sha256:bc9261bd7205b13a8d1f74f17399694d9ba9f6d2fc05c50415b297a671ff9143","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # setup() uses nequip's loader internals and keys its compile cache on the\n# # nequip version, so stay within one minor release.\n# \"nequip>=0.19,<0.20\",\n# \"nequip-allegro>=0.8.3,<0.9\",\n# \"ase>=3.23\",\n# # AOTInductor compilation needs torch>=2.6. Listed here so the explicit\n# # cu128 index below routes it: PyPI's recent torch wheels bundle a CUDA\n# # newer than many cluster drivers support.\n# \"torch>=2.6,<2.12\",\n# # CUDA headers for compiling on nodes with no CUDA toolkit loaded (the\n# # runtime headers come with torch); matches the cu128 torch build.\n# \"cuda-toolkit[cccl,nvcc]>=12.8,<12.9\",\n# # AOTInductor's C++ builder imports torch.utils.cpp_extension, which\n# # imports setuptools at module level; uv venvs do not seed it.\n# \"setuptools>=70\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"Allegro env — hosts the Allegro foundation models from nequip.net.\n\nAllegro is a strictly local E(3)-equivariant interatomic potential from the\nKozinsky group (Harvard), built on the NequIP framework; its locality is what\nlets it scale across many GPUs in LAMMPS. Allegro-OAM-L is trained on OMat24\nand fine-tuned on sAlex + MPtrj; Allegro-MP-L is trained on MPtrj alone. All\nare PBE-level and licensed CC-BY-4.0. Versions are pinned to fixed Zenodo\nartifacts (md5-checked) so rebuilds serve the same weights; the package file\nis downloaded into the shared model cache at `rootstock add` time and loaded\nfrom there, so serving needs no network.\n\nCompilation: on GPU, models run through NequIP's AOTInductor path\n(``nequip-compile --mode aotinductor --target ase``), the path NequIP\nrecommends for inference. A compiled model only runs on the hardware and\nsoftware it was compiled with, so setup() compiles once per (model, GPU\narchitecture, torch version, nequip version) and caches the result: in the\nshared model cache when the caller can write there, otherwise in the caller's\nTorchInductor cache directory. The first GPU load of a model on a new GPU type\ntherefore compiles first: on an 8-core node with an A10G, about 2 minutes for\nAllegro-OAM-L, 5 for NequIP-OAM-L and 6-7 for NequIP-OAM-XL. If that load runs\ninto the worker startup timeout, raise it (``timeout=`` on\nRootstockCalculator, ``--verify-timeout`` for verification); an interrupted\ncompile still finishes in the background and later loads reuse the compiled\nfile. AOTInductor needs a C++17 host compiler (``c++``/``gcc`` >= 8, or\nwhatever ``CXX`` names) at compile time; CUDA headers come from the env's own\nNVIDIA wheels when no CUDA toolkit is loaded. Pass\n``setup_kwargs={\"compile\": False}`` to run eager instead. CPU loads always run\neager.\n\nCustom weights: ``allegro:custom`` with weights= accepts a packaged model\n(``.nequip.zip``), a training checkpoint (``.ckpt``), or a model already\ncompiled for ASE on this hardware (``.nequip.pt2``). If your model's type\nnames are not chemical symbols, pass\n``setup_kwargs={\"chemical_species_to_atom_type_map\": {\"H\": \"\", ...}}``.\n\"\"\"\n\nimport hashlib\nimport os\nimport shutil\nimport subprocess\nimport sys\nimport tempfile\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"allegro-oam-l\": \"Allegro-OAM-L-0.1\",\n \"allegro-mp-l\": \"Allegro-MP-L-0.1\",\n # Your own Allegro model: pair with weights= (loaded via setup_from_path).\n \"allegro:custom\": None,\n}\n\n# Explicit scope: both NVIDIA machines of the shared ALCF install.\nCLUSTERS = [\"sophia\", \"polaris\"]\n\n# Zenodo record and md5 of each pinned package file (the artifacts nequip.net\n# resolves mir-group/:0.1 to).\n_ARTIFACTS = {\n \"Allegro-OAM-L-0.1\": (\"16980200\", \"0db7f9b3c3a62e74d78b3fcf2973c462\"),\n \"Allegro-MP-L-0.1\": (\"16980200\", \"c68d5345e925ee444543d6d4d968161f\"),\n}\n\n\ndef _shared_cache() -> Path:\n return Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\") / \"nequip\"\n\n\ndef _package_path(name: str) -> Path:\n \"\"\"The pinned package file in the shared cache, downloading it if missing.\"\"\"\n path = _shared_cache() / \"packages\" / f\"{name}.nequip.zip\"\n if path.is_file():\n return path\n record, md5 = _ARTIFACTS[name]\n url = f\"https://zenodo.org/api/records/{record}/files/{name}.nequip.zip/content\"\n path.parent.mkdir(parents=True, exist_ok=True)\n tmp = path.with_name(f\".{path.name}.{os.getpid()}.part\")\n try:\n digest = hashlib.md5()\n with urllib.request.urlopen(url) as response, open(tmp, \"wb\") as f:\n while chunk := response.read(1 << 20):\n f.write(chunk)\n digest.update(chunk)\n if digest.hexdigest() != md5:\n raise RuntimeError(f\"{url}: md5 {digest.hexdigest()} does not match {md5}\")\n os.replace(tmp, path)\n finally:\n tmp.unlink(missing_ok=True)\n return path\n\n\ndef _compiled_dirs() -> list[Path]:\n \"\"\"Where compiled models are looked up, in order; the first writable one\n receives new compilations.\"\"\"\n dirs = [_shared_cache() / \"compiled\"]\n inductor = os.environ.get(\"TORCHINDUCTOR_CACHE_DIR\")\n if inductor:\n dirs.append(Path(inductor) / \"nequip\")\n else:\n dirs.append(Path.home() / \".cache\" / \"nequip\" / \"compiled\")\n return dirs\n\n\ndef _writable(directory: Path) -> bool:\n probe = directory\n while not probe.exists():\n probe = probe.parent\n return os.access(probe, os.W_OK | os.X_OK)\n\n\ndef _hardware_tag(device: str) -> str:\n import torch\n\n major, minor = torch.cuda.get_device_capability(torch.device(device))\n return f\"sm{major}{minor}\"\n\n\ndef _toolkit_env(env: dict, link_dir: Path) -> None:\n \"\"\"Point AOTInductor at CUDA headers and libcuda.\n\n The C++ wrapper AOTInductor builds needs ``$CUDA_HOME/include`` and links\n ``-lcuda``. torch derives CUDA_HOME from ``nvcc`` on PATH when unset, which\n goes wrong when that nvcc is a compiler-suite wrapper (e.g. the NVIDIA HPC\n SDK's, whose include dir has its own intrinsics headers and no cuda.h).\n So use CUDA_HOME, CUDA_PATH or nvcc's toolkit only if it really ships\n ``include/cuda.h``; otherwise fall back to the headers in NVIDIA's CUDA\n wheels (``nvidia/cuda_runtime/``, ``cuda_nvcc/`` and ``cuda_cccl/``) and\n link against the driver library this process has loaded.\n \"\"\"\n candidates = [env.get(\"CUDA_HOME\"), env.get(\"CUDA_PATH\")]\n nvcc = shutil.which(\"nvcc\", path=env.get(\"PATH\"))\n if nvcc:\n candidates.append(str(Path(nvcc).resolve().parent.parent))\n for home in filter(None, candidates):\n if (Path(home) / \"include\" / \"cuda.h\").is_file():\n env[\"CUDA_HOME\"] = home\n return\n env.pop(\"CUDA_PATH\", None)\n import nvidia\n\n headers = next(\n (h for base in nvidia.__path__ for h in Path(base).glob(\"*/include/cuda.h\")), None\n )\n if headers is None:\n return\n env[\"CUDA_HOME\"] = str(headers.parent.parent)\n extra = [\n str(h.parent.parent)\n for base in nvidia.__path__\n for marker in (\"crt/host_defines.h\", \"nv/target\")\n for h in Path(base).glob(f\"*/include/{marker}\")\n ]\n env[\"CPATH\"] = os.pathsep.join(filter(None, [*extra, env.get(\"CPATH\")]))\n with open(\"/proc/self/maps\") as maps:\n libcuda = next((line.split()[-1] for line in maps if \"/libcuda.so\" in line), None)\n if libcuda is not None:\n (link_dir / \"libcuda.so\").symlink_to(libcuda)\n env[\"LIBRARY_PATH\"] = os.pathsep.join(\n filter(None, [str(link_dir), env.get(\"LIBRARY_PATH\")])\n )\n\n\n# Host compilers tried, in order, when CXX is unset or not GNU-compatible.\n_CXX_CANDIDATES = (\"g++-14\", \"g++-13\", \"g++-12\", \"g++-11\", \"g++-10\", \"g++-9\", \"g++\", \"clang++\")\n\n\ndef _gnu_major(cxx: str) -> int | None:\n \"\"\"Major version of a GNU-compatible (gcc or clang) compiler, else None.\n\n AOTInductor passes gcc-style flags, which other compilers on HPC systems\n (e.g. NVIDIA's nvc++) reject.\n \"\"\"\n try:\n banner = subprocess.run([cxx, \"--version\"], capture_output=True, text=True).stdout\n version = subprocess.run([cxx, \"-dumpversion\"], capture_output=True, text=True).stdout\n except OSError:\n return None\n if \"Free Software Foundation\" not in banner and \"clang version\" not in banner:\n return None\n try:\n return int(version.strip().split(\".\")[0])\n except ValueError:\n return None\n\n\ndef _compiler_env(env: dict) -> None:\n \"\"\"Point CXX/CC at a GNU-compatible compiler new enough for C++17.\n\n Cluster compiler modules often export CXX/CC for a vendor compiler (e.g.\n nvc++/nvc). Keep the caller's CXX when it is gcc >= 9 or clang;\n otherwise use the first suitable candidate on PATH, with its matching C\n compiler (Triton builds a small C helper with CC during compilation).\n \"\"\"\n path = os.pathsep.join(filter(None, [env.get(\"PATH\"), \"/usr/bin\"]))\n for name in filter(None, [env.get(\"CXX\"), *_CXX_CANDIDATES]):\n cxx = shutil.which(name, path=path)\n major = _gnu_major(cxx) if cxx else None\n if major is None or (\"clang\" not in Path(cxx).name and major < 9):\n continue\n env[\"CXX\"] = cxx\n cc = shutil.which(\n Path(cxx).name.replace(\"clang++\", \"clang\").replace(\"g++\", \"gcc\"), path=path\n )\n if cc:\n env[\"CC\"] = cc\n return\n\n\ndef _compiled_model(source: Path, key: str, device: str) -> Path:\n \"\"\"Path to ``source`` compiled for ASE on ``device``, compiling on a miss.\"\"\"\n import nequip\n import torch\n\n tag = f\"{key}-{_hardware_tag(device)}-torch{torch.__version__}-nequip{nequip.__version__}\"\n name = f\"{tag.replace('+', '_')}.nequip.pt2\"\n dirs = _compiled_dirs()\n for directory in dirs:\n if (directory / name).is_file():\n return directory / name\n\n target_dir = next((d for d in dirs if _writable(d)), None)\n if target_dir is None:\n raise RuntimeError(f\"no writable directory to cache the compiled model in: {dirs}\")\n target_dir.mkdir(parents=True, exist_ok=True)\n out = target_dir / name\n tmp = target_dir / f\".{tag}.{os.getpid()}.nequip.pt2\"\n\n # Compile in a child process pinned to the same GPU, so the compiler's\n # memory is released before the model is loaded. The child moves its\n # output into place itself: if this worker is killed mid-compile (e.g. by\n # a startup timeout), the compile still completes and lands in the cache.\n env = dict(os.environ)\n if \":\" in device:\n index = int(device.split(\":\", 1)[1])\n visible = env.get(\"CUDA_VISIBLE_DEVICES\")\n env[\"CUDA_VISIBLE_DEVICES\"] = visible.split(\",\")[index] if visible else str(index)\n compile_and_move = (\n \"import os, sys; from nequip.scripts.compile import main; \"\n \"main(sys.argv[1:-1]); os.replace(sys.argv[2], sys.argv[-1])\"\n )\n cmd = [\n sys.executable,\n \"-c\",\n compile_and_move,\n str(source),\n str(tmp),\n \"--mode\",\n \"aotinductor\",\n \"--device\",\n \"cuda\",\n \"--target\",\n \"ase\",\n str(out),\n ]\n print(f\"nequip: compiling {source.name} for {tag} (one-time, several minutes)\", file=sys.stderr)\n try:\n with tempfile.TemporaryDirectory() as link_dir:\n _toolkit_env(env, Path(link_dir))\n _compiler_env(env)\n result = subprocess.run(cmd, env=env, stdout=sys.stderr, stderr=sys.stderr)\n if result.returncode != 0:\n raise RuntimeError(\n f\"nequip-compile failed (exit {result.returncode}); see the log above. \"\n \"It needs gcc >= 9 or clang (set CXX if neither is on PATH). \"\n 'Pass setup_kwargs={\"compile\": False} to run eager instead.'\n )\n finally:\n tmp.unlink(missing_ok=True)\n return out\n\n\ndef _load(source: Path, key: str, device: str, compile: bool, kwargs: dict):\n # torch's AOTInductor package loader reaches torch._inductor.codecache\n # without importing it, which fails in a fresh process.\n import torch._inductor.codecache # noqa: F401\n from nequip.integrations.ase import NequIPCalculator\n\n if source.name.endswith(\".nequip.pt2\"):\n return NequIPCalculator.from_compiled_model(str(source), device=device, **kwargs)\n if compile and device != \"cpu\":\n compiled = _compiled_model(source, key, device)\n return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs)\n # Eager load of a package or checkpoint. nequip exposes this only as a\n # private classmethod (its own tests use it); the public loader takes\n # compiled models only.\n return NequIPCalculator._from_saved_model(str(source), device=device, **kwargs)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", compile: bool = True, **kwargs):\n \"\"\"\n Load an Allegro calculator for a pinned nequip.net foundation model.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\").\n compile: Run the AOTInductor-compiled model on GPU (compiling and\n caching it on first use). False runs eager PyTorch.\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n name = CHECKPOINTS[checkpoint]\n kwargs.setdefault(\"chemical_species_to_atom_type_map\", True)\n return _load(_package_path(name), name, device, compile, kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", compile: bool = True, **kwargs):\n # Custom checkpoints (`:custom` ids with user weights): an Allegro .nequip.zip\n # package or .ckpt checkpoint (compiled on first GPU use, cached under a\n # key that includes the file's hash) or a ready .nequip.pt2.\n source = Path(path)\n key = source.name\n if not key.endswith(\".nequip.pt2\") and compile and device != \"cpu\":\n digest = hashlib.sha256()\n with open(source, \"rb\") as f:\n while chunk := f.read(1 << 20):\n digest.update(chunk)\n key = f\"custom-{digest.hexdigest()[:16]}\"\n return _load(source, key, device, compile, kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","cuda-toolkit":"12.8.2.0","nequip":"0.19.1","nequip-allegro":"0.8.3","rootstock":"1.6.4","setuptools":"84.0.0","torch":"2.10.0+cu128"},"lock_hash":"sha256:c022acfd951ab8b9fe6e1894c748a5f576b867ef693eb05d7eca18adcfae2867","checkpoints":{"allegro-oam-l":{"fetched_at":"2026-10-05T21:00:14.206502+00:00","verified_at":"2026-10-25T19:59:55.999527+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/nequip/compiled/Allegro-OAM-L-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":42039109}],"weights_recorded_at":"2026-10-25T19:59:55.999530+00:00"},"allegro-mp-l":{"fetched_at":"2026-10-05T21:00:21.906775+00:00","verified_at":"2026-10-25T19:59:55.999533+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/nequip/compiled/Allegro-MP-L-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":78600784}],"weights_recorded_at":"2026-10-25T19:59:55.999535+00:00"},"allegro:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:59:55.999585+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"nequip":{"built_at":"2026-10-16T21:43:51.189255+00:00","source_hash":"sha256:eada3a775ec8af9ad476154b832df248cbd29157d1b14f7fad2453492e187cb7","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # setup() uses nequip's loader internals and keys its compile cache on the\n# # nequip version, so stay within one minor release.\n# \"nequip>=0.19,<0.20\",\n# \"ase>=3.23\",\n# # AOTInductor compilation needs torch>=2.6. Listed here so the explicit\n# # cu128 index below routes it: PyPI's recent torch wheels bundle a CUDA\n# # newer than many cluster drivers support.\n# \"torch>=2.6,<2.12\",\n# # CUDA headers for compiling on nodes with no CUDA toolkit loaded (the\n# # runtime headers come with torch); matches the cu128 torch build.\n# \"cuda-toolkit[cccl,nvcc]>=12.8,<12.9\",\n# # AOTInductor's C++ builder imports torch.utils.cpp_extension, which\n# # imports setuptools at module level; uv venvs do not seed it.\n# \"setuptools>=70\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"NequIP env — hosts the NequIP-OAM foundation models from nequip.net.\n\nNequIP is an E(3)-equivariant message-passing interatomic potential from the\nKozinsky group (Harvard). The OAM models are trained on OMat24 and fine-tuned\non sAlex + MPtrj; NequIP-MP-L is trained on MPtrj alone. All are PBE-level and\nlicensed CC-BY-4.0. Versions are pinned to fixed Zenodo artifacts\n(md5-checked) so rebuilds serve the same weights; the package file is\ndownloaded into the shared model cache at `rootstock add` time and loaded from\nthere, so serving needs no network.\n\nCompilation: on GPU, models run through nequip's AOTInductor path\n(``nequip-compile --mode aotinductor --target ase``), the path NequIP\nrecommends for inference. A compiled model only runs on the hardware and\nsoftware it was compiled with, so setup() compiles once per (model, GPU\narchitecture, torch version, nequip version) and caches the result: in the\nshared model cache when the caller can write there, otherwise in the caller's\nTorchInductor cache directory. The first GPU load of a model on a new GPU type\ntherefore compiles first: on an 8-core node with an A10G, about 2 minutes for\nAllegro-OAM-L, 5 for NequIP-OAM-L and 6-7 for NequIP-OAM-XL. If that load runs\ninto the worker startup timeout, raise it (``timeout=`` on\nRootstockCalculator, ``--verify-timeout`` for verification); an interrupted\ncompile still finishes in the background and later loads reuse the compiled\nfile. AOTInductor needs a C++17 host compiler (``c++``/``gcc`` >= 8, or\nwhatever ``CXX`` names) at compile time; CUDA headers come from the env's own\nNVIDIA wheels when no CUDA toolkit is loaded. Pass\n``setup_kwargs={\"compile\": False}`` to run eager instead. CPU loads always run\neager.\n\nCustom weights: ``nequip:custom`` with weights= accepts a packaged model\n(``.nequip.zip``), a training checkpoint (``.ckpt``), or a model already\ncompiled for ASE on this hardware (``.nequip.pt2``). If your model's type\nnames are not chemical symbols, pass\n``setup_kwargs={\"chemical_species_to_atom_type_map\": {\"H\": \"\", ...}}``.\n\"\"\"\n\nimport hashlib\nimport os\nimport shutil\nimport subprocess\nimport sys\nimport tempfile\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"nequip-oam-s\": \"NequIP-OAM-S-0.1\",\n \"nequip-oam-m\": \"NequIP-OAM-M-0.1\",\n \"nequip-oam-l\": \"NequIP-OAM-L-0.1\",\n \"nequip-oam-xl\": \"NequIP-OAM-XL-0.1\",\n \"nequip-mp-l\": \"NequIP-MP-L-0.1\",\n # Your own NequIP model: pair with weights= (loaded via setup_from_path).\n \"nequip:custom\": None,\n}\n\n# Explicit scope: both NVIDIA machines of the shared ALCF install.\nCLUSTERS = [\"sophia\", \"polaris\"]\n\n# Zenodo record and md5 of each pinned package file (the artifacts nequip.net\n# resolves mir-group/:0.1 to).\n_ARTIFACTS = {\n \"NequIP-OAM-S-0.1\": (\"18775904\", \"399a98bf36fc4550bc85d48f35451c6c\"),\n \"NequIP-OAM-M-0.1\": (\"18775904\", \"e5b44583f4421be702b35d76503754cd\"),\n \"NequIP-OAM-L-0.1\": (\"18775904\", \"67144367c710a70a53a8e21acf331980\"),\n \"NequIP-OAM-XL-0.1\": (\"18775904\", \"3d2369c7238eb83a23141abdcb055a8f\"),\n \"NequIP-MP-L-0.1\": (\"18775904\", \"624a01d4cbbae2d3bb75dc7074532d13\"),\n}\n\n\ndef _shared_cache() -> Path:\n return Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\") / \"nequip\"\n\n\ndef _package_path(name: str) -> Path:\n \"\"\"The pinned package file in the shared cache, downloading it if missing.\"\"\"\n path = _shared_cache() / \"packages\" / f\"{name}.nequip.zip\"\n if path.is_file():\n return path\n record, md5 = _ARTIFACTS[name]\n url = f\"https://zenodo.org/api/records/{record}/files/{name}.nequip.zip/content\"\n path.parent.mkdir(parents=True, exist_ok=True)\n tmp = path.with_name(f\".{path.name}.{os.getpid()}.part\")\n try:\n digest = hashlib.md5()\n with urllib.request.urlopen(url) as response, open(tmp, \"wb\") as f:\n while chunk := response.read(1 << 20):\n f.write(chunk)\n digest.update(chunk)\n if digest.hexdigest() != md5:\n raise RuntimeError(f\"{url}: md5 {digest.hexdigest()} does not match {md5}\")\n os.replace(tmp, path)\n finally:\n tmp.unlink(missing_ok=True)\n return path\n\n\ndef _compiled_dirs() -> list[Path]:\n \"\"\"Where compiled models are looked up, in order; the first writable one\n receives new compilations.\"\"\"\n dirs = [_shared_cache() / \"compiled\"]\n inductor = os.environ.get(\"TORCHINDUCTOR_CACHE_DIR\")\n if inductor:\n dirs.append(Path(inductor) / \"nequip\")\n else:\n dirs.append(Path.home() / \".cache\" / \"nequip\" / \"compiled\")\n return dirs\n\n\ndef _writable(directory: Path) -> bool:\n probe = directory\n while not probe.exists():\n probe = probe.parent\n return os.access(probe, os.W_OK | os.X_OK)\n\n\ndef _hardware_tag(device: str) -> str:\n import torch\n\n major, minor = torch.cuda.get_device_capability(torch.device(device))\n return f\"sm{major}{minor}\"\n\n\ndef _toolkit_env(env: dict, link_dir: Path) -> None:\n \"\"\"Point AOTInductor at CUDA headers and libcuda.\n\n The C++ wrapper AOTInductor builds needs ``$CUDA_HOME/include`` and links\n ``-lcuda``. torch derives CUDA_HOME from ``nvcc`` on PATH when unset, which\n goes wrong when that nvcc is a compiler-suite wrapper (e.g. the NVIDIA HPC\n SDK's, whose include dir has its own intrinsics headers and no cuda.h).\n So use CUDA_HOME, CUDA_PATH or nvcc's toolkit only if it really ships\n ``include/cuda.h``; otherwise fall back to the headers in NVIDIA's CUDA\n wheels (``nvidia/cuda_runtime/``, ``cuda_nvcc/`` and ``cuda_cccl/``) and\n link against the driver library this process has loaded.\n \"\"\"\n candidates = [env.get(\"CUDA_HOME\"), env.get(\"CUDA_PATH\")]\n nvcc = shutil.which(\"nvcc\", path=env.get(\"PATH\"))\n if nvcc:\n candidates.append(str(Path(nvcc).resolve().parent.parent))\n for home in filter(None, candidates):\n if (Path(home) / \"include\" / \"cuda.h\").is_file():\n env[\"CUDA_HOME\"] = home\n return\n env.pop(\"CUDA_PATH\", None)\n import nvidia\n\n headers = next(\n (h for base in nvidia.__path__ for h in Path(base).glob(\"*/include/cuda.h\")), None\n )\n if headers is None:\n return\n env[\"CUDA_HOME\"] = str(headers.parent.parent)\n extra = [\n str(h.parent.parent)\n for base in nvidia.__path__\n for marker in (\"crt/host_defines.h\", \"nv/target\")\n for h in Path(base).glob(f\"*/include/{marker}\")\n ]\n env[\"CPATH\"] = os.pathsep.join(filter(None, [*extra, env.get(\"CPATH\")]))\n with open(\"/proc/self/maps\") as maps:\n libcuda = next((line.split()[-1] for line in maps if \"/libcuda.so\" in line), None)\n if libcuda is not None:\n (link_dir / \"libcuda.so\").symlink_to(libcuda)\n env[\"LIBRARY_PATH\"] = os.pathsep.join(\n filter(None, [str(link_dir), env.get(\"LIBRARY_PATH\")])\n )\n\n\n# Host compilers tried, in order, when CXX is unset or not GNU-compatible.\n_CXX_CANDIDATES = (\"g++-14\", \"g++-13\", \"g++-12\", \"g++-11\", \"g++-10\", \"g++-9\", \"g++\", \"clang++\")\n\n\ndef _gnu_major(cxx: str) -> int | None:\n \"\"\"Major version of a GNU-compatible (gcc or clang) compiler, else None.\n\n AOTInductor passes gcc-style flags, which other compilers on HPC systems\n (e.g. NVIDIA's nvc++) reject.\n \"\"\"\n try:\n banner = subprocess.run([cxx, \"--version\"], capture_output=True, text=True).stdout\n version = subprocess.run([cxx, \"-dumpversion\"], capture_output=True, text=True).stdout\n except OSError:\n return None\n if \"Free Software Foundation\" not in banner and \"clang version\" not in banner:\n return None\n try:\n return int(version.strip().split(\".\")[0])\n except ValueError:\n return None\n\n\ndef _compiler_env(env: dict) -> None:\n \"\"\"Point CXX/CC at a GNU-compatible compiler new enough for C++17.\n\n Cluster compiler modules often export CXX/CC for a vendor compiler (e.g.\n nvc++/nvc). Keep the caller's CXX when it is gcc >= 9 or clang;\n otherwise use the first suitable candidate on PATH, with its matching C\n compiler (Triton builds a small C helper with CC during compilation).\n \"\"\"\n path = os.pathsep.join(filter(None, [env.get(\"PATH\"), \"/usr/bin\"]))\n for name in filter(None, [env.get(\"CXX\"), *_CXX_CANDIDATES]):\n cxx = shutil.which(name, path=path)\n major = _gnu_major(cxx) if cxx else None\n if major is None or (\"clang\" not in Path(cxx).name and major < 9):\n continue\n env[\"CXX\"] = cxx\n cc = shutil.which(\n Path(cxx).name.replace(\"clang++\", \"clang\").replace(\"g++\", \"gcc\"), path=path\n )\n if cc:\n env[\"CC\"] = cc\n return\n\n\ndef _compiled_model(source: Path, key: str, device: str) -> Path:\n \"\"\"Path to ``source`` compiled for ASE on ``device``, compiling on a miss.\"\"\"\n import nequip\n import torch\n\n tag = f\"{key}-{_hardware_tag(device)}-torch{torch.__version__}-nequip{nequip.__version__}\"\n name = f\"{tag.replace('+', '_')}.nequip.pt2\"\n dirs = _compiled_dirs()\n for directory in dirs:\n if (directory / name).is_file():\n return directory / name\n\n target_dir = next((d for d in dirs if _writable(d)), None)\n if target_dir is None:\n raise RuntimeError(f\"no writable directory to cache the compiled model in: {dirs}\")\n target_dir.mkdir(parents=True, exist_ok=True)\n out = target_dir / name\n tmp = target_dir / f\".{tag}.{os.getpid()}.nequip.pt2\"\n\n # Compile in a child process pinned to the same GPU, so the compiler's\n # memory is released before the model is loaded. The child moves its\n # output into place itself: if this worker is killed mid-compile (e.g. by\n # a startup timeout), the compile still completes and lands in the cache.\n env = dict(os.environ)\n if \":\" in device:\n index = int(device.split(\":\", 1)[1])\n visible = env.get(\"CUDA_VISIBLE_DEVICES\")\n env[\"CUDA_VISIBLE_DEVICES\"] = visible.split(\",\")[index] if visible else str(index)\n compile_and_move = (\n \"import os, sys; from nequip.scripts.compile import main; \"\n \"main(sys.argv[1:-1]); os.replace(sys.argv[2], sys.argv[-1])\"\n )\n cmd = [\n sys.executable,\n \"-c\",\n compile_and_move,\n str(source),\n str(tmp),\n \"--mode\",\n \"aotinductor\",\n \"--device\",\n \"cuda\",\n \"--target\",\n \"ase\",\n str(out),\n ]\n print(f\"nequip: compiling {source.name} for {tag} (one-time, several minutes)\", file=sys.stderr)\n try:\n with tempfile.TemporaryDirectory() as link_dir:\n _toolkit_env(env, Path(link_dir))\n _compiler_env(env)\n result = subprocess.run(cmd, env=env, stdout=sys.stderr, stderr=sys.stderr)\n if result.returncode != 0:\n raise RuntimeError(\n f\"nequip-compile failed (exit {result.returncode}); see the log above. \"\n \"It needs gcc >= 9 or clang (set CXX if neither is on PATH). \"\n 'Pass setup_kwargs={\"compile\": False} to run eager instead.'\n )\n finally:\n tmp.unlink(missing_ok=True)\n return out\n\n\ndef _load(source: Path, key: str, device: str, compile: bool, kwargs: dict):\n # torch's AOTInductor package loader reaches torch._inductor.codecache\n # without importing it, which fails in a fresh process.\n import torch._inductor.codecache # noqa: F401\n from nequip.integrations.ase import NequIPCalculator\n\n if source.name.endswith(\".nequip.pt2\"):\n return NequIPCalculator.from_compiled_model(str(source), device=device, **kwargs)\n if compile and device != \"cpu\":\n compiled = _compiled_model(source, key, device)\n return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs)\n # Eager load of a package or checkpoint. nequip exposes this only as a\n # private classmethod (its own tests use it); the public loader takes\n # compiled models only.\n return NequIPCalculator._from_saved_model(str(source), device=device, **kwargs)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", compile: bool = True, **kwargs):\n \"\"\"\n Load a NequIP calculator for a pinned nequip.net foundation model.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\").\n compile: Run the AOTInductor-compiled model on GPU (compiling and\n caching it on first use). False runs eager PyTorch.\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n name = CHECKPOINTS[checkpoint]\n kwargs.setdefault(\"chemical_species_to_atom_type_map\", True)\n return _load(_package_path(name), name, device, compile, kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", compile: bool = True, **kwargs):\n # Custom checkpoints (`:custom` ids with user weights): a .nequip.zip\n # package or .ckpt checkpoint (compiled on first GPU use, cached under a\n # key that includes the file's hash) or a ready .nequip.pt2.\n source = Path(path)\n key = source.name\n if not key.endswith(\".nequip.pt2\") and compile and device != \"cpu\":\n digest = hashlib.sha256()\n with open(source, \"rb\") as f:\n while chunk := f.read(1 << 20):\n digest.update(chunk)\n key = f\"custom-{digest.hexdigest()[:16]}\"\n return _load(source, key, device, compile, kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","cuda-toolkit":"12.8.2.0","nequip":"0.19.1","rootstock":"1.6.4","setuptools":"84.0.0","torch":"2.10.0+cu128"},"lock_hash":"sha256:52b324a75bb7630844f3a1faad8974c4e9aef36eff2b400477e4c96004d6811c","checkpoints":{"nequip-mp-l":{"fetched_at":"2026-10-05T21:00:13.864265+00:00","verified_at":"2026-10-25T19:59:55.999558+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/nequip/compiled/NequIP-MP-L-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":46051639}],"weights_recorded_at":"2026-10-25T19:59:55.999561+00:00"},"nequip-oam-l":{"fetched_at":"2026-10-05T21:00:14.053978+00:00","verified_at":"2026-10-25T19:59:55.999548+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/nequip/compiled/NequIP-OAM-L-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":45969552}],"weights_recorded_at":"2026-10-25T19:59:55.999551+00:00"},"nequip-oam-s":{"fetched_at":"2026-10-05T21:00:31.367040+00:00","verified_at":"2026-10-25T19:59:55.999538+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/nequip/compiled/NequIP-OAM-S-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":4277568}],"weights_recorded_at":"2026-10-25T19:59:55.999540+00:00"},"nequip-oam-m":{"fetched_at":"2026-10-05T21:00:37.557331+00:00","verified_at":"2026-10-25T19:59:55.999543+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/nequip/compiled/NequIP-OAM-M-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":17148312}],"weights_recorded_at":"2026-10-25T19:59:55.999546+00:00"},"nequip-oam-xl":{"fetched_at":"2026-10-05T21:01:39.674414+00:00","verified_at":"2026-10-25T19:59:55.999553+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/nequip/compiled/NequIP-OAM-XL-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":140072170}],"weights_recorded_at":"2026-10-25T19:59:55.999556+00:00"},"nequip:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:59:55.999587+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}}}},{"schema_version":5,"cluster":"sophia","root":"/eagle/projects/Rootstock/rootstock","maintainer":{"name":"Owen Price Skelly","email":"OwenPriceSkelly@uchicago.edu"},"rootstock_version":"1.6.4","python_version":"unknown","last_updated":"2026-10-25T20:00:19.052187+00:00","environments":{"allscaip":{"built_at":"2026-08-06T19:38:53.967540+00:00","source_hash":"sha256:bde4548ac275b9af7d17c3103e3da8008038f306ab5cb03e898c95ce6ad994d9","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"AllScAIP env — FAIRChem scalable attention MLIP trained on OMol25.\n\nallscaip-md-conserving-all-omol is an energy-conserving, all-to-all node\nattention model served through fairchem-core's get_predict_unit — the same\nAPI as eSEN. fairchem v2 carries the architecture in-package, so no\nflash-attention or custom CUDA kernels are needed.\n\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"allscaip-md-conserving-all-omol\": \"allscaip-md-conserving-all-omol\",\n \"allscaip-md-direct-all-omol\": \"allscaip-md-direct-all-omol\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"allscaip:custom\": None,\n}\n\n\ndef _fairchem_device(device: str) -> str:\n \"\"\"Translate an indexed device (\"cuda:2\") into what fairchem v2 accepts.\n\n MLIPPredictUnit._setup_device asserts `device in [\"cpu\", \"cuda\"]` and then\n resolves the real GPU itself via get_device_for_local_rank(), which returns\n f\"cuda:{torch.cuda.current_device()}\". So an index has to travel through\n torch's current-device state, not the argument. Verifying several\n checkpoints at once on a multi-GPU node hands each worker \"cuda:N\" — that\n killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync\n (4x A100, VERIFY_JOBS=4), while single-GPU Sophia never hit it.\n \"\"\"\n if device.startswith(\"cuda:\"):\n import torch\n\n torch.cuda.set_device(int(device.split(\":\", 1)[1]))\n return \"cuda\"\n return device\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(\n CHECKPOINTS[checkpoint], device=_fairchem_device(device)\n )\n return FAIRChemCalculator(predictor)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=_fairchem_device(device))\n return FAIRChemCalculator(predictor)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.6.0","torch":"2.8.0"},"lock_hash":"sha256:c99d413acdd9b5b2cdd0241ba29966574fe6800af2b63376abacd0ad28fed8b0","checkpoints":{"allscaip-md-direct-all-omol":{"fetched_at":"2026-08-03T17:03:56.967701+00:00","verified_at":"2026-10-25T19:51:31.848362+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/5517360cf5a23ca141f1b51b708d131289e35080bfa79f27b990695e3e78d18e","size":694955368},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.994651+00:00"},"allscaip-md-conserving-all-omol":{"fetched_at":"2026-08-03T17:03:57.183635+00:00","verified_at":"2026-10-25T19:51:31.848326+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/14fe8626e9e75337769ecf98a765cdc02667763ec4dc1103d958e1c9f56b26e9","size":688432239},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.993953+00:00"},"allscaip:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:51:31.852988+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"ani":{"built_at":"2026-07-31T18:31:40.101251+00:00","source_hash":"sha256:16ebe0eda0074fa275e0ce2c2d68c84dc10d2069ef0290d19802bd46bcf96d0e","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"torchani>=2.2\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n# ///\n\"\"\"\nANI-2x environment for Rootstock.\n\nANI-2x is a neural network potential for organic molecules containing\nH, C, N, O, F, S, Cl. It is not a universal potential — do not use it\nfor inorganic or periodic systems.\n\nModels:\n - \"ANI2x\": ANI-2x ensemble (default, 8 networks)\n - \"ANI1ccx\": ANI-1ccx, trained on CCSD(T)/CBS data (H, C, N, O only)\n - \"ANI1x\": ANI-1x (H, C, N, O only)\n\"\"\"\n\nCHECKPOINTS = {\n \"ani-2x\": \"ANI2x\",\n \"ani-1ccx\": \"ANI1ccx\",\n \"ani-1x\": \"ANI1x\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an ANI calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n import torchani\n\n model_map = {\n \"ANI2x\": torchani.models.ANI2x,\n \"ANI1ccx\": torchani.models.ANI1ccx,\n \"ANI1x\": torchani.models.ANI1x,\n }\n model = CHECKPOINTS[checkpoint]\n\n return model_map[model](periodic_table_index=True).to(device).ase()\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.3.0","torch":"2.13.0","torchani":"2.8.4"},"lock_hash":"sha256:8864bbf63ae2a9aa45ee0fc265b0c30e3c4d7b67fb9631b3efd90b5b7f774934","checkpoints":{"ani-1x":{"fetched_at":"2026-08-03T17:03:11.273898+00:00","verified_at":"2026-10-25T19:51:31.848396+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.local/share/Torchani/StateDicts/.cache/huggingface/download/ani1x_state_dict.pt.metadata","size":124},{"path":"home/.local/share/Torchani/StateDicts/ani1x_state_dict.pt","size":10546029}],"weights_recorded_at":"2026-10-25T19:51:31.848401+00:00"},"ani-1ccx":{"fetched_at":"2026-08-03T17:03:11.776896+00:00","verified_at":"2026-10-25T19:51:31.848387+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.local/share/Torchani/StateDicts/.cache/huggingface/download/ani1ccx_state_dict.pt.metadata","size":125},{"path":"home/.local/share/Torchani/StateDicts/ani1ccx_state_dict.pt","size":10546029}],"weights_recorded_at":"2026-10-25T19:51:31.848392+00:00"},"ani-2x":{"fetched_at":"2026-08-03T17:04:29.856289+00:00","verified_at":"2026-10-25T19:51:31.848376+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.local/share/Torchani/StateDicts/.cache/huggingface/download/ani2x_state_dict.pt.metadata","size":125},{"path":"home/.local/share/Torchani/StateDicts/ani2x_state_dict.pt","size":54982957}],"weights_recorded_at":"2026-10-25T19:51:31.848382+00:00"}}},"chgnet":{"built_at":"2026-07-31T18:31:25+00:00","source_hash":"sha256:b6ea6ac08a17cb3dc17d5455571dc5c21d513cd802916015de7f272e83eefe37","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"chgnet>=0.3.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"CHGNet env — hosts pretrained charge-informed universal potentials.\"\"\"\n\nCHECKPOINTS = {\n \"chgnet-default\": \"chgnet-default\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"chgnet:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a CHGNet calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\")\n\n Returns:\n ASE-compatible calculator\n \"\"\"\n from chgnet.model import CHGNet, CHGNetCalculator\n\n model_name = CHECKPOINTS[checkpoint]\n model = CHGNet.load() if model_name == \"chgnet-default\" else CHGNet.load(model_name)\n return CHGNetCalculator(model=model, use_device=device)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # CHGNet.from_file, not the named-model CHGNet.load() setup() uses.\n from chgnet.model import CHGNet, CHGNetCalculator\n\n model = CHGNet.from_file(path)\n return CHGNetCalculator(model=model, use_device=device)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","chgnet":"0.4.2","rootstock":"1.3.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:b5bc1f87f4c596d669c56f6a5053d21b488b1e141dc8e38e7a7a964b887fcc73","checkpoints":{"chgnet-default":{"fetched_at":"2026-08-03T17:04:29.877040+00:00","verified_at":"2026-10-25T19:51:31.848405+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"equiformer":{"built_at":"2026-07-31T18:31:42.640544+00:00","source_hash":"sha256:ac01aec3be646cbd2c63a2e9f6a0cc8308aa03f5d8528150c5f664b2aa06426f","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\nEquiformerV2 environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 EquiformerV2 checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"EquiformerV2-153M-S2EF-OC20-All+MD\": default\n - \"EquiformerV2-31M-S2EF-OC20-All+MD\"\n - \"EquiformerV2-83M-S2EF-OC20-2M\"\n\"\"\"\n\nCHECKPOINTS = {\n \"equiformer-v2-153m-s2ef-oc20-all-md\": \"EquiformerV2-153M-S2EF-OC20-All+MD\",\n \"equiformer-v2-31m-s2ef-oc20-all-md\": \"EquiformerV2-31M-S2EF-OC20-All+MD\",\n \"equiformer-v2-83m-s2ef-oc20-2m\": \"EquiformerV2-83M-S2EF-OC20-2M\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"equiformer:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an EquiformerV2 OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): OCPCalculator loads a\n # checkpoint file natively — this is setup() minus the registry download.\n from fairchem.core import OCPCalculator\n\n return OCPCalculator(checkpoint_path=path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.3.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"equiformer-v2-31m-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:07:10.587007+00:00","verified_at":"2026-10-25T19:51:31.848419+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_31M_ec4_allmd.pt","size":125685325}],"weights_recorded_at":"2026-10-25T19:59:55.994689+00:00"},"equiformer-v2-83m-s2ef-oc20-2m":{"fetched_at":"2026-08-03T17:07:10.622992+00:00","verified_at":"2026-10-25T19:51:31.848427+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_83M_2M.pt","size":336843381}],"weights_recorded_at":"2026-10-25T19:59:55.994695+00:00"},"equiformer-v2-153m-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:07:12.126551+00:00","verified_at":"2026-10-25T19:51:31.848410+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_153M_ec4_allmd.pt","size":619855697}],"weights_recorded_at":"2026-10-25T19:59:55.994683+00:00"},"equiformer:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: weights= run diverges from equiformer-v2-153m-s2ef-oc20-all-md: energy differs by 1.397e-04 eV (tolerance 1e-04)","weight_files":null,"weights_recorded_at":null}}},"escn":{"built_at":"2026-07-31T18:32:58.772603+00:00","source_hash":"sha256:e1301f3e20ac38e355e0b36e2ab9fb0f0a78f58059b5e2c57ab280aa9dad7ff5","source":"# /// script\n# requires-python = \">=3.11,<3.12\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# # scipy.special.sph_harm was removed in scipy 1.17 and fairchem-core 1.x\n# # still imports it — an uncapped rebuild breaks at import (Delta, 2026-07-18).\n# \"scipy<1.17\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\neSCN environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 eSCN checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"eSCN-L6-M2-Lay12-S2EF-OC20-All+MD\": default\n - \"eSCN-L6-M3-Lay20-S2EF-OC20-All+MD\"\n - \"eSCN-L6-M2-Lay12-S2EF-OC20-2M\"\n - \"eSCN-L4-M2-Lay12-S2EF-OC20-2M\"\n\"\"\"\n\nCHECKPOINTS = {\n \"escn-l6-m2-lay12-s2ef-oc20-all-md\": \"eSCN-L6-M2-Lay12-S2EF-OC20-All+MD\",\n \"escn-l6-m3-lay20-s2ef-oc20-all-md\": \"eSCN-L6-M3-Lay20-S2EF-OC20-All+MD\",\n \"escn-l6-m2-lay12-s2ef-oc20-2m\": \"eSCN-L6-M2-Lay12-S2EF-OC20-2M\",\n \"escn-l4-m2-lay12-s2ef-oc20-2m\": \"eSCN-L4-M2-Lay12-S2EF-OC20-2M\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"escn:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an eSCN OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): OCPCalculator loads a\n # checkpoint file natively — this is setup() minus the registry download.\n from fairchem.core import OCPCalculator\n\n return OCPCalculator(checkpoint_path=path, cpu=(device == \"cpu\"))\n","python_requires":">=3.11,<3.12","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"1.3.0","scipy":"1.16.3","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"lock_hash":null,"checkpoints":{"escn-l4-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-08-03T17:07:19.768803+00:00","verified_at":"2026-10-25T19:51:31.848458+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l4_m2_lay12_2M_s2ef.pt","size":144761989}],"weights_recorded_at":"2026-10-25T19:59:55.995306+00:00"},"escn-l6-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-08-03T17:08:03.373277+00:00","verified_at":"2026-10-25T19:51:31.848451+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l6_m2_lay12_2M_s2ef.pt","size":207735045}],"weights_recorded_at":"2026-10-25T19:59:55.995300+00:00"},"escn-l6-m2-lay12-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:08:03.392125+00:00","verified_at":"2026-10-25T19:51:31.848435+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l6_m2_lay12_all_md_s2ef.pt","size":207735109}],"weights_recorded_at":"2026-10-25T19:59:55.994700+00:00"},"escn-l6-m3-lay20-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:08:14.400434+00:00","verified_at":"2026-10-25T19:51:31.848443+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/escn_l6_m3_lay20_all_md_s2ef.pt","size":801642787}],"weights_recorded_at":"2026-10-25T19:59:55.995293+00:00"},"escn:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"smoke-test: weights= run diverges from escn-l6-m2-lay12-s2ef-oc20-all-md: energy differs by 1.121e-02 eV (tolerance 1e-04)","weight_files":null,"weights_recorded_at":null}}},"esen":{"built_at":"2026-08-06T19:39:02.621368+00:00","source_hash":"sha256:d07f4d08c7ce3ed5aa36b795281d2be931b17cd6d342e2b569ef0fd1f53e59db","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"eSEN env — hosts FAIRChem eSEN single-task checkpoints.\n\nfairchem-core v2 dropped the torch-geometric / pyg-find-links install dance\n(simplified here 2026-07-30 to match uma.py/allscaip.py — the find-links were\nalso blocking universal lockfile resolution).\n\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"esen-md-direct-all-omol\": \"esen-md-direct-all-omol\",\n \"esen-sm-conserving-all-omol\": \"esen-sm-conserving-all-omol\",\n \"esen-sm-direct-all-omol\": \"esen-sm-direct-all-omol\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"esen:custom\": None,\n}\n\n\ndef _fairchem_device(device: str) -> str:\n \"\"\"Translate an indexed device (\"cuda:2\") into what fairchem v2 accepts.\n\n MLIPPredictUnit._setup_device asserts `device in [\"cpu\", \"cuda\"]` and then\n resolves the real GPU itself via get_device_for_local_rank(), which returns\n f\"cuda:{torch.cuda.current_device()}\". So an index has to travel through\n torch's current-device state, not the argument. Verifying several\n checkpoints at once on a multi-GPU node hands each worker \"cuda:N\" — that\n killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync\n (4x A100, VERIFY_JOBS=4), while single-GPU Sophia never hit it.\n \"\"\"\n if device.startswith(\"cuda:\"):\n import torch\n\n torch.cuda.set_device(int(device.split(\":\", 1)[1]))\n return \"cuda\"\n return device\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(\n CHECKPOINTS[checkpoint], device=_fairchem_device(device)\n )\n return FAIRChemCalculator(predictor)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=_fairchem_device(device))\n return FAIRChemCalculator(predictor)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.6.0","torch":"2.8.0"},"lock_hash":"sha256:c99d413acdd9b5b2cdd0241ba29966574fe6800af2b63376abacd0ad28fed8b0","checkpoints":{"esen-sm-direct-all-omol":{"fetched_at":"2026-08-03T17:11:00.904446+00:00","verified_at":"2026-10-25T19:51:31.848484+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/a1b7109bcca6e1288c5f17c68776d3c6a1ab0766f6d7fdd45f901605355d2831","size":50960291},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.995326+00:00"},"esen-sm-conserving-all-omol":{"fetched_at":"2026-08-03T17:11:00.922460+00:00","verified_at":"2026-10-25T19:51:31.848474+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/01f63da2d071e39fc46a5f22f8369d0fc9de317ab2ef361a76603f5661238025","size":50958067},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.995320+00:00"},"esen-md-direct-all-omol":{"fetched_at":"2026-08-03T17:11:00.992499+00:00","verified_at":"2026-10-25T19:51:31.848466+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--OMol25/blobs/e0cd2ac37d2992b385a39778469c443ccbc679c7c4fb3b40b38aa2ee23c29b1d","size":405789591},{"path":"home/.cache/fairchem/models--facebook--OMol25/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.995313+00:00"},"esen:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:51:31.852996+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"grace":{"built_at":"2026-07-31T18:34:21.902859+00:00","source_hash":"sha256:f46207cbbd702910ea889c09cae3f9933de6c48baf22d75c4039caf02467b635","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"tensorpotential>=0.6.0\",\n# \"ase>=3.22\",\n# ]\n# ///\n\"\"\"GRACE env — hosts GRACE foundation checkpoints via tensorpotential.\n\nGRACE runs on TensorFlow, not torch. TPCalculator has no device argument:\nTF grabs whatever GPU it sees, so device selection happens via\nCUDA_VISIBLE_DEVICES, and both it and TF_USE_LEGACY_KERAS must be set\nbefore the first TF import. The first calculation triggers an XLA compile —\na slow first step is expected.\n\"\"\"\n\nCHECKPOINTS = {\n \"grace-2l-smax-omat-large\": \"GRACE-2L-SMAX-OMAT-large\",\n \"grace-3l-omat-large-ft-am\": \"GRACE-3L-OMAT-large-ft-AM\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n import os\n\n os.environ[\"TF_USE_LEGACY_KERAS\"] = \"1\"\n if device == \"cpu\":\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n elif device.startswith(\"cuda:\"):\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = device.split(\":\", 1)[1]\n\n from tensorpotential.calculator import grace_fm\n\n return grace_fm(CHECKPOINTS[checkpoint])\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.3.0","tensorpotential":"0.6.0"},"lock_hash":"sha256:1ea87e2d3e8620111788e4f3804932ac53f399ab75df08ee6b2b38195a82d50e","checkpoints":{"grace-2l-smax-omat-large":{"fetched_at":"2026-08-03T17:10:49.308918+00:00","verified_at":"2026-10-25T19:51:31.848492+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"grace-3l-omat-large-ft-am":{"fetched_at":"2026-08-03T17:11:39.418528+00:00","verified_at":"2026-10-25T19:51:31.848497+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"mace":{"built_at":"2026-08-03T17:00:55.121628+00:00","source_hash":"sha256:eb8e9790141ad43e2735656a942a3b3af3a551ae6e89aef0b99c72ff89f80360","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # 0.3.15+ needed for the mh-1 registry entry (matpes needs 0.3.13,\n# # omol needs 0.3.14, mpa-0 needs 0.3.10).\n# \"mace-torch>=0.3.15\",\n# \"ase>=3.22\",\n# # 2.4.1 is explicitly unsupported by mace-torch.\n# \"torch>=2.4.0,!=2.4.1,<2.10\",\n# ]\n# ///\n\"\"\"MACE env — hosts MACE-MP-0, MACE-OFF23, MPA-0, MATPES, MH-1, and OMOL checkpoints.\n\nAll ship in the same `mace-torch` package, so they share an environment.\nUpstream-string routing in CHECKPOINTS: an `off:` prefix routes to mace_off()\nand an `omol:` prefix to mace_omol() (float64, molecules only); an `mh:`\nprefix marks a multi-head model (float64, per the MACE-MH-1 model card).\n\nMulti-head checkpoints select a head via the `head` kwarg on setup()\n(setup_kwargs={\"head\": ...} / --kwarg head=...), named by upstream's training\ncorpus — see MH1_HEADS; omat_pbe is the default.\n\nThe OMOL checkpoint expects `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-mp-0-small\": \"small\",\n \"mace-mp-0-medium\": \"medium\",\n \"mace-mp-0-large\": \"large\",\n \"mace-off23-small\": \"off:small\",\n \"mace-off23-medium\": \"off:medium\",\n \"mace-off23-large\": \"off:large\",\n # Only a medium MPA-0 has been released, but upstream names the weights\n # file mace-mpa-0-medium.model — keep the size explicit like mace-mp-0.\n \"mace-mpa-0-medium\": \"medium-mpa-0\",\n \"mace-matpes-r2scan-0\": \"mace-matpes-r2scan-0\",\n # One entry per weights file: MH-1's heads are selected by setup(head=...).\n \"mace-mh-1\": \"mh:mh-1\",\n # Only the extra-large OMOL model has been released.\n \"mace-omol-0-extra-large\": \"omol:extra_large\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace:custom\": None,\n}\n\n# MH-1's heads, named by training corpus. The released weights file is the\n# authority (mace_select_head --list_heads; the model card also lists a\n# rgd1_b3lyp head, but that shipped only in mh-0 — ACEsuit/mace#1462).\n# Validated here because upstream only warns on an unknown head and silently\n# falls back to the last one.\nMH1_HEADS = (\n \"omat_pbe\",\n \"omol\",\n \"spice_wB97M\",\n \"oc20_usemppbe\",\n \"mp_pbe_refit_add\",\n \"matpes_r2scan\",\n)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", head: str | None = None):\n arg = CHECKPOINTS[checkpoint]\n if arg.startswith(\"mh:\"):\n head = head or \"omat_pbe\"\n if head not in MH1_HEADS:\n raise ValueError(f\"unknown head {head!r}; expected one of {', '.join(MH1_HEADS)}\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg[3:], device=device, default_dtype=\"float64\", head=head)\n if head is not None:\n raise ValueError(f\"'head' selects a head of a multi-head model; {checkpoint} has one head\")\n if arg.startswith(\"off:\"):\n from mace.calculators import mace_off\n\n return mace_off(model=arg[4:], device=device, default_dtype=\"float32\")\n if arg.startswith(\"omol:\"):\n from mace.calculators import mace_omol\n\n return mace_omol(model=arg[5:], device=device, default_dtype=\"float64\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg, device=device, default_dtype=\"float32\")\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", default_dtype: str = \"float32\", head: str | None = None):\n # Custom checkpoints (`:custom` ids with user weights): fine-tunes load through\n # MACECalculator directly — the mp/off dispatch in setup() only exists\n # to pick which pretrained file to download. `head` is for fine-tunes that\n # keep multiple heads; single-head weights load without it.\n from mace.calculators import MACECalculator\n\n return MACECalculator(model_paths=path, device=device, default_dtype=default_dtype, head=head)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"1.4.3","torch":"2.9.1"},"lock_hash":"sha256:2ef43c5e21b8c2dc7b3830833e9ae38555d87ade585a064e9fba3d63ca792931","checkpoints":{"mace-mp-0-medium":{"fetched_at":"2026-08-03T17:12:25.735382+00:00","verified_at":"2026-10-25T19:51:31.848509+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231203mace128L1_epoch199model","size":44422970}],"weights_recorded_at":"2026-10-25T19:59:55.995344+00:00"},"mace-mh-1":{"fetched_at":"2026-08-03T17:12:25.771339+00:00","verified_at":"2026-10-25T19:51:31.848565+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macemh1model","size":59208139}],"weights_recorded_at":"2026-10-25T19:59:55.995382+00:00"},"mace-matpes-r2scan-0":{"fetched_at":"2026-08-03T17:12:26.162740+00:00","verified_at":"2026-10-25T19:51:31.848557+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEmatpesr2scanomatftmodel","size":79470738}],"weights_recorded_at":"2026-10-25T19:59:55.995377+00:00"},"mace-mp-0-large":{"fetched_at":"2026-08-03T17:12:26.906848+00:00","verified_at":"2026-10-25T19:51:31.848517+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE_MPtrj_20229model","size":133803220}],"weights_recorded_at":"2026-10-25T19:59:55.995349+00:00"},"mace-off23-medium":{"fetched_at":"2026-08-03T17:12:44.272336+00:00","verified_at":"2026-10-25T19:51:31.848532+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_medium.model","size":18350596}],"weights_recorded_at":"2026-10-25T19:59:55.995361+00:00"},"mace-mp-0-small":{"fetched_at":"2026-08-03T17:12:44.380401+00:00","verified_at":"2026-10-25T19:51:31.848501+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231210mace128L0_energy_epoch249model","size":32581838}],"weights_recorded_at":"2026-10-25T19:59:55.995338+00:00"},"mace-mpa-0-medium":{"fetched_at":"2026-08-03T17:12:44.712260+00:00","verified_at":"2026-10-25T19:51:31.848548+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macempa0mediummodel","size":79462305}],"weights_recorded_at":"2026-10-25T19:59:55.995371+00:00"},"mace-off23-large":{"fetched_at":"2026-08-03T17:12:44.790639+00:00","verified_at":"2026-10-25T19:51:31.848540+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_large.model","size":55492786}],"weights_recorded_at":"2026-10-25T19:59:55.995366+00:00"},"mace-off23-small":{"fetched_at":"2026-08-03T17:13:00.903240+00:00","verified_at":"2026-10-25T19:51:31.848524+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_small.model","size":7347350}],"weights_recorded_at":"2026-10-25T19:59:55.995354+00:00"},"mace-omol-0-extra-large":{"fetched_at":"2026-08-03T17:13:02.181806+00:00","verified_at":"2026-10-25T19:51:31.848573+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-omol-0-extra-large-1024.model","size":422242640}],"weights_recorded_at":"2026-10-25T19:59:55.995387+00:00"},"mace:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:51:31.853000+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"mace_polar":{"built_at":"2026-07-31T18:34:29.831867+00:00","source_hash":"sha256:8db4cfd29a59c415e3762b04195e30b9e08d5f044ff6aa21ebcf01bcb78bc1ef","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"ase>=3.22\",\n# \"torch>=2.4.0,<2.10\",\n# # 0.3.16 is the first PyPI release with mace_polar().\n# \"mace-torch>=0.3.16\",\n# # PolarMACE imports graph_longrange at runtime; the distribution is named\n# # graph-longrange and exists only as this git repo (no PyPI release).\n# \"graph-longrange @ git+https://github.com/WillBaldwin0/graph_electrostatics.git\",\n# ]\n# ///\n\"\"\"MACE-POLAR env — electrostatic/polarizable MACE foundation models (OMol25).\n\nKept separate from the stable `mace` env because of the extra git-only\ngraph-longrange dependency.\n\nPOLAR checkpoints expect `charge`, `spin`, and `external_field` in atoms.info.\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-polar-1-s\": \"polar-1-s\",\n \"mace-polar-1-m\": \"polar-1-m\",\n \"mace-polar-1-l\": \"polar-1-l\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"mace-polar:custom\": None,\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from mace.calculators import mace_polar\n\n return mace_polar(model=CHECKPOINTS[checkpoint], device=device, default_dtype=\"float32\")\n\n\ndef setup_from_path(path: str, device: str = \"cuda\"):\n # Custom checkpoints (`:custom` ids with user weights): mace_polar() accepts a\n # weights file directly, keeping the PolarMACE model-type wiring.\n from mace.calculators import mace_polar\n\n return mace_polar(model=path, device=device, default_dtype=\"float32\")\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"1.3.0","torch":"2.9.1"},"lock_hash":"sha256:bf60a2e3dd4771b8eb5b17a99053c1bf90448615005aab598328165f8da4f5aa","checkpoints":{"mace-polar-1-s":{"fetched_at":"2026-08-03T17:15:38.414289+00:00","verified_at":"2026-10-25T19:51:31.848580+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Smodel","size":33375439}],"weights_recorded_at":"2026-10-25T19:59:55.995393+00:00"},"mace-polar-1-m":{"fetched_at":"2026-08-03T17:15:38.434404+00:00","verified_at":"2026-10-25T19:51:31.848588+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Mmodel","size":68133235}],"weights_recorded_at":"2026-10-25T19:59:55.995398+00:00"},"mace-polar-1-l":{"fetched_at":"2026-08-03T17:15:38.454464+00:00","verified_at":"2026-10-25T19:51:31.848595+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Lmodel","size":129808715}],"weights_recorded_at":"2026-10-25T19:59:55.995404+00:00"},"mace-polar:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:51:31.853003+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"orb_v3":{"built_at":"2026-07-31T18:36:13.727862+00:00","source_hash":"sha256:1f63f36fe8c791313c04bbb5bf534a1a7168c327c1a35f98ecd3d80e3ff74122","source":"# /// script\n# # <3.13: orb-models pins dm-tree==0.1.8, which has no cp313 wheel and whose\n# # sdist doesn't compile against modern GCC (vendored abseil).\n# requires-python = \">=3.12,<3.13\"\n# dependencies = [\n# \"orb-models>=0.6.2\",\n# \"ase>=3.25\",\n# \"torch>=2.8\",\n# # Not imported here — constrains orb-models' transitive dep. setup()'s\n# # no-lock serve path relies on cached_path returning local files without\n# # locking or writing, verified against exactly this version (#67).\n# \"cached_path==1.8.10\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"Orb v3 env — the primary orb env, Orbital Materials' Orb v3 potentials.\n\nSeparate from orb.py because the v3 loaders need orb-models>=0.6.2, which\nbumped the Python floor to 3.12 and torch to 2.8; the v3 loader API also\ndiffers (returns a tuple, requires `atoms_adapter` on ORBCalculator, imports\nthe calculator from forcefield.inference). orb.py (v2) survives only for\norb-d3-v2, the dispersion-corrected variant with no v3 equivalent.\n\"\"\"\n\nimport os\nimport shutil\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"orb-v3-conservative-inf-omat\": \"orb-v3-conservative-inf-omat\",\n \"orb-v3-conservative-20-omat\": \"orb-v3-conservative-20-omat\",\n \"orb-v3-direct-inf-omat\": \"orb-v3-direct-inf-omat\",\n \"orb-v3-direct-20-omat\": \"orb-v3-direct-20-omat\",\n \"orb-v3-conservative-inf-mpa\": \"orb-v3-conservative-inf-mpa\",\n \"orb-v3-conservative-20-mpa\": \"orb-v3-conservative-20-mpa\",\n \"orb-v3-direct-inf-mpa\": \"orb-v3-direct-inf-mpa\",\n \"orb-v3-direct-20-mpa\": \"orb-v3-direct-20-mpa\",\n # The omol ids (orb-v3-{conservative,direct}-omol) are dropped from the\n # catalog 2026-07-30: they had been failing verify on every cluster since\n # 2026-05. Re-add once the failure is understood.\n # Your own fine-tuned v3 weights: pair with weights= (loaded via\n # setup_from_path).\n \"orb-v3:custom\": None,\n}\n\n\ndef _default_weights_url(load_fn) -> str:\n \"\"\"The upstream URL baked into the loader's ``weights_path`` default.\"\"\"\n import inspect\n\n default = inspect.signature(load_fn).parameters[\"weights_path\"].default\n if not isinstance(default, str) or not default.startswith((\"http://\", \"https://\")):\n raise RuntimeError(\n f\"{load_fn.__name__} has no URL default for weights_path \"\n f\"(got {default!r}); update this env file for the installed orb-models\"\n )\n return default\n\n\ndef _local_weights_path(url: str) -> Path:\n \"\"\"Where the checkpoint lives in the shared model cache.\"\"\"\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"orb\" / os.path.basename(url)\n\n\ndef _fetch(url: str, dest: Path) -> None:\n \"\"\"Download ``url`` to ``dest`` atomically (tmp file + rename).\"\"\"\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp = dest.with_name(f\"{dest.name}.tmp.{os.getpid()}\")\n try:\n with urllib.request.urlopen(url) as resp, open(tmp, \"wb\") as out:\n shutil.copyfileobj(resp, out)\n os.replace(tmp, dest)\n finally:\n tmp.unlink(missing_ok=True)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", precision: str = \"float32-high\"):\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.inference.calculator import ORBCalculator\n\n fn_name = CHECKPOINTS[checkpoint].replace(\"-\", \"_\")\n load_fn = getattr(pretrained, fn_name)\n\n # orb-models resolves its default weights URL through `cached_path`, which\n # write-locks its cache dir even on warm hits — EACCES for anyone who can't\n # write the shared install (Garden-AI/rootstock#67). Handed a *local* path\n # instead, cached_path returns it without locking. So the weights are\n # pre-fetched into the shared model cache at `rootstock add` time\n # (maintainer, cache writable) and every later serve loads that file.\n url = _default_weights_url(load_fn)\n weights = _local_weights_path(url)\n if not weights.exists():\n _fetch(url, weights)\n\n orbff, atoms_adapter = load_fn(\n weights_path=str(weights), device=torch.device(device), precision=precision\n )\n return ORBCalculator(orbff, atoms_adapter=atoms_adapter, device=torch.device(device))\n\n\ndef setup_from_path(\n path: str,\n device: str = \"cuda\",\n arch: str = \"orb-v3-conservative-inf-omat\",\n precision: str = \"float32-high\",\n):\n # Custom checkpoints (`:custom` ids with user weights). A weights file doesn't say\n # which orb architecture produced it, so `arch` names the pretrained\n # loader to instantiate — pass the right one at call time\n # (setup_kwargs={\"arch\": ...} / --kwarg arch=...). Handing the loader a\n # local path also means no network and no cached_path locking (see setup()).\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.inference.calculator import ORBCalculator\n\n fn_name = arch.replace(\"-\", \"_\")\n try:\n load_fn = getattr(pretrained, fn_name)\n except AttributeError:\n raise ValueError(\n f\"unknown orb architecture {arch!r}; expected a loader name from \"\n f\"orb_models.forcefield.pretrained, e.g. orb-v3-conservative-inf-omat\"\n ) from None\n\n orbff, atoms_adapter = load_fn(\n weights_path=path, device=torch.device(device), precision=precision\n )\n return ORBCalculator(orbff, atoms_adapter=atoms_adapter, device=torch.device(device))\n","python_requires":">=3.12,<3.13","dependencies":{"ase":"3.29.0","cached-path":"1.8.10","orb-models":"0.7.0","rootstock":"1.3.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:cf6466e7518601cbbb8b296939d1804006146ec5a541c64179818df29b2ecda4","checkpoints":{"orb-v3-conservative-20-mpa":{"fetched_at":"2026-08-03T17:18:18.499544+00:00","verified_at":"2026-10-25T19:51:31.848640+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-20-mpa-20250404.ckpt","size":102097380}],"weights_recorded_at":"2026-10-25T19:59:55.995435+00:00"},"orb-v3-conservative-20-omat":{"fetched_at":"2026-08-03T17:18:18.543840+00:00","verified_at":"2026-10-25T19:51:31.848611+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-20-omat-20250404.ckpt","size":102097517}],"weights_recorded_at":"2026-10-25T19:59:55.995414+00:00"},"orb-v3-direct-20-mpa":{"fetched_at":"2026-08-03T17:18:51.755916+00:00","verified_at":"2026-10-25T19:51:31.848655+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-20-mpa-20250404.ckpt","size":102636362}],"weights_recorded_at":"2026-10-25T19:59:55.995447+00:00"},"orb-v3-conservative-inf-omat":{"fetched_at":"2026-08-03T17:18:52.017384+00:00","verified_at":"2026-10-25T19:51:31.848603+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-inf-omat-20250404.ckpt","size":102097517}],"weights_recorded_at":"2026-10-25T19:59:55.995409+00:00"},"orb-v3-conservative-inf-mpa":{"fetched_at":"2026-08-03T17:18:52.075825+00:00","verified_at":"2026-10-25T19:51:31.848633+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-conservative-inf-mpa-20250404.ckpt","size":102097517}],"weights_recorded_at":"2026-10-25T19:59:55.995430+00:00"},"orb-v3-direct-20-omat":{"fetched_at":"2026-08-03T17:18:53.566555+00:00","verified_at":"2026-10-25T19:51:31.848626+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-20-omat-20250404.ckpt","size":102636510}],"weights_recorded_at":"2026-10-25T19:59:55.995425+00:00"},"orb-v3-direct-inf-omat":{"fetched_at":"2026-08-03T17:19:14.412381+00:00","verified_at":"2026-10-25T19:51:31.848618+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-inf-omat-20250404.ckpt","size":102636658}],"weights_recorded_at":"2026-10-25T19:59:55.995420+00:00"},"orb-v3-direct-inf-mpa":{"fetched_at":"2026-08-03T17:19:16.914781+00:00","verified_at":"2026-10-25T19:51:31.848648+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-v3-direct-inf-mpa-20250404.ckpt","size":102636510}],"weights_recorded_at":"2026-10-25T19:59:55.995441+00:00"},"orb-v3:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:51:31.853006+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"orb":{"built_at":"2026-07-31T18:36:40.771968+00:00","source_hash":"sha256:4ae1362a989170b29d5082a23f9586c50d4f04b700bd29b8c65609efcbe52bef","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # >=0.5,<0.6: 0.5.5 is what the verified Delta env resolved — the v2\n# # loaders keep their single-return API through 0.5.x. 0.4.x is broken\n# # for us: it imports pynanoflann, which is git-only and undeclared, so\n# # a fresh build dies at import (Delta, 2026-07-31). 0.6 raises the\n# # Python floor to 3.12 — that line lives in orb_v3.py.\n# \"orb-models>=0.5,<0.6\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# # Not imported here — constrains orb-models' transitive dep. setup()'s\n# # no-lock serve path relies on cached_path returning local files without\n# # locking or writing, verified against exactly this version (#67).\n# \"cached_path==1.8.10\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"Orb v2 env — kept only for the built-in-D3 dispersion variant.\n\nOrb v3 (orb_v3.py) is the primary orb env; it supersedes the v2 checkpoints\nexcept orb-d3-v2, which has no v3 equivalent (v3 ships no dispersion-corrected\nmodel). Catalog trimmed to that one id 2026-07-30. The two lines can't share\nan env: the v3 loaders need orb-models>=0.6, which raises the Python floor\nto 3.12 and torch to 2.8 (the v2 loaders here are fine through 0.5.x).\n\"\"\"\n\nimport os\nimport shutil\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"orb-d3-v2\": \"orb-d3-v2\",\n # Your own fine-tuned v2-architecture weights: pair with weights=\n # (loaded via setup_from_path). v3 fine-tunes go to orb-v3:custom.\n \"orb-v2:custom\": None,\n}\n\n\ndef _default_weights_url(load_fn) -> str:\n \"\"\"The upstream URL baked into the loader's ``weights_path`` default.\"\"\"\n import inspect\n\n default = inspect.signature(load_fn).parameters[\"weights_path\"].default\n if not isinstance(default, str) or not default.startswith((\"http://\", \"https://\")):\n raise RuntimeError(\n f\"{load_fn.__name__} has no URL default for weights_path \"\n f\"(got {default!r}); update this env file for the installed orb-models\"\n )\n return default\n\n\ndef _local_weights_path(url: str) -> Path:\n \"\"\"Where the checkpoint lives in the shared model cache.\"\"\"\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"orb\" / os.path.basename(url)\n\n\ndef _fetch(url: str, dest: Path) -> None:\n \"\"\"Download ``url`` to ``dest`` atomically (tmp file + rename).\"\"\"\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp = dest.with_name(f\"{dest.name}.tmp.{os.getpid()}\")\n try:\n with urllib.request.urlopen(url) as resp, open(tmp, \"wb\") as out:\n shutil.copyfileobj(resp, out)\n os.replace(tmp, dest)\n finally:\n tmp.unlink(missing_ok=True)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.calculator import ORBCalculator\n\n # orb-models exposes one function per checkpoint, e.g. pretrained.orb_v2().\n fn_name = CHECKPOINTS[checkpoint].replace(\"-\", \"_\")\n load_fn = getattr(pretrained, fn_name)\n\n # orb-models resolves its default weights URL through `cached_path`, which\n # write-locks its cache dir even on warm hits — EACCES for anyone who can't\n # write the shared install (Garden-AI/rootstock#67). Handed a *local* path\n # instead, cached_path returns it without locking. So the weights are\n # pre-fetched into the shared model cache at `rootstock add` time\n # (maintainer, cache writable) and every later serve loads that file.\n url = _default_weights_url(load_fn)\n weights = _local_weights_path(url)\n if not weights.exists():\n _fetch(url, weights)\n\n orbff = load_fn(weights_path=str(weights), device=torch.device(device))\n return ORBCalculator(orbff, device=torch.device(device))\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", arch: str = \"orb-v2\"):\n # Custom checkpoints (`:custom` ids with user weights). A weights file doesn't say\n # which orb architecture produced it, so `arch` names the pretrained\n # loader to instantiate — pass the right one at call time\n # (setup_kwargs={\"arch\": ...} / --kwarg arch=...). Handing the loader a local path also means no\n # network and no cached_path locking (see setup()).\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.calculator import ORBCalculator\n\n fn_name = arch.replace(\"-\", \"_\")\n try:\n load_fn = getattr(pretrained, fn_name)\n except AttributeError:\n raise ValueError(\n f\"unknown orb architecture {arch!r}; expected a loader name from \"\n f\"orb_models.forcefield.pretrained, e.g. orb-v2, orb-d3-v2\"\n ) from None\n\n orbff = load_fn(weights_path=path, device=torch.device(device))\n return ORBCalculator(orbff, device=torch.device(device))\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","cached-path":"1.8.10","orb-models":"0.5.5","rootstock":"1.3.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:58ce179e368a874745398cf83be4cf4a5a96dc6e28669213de8e201b67be49d3","checkpoints":{"orb-d3-v2":{"fetched_at":"2026-08-03T17:18:30.895591+00:00","verified_at":"2026-10-25T19:51:31.848663+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-d3-v2-20241011.ckpt","size":100967743}],"weights_recorded_at":"2026-10-25T19:59:55.995452+00:00"}}},"pet":{"built_at":"2026-08-03T20:28:04.390522+00:00","source_hash":"sha256:2ed7f37a981d7b622c8f25c58a80b4269a919554a0fd61e36ddee4a7678654fd","source":"# /// script\n# requires-python = \">=3.11,<3.15\"\n# dependencies = [\n# \"upet>=0.2.6\",\n# \"ase>=3.22\",\n# # upet pulls nvalchemi-toolkit-ops unpinned; 0.4+ needs torch>=2.8 at\n# # runtime but only declares the constraint on its extras, so the\n# # resolver won't catch it (same trap as the tensornet env).\n# \"torch>=2.8,<2.14\",\n# ]\n# ///\n\"\"\"PET env — hosts lab-cosmo's UPET foundation checkpoints (PET-MAD successor).\n\nThe upstream string encodes model@version; versions are pinned rather than\n\"latest\" so rebuilds serve the same weights. pet-omatpes-l is trained at the\nr2SCAN level of theory — its energies are not comparable to the PBE-level\npet-oam models.\n\"\"\"\n\nCHECKPOINTS = {\n \"pet-oam-xl\": \"pet-oam-xl@1.0.0\",\n \"pet-omatpes-l\": \"pet-omatpes-l@0.1.0\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from huggingface_hub import hf_hub_download\n from upet.calculator import UPETCalculator\n\n # metatomic-ase 0.1.2's nvalchemi fast path passes a float max_neighbors\n # (len * max(128, cutoff**3)) into torch.full — TypeError for cutoffs > ~5 Å.\n # Force the vesin fallback; the flag is read at call time.\n import metatomic_ase._neighbors as _mta_n\n _mta_n.HAS_NVALCHEMIOPS = False\n\n # Passing model=/version= makes UPETCalculator resolve the name by listing\n # the hub repo — an uncached API call that fails on workers, which run\n # with HF_HUB_OFFLINE=1 (and on any node without internet). Fetch the\n # pinned file ourselves — a cache hit needs no network even offline — and\n # hand it over as checkpoint_path, which skips the resolve entirely.\n model, version = CHECKPOINTS[checkpoint].split(\"@\", 1)\n path = hf_hub_download(\n repo_id=\"lab-cosmo/upet\",\n filename=f\"{model}-v{version}.ckpt\",\n subfolder=\"models\",\n )\n return UPETCalculator(checkpoint_path=path, device=device)\n","python_requires":">=3.11,<3.15","dependencies":{"ase":"3.29.0","rootstock":"1.4.3","torch":"2.13.0","upet":"0.2.6"},"lock_hash":"sha256:9809737b1bd5fea93de01e7f7340e8140d88dbda6fa66c09ef1d7144435a6f5c","checkpoints":{"pet-omatpes-l":{"fetched_at":"2026-08-03T17:21:45.117310+00:00","verified_at":"2026-10-25T19:51:31.852904+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/blobs/879b1045391d88869522605a8b8b3cedeed74668e7062fdd7487548ab7b08004","size":771692879},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:51:31.852912+00:00"},"pet-oam-xl":{"fetched_at":"2026-08-03T17:21:45.473434+00:00","verified_at":"2026-10-25T19:51:31.848670+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/blobs/c3a67cd019969dfd4dcabe9574682fe035f861d3f1c10190989b36c983699409","size":2920687712},{"path":"cache/huggingface/hub/models--lab-cosmo--upet/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:51:31.848676+00:00"}}},"sevennet":{"built_at":"2026-07-31T18:37:57.316160+00:00","source_hash":"sha256:560f24594f3048803e4527d1a3761784e11c9a1d198c14781697f323a1384057","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"sevenn>=0.10.0\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"SevenNet env — hosts pretrained SevenNet universal potentials.\n\nSevenNet (SCalable EquiVariance-Enabled Neural Network) ships several\npretrained models, loaded by keyword through ``SevenNetCalculator``.\n\nMulti-fidelity models (``7net-omni``, ``7net-mf-ompa``) take a ``modal``\nargument selecting the training fidelity (e.g. ``\"mpa\"`` or ``\"omat24\"``),\nand sevenn releases after mid-2026 make it mandatory (older ones defaulted\nsilently). When not given, setup() defaults them to ``\"mpa\"`` — the\nMPtrj+sAlex fidelity, consistent with the lineage of the other checkpoints\nhere — so no-kwarg paths (smoke-test, plain ``rootstock add``) keep working.\nOverride at runtime via ``setup_kwargs={\"modal\": ...}`` on\nRootstockCalculator (or ``--kwarg modal=...`` for ``rootstock add``);\nsingle-fidelity models ignore it.\n\"\"\"\n\nCHECKPOINTS = {\n \"sevennet-0\": \"7net-0\",\n \"sevennet-l3i5\": \"7net-l3i5\",\n \"sevennet-omat\": \"7net-omat\",\n \"sevennet-mf-ompa\": \"7net-mf-ompa\",\n \"sevennet-omni\": \"7net-omni\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"sevennet:custom\": None,\n}\n\n\n# Multi-fidelity models and the fidelity used when modal isn't specified.\nMULTI_FIDELITY_DEFAULT_MODAL = {\n \"sevennet-mf-ompa\": \"mpa\",\n \"sevennet-omni\": \"mpa\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", modal: str | None = None):\n \"\"\"\n Load a SevenNet calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\").\n modal: Fidelity selector for multi-fidelity models (e.g. \"mpa\",\n \"omat24\"). Defaults to \"mpa\" for 7net-omni / 7net-mf-ompa;\n ignored by the single-fidelity models.\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n from sevenn.calculator import SevenNetCalculator\n\n if modal is None:\n modal = MULTI_FIDELITY_DEFAULT_MODAL.get(checkpoint)\n kwargs = {\"modal\": modal} if modal is not None else {}\n return SevenNetCalculator(model=CHECKPOINTS[checkpoint], device=device, **kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", modal: str | None = None):\n # Custom checkpoints (`:custom` ids with user weights): SevenNetCalculator loads a\n # checkpoint file directly. A multi-fidelity fine-tune must pass its\n # fidelity (setup_kwargs={\"modal\": \"mpa\"} / --kwarg modal=mpa);\n # single-fidelity ones need no kwargs.\n from sevenn.calculator import SevenNetCalculator\n\n kwargs = {\"modal\": modal} if modal is not None else {}\n return SevenNetCalculator(model=path, device=device, **kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","rootstock":"1.3.0","sevenn":"0.13.0","torch":"2.11.0+cu128"},"lock_hash":"sha256:5e341191eeda23990bfa2217b40a922d8c0e31eb5361ebb1233fc10c4edced4d","checkpoints":{"sevennet-0":{"fetched_at":"2026-08-03T17:21:35.263018+00:00","verified_at":"2026-10-25T19:51:31.852916+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-l3i5":{"fetched_at":"2026-08-03T17:21:35.332510+00:00","verified_at":"2026-10-25T19:51:31.852920+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-omat":{"fetched_at":"2026-08-03T17:22:03.544304+00:00","verified_at":"2026-10-25T19:51:31.852924+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-mf-ompa":{"fetched_at":"2026-08-03T17:22:04.480485+00:00","verified_at":"2026-10-25T19:51:31.852928+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-omni":{"fetched_at":"2026-08-03T17:22:06.723092+00:00","verified_at":"2026-10-25T19:51:31.852932+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"uma":{"built_at":"2026-08-06T19:38:53.367289+00:00","source_hash":"sha256:a21a78ddc6bc95aaabc4e17428d768b8d713671906246b4baf30a4e965b7d1f5","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"UMA env — hosts Meta's UMA foundation model via FAIRChem.\n\nfairchem-core v2 dropped the torch-geometric / pyg-find-links install dance, so\nthis env is a plain PyPI install. The original uma-s-1 had an extensivity bug\nand was removed from the fairchem 2.20 registry — use uma-s-1p1 or uma-s-1p2p1.\n\"\"\"\n\nCHECKPOINTS = {\n \"uma-s-1p1\": \"uma-s-1p1\",\n # uma-s-1p2 has a known major bug; uma-s-1p2p1 fixes it and is the\n # upstream-recommended small model. 1p2 stays listed for reproducibility\n # of existing runs.\n \"uma-s-1p2\": \"uma-s-1p2\",\n # uma-s-1p2p1 is in fairchem's registry on git main but NOT in any\n # release yet (latest fairchem-core 2.21.0, 2026-06-08, lacks it — the\n # 2026-07-30 sync failed on exactly this). Re-add when the next\n # fairchem-core ships, and bump the dependency floor to that version.\n # \"uma-s-1p2p1\": \"uma-s-1p2p1\",\n \"uma-m-1p1\": \"uma-m-1p1\",\n # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).\n \"uma:custom\": None,\n}\n\n\ndef _fairchem_device(device: str) -> str:\n \"\"\"Translate an indexed device (\"cuda:2\") into what fairchem v2 accepts.\n\n MLIPPredictUnit._setup_device asserts `device in [\"cpu\", \"cuda\"]` and then\n resolves the real GPU itself via get_device_for_local_rank(), which returns\n f\"cuda:{torch.cuda.current_device()}\". So an index has to travel through\n torch's current-device state, not the argument. Verifying several\n checkpoints at once on a multi-GPU node hands each worker \"cuda:N\" — that\n killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync\n (4x A100, VERIFY_JOBS=4), while single-GPU Sophia never hit it.\n \"\"\"\n if device.startswith(\"cuda:\"):\n import torch\n\n torch.cuda.set_device(int(device.split(\":\", 1)[1]))\n return \"cuda\"\n return device\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", task: str = \"omat\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(\n CHECKPOINTS[checkpoint], device=_fairchem_device(device)\n )\n return FAIRChemCalculator(predictor, task_name=task)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", task: str = \"omat\"):\n # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through\n # load_predict_unit, not the registry-name lookup setup() uses.\n from fairchem.core import FAIRChemCalculator\n from fairchem.core.units.mlip_unit import load_predict_unit\n\n predictor = load_predict_unit(path, device=_fairchem_device(device))\n return FAIRChemCalculator(predictor, task_name=task)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"1.6.0","torch":"2.8.0"},"lock_hash":"sha256:c99d413acdd9b5b2cdd0241ba29966574fe6800af2b63376abacd0ad28fed8b0","checkpoints":{"uma-s-1p1":{"fetched_at":"2026-08-03T17:25:15.360399+00:00","verified_at":"2026-10-25T19:51:31.852936+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/07068e9c76702ca173d13155095f2117c1b327ec228557e64cd2709c777b824a","size":1173759575},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.999502+00:00"},"uma-m-1p1":{"fetched_at":"2026-08-03T17:25:21.596525+00:00","verified_at":"2026-10-25T19:51:31.852955+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/c30034edbf2e127f703f814cacb632661767da99b3e71c5b2ee5290510a52d68","size":11174706771},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.999516+00:00"},"uma-s-1p2":{"fetched_at":"2026-08-03T17:26:22.685148+00:00","verified_at":"2026-10-25T19:51:31.852945+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/856080f16b62f6a42a1accd431308de37c6ed8c6","size":9000},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/ba5c0d912efa22dc238e5fb1b5b7f66ee2e68c48c1b95b7cfd7fe1da5938398b","size":2333393167},{"path":"home/.cache/fairchem/models--facebook--UMA/blobs/dfe22a26592f112ff3322ada76960ca1576afe90","size":11770},{"path":"home/.cache/fairchem/models--facebook--UMA/refs/main","size":40}],"weights_recorded_at":"2026-10-25T19:59:55.999509+00:00"},"uma:custom":{"fetched_at":null,"verified_at":"2026-10-25T19:51:31.853009+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"tensornet":{"built_at":"2026-08-03T20:28:00.083517+00:00","source_hash":"sha256:659d5d93612741c0a62befa641f1f5d2259ab760837e261d8c84ecf5d4564014","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"torch>=2.4.0,<2.5\",\n# \"ase>=3.22\",\n# \"huggingface_hub\",\n# \"matgl\",\n# # 0.4+ needs torch>=2.8 at runtime (custom-op registration uses string\n# # annotations infer_schema can't parse on older torch) but only declares\n# # the constraint on its extras, so the resolver won't catch it.\n# # \"nvalchemi-toolkit-ops<0.4\",\n# \"pymatgen\",\n# \"monty\",\n# \"ruamel.yaml\",\n# \"scipy\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# \"torch-spline-conv\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n#\n# [tool.uv.sources]\n# matgl = { git = \"https://github.com/materialsvirtuallab/matgl.git\", rev = \"35f3ea2d2028febfeb99bb6707f4c6f5223a0851\" }\n# ///\n\"\"\"TensorNet env — hosts MatPES TensorNet checkpoints via MatGL.\"\"\"\n\nCHECKPOINTS = {\n \"tensornet-matpes-pbe-2025-2\": \"materialyze/TensorNet-PES-MatPES-PBE-2025.2\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n # matgl 1.0.0 imports ExpCellFilter from ase.constraints, but it moved to\n # ase.filters in ASE 3.23. Patch it in before matgl imports.\n import ase.constraints\n\n if not hasattr(ase.constraints, \"ExpCellFilter\"):\n from ase.filters import ExpCellFilter\n\n ase.constraints.ExpCellFilter = ExpCellFilter\n\n # DGL 2.x graphbolt imports torchdata submodules removed in torchdata>=0.7.\n # Stub the entire graphbolt subpackage before `import dgl` runs; DGL's\n # __init__ will use our empty stub and skip the real graphbolt initialisation.\n # matgl only uses DGL for graph construction — graphbolt is never called.\n import sys, types\n\n for _name in [\n \"dgl.graphbolt\",\n \"dgl.graphbolt.base\",\n \"dgl.graphbolt.dataloader\",\n \"dgl.graphbolt.feature_fetcher\",\n \"dgl.graphbolt.minibatch_transformer\",\n ]:\n if _name not in sys.modules:\n sys.modules[_name] = types.ModuleType(_name)\n\n from huggingface_hub import snapshot_download\n\n import matgl\n from matgl.ext.ase import PESCalculator\n\n # matgl 1.0.0 load_model only checks the GitHub manifest; HF models must\n # be downloaded explicitly and passed as a local path.\n local_path = snapshot_download(repo_id=CHECKPOINTS[checkpoint])\n pot = matgl.load_model(local_path).to(device)\n return PESCalculator(potential=pot)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","huggingface-hub":"1.26.0","matgl":"4.0.3","monty":"2026.7.16","pymatgen":"2026.5.4","rootstock":"1.4.3","ruamel-yaml":"0.19.1","scipy":"1.17.1","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0.post1","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121","torch-spline-conv":"1.2.2+pt24cu121"},"lock_hash":null,"checkpoints":{"tensornet-matpes-pbe-2025-2":{"fetched_at":"2026-08-03T19:34:11.158461+00:00","verified_at":"2026-10-25T19:51:31.852964+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/blobs/19a0a506a58ea454e407040d55c393a174a5165834f22ea95c4c923ad9a406cc","size":4495},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/blobs/544966013206da7faa1966ffa360a82fb2d31c3e","size":5298},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/blobs/b9efa8c708014dd6c37e348f65dc1d8ae5c7bd5c7290de4d9f210e16c7fc3766","size":3402824},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/refs/main","size":40},{"path":"cache/huggingface/hub/models--materialyze--TensorNet-PES-MatPES-PBE-2025.2/trees/7ac2b90400a87eca7f44c1e7956dca6b1c85e418.json","size":906}],"weights_recorded_at":"2026-10-25T19:59:55.999524+00:00"}}},"allegro":{"built_at":"2026-10-16T21:43:39.464385+00:00","source_hash":"sha256:bc9261bd7205b13a8d1f74f17399694d9ba9f6d2fc05c50415b297a671ff9143","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # setup() uses nequip's loader internals and keys its compile cache on the\n# # nequip version, so stay within one minor release.\n# \"nequip>=0.19,<0.20\",\n# \"nequip-allegro>=0.8.3,<0.9\",\n# \"ase>=3.23\",\n# # AOTInductor compilation needs torch>=2.6. Listed here so the explicit\n# # cu128 index below routes it: PyPI's recent torch wheels bundle a CUDA\n# # newer than many cluster drivers support.\n# \"torch>=2.6,<2.12\",\n# # CUDA headers for compiling on nodes with no CUDA toolkit loaded (the\n# # runtime headers come with torch); matches the cu128 torch build.\n# \"cuda-toolkit[cccl,nvcc]>=12.8,<12.9\",\n# # AOTInductor's C++ builder imports torch.utils.cpp_extension, which\n# # imports setuptools at module level; uv venvs do not seed it.\n# \"setuptools>=70\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"Allegro env — hosts the Allegro foundation models from nequip.net.\n\nAllegro is a strictly local E(3)-equivariant interatomic potential from the\nKozinsky group (Harvard), built on the NequIP framework; its locality is what\nlets it scale across many GPUs in LAMMPS. Allegro-OAM-L is trained on OMat24\nand fine-tuned on sAlex + MPtrj; Allegro-MP-L is trained on MPtrj alone. All\nare PBE-level and licensed CC-BY-4.0. Versions are pinned to fixed Zenodo\nartifacts (md5-checked) so rebuilds serve the same weights; the package file\nis downloaded into the shared model cache at `rootstock add` time and loaded\nfrom there, so serving needs no network.\n\nCompilation: on GPU, models run through NequIP's AOTInductor path\n(``nequip-compile --mode aotinductor --target ase``), the path NequIP\nrecommends for inference. A compiled model only runs on the hardware and\nsoftware it was compiled with, so setup() compiles once per (model, GPU\narchitecture, torch version, nequip version) and caches the result: in the\nshared model cache when the caller can write there, otherwise in the caller's\nTorchInductor cache directory. The first GPU load of a model on a new GPU type\ntherefore compiles first: on an 8-core node with an A10G, about 2 minutes for\nAllegro-OAM-L, 5 for NequIP-OAM-L and 6-7 for NequIP-OAM-XL. If that load runs\ninto the worker startup timeout, raise it (``timeout=`` on\nRootstockCalculator, ``--verify-timeout`` for verification); an interrupted\ncompile still finishes in the background and later loads reuse the compiled\nfile. AOTInductor needs a C++17 host compiler (``c++``/``gcc`` >= 8, or\nwhatever ``CXX`` names) at compile time; CUDA headers come from the env's own\nNVIDIA wheels when no CUDA toolkit is loaded. Pass\n``setup_kwargs={\"compile\": False}`` to run eager instead. CPU loads always run\neager.\n\nCustom weights: ``allegro:custom`` with weights= accepts a packaged model\n(``.nequip.zip``), a training checkpoint (``.ckpt``), or a model already\ncompiled for ASE on this hardware (``.nequip.pt2``). If your model's type\nnames are not chemical symbols, pass\n``setup_kwargs={\"chemical_species_to_atom_type_map\": {\"H\": \"\", ...}}``.\n\"\"\"\n\nimport hashlib\nimport os\nimport shutil\nimport subprocess\nimport sys\nimport tempfile\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"allegro-oam-l\": \"Allegro-OAM-L-0.1\",\n \"allegro-mp-l\": \"Allegro-MP-L-0.1\",\n # Your own Allegro model: pair with weights= (loaded via setup_from_path).\n \"allegro:custom\": None,\n}\n\n# Explicit scope: both NVIDIA machines of the shared ALCF install.\nCLUSTERS = [\"sophia\", \"polaris\"]\n\n# Zenodo record and md5 of each pinned package file (the artifacts nequip.net\n# resolves mir-group/:0.1 to).\n_ARTIFACTS = {\n \"Allegro-OAM-L-0.1\": (\"16980200\", \"0db7f9b3c3a62e74d78b3fcf2973c462\"),\n \"Allegro-MP-L-0.1\": (\"16980200\", \"c68d5345e925ee444543d6d4d968161f\"),\n}\n\n\ndef _shared_cache() -> Path:\n return Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\") / \"nequip\"\n\n\ndef _package_path(name: str) -> Path:\n \"\"\"The pinned package file in the shared cache, downloading it if missing.\"\"\"\n path = _shared_cache() / \"packages\" / f\"{name}.nequip.zip\"\n if path.is_file():\n return path\n record, md5 = _ARTIFACTS[name]\n url = f\"https://zenodo.org/api/records/{record}/files/{name}.nequip.zip/content\"\n path.parent.mkdir(parents=True, exist_ok=True)\n tmp = path.with_name(f\".{path.name}.{os.getpid()}.part\")\n try:\n digest = hashlib.md5()\n with urllib.request.urlopen(url) as response, open(tmp, \"wb\") as f:\n while chunk := response.read(1 << 20):\n f.write(chunk)\n digest.update(chunk)\n if digest.hexdigest() != md5:\n raise RuntimeError(f\"{url}: md5 {digest.hexdigest()} does not match {md5}\")\n os.replace(tmp, path)\n finally:\n tmp.unlink(missing_ok=True)\n return path\n\n\ndef _compiled_dirs() -> list[Path]:\n \"\"\"Where compiled models are looked up, in order; the first writable one\n receives new compilations.\"\"\"\n dirs = [_shared_cache() / \"compiled\"]\n inductor = os.environ.get(\"TORCHINDUCTOR_CACHE_DIR\")\n if inductor:\n dirs.append(Path(inductor) / \"nequip\")\n else:\n dirs.append(Path.home() / \".cache\" / \"nequip\" / \"compiled\")\n return dirs\n\n\ndef _writable(directory: Path) -> bool:\n probe = directory\n while not probe.exists():\n probe = probe.parent\n return os.access(probe, os.W_OK | os.X_OK)\n\n\ndef _hardware_tag(device: str) -> str:\n import torch\n\n major, minor = torch.cuda.get_device_capability(torch.device(device))\n return f\"sm{major}{minor}\"\n\n\ndef _toolkit_env(env: dict, link_dir: Path) -> None:\n \"\"\"Point AOTInductor at CUDA headers and libcuda.\n\n The C++ wrapper AOTInductor builds needs ``$CUDA_HOME/include`` and links\n ``-lcuda``. torch derives CUDA_HOME from ``nvcc`` on PATH when unset, which\n goes wrong when that nvcc is a compiler-suite wrapper (e.g. the NVIDIA HPC\n SDK's, whose include dir has its own intrinsics headers and no cuda.h).\n So use CUDA_HOME, CUDA_PATH or nvcc's toolkit only if it really ships\n ``include/cuda.h``; otherwise fall back to the headers in NVIDIA's CUDA\n wheels (``nvidia/cuda_runtime/``, ``cuda_nvcc/`` and ``cuda_cccl/``) and\n link against the driver library this process has loaded.\n \"\"\"\n candidates = [env.get(\"CUDA_HOME\"), env.get(\"CUDA_PATH\")]\n nvcc = shutil.which(\"nvcc\", path=env.get(\"PATH\"))\n if nvcc:\n candidates.append(str(Path(nvcc).resolve().parent.parent))\n for home in filter(None, candidates):\n if (Path(home) / \"include\" / \"cuda.h\").is_file():\n env[\"CUDA_HOME\"] = home\n return\n env.pop(\"CUDA_PATH\", None)\n import nvidia\n\n headers = next(\n (h for base in nvidia.__path__ for h in Path(base).glob(\"*/include/cuda.h\")), None\n )\n if headers is None:\n return\n env[\"CUDA_HOME\"] = str(headers.parent.parent)\n extra = [\n str(h.parent.parent)\n for base in nvidia.__path__\n for marker in (\"crt/host_defines.h\", \"nv/target\")\n for h in Path(base).glob(f\"*/include/{marker}\")\n ]\n env[\"CPATH\"] = os.pathsep.join(filter(None, [*extra, env.get(\"CPATH\")]))\n with open(\"/proc/self/maps\") as maps:\n libcuda = next((line.split()[-1] for line in maps if \"/libcuda.so\" in line), None)\n if libcuda is not None:\n (link_dir / \"libcuda.so\").symlink_to(libcuda)\n env[\"LIBRARY_PATH\"] = os.pathsep.join(\n filter(None, [str(link_dir), env.get(\"LIBRARY_PATH\")])\n )\n\n\n# Host compilers tried, in order, when CXX is unset or not GNU-compatible.\n_CXX_CANDIDATES = (\"g++-14\", \"g++-13\", \"g++-12\", \"g++-11\", \"g++-10\", \"g++-9\", \"g++\", \"clang++\")\n\n\ndef _gnu_major(cxx: str) -> int | None:\n \"\"\"Major version of a GNU-compatible (gcc or clang) compiler, else None.\n\n AOTInductor passes gcc-style flags, which other compilers on HPC systems\n (e.g. NVIDIA's nvc++) reject.\n \"\"\"\n try:\n banner = subprocess.run([cxx, \"--version\"], capture_output=True, text=True).stdout\n version = subprocess.run([cxx, \"-dumpversion\"], capture_output=True, text=True).stdout\n except OSError:\n return None\n if \"Free Software Foundation\" not in banner and \"clang version\" not in banner:\n return None\n try:\n return int(version.strip().split(\".\")[0])\n except ValueError:\n return None\n\n\ndef _compiler_env(env: dict) -> None:\n \"\"\"Point CXX/CC at a GNU-compatible compiler new enough for C++17.\n\n Cluster compiler modules often export CXX/CC for a vendor compiler (e.g.\n nvc++/nvc). Keep the caller's CXX when it is gcc >= 9 or clang;\n otherwise use the first suitable candidate on PATH, with its matching C\n compiler (Triton builds a small C helper with CC during compilation).\n \"\"\"\n path = os.pathsep.join(filter(None, [env.get(\"PATH\"), \"/usr/bin\"]))\n for name in filter(None, [env.get(\"CXX\"), *_CXX_CANDIDATES]):\n cxx = shutil.which(name, path=path)\n major = _gnu_major(cxx) if cxx else None\n if major is None or (\"clang\" not in Path(cxx).name and major < 9):\n continue\n env[\"CXX\"] = cxx\n cc = shutil.which(\n Path(cxx).name.replace(\"clang++\", \"clang\").replace(\"g++\", \"gcc\"), path=path\n )\n if cc:\n env[\"CC\"] = cc\n return\n\n\ndef _compiled_model(source: Path, key: str, device: str) -> Path:\n \"\"\"Path to ``source`` compiled for ASE on ``device``, compiling on a miss.\"\"\"\n import nequip\n import torch\n\n tag = f\"{key}-{_hardware_tag(device)}-torch{torch.__version__}-nequip{nequip.__version__}\"\n name = f\"{tag.replace('+', '_')}.nequip.pt2\"\n dirs = _compiled_dirs()\n for directory in dirs:\n if (directory / name).is_file():\n return directory / name\n\n target_dir = next((d for d in dirs if _writable(d)), None)\n if target_dir is None:\n raise RuntimeError(f\"no writable directory to cache the compiled model in: {dirs}\")\n target_dir.mkdir(parents=True, exist_ok=True)\n out = target_dir / name\n tmp = target_dir / f\".{tag}.{os.getpid()}.nequip.pt2\"\n\n # Compile in a child process pinned to the same GPU, so the compiler's\n # memory is released before the model is loaded. The child moves its\n # output into place itself: if this worker is killed mid-compile (e.g. by\n # a startup timeout), the compile still completes and lands in the cache.\n env = dict(os.environ)\n if \":\" in device:\n index = int(device.split(\":\", 1)[1])\n visible = env.get(\"CUDA_VISIBLE_DEVICES\")\n env[\"CUDA_VISIBLE_DEVICES\"] = visible.split(\",\")[index] if visible else str(index)\n compile_and_move = (\n \"import os, sys; from nequip.scripts.compile import main; \"\n \"main(sys.argv[1:-1]); os.replace(sys.argv[2], sys.argv[-1])\"\n )\n cmd = [\n sys.executable,\n \"-c\",\n compile_and_move,\n str(source),\n str(tmp),\n \"--mode\",\n \"aotinductor\",\n \"--device\",\n \"cuda\",\n \"--target\",\n \"ase\",\n str(out),\n ]\n print(f\"nequip: compiling {source.name} for {tag} (one-time, several minutes)\", file=sys.stderr)\n try:\n with tempfile.TemporaryDirectory() as link_dir:\n _toolkit_env(env, Path(link_dir))\n _compiler_env(env)\n result = subprocess.run(cmd, env=env, stdout=sys.stderr, stderr=sys.stderr)\n if result.returncode != 0:\n raise RuntimeError(\n f\"nequip-compile failed (exit {result.returncode}); see the log above. \"\n \"It needs gcc >= 9 or clang (set CXX if neither is on PATH). \"\n 'Pass setup_kwargs={\"compile\": False} to run eager instead.'\n )\n finally:\n tmp.unlink(missing_ok=True)\n return out\n\n\ndef _load(source: Path, key: str, device: str, compile: bool, kwargs: dict):\n # torch's AOTInductor package loader reaches torch._inductor.codecache\n # without importing it, which fails in a fresh process.\n import torch._inductor.codecache # noqa: F401\n from nequip.integrations.ase import NequIPCalculator\n\n if source.name.endswith(\".nequip.pt2\"):\n return NequIPCalculator.from_compiled_model(str(source), device=device, **kwargs)\n if compile and device != \"cpu\":\n compiled = _compiled_model(source, key, device)\n return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs)\n # Eager load of a package or checkpoint. nequip exposes this only as a\n # private classmethod (its own tests use it); the public loader takes\n # compiled models only.\n return NequIPCalculator._from_saved_model(str(source), device=device, **kwargs)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", compile: bool = True, **kwargs):\n \"\"\"\n Load an Allegro calculator for a pinned nequip.net foundation model.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\").\n compile: Run the AOTInductor-compiled model on GPU (compiling and\n caching it on first use). False runs eager PyTorch.\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n name = CHECKPOINTS[checkpoint]\n kwargs.setdefault(\"chemical_species_to_atom_type_map\", True)\n return _load(_package_path(name), name, device, compile, kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", compile: bool = True, **kwargs):\n # Custom checkpoints (`:custom` ids with user weights): an Allegro .nequip.zip\n # package or .ckpt checkpoint (compiled on first GPU use, cached under a\n # key that includes the file's hash) or a ready .nequip.pt2.\n source = Path(path)\n key = source.name\n if not key.endswith(\".nequip.pt2\") and compile and device != \"cpu\":\n digest = hashlib.sha256()\n with open(source, \"rb\") as f:\n while chunk := f.read(1 << 20):\n digest.update(chunk)\n key = f\"custom-{digest.hexdigest()[:16]}\"\n return _load(source, key, device, compile, kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","cuda-toolkit":"12.8.2.0","nequip":"0.19.1","nequip-allegro":"0.8.3","rootstock":"1.6.4","setuptools":"84.0.0","torch":"2.10.0+cu128"},"lock_hash":"sha256:c022acfd951ab8b9fe6e1894c748a5f576b867ef693eb05d7eca18adcfae2867","checkpoints":{"allegro-oam-l":{"fetched_at":"2026-10-05T21:00:14.206502+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker process died before connecting: worker process exited with code 1.\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 34540 files, 8.9 GB in 65.1s; weights: 1.1 GB (heuristic)\nW1025 19:40:07.471000 2561461 lib/python3.11/site-packages/torch/_inductor/package/package.py:120] Loading outdated pt2 file. Please regenerate your package.\nTraceback (most recent call last):\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_xsjee7n7/wrapper.py\", line 46, in \n run_worker(\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/lib/python3.11/site-packages/rootstock/worker.py\", line 378, in run_worker\n calculator = setup_fn(checkpoint, device, **setup_kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_xsjee7n7/weights_capture.py\", line 192, in _capturing_setup\n calculator = setup_fn(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/env_source.py\", line 331, in setup\n return _load(_package_path(name), name, device, compile, kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/env_source.py\", line 309, in _load\n return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/lib/python3.11/site-packages/nequip/integrations/mixins.py\", line 52, in from_compiled_model\n model, metadata = load_compiled_model(\n ^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/lib/python3.11/site-packages/nequip/model/inference_models/compiled.py\", line 40, in load_compiled_model\n model, metadata = load_aotinductor_model(\n ^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/lib/python3.11/site-packages/nequip/model/inference_models/aotinductor.py\", line 82, in load_aotinductor_model\n compiled_model = torch._inductor.aoti_load_package(compile_path)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/lib/python3.11/site-packages/torch/_inductor/__init__.py\", line 268, in aoti_load_package\n return load_package(\n ^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/lib/python3.11/site-packages/torch/_inductor/package/package.py\", line 135, in load_package\n loader = torch._C._aoti.AOTIModelPackageLoader(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nRuntimeError: Error in dlopen: /lib64/libc.so.6: version `GLIBC_2.38' not found (required by /tmp/FaDyaO/.Allegro-OAM-L-0.1-sm80-torch2.10.0+cu128-nequip0.19.1.4062231.nequip/data/aotinductor/model/cccpe3z72wgcinzqephtwfif7gxl5xc7nnd5gpledt5eqpoev63o.wrapper.so)","weight_files":[{"path":"cache/nequip/compiled/Allegro-OAM-L-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":42039109}],"weights_recorded_at":"2026-10-25T19:59:55.999530+00:00"},"allegro-mp-l":{"fetched_at":"2026-10-05T21:00:21.906775+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker process died before connecting: worker process exited with code 1.\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 34540 files, 8.9 GB in 27.1s; weights: 1.1 GB (heuristic)\nW1025 19:40:47.824000 2563537 lib/python3.11/site-packages/torch/_inductor/package/package.py:120] Loading outdated pt2 file. Please regenerate your package.\nTraceback (most recent call last):\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_mirq_uiz/wrapper.py\", line 46, in \n run_worker(\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/lib/python3.11/site-packages/rootstock/worker.py\", line 378, in run_worker\n calculator = setup_fn(checkpoint, device, **setup_kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_mirq_uiz/weights_capture.py\", line 192, in _capturing_setup\n calculator = setup_fn(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/env_source.py\", line 331, in setup\n return _load(_package_path(name), name, device, compile, kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/env_source.py\", line 309, in _load\n return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/lib/python3.11/site-packages/nequip/integrations/mixins.py\", line 52, in from_compiled_model\n model, metadata = load_compiled_model(\n ^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/lib/python3.11/site-packages/nequip/model/inference_models/compiled.py\", line 40, in load_compiled_model\n model, metadata = load_aotinductor_model(\n ^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/lib/python3.11/site-packages/nequip/model/inference_models/aotinductor.py\", line 82, in load_aotinductor_model\n compiled_model = torch._inductor.aoti_load_package(compile_path)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/lib/python3.11/site-packages/torch/_inductor/__init__.py\", line 268, in aoti_load_package\n return load_package(\n ^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/allegro/lib/python3.11/site-packages/torch/_inductor/package/package.py\", line 135, in load_package\n loader = torch._C._aoti.AOTIModelPackageLoader(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nRuntimeError: Error in dlopen: /lib64/libc.so.6: version `GLIBC_2.38' not found (required by /tmp/Ka71wm/.Allegro-MP-L-0.1-sm80-torch2.10.0+cu128-nequip0.19.1.4062229.nequip/data/aotinductor/model/c7xppskofcwuv34gtpr47tgdduf3kn3ocdl4jyccpgyx7i7zgysz.wrapper.so)","weight_files":[{"path":"cache/nequip/compiled/Allegro-MP-L-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":78600784}],"weights_recorded_at":"2026-10-25T19:59:55.999535+00:00"},"allegro:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":null,"weight_files":null,"weights_recorded_at":null}}},"nequip":{"built_at":"2026-10-16T21:43:51.189255+00:00","source_hash":"sha256:eada3a775ec8af9ad476154b832df248cbd29157d1b14f7fad2453492e187cb7","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# # setup() uses nequip's loader internals and keys its compile cache on the\n# # nequip version, so stay within one minor release.\n# \"nequip>=0.19,<0.20\",\n# \"ase>=3.23\",\n# # AOTInductor compilation needs torch>=2.6. Listed here so the explicit\n# # cu128 index below routes it: PyPI's recent torch wheels bundle a CUDA\n# # newer than many cluster drivers support.\n# \"torch>=2.6,<2.12\",\n# # CUDA headers for compiling on nodes with no CUDA toolkit loaded (the\n# # runtime headers come with torch); matches the cu128 torch build.\n# \"cuda-toolkit[cccl,nvcc]>=12.8,<12.9\",\n# # AOTInductor's C++ builder imports torch.utils.cpp_extension, which\n# # imports setuptools at module level; uv venvs do not seed it.\n# \"setuptools>=70\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"NequIP env — hosts the NequIP-OAM foundation models from nequip.net.\n\nNequIP is an E(3)-equivariant message-passing interatomic potential from the\nKozinsky group (Harvard). The OAM models are trained on OMat24 and fine-tuned\non sAlex + MPtrj; NequIP-MP-L is trained on MPtrj alone. All are PBE-level and\nlicensed CC-BY-4.0. Versions are pinned to fixed Zenodo artifacts\n(md5-checked) so rebuilds serve the same weights; the package file is\ndownloaded into the shared model cache at `rootstock add` time and loaded from\nthere, so serving needs no network.\n\nCompilation: on GPU, models run through nequip's AOTInductor path\n(``nequip-compile --mode aotinductor --target ase``), the path NequIP\nrecommends for inference. A compiled model only runs on the hardware and\nsoftware it was compiled with, so setup() compiles once per (model, GPU\narchitecture, torch version, nequip version) and caches the result: in the\nshared model cache when the caller can write there, otherwise in the caller's\nTorchInductor cache directory. The first GPU load of a model on a new GPU type\ntherefore compiles first: on an 8-core node with an A10G, about 2 minutes for\nAllegro-OAM-L, 5 for NequIP-OAM-L and 6-7 for NequIP-OAM-XL. If that load runs\ninto the worker startup timeout, raise it (``timeout=`` on\nRootstockCalculator, ``--verify-timeout`` for verification); an interrupted\ncompile still finishes in the background and later loads reuse the compiled\nfile. AOTInductor needs a C++17 host compiler (``c++``/``gcc`` >= 8, or\nwhatever ``CXX`` names) at compile time; CUDA headers come from the env's own\nNVIDIA wheels when no CUDA toolkit is loaded. Pass\n``setup_kwargs={\"compile\": False}`` to run eager instead. CPU loads always run\neager.\n\nCustom weights: ``nequip:custom`` with weights= accepts a packaged model\n(``.nequip.zip``), a training checkpoint (``.ckpt``), or a model already\ncompiled for ASE on this hardware (``.nequip.pt2``). If your model's type\nnames are not chemical symbols, pass\n``setup_kwargs={\"chemical_species_to_atom_type_map\": {\"H\": \"\", ...}}``.\n\"\"\"\n\nimport hashlib\nimport os\nimport shutil\nimport subprocess\nimport sys\nimport tempfile\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"nequip-oam-s\": \"NequIP-OAM-S-0.1\",\n \"nequip-oam-m\": \"NequIP-OAM-M-0.1\",\n \"nequip-oam-l\": \"NequIP-OAM-L-0.1\",\n \"nequip-oam-xl\": \"NequIP-OAM-XL-0.1\",\n \"nequip-mp-l\": \"NequIP-MP-L-0.1\",\n # Your own NequIP model: pair with weights= (loaded via setup_from_path).\n \"nequip:custom\": None,\n}\n\n# Explicit scope: both NVIDIA machines of the shared ALCF install.\nCLUSTERS = [\"sophia\", \"polaris\"]\n\n# Zenodo record and md5 of each pinned package file (the artifacts nequip.net\n# resolves mir-group/:0.1 to).\n_ARTIFACTS = {\n \"NequIP-OAM-S-0.1\": (\"18775904\", \"399a98bf36fc4550bc85d48f35451c6c\"),\n \"NequIP-OAM-M-0.1\": (\"18775904\", \"e5b44583f4421be702b35d76503754cd\"),\n \"NequIP-OAM-L-0.1\": (\"18775904\", \"67144367c710a70a53a8e21acf331980\"),\n \"NequIP-OAM-XL-0.1\": (\"18775904\", \"3d2369c7238eb83a23141abdcb055a8f\"),\n \"NequIP-MP-L-0.1\": (\"18775904\", \"624a01d4cbbae2d3bb75dc7074532d13\"),\n}\n\n\ndef _shared_cache() -> Path:\n return Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\") / \"nequip\"\n\n\ndef _package_path(name: str) -> Path:\n \"\"\"The pinned package file in the shared cache, downloading it if missing.\"\"\"\n path = _shared_cache() / \"packages\" / f\"{name}.nequip.zip\"\n if path.is_file():\n return path\n record, md5 = _ARTIFACTS[name]\n url = f\"https://zenodo.org/api/records/{record}/files/{name}.nequip.zip/content\"\n path.parent.mkdir(parents=True, exist_ok=True)\n tmp = path.with_name(f\".{path.name}.{os.getpid()}.part\")\n try:\n digest = hashlib.md5()\n with urllib.request.urlopen(url) as response, open(tmp, \"wb\") as f:\n while chunk := response.read(1 << 20):\n f.write(chunk)\n digest.update(chunk)\n if digest.hexdigest() != md5:\n raise RuntimeError(f\"{url}: md5 {digest.hexdigest()} does not match {md5}\")\n os.replace(tmp, path)\n finally:\n tmp.unlink(missing_ok=True)\n return path\n\n\ndef _compiled_dirs() -> list[Path]:\n \"\"\"Where compiled models are looked up, in order; the first writable one\n receives new compilations.\"\"\"\n dirs = [_shared_cache() / \"compiled\"]\n inductor = os.environ.get(\"TORCHINDUCTOR_CACHE_DIR\")\n if inductor:\n dirs.append(Path(inductor) / \"nequip\")\n else:\n dirs.append(Path.home() / \".cache\" / \"nequip\" / \"compiled\")\n return dirs\n\n\ndef _writable(directory: Path) -> bool:\n probe = directory\n while not probe.exists():\n probe = probe.parent\n return os.access(probe, os.W_OK | os.X_OK)\n\n\ndef _hardware_tag(device: str) -> str:\n import torch\n\n major, minor = torch.cuda.get_device_capability(torch.device(device))\n return f\"sm{major}{minor}\"\n\n\ndef _toolkit_env(env: dict, link_dir: Path) -> None:\n \"\"\"Point AOTInductor at CUDA headers and libcuda.\n\n The C++ wrapper AOTInductor builds needs ``$CUDA_HOME/include`` and links\n ``-lcuda``. torch derives CUDA_HOME from ``nvcc`` on PATH when unset, which\n goes wrong when that nvcc is a compiler-suite wrapper (e.g. the NVIDIA HPC\n SDK's, whose include dir has its own intrinsics headers and no cuda.h).\n So use CUDA_HOME, CUDA_PATH or nvcc's toolkit only if it really ships\n ``include/cuda.h``; otherwise fall back to the headers in NVIDIA's CUDA\n wheels (``nvidia/cuda_runtime/``, ``cuda_nvcc/`` and ``cuda_cccl/``) and\n link against the driver library this process has loaded.\n \"\"\"\n candidates = [env.get(\"CUDA_HOME\"), env.get(\"CUDA_PATH\")]\n nvcc = shutil.which(\"nvcc\", path=env.get(\"PATH\"))\n if nvcc:\n candidates.append(str(Path(nvcc).resolve().parent.parent))\n for home in filter(None, candidates):\n if (Path(home) / \"include\" / \"cuda.h\").is_file():\n env[\"CUDA_HOME\"] = home\n return\n env.pop(\"CUDA_PATH\", None)\n import nvidia\n\n headers = next(\n (h for base in nvidia.__path__ for h in Path(base).glob(\"*/include/cuda.h\")), None\n )\n if headers is None:\n return\n env[\"CUDA_HOME\"] = str(headers.parent.parent)\n extra = [\n str(h.parent.parent)\n for base in nvidia.__path__\n for marker in (\"crt/host_defines.h\", \"nv/target\")\n for h in Path(base).glob(f\"*/include/{marker}\")\n ]\n env[\"CPATH\"] = os.pathsep.join(filter(None, [*extra, env.get(\"CPATH\")]))\n with open(\"/proc/self/maps\") as maps:\n libcuda = next((line.split()[-1] for line in maps if \"/libcuda.so\" in line), None)\n if libcuda is not None:\n (link_dir / \"libcuda.so\").symlink_to(libcuda)\n env[\"LIBRARY_PATH\"] = os.pathsep.join(\n filter(None, [str(link_dir), env.get(\"LIBRARY_PATH\")])\n )\n\n\n# Host compilers tried, in order, when CXX is unset or not GNU-compatible.\n_CXX_CANDIDATES = (\"g++-14\", \"g++-13\", \"g++-12\", \"g++-11\", \"g++-10\", \"g++-9\", \"g++\", \"clang++\")\n\n\ndef _gnu_major(cxx: str) -> int | None:\n \"\"\"Major version of a GNU-compatible (gcc or clang) compiler, else None.\n\n AOTInductor passes gcc-style flags, which other compilers on HPC systems\n (e.g. NVIDIA's nvc++) reject.\n \"\"\"\n try:\n banner = subprocess.run([cxx, \"--version\"], capture_output=True, text=True).stdout\n version = subprocess.run([cxx, \"-dumpversion\"], capture_output=True, text=True).stdout\n except OSError:\n return None\n if \"Free Software Foundation\" not in banner and \"clang version\" not in banner:\n return None\n try:\n return int(version.strip().split(\".\")[0])\n except ValueError:\n return None\n\n\ndef _compiler_env(env: dict) -> None:\n \"\"\"Point CXX/CC at a GNU-compatible compiler new enough for C++17.\n\n Cluster compiler modules often export CXX/CC for a vendor compiler (e.g.\n nvc++/nvc). Keep the caller's CXX when it is gcc >= 9 or clang;\n otherwise use the first suitable candidate on PATH, with its matching C\n compiler (Triton builds a small C helper with CC during compilation).\n \"\"\"\n path = os.pathsep.join(filter(None, [env.get(\"PATH\"), \"/usr/bin\"]))\n for name in filter(None, [env.get(\"CXX\"), *_CXX_CANDIDATES]):\n cxx = shutil.which(name, path=path)\n major = _gnu_major(cxx) if cxx else None\n if major is None or (\"clang\" not in Path(cxx).name and major < 9):\n continue\n env[\"CXX\"] = cxx\n cc = shutil.which(\n Path(cxx).name.replace(\"clang++\", \"clang\").replace(\"g++\", \"gcc\"), path=path\n )\n if cc:\n env[\"CC\"] = cc\n return\n\n\ndef _compiled_model(source: Path, key: str, device: str) -> Path:\n \"\"\"Path to ``source`` compiled for ASE on ``device``, compiling on a miss.\"\"\"\n import nequip\n import torch\n\n tag = f\"{key}-{_hardware_tag(device)}-torch{torch.__version__}-nequip{nequip.__version__}\"\n name = f\"{tag.replace('+', '_')}.nequip.pt2\"\n dirs = _compiled_dirs()\n for directory in dirs:\n if (directory / name).is_file():\n return directory / name\n\n target_dir = next((d for d in dirs if _writable(d)), None)\n if target_dir is None:\n raise RuntimeError(f\"no writable directory to cache the compiled model in: {dirs}\")\n target_dir.mkdir(parents=True, exist_ok=True)\n out = target_dir / name\n tmp = target_dir / f\".{tag}.{os.getpid()}.nequip.pt2\"\n\n # Compile in a child process pinned to the same GPU, so the compiler's\n # memory is released before the model is loaded. The child moves its\n # output into place itself: if this worker is killed mid-compile (e.g. by\n # a startup timeout), the compile still completes and lands in the cache.\n env = dict(os.environ)\n if \":\" in device:\n index = int(device.split(\":\", 1)[1])\n visible = env.get(\"CUDA_VISIBLE_DEVICES\")\n env[\"CUDA_VISIBLE_DEVICES\"] = visible.split(\",\")[index] if visible else str(index)\n compile_and_move = (\n \"import os, sys; from nequip.scripts.compile import main; \"\n \"main(sys.argv[1:-1]); os.replace(sys.argv[2], sys.argv[-1])\"\n )\n cmd = [\n sys.executable,\n \"-c\",\n compile_and_move,\n str(source),\n str(tmp),\n \"--mode\",\n \"aotinductor\",\n \"--device\",\n \"cuda\",\n \"--target\",\n \"ase\",\n str(out),\n ]\n print(f\"nequip: compiling {source.name} for {tag} (one-time, several minutes)\", file=sys.stderr)\n try:\n with tempfile.TemporaryDirectory() as link_dir:\n _toolkit_env(env, Path(link_dir))\n _compiler_env(env)\n result = subprocess.run(cmd, env=env, stdout=sys.stderr, stderr=sys.stderr)\n if result.returncode != 0:\n raise RuntimeError(\n f\"nequip-compile failed (exit {result.returncode}); see the log above. \"\n \"It needs gcc >= 9 or clang (set CXX if neither is on PATH). \"\n 'Pass setup_kwargs={\"compile\": False} to run eager instead.'\n )\n finally:\n tmp.unlink(missing_ok=True)\n return out\n\n\ndef _load(source: Path, key: str, device: str, compile: bool, kwargs: dict):\n # torch's AOTInductor package loader reaches torch._inductor.codecache\n # without importing it, which fails in a fresh process.\n import torch._inductor.codecache # noqa: F401\n from nequip.integrations.ase import NequIPCalculator\n\n if source.name.endswith(\".nequip.pt2\"):\n return NequIPCalculator.from_compiled_model(str(source), device=device, **kwargs)\n if compile and device != \"cpu\":\n compiled = _compiled_model(source, key, device)\n return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs)\n # Eager load of a package or checkpoint. nequip exposes this only as a\n # private classmethod (its own tests use it); the public loader takes\n # compiled models only.\n return NequIPCalculator._from_saved_model(str(source), device=device, **kwargs)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", compile: bool = True, **kwargs):\n \"\"\"\n Load a NequIP calculator for a pinned nequip.net foundation model.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cuda:0\", \"cpu\").\n compile: Run the AOTInductor-compiled model on GPU (compiling and\n caching it on first use). False runs eager PyTorch.\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n name = CHECKPOINTS[checkpoint]\n kwargs.setdefault(\"chemical_species_to_atom_type_map\", True)\n return _load(_package_path(name), name, device, compile, kwargs)\n\n\ndef setup_from_path(path: str, device: str = \"cuda\", compile: bool = True, **kwargs):\n # Custom checkpoints (`:custom` ids with user weights): a .nequip.zip\n # package or .ckpt checkpoint (compiled on first GPU use, cached under a\n # key that includes the file's hash) or a ready .nequip.pt2.\n source = Path(path)\n key = source.name\n if not key.endswith(\".nequip.pt2\") and compile and device != \"cpu\":\n digest = hashlib.sha256()\n with open(source, \"rb\") as f:\n while chunk := f.read(1 << 20):\n digest.update(chunk)\n key = f\"custom-{digest.hexdigest()[:16]}\"\n return _load(source, key, device, compile, kwargs)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","cuda-toolkit":"12.8.2.0","nequip":"0.19.1","rootstock":"1.6.4","setuptools":"84.0.0","torch":"2.10.0+cu128"},"lock_hash":"sha256:52b324a75bb7630844f3a1faad8974c4e9aef36eff2b400477e4c96004d6811c","checkpoints":{"nequip-mp-l":{"fetched_at":"2026-10-05T21:00:13.864265+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker process died before connecting: worker process exited with code 1.\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 34488 files, 8.9 GB in 15.4s; weights: 1.1 GB (heuristic)\nW1025 19:44:35.410000 2569762 lib/python3.11/site-packages/torch/_inductor/package/package.py:120] Loading outdated pt2 file. Please regenerate your package.\nTraceback (most recent call last):\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_4nb7zgq0/wrapper.py\", line 46, in \n run_worker(\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/rootstock/worker.py\", line 378, in run_worker\n calculator = setup_fn(checkpoint, device, **setup_kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_4nb7zgq0/weights_capture.py\", line 192, in _capturing_setup\n calculator = setup_fn(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/env_source.py\", line 335, in setup\n return _load(_package_path(name), name, device, compile, kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/env_source.py\", line 313, in _load\n return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/integrations/mixins.py\", line 52, in from_compiled_model\n model, metadata = load_compiled_model(\n ^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/model/inference_models/compiled.py\", line 40, in load_compiled_model\n model, metadata = load_aotinductor_model(\n ^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/model/inference_models/aotinductor.py\", line 82, in load_aotinductor_model\n compiled_model = torch._inductor.aoti_load_package(compile_path)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/torch/_inductor/__init__.py\", line 268, in aoti_load_package\n return load_package(\n ^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/torch/_inductor/package/package.py\", line 135, in load_package\n loader = torch._C._aoti.AOTIModelPackageLoader(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nRuntimeError: Error in dlopen: /lib64/libc.so.6: version `GLIBC_2.38' not found (required by /tmp/ezbudI/.NequIP-MP-L-0.1-sm80-torch2.10.0+cu128-nequip0.19.1.4062228.nequip/data/aotinductor/model/cjuanxdtswaxorz54qfieopfcrivzctvazcejcmnvvihc6arlqke.wrapper.so)","weight_files":[{"path":"cache/nequip/compiled/NequIP-MP-L-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":46051639}],"weights_recorded_at":"2026-10-25T19:59:55.999561+00:00"},"nequip-oam-l":{"fetched_at":"2026-10-05T21:00:14.053978+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker process died before connecting: worker process exited with code 1.\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 34488 files, 8.9 GB in 12.5s; weights: 1.1 GB (heuristic)\nW1025 19:43:32.413000 2568613 lib/python3.11/site-packages/torch/_inductor/package/package.py:120] Loading outdated pt2 file. Please regenerate your package.\nTraceback (most recent call last):\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_qyrlc0_h/wrapper.py\", line 46, in \n run_worker(\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/rootstock/worker.py\", line 378, in run_worker\n calculator = setup_fn(checkpoint, device, **setup_kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_qyrlc0_h/weights_capture.py\", line 192, in _capturing_setup\n calculator = setup_fn(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/env_source.py\", line 335, in setup\n return _load(_package_path(name), name, device, compile, kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/env_source.py\", line 313, in _load\n return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/integrations/mixins.py\", line 52, in from_compiled_model\n model, metadata = load_compiled_model(\n ^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/model/inference_models/compiled.py\", line 40, in load_compiled_model\n model, metadata = load_aotinductor_model(\n ^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/model/inference_models/aotinductor.py\", line 82, in load_aotinductor_model\n compiled_model = torch._inductor.aoti_load_package(compile_path)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/torch/_inductor/__init__.py\", line 268, in aoti_load_package\n return load_package(\n ^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/torch/_inductor/package/package.py\", line 135, in load_package\n loader = torch._C._aoti.AOTIModelPackageLoader(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nRuntimeError: Error in dlopen: /lib64/libc.so.6: version `GLIBC_2.38' not found (required by /tmp/dJ0UOW/.NequIP-OAM-L-0.1-sm80-torch2.10.0+cu128-nequip0.19.1.4062230.nequip/data/aotinductor/model/caf3wvfkfgcij4ki2ajqaozm72iisambhbq4ldgdznddd5wb5eg4.wrapper.so)","weight_files":[{"path":"cache/nequip/compiled/NequIP-OAM-L-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":45969552}],"weights_recorded_at":"2026-10-25T19:59:55.999551+00:00"},"nequip-oam-s":{"fetched_at":"2026-10-05T21:00:31.367040+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker process died before connecting: worker process exited with code 1.\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 34488 files, 8.9 GB in 64.2s; weights: 1.1 GB (heuristic)\nW1025 19:42:27.966000 2564327 lib/python3.11/site-packages/torch/_inductor/package/package.py:120] Loading outdated pt2 file. Please regenerate your package.\nTraceback (most recent call last):\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_6c01orlk/wrapper.py\", line 46, in \n run_worker(\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/rootstock/worker.py\", line 378, in run_worker\n calculator = setup_fn(checkpoint, device, **setup_kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_6c01orlk/weights_capture.py\", line 192, in _capturing_setup\n calculator = setup_fn(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/env_source.py\", line 335, in setup\n return _load(_package_path(name), name, device, compile, kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/env_source.py\", line 313, in _load\n return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/integrations/mixins.py\", line 52, in from_compiled_model\n model, metadata = load_compiled_model(\n ^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/model/inference_models/compiled.py\", line 40, in load_compiled_model\n model, metadata = load_aotinductor_model(\n ^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/model/inference_models/aotinductor.py\", line 82, in load_aotinductor_model\n compiled_model = torch._inductor.aoti_load_package(compile_path)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/torch/_inductor/__init__.py\", line 268, in aoti_load_package\n return load_package(\n ^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/torch/_inductor/package/package.py\", line 135, in load_package\n loader = torch._C._aoti.AOTIModelPackageLoader(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nRuntimeError: Error in dlopen: /lib64/libc.so.6: version `GLIBC_2.38' not found (required by /tmp/AZhYv9/.NequIP-OAM-S-0.1-sm80-torch2.10.0+cu128-nequip0.19.1.4064931.nequip/data/aotinductor/model/cnotc2zbk27tvdwv22ckbmvyjljvxg2sfm2n57gkvdumunnmgvcl.wrapper.so)","weight_files":[{"path":"cache/nequip/compiled/NequIP-OAM-S-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":4277568}],"weights_recorded_at":"2026-10-25T19:59:55.999540+00:00"},"nequip-oam-m":{"fetched_at":"2026-10-05T21:00:37.557331+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker process died before connecting: worker process exited with code 1.\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 34488 files, 8.9 GB in 29.1s; weights: 1.1 GB (heuristic)\nW1025 19:43:10.141000 2567432 lib/python3.11/site-packages/torch/_inductor/package/package.py:120] Loading outdated pt2 file. Please regenerate your package.\nTraceback (most recent call last):\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_tx9rkbsk/wrapper.py\", line 46, in \n run_worker(\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/rootstock/worker.py\", line 378, in run_worker\n calculator = setup_fn(checkpoint, device, **setup_kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_tx9rkbsk/weights_capture.py\", line 192, in _capturing_setup\n calculator = setup_fn(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/env_source.py\", line 335, in setup\n return _load(_package_path(name), name, device, compile, kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/env_source.py\", line 313, in _load\n return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/integrations/mixins.py\", line 52, in from_compiled_model\n model, metadata = load_compiled_model(\n ^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/model/inference_models/compiled.py\", line 40, in load_compiled_model\n model, metadata = load_aotinductor_model(\n ^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/model/inference_models/aotinductor.py\", line 82, in load_aotinductor_model\n compiled_model = torch._inductor.aoti_load_package(compile_path)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/torch/_inductor/__init__.py\", line 268, in aoti_load_package\n return load_package(\n ^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/torch/_inductor/package/package.py\", line 135, in load_package\n loader = torch._C._aoti.AOTIModelPackageLoader(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nRuntimeError: Error in dlopen: /lib64/libc.so.6: version `GLIBC_2.38' not found (required by /tmp/tHY4bG/.NequIP-OAM-M-0.1-sm80-torch2.10.0+cu128-nequip0.19.1.4064927.nequip/data/aotinductor/model/cz3tmbbrqp3okmcfal5jw5zdfxbtwz4vldyahb7fx5fldz5xtzv5.wrapper.so)","weight_files":[{"path":"cache/nequip/compiled/NequIP-OAM-M-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":17148312}],"weights_recorded_at":"2026-10-25T19:59:55.999546+00:00"},"nequip-oam-xl":{"fetched_at":"2026-10-05T21:01:39.674414+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: WorkerDiedError: Worker process died before connecting: worker process exited with code 1.\n--- worker stderr (tail) ---\n[Worker] Prewarmed page cache: 34488 files, 8.9 GB in 25.1s; weights: 1.1 GB (heuristic)\nW1025 19:44:09.814000 2568855 lib/python3.11/site-packages/torch/_inductor/package/package.py:120] Loading outdated pt2 file. Please regenerate your package.\nTraceback (most recent call last):\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_cqr3doz5/wrapper.py\", line 46, in \n run_worker(\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/rootstock/worker.py\", line 378, in run_worker\n calculator = setup_fn(checkpoint, device, **setup_kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/var/tmp/pbs.196937.sophia-pbs-01.lab.alcf.anl.gov/rootstock_spawn_cqr3doz5/weights_capture.py\", line 192, in _capturing_setup\n calculator = setup_fn(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/env_source.py\", line 335, in setup\n return _load(_package_path(name), name, device, compile, kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/env_source.py\", line 313, in _load\n return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/integrations/mixins.py\", line 52, in from_compiled_model\n model, metadata = load_compiled_model(\n ^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/model/inference_models/compiled.py\", line 40, in load_compiled_model\n model, metadata = load_aotinductor_model(\n ^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/nequip/model/inference_models/aotinductor.py\", line 82, in load_aotinductor_model\n compiled_model = torch._inductor.aoti_load_package(compile_path)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/torch/_inductor/__init__.py\", line 268, in aoti_load_package\n return load_package(\n ^^^^^^^^^^^^^\n File \"/lus/eagle/projects/Garden-Ai/rootstock/envs/nequip/lib/python3.11/site-packages/torch/_inductor/package/package.py\", line 135, in load_package\n loader = torch._C._aoti.AOTIModelPackageLoader(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nRuntimeError: Error in dlopen: /lib64/libc.so.6: version `GLIBC_2.38' not found (required by /tmp/vetSCJ/.NequIP-OAM-XL-0.1-sm80-torch2.10.0+cu128-nequip0.19.1.4066696.nequip/data/aotinductor/model/cvba24m2ccwzdtcw4shay5wwzwdbnc4gyl3fjfjuhhel7m7xezau.wrapper.so)","weight_files":[{"path":"cache/nequip/compiled/NequIP-OAM-XL-0.1-sm80-torch2.10.0_cu128-nequip0.19.1.nequip.pt2","size":140072170}],"weights_recorded_at":"2026-10-25T19:59:55.999556+00:00"},"nequip:custom":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":null,"weight_files":null,"weights_recorded_at":null}}}}},{"schema_version":3,"cluster":"unknown","root":"/lus/eagle/projects/Garden-Ai/rootstock","maintainer":{"name":"Owen Price Skelly","email":"OwenPriceSkelly@uchicago.edu"},"rootstock_version":"0.9.5","python_version":"3.10","last_updated":"2026-07-27T22:40:47.449765+00:00","environments":{"ani":{"status":"ready","built_at":"2026-07-11T05:31:39.473889+00:00","source_hash":"sha256:4519d6f76936f874039d5a402b9a8470d2432736a839d84163af67ce1a20d3e7","source":"# /// script\n# requires-python = \">=3.10\"\n# dependencies = [\n# \"torchani>=2.2\",\n# \"ase>=3.22\",\n# \"torch>=2.0\",\n# ]\n# ///\n\"\"\"\nANI-2x environment for Rootstock.\n\nANI-2x is a neural network potential for organic molecules containing\nH, C, N, O, F, S, Cl. It is not a universal potential — do not use it\nfor inorganic or periodic systems.\n\nModels:\n - \"ANI2x\": ANI-2x ensemble (default, 8 networks)\n - \"ANI1ccx\": ANI-1ccx, trained on CCSD(T)/CBS data (H, C, N, O only)\n - \"ANI1x\": ANI-1x (H, C, N, O only)\n\"\"\"\n\nCHECKPOINTS = {\n \"ani-2x\": \"ANI2x\",\n \"ani-1ccx\": \"ANI1ccx\",\n \"ani-1x\": \"ANI1x\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an ANI calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n import torchani\n\n model_map = {\n \"ANI2x\": torchani.models.ANI2x,\n \"ANI1ccx\": torchani.models.ANI1ccx,\n \"ANI1x\": torchani.models.ANI1x,\n }\n model = CHECKPOINTS[checkpoint]\n\n return model_map[model](periodic_table_index=True).to(device).ase()\n","python_requires":">=3.10","dependencies":{"ase":"3.29.0","rootstock":"0.9.5","torch":"2.8.0","torchani":"2.8.2"},"checkpoints":{"ani-2x":{"fetched_at":"2026-07-11T05:31:49.196284+00:00","verified_at":"2026-07-27T22:34:04.937177+00:00","verified_device":"cuda","last_error":null},"ani-1ccx":{"fetched_at":"2026-07-11T05:32:06.608687+00:00","verified_at":"2026-07-27T22:34:08.583203+00:00","verified_device":"cuda","last_error":null},"ani-1x":{"fetched_at":"2026-07-11T05:32:17.424680+00:00","verified_at":"2026-07-27T22:34:12.313409+00:00","verified_device":"cuda","last_error":null}}},"equiformer":{"status":"ready","built_at":"2026-07-11T05:31:54.529081+00:00","source_hash":"sha256:8400e3742fb766020bbba69feed4de61d113b9a302b63bdf3ffed24807de5558","source":"# /// script\n# requires-python = \">=3.10,<3.11\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\nEquiformerV2 environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 EquiformerV2 checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"EquiformerV2-153M-S2EF-OC20-All+MD\": default\n - \"EquiformerV2-31M-S2EF-OC20-All+MD\"\n - \"EquiformerV2-83M-S2EF-OC20-2M\"\n\"\"\"\n\nCHECKPOINTS = {\n \"equiformer-v2-153m-s2ef-oc20-all-md\": \"EquiformerV2-153M-S2EF-OC20-All+MD\",\n \"equiformer-v2-31m-s2ef-oc20-all-md\": \"EquiformerV2-31M-S2EF-OC20-All+MD\",\n \"equiformer-v2-83m-s2ef-oc20-2m\": \"EquiformerV2-83M-S2EF-OC20-2M\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an EquiformerV2 OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n","python_requires":">=3.10,<3.11","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"0.9.5","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"checkpoints":{"equiformer-v2-153m-s2ef-oc20-all-md":{"fetched_at":"2026-07-11T05:33:01.976220+00:00","verified_at":"2026-07-27T22:35:06.297876+00:00","verified_device":"cuda","last_error":null},"equiformer-v2-31m-s2ef-oc20-all-md":{"fetched_at":"2026-07-11T05:33:39.405390+00:00","verified_at":"2026-07-27T22:35:18.361081+00:00","verified_device":"cuda","last_error":null},"equiformer-v2-83m-s2ef-oc20-2m":{"fetched_at":"2026-07-11T05:34:20.810222+00:00","verified_at":"2026-07-27T22:35:34.720037+00:00","verified_device":"cuda","last_error":null}}},"escn":{"status":"ready","built_at":"2026-07-11T05:31:54.698039+00:00","source_hash":"sha256:3ab4edb97bb0536fe61f79e84262eebe56b25d83aab9e073663de1a3aadaecec","source":"# /// script\n# requires-python = \">=3.10,<3.11\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=1.0.0,<2.0.0\",\n# \"ase>=3.22\",\n# \"torch-geometric\",\n# \"torch-scatter\",\n# \"torch-sparse\",\n# \"torch-cluster\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"\neSCN environment for Rootstock.\n\nUses fairchem-core 1.x to access legacy OC20 eSCN checkpoints via\nOCPCalculator. These checkpoints are optimized for catalysis systems\n(slabs + adsorbates).\n\nModels:\n - \"eSCN-L6-M2-Lay12-S2EF-OC20-All+MD\": default\n - \"eSCN-L6-M3-Lay20-S2EF-OC20-All+MD\"\n - \"eSCN-L6-M2-Lay12-S2EF-OC20-2M\"\n - \"eSCN-L4-M2-Lay12-S2EF-OC20-2M\"\n\"\"\"\n\nCHECKPOINTS = {\n \"escn-l6-m2-lay12-s2ef-oc20-all-md\": \"eSCN-L6-M2-Lay12-S2EF-OC20-All+MD\",\n \"escn-l6-m3-lay20-s2ef-oc20-all-md\": \"eSCN-L6-M3-Lay20-S2EF-OC20-All+MD\",\n \"escn-l6-m2-lay12-s2ef-oc20-2m\": \"eSCN-L6-M2-Lay12-S2EF-OC20-2M\",\n \"escn-l4-m2-lay12-s2ef-oc20-2m\": \"eSCN-L4-M2-Lay12-S2EF-OC20-2M\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load an eSCN OC20 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible OCPCalculator.\n \"\"\"\n import os\n from fairchem.core import OCPCalculator\n from fairchem.core.models.model_registry import model_name_to_local_file\n\n cache_dir = os.environ.get(\"XDG_CACHE_HOME\", os.path.expanduser(\"~/.cache\"))\n local_path = model_name_to_local_file(CHECKPOINTS[checkpoint], local_cache=cache_dir)\n return OCPCalculator(checkpoint_path=local_path, cpu=(device == \"cpu\"))\n","python_requires":">=3.10,<3.11","dependencies":{"ase":"3.29.0","fairchem-core":"1.10.0","rootstock":"0.9.5","torch":"2.4.1","torch-cluster":"1.6.3+pt24cu121","torch-geometric":"2.8.0","torch-scatter":"2.1.2+pt24cu121","torch-sparse":"0.6.18+pt24cu121"},"checkpoints":{"escn-l6-m2-lay12-s2ef-oc20-all-md":{"fetched_at":"2026-07-11T05:35:13.166232+00:00","verified_at":"2026-07-27T22:35:50.609733+00:00","verified_device":"cuda","last_error":null},"escn-l6-m3-lay20-s2ef-oc20-all-md":{"fetched_at":"2026-07-11T05:36:05.078389+00:00","verified_at":"2026-07-27T22:36:06.392224+00:00","verified_device":"cuda","last_error":null},"escn-l6-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-07-11T05:36:39.770642+00:00","verified_at":"2026-07-27T22:36:18.135903+00:00","verified_device":"cuda","last_error":null},"escn-l4-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-07-11T05:37:03.949516+00:00","verified_at":"2026-07-27T22:36:29.085140+00:00","verified_device":"cuda","last_error":null}}},"esen":{"status":"ready","built_at":"2026-07-11T05:31:54.877600+00:00","source_hash":"sha256:5791286b7d2ede75fa6e96e4e3d83ebecfe56f427a86680329b5010703f4f492","source":"# /// script\n# requires-python = \">=3.10,<3.11\"\n# dependencies = [\n# \"torch>=2.4.0\",\n# \"fairchem-core>=2.0.0\",\n# \"ase>=3.22\",\n# \"torch-geometric\",\n# ]\n#\n# [tool.uv]\n# find-links = [\"https://data.pyg.org/whl/torch-2.4.0+cu121.html\"]\n# ///\n\"\"\"eSEN env — hosts FAIRChem eSEN single-task checkpoints.\n\nOMol checkpoints expect `charge` and `spin` in `atoms.info`.\n\"\"\"\n\nCHECKPOINTS = {\n \"esen-md-direct-all-omol\": \"esen-md-direct-all-omol\",\n \"esen-sm-conserving-all-omol\": \"esen-sm-conserving-all-omol\",\n \"esen-sm-direct-all-omol\": \"esen-sm-direct-all-omol\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(CHECKPOINTS[checkpoint], device=device)\n return FAIRChemCalculator(predictor)\n","python_requires":">=3.10,<3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.14.0","rootstock":"0.9.5","torch":"2.8.0","torch-geometric":"2.8.0"},"checkpoints":{"esen-md-direct-all-omol":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"download: Access to model facebook/OMol25 is restricted. You must have access to it and be authenticated to access it. Please log in."},"esen-sm-conserving-all-omol":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"download: Access to model facebook/OMol25 is restricted. You must have access to it and be authenticated to access it. Please log in."},"esen-sm-direct-all-omol":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"download: Access to model facebook/OMol25 is restricted. You must have access to it and be authenticated to access it. Please log in."}}},"mace":{"status":"ready","built_at":"2026-07-11T05:31:55.025303+00:00","source_hash":"sha256:179de9d31a8c2763021ce581ca3e39d854eaea1caeea4dd4e80e1394d4e76c17","source":"# /// script\n# requires-python = \">=3.10\"\n# dependencies = [\n# \"mace-torch>=0.3.0\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0,<2.10\",\n# ]\n# ///\n\"\"\"MACE env — hosts MACE-MP-0 and MACE-OFF23 checkpoints.\n\nBoth ship in the same `mace-torch` package, so they share an environment.\nThe `off:` prefix on the upstream string in CHECKPOINTS routes to mace_off()\ninstead of mace_mp().\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-mp-0-small\": \"small\",\n \"mace-mp-0-medium\": \"medium\",\n \"mace-mp-0-large\": \"large\",\n \"mace-off23-small\": \"off:small\",\n \"mace-off23-medium\": \"off:medium\",\n \"mace-off23-large\": \"off:large\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n arg = CHECKPOINTS[checkpoint]\n if arg.startswith(\"off:\"):\n from mace.calculators import mace_off\n\n return mace_off(model=arg[4:], device=device, default_dtype=\"float32\")\n from mace.calculators import mace_mp\n\n return mace_mp(model=arg, device=device, default_dtype=\"float32\")\n","python_requires":">=3.10","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"0.9.5","torch":"2.9.1"},"checkpoints":{"mace-mp-0-small":{"fetched_at":"2026-07-11T05:38:04.014483+00:00","verified_at":"2026-07-27T22:36:54.322994+00:00","verified_device":"cuda","last_error":null},"mace-mp-0-medium":{"fetched_at":"2026-07-11T05:38:21.195651+00:00","verified_at":"2026-07-27T22:37:01.557409+00:00","verified_device":"cuda","last_error":null},"mace-mp-0-large":{"fetched_at":"2026-07-11T05:39:14.473851+00:00","verified_at":"2026-07-27T22:37:08.932359+00:00","verified_device":"cuda","last_error":null},"mace-off23-small":{"fetched_at":"2026-07-11T05:39:28.975081+00:00","verified_at":"2026-07-27T22:37:15.872033+00:00","verified_device":"cuda","last_error":null},"mace-off23-medium":{"fetched_at":"2026-07-11T05:39:43.433571+00:00","verified_at":"2026-07-27T22:37:22.646981+00:00","verified_device":"cuda","last_error":null},"mace-off23-large":{"fetched_at":"2026-07-11T05:39:58.220376+00:00","verified_at":"2026-07-27T22:37:28.954859+00:00","verified_device":"cuda","last_error":null}}},"mace_off23":{"status":"ready","built_at":"2026-07-11T05:31:55.174555+00:00","source_hash":"sha256:37d632e7045d40c8b9b7461e6c8e024c6c9f47699152adbf311ad05b58f11240","source":"# /// script\n# requires-python = \">=3.10\"\n# dependencies = [\n# \"mace-torch>=0.3.0\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0,<2.10\",\n# ]\n# ///\n\"\"\"\nMACE-OFF23 environment for Rootstock.\n\nMACE-OFF23 is a transferable force field for organic molecules, distinct from\nMACE-MP-0 (which targets inorganic materials). Use this env for molecular\ndynamics and geometry optimisation of drug-like and organic systems.\n\nModels:\n - \"small\": ~4M params, fastest\n - \"medium\": ~10M params, balanced (default)\n - \"large\": ~28M params, most accurate\n\"\"\"\n\nCHECKPOINTS = {\n \"mace-off23-small\": \"small\",\n \"mace-off23-medium\": \"medium\",\n \"mace-off23-large\": \"large\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\"):\n \"\"\"\n Load a MACE-OFF23 calculator.\n\n Args:\n checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS.\n device: PyTorch device string (e.g., \"cuda\", \"cpu\").\n\n Returns:\n ASE-compatible calculator.\n \"\"\"\n from mace.calculators import mace_off\n\n return mace_off(model=CHECKPOINTS[checkpoint], device=device, default_dtype=\"float32\")\n","python_requires":">=3.10","dependencies":{"ase":"3.29.0","mace-torch":"0.3.16","rootstock":"0.9.5","torch":"2.9.1"},"checkpoints":{}},"orb_v3":{"status":"ready","built_at":"2026-07-11T05:31:55.324768+00:00","source_hash":"sha256:305f3d2e8183d7fb0912e6d4bf3a737c0f42e383cf12e4fec4289a67c00be837","source":"# /// script\n# requires-python = \">=3.12\"\n# dependencies = [\n# \"orb-models>=0.6.2\",\n# \"ase>=3.25\",\n# \"torch>=2.8\",\n# # Not imported here — constrains orb-models' transitive dep. setup()'s\n# # no-lock serve path relies on cached_path returning local files without\n# # locking or writing, verified against exactly this version (#67).\n# \"cached_path==1.8.10\",\n# ]\n#\n# [tool.uv.sources]\n# torch = { index = \"pytorch-cu128\" }\n#\n# [[tool.uv.index]]\n# name = \"pytorch-cu128\"\n# url = \"https://download.pytorch.org/whl/cu128\"\n# explicit = true\n# ///\n\"\"\"Orb v3 env — Orbital Materials' Orb v3 universal potentials.\n\nSeparate from orb.py because orb-models>=0.5 changed the loader API\n(returns a tuple, requires `atoms_adapter` on ORBCalculator, moved calculator\nimport path) and 0.6.x bumped the Python floor to 3.12 and torch to 2.8.\n\"\"\"\n\nimport os\nimport shutil\nimport urllib.request\nfrom pathlib import Path\n\nCHECKPOINTS = {\n \"orb-v3-conservative-inf-omat\": \"orb-v3-conservative-inf-omat\",\n \"orb-v3-conservative-20-omat\": \"orb-v3-conservative-20-omat\",\n \"orb-v3-direct-inf-omat\": \"orb-v3-direct-inf-omat\",\n \"orb-v3-direct-20-omat\": \"orb-v3-direct-20-omat\",\n \"orb-v3-conservative-inf-mpa\": \"orb-v3-conservative-inf-mpa\",\n \"orb-v3-conservative-20-mpa\": \"orb-v3-conservative-20-mpa\",\n \"orb-v3-direct-inf-mpa\": \"orb-v3-direct-inf-mpa\",\n \"orb-v3-direct-20-mpa\": \"orb-v3-direct-20-mpa\",\n \"orb-v3-conservative-omol\": \"orb-v3-conservative-omol\",\n \"orb-v3-direct-omol\": \"orb-v3-direct-omol\",\n}\n\n\ndef _default_weights_url(load_fn) -> str:\n \"\"\"The upstream URL baked into the loader's ``weights_path`` default.\"\"\"\n import inspect\n\n default = inspect.signature(load_fn).parameters[\"weights_path\"].default\n if not isinstance(default, str) or not default.startswith((\"http://\", \"https://\")):\n raise RuntimeError(\n f\"{load_fn.__name__} has no URL default for weights_path \"\n f\"(got {default!r}); update this env file for the installed orb-models\"\n )\n return default\n\n\ndef _local_weights_path(url: str) -> Path:\n \"\"\"Where the checkpoint lives in the shared model cache.\"\"\"\n cache = Path(os.environ.get(\"XDG_CACHE_HOME\") or Path.home() / \".cache\")\n return cache / \"orb\" / os.path.basename(url)\n\n\ndef _fetch(url: str, dest: Path) -> None:\n \"\"\"Download ``url`` to ``dest`` atomically (tmp file + rename).\"\"\"\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp = dest.with_name(f\"{dest.name}.tmp.{os.getpid()}\")\n try:\n with urllib.request.urlopen(url) as resp, open(tmp, \"wb\") as out:\n shutil.copyfileobj(resp, out)\n os.replace(tmp, dest)\n finally:\n tmp.unlink(missing_ok=True)\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", precision: str = \"float32-high\"):\n import torch\n from orb_models.forcefield import pretrained\n from orb_models.forcefield.inference.calculator import ORBCalculator\n\n fn_name = CHECKPOINTS[checkpoint].replace(\"-\", \"_\")\n load_fn = getattr(pretrained, fn_name)\n\n # orb-models resolves its default weights URL through `cached_path`, which\n # write-locks its cache dir even on warm hits — EACCES for anyone who can't\n # write the shared install (Garden-AI/rootstock#67). Handed a *local* path\n # instead, cached_path returns it without locking. So the weights are\n # pre-fetched into the shared model cache at `rootstock add` time\n # (maintainer, cache writable) and every later serve loads that file.\n url = _default_weights_url(load_fn)\n weights = _local_weights_path(url)\n if not weights.exists():\n _fetch(url, weights)\n\n orbff, atoms_adapter = load_fn(\n weights_path=str(weights), device=torch.device(device), precision=precision\n )\n return ORBCalculator(orbff, atoms_adapter=atoms_adapter, device=torch.device(device))\n","python_requires":">=3.12","dependencies":{"ase":"3.29.0","cached-path":"1.8.10","orb-models":"0.7.0","rootstock":"0.9.5","torch":"2.11.0+cu128"},"checkpoints":{"orb-v3-conservative-inf-omat":{"fetched_at":"2026-07-11T05:40:59.847916+00:00","verified_at":"2026-07-27T22:38:50.668616+00:00","verified_device":"cuda","last_error":null},"orb-v3-conservative-20-omat":{"fetched_at":"2026-07-11T05:42:12.601799+00:00","verified_at":"2026-07-27T22:39:12.262785+00:00","verified_device":"cuda","last_error":null},"orb-v3-direct-inf-omat":{"fetched_at":"2026-07-11T05:42:49.249049+00:00","verified_at":"2026-07-27T22:39:22.259054+00:00","verified_device":"cuda","last_error":null},"orb-v3-direct-20-omat":{"fetched_at":"2026-07-11T05:43:13.274625+00:00","verified_at":"2026-07-27T22:39:31.529740+00:00","verified_device":"cuda","last_error":null},"orb-v3-conservative-inf-mpa":{"fetched_at":"2026-07-11T05:43:36.040425+00:00","verified_at":"2026-07-27T22:39:52.500496+00:00","verified_device":"cuda","last_error":null},"orb-v3-conservative-20-mpa":{"fetched_at":"2026-07-11T05:44:11.758182+00:00","verified_at":"2026-07-27T22:40:13.437050+00:00","verified_device":"cuda","last_error":null},"orb-v3-direct-inf-mpa":{"fetched_at":"2026-07-11T05:44:49.463452+00:00","verified_at":"2026-07-27T22:40:22.190201+00:00","verified_device":"cuda","last_error":null},"orb-v3-direct-20-mpa":{"fetched_at":"2026-07-11T05:45:12.852572+00:00","verified_at":"2026-07-27T22:40:30.961710+00:00","verified_device":"cuda","last_error":null},"orb-v3-conservative-omol":{"fetched_at":"2026-07-11T05:45:35.440591+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: ConnectionResetError: [Errno 104] Connection reset by peer"},"orb-v3-direct-omol":{"fetched_at":"2026-07-11T05:45:56.063427+00:00","verified_at":null,"verified_device":null,"last_error":"smoke-test: ConnectionResetError: [Errno 104] Connection reset by peer"}}},"uma":{"status":"ready","built_at":"2026-07-11T05:31:55.492450+00:00","source_hash":"sha256:90fa56797dbc112f801242475202c036de053f8cc9b27d7d6aeb2410a37b8149","source":"# /// script\n# requires-python = \">=3.11\"\n# dependencies = [\n# \"fairchem-core>=2.20\",\n# \"ase>=3.22\",\n# \"torch>=2.4.0\",\n# ]\n# ///\n\"\"\"UMA env — hosts Meta's UMA foundation model via FAIRChem.\n\nfairchem-core v2 dropped the torch-geometric / pyg-find-links install dance, so\nthis env is a plain PyPI install. The original uma-s-1 had an extensivity bug\nand was removed from the fairchem 2.20 registry — use uma-s-1p1 or uma-s-1p2.\n\"\"\"\n\nCHECKPOINTS = {\n \"uma-s-1p1\": \"uma-s-1p1\",\n \"uma-s-1p2\": \"uma-s-1p2\",\n \"uma-m-1p1\": \"uma-m-1p1\",\n}\n\n\ndef setup(checkpoint: str, device: str = \"cuda\", task: str = \"omat\"):\n from fairchem.core import FAIRChemCalculator, pretrained_mlip\n\n predictor = pretrained_mlip.get_predict_unit(CHECKPOINTS[checkpoint], device=device)\n return FAIRChemCalculator(predictor, task_name=task)\n","python_requires":">=3.11","dependencies":{"ase":"3.29.0","fairchem-core":"2.21.0","rootstock":"0.9.5","torch":"2.8.0"},"checkpoints":{"uma-s-1p1":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"download: Access to model facebook/UMA is restricted. You must have access to it and be authenticated to access it. Please log in."},"uma-s-1p2":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"download: Access to model facebook/UMA is restricted. You must have access to it and be authenticated to access it. Please log in."},"uma-m-1p1":{"fetched_at":null,"verified_at":null,"verified_device":null,"last_error":"download: Access to model facebook/UMA is restricted. You must have access to it and be authenticated to access it. Please log in."}}}}}]}