{"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-09-04T21:16:34.821419+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":"2026-07-18T03:53:42.849047+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"ani-1ccx":{"fetched_at":"2026-07-18T01:07:52.739750+00:00","verified_at":"2026-07-18T03:53:42.849056+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"ani-1x":{"fetched_at":"2026-07-18T01:08:01.222324+00:00","verified_at":"2026-07-18T03:53:42.849060+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"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":"2026-07-31T00:24:01.894144+00:00","verified_device":"cuda","last_error":null,"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":"2026-07-20T20:07:22.716394+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"dimenet-plus-plus-s2ef-oc20-20m":{"fetched_at":"2026-07-20T20:08:11.040759+00:00","verified_at":"2026-07-20T20:08:31.773415+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"dimenet-plus-plus-s2ef-oc20-2m":{"fetched_at":"2026-07-20T20:08:53.776120+00:00","verified_at":"2026-07-20T20:09:14.588379+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"dimenet-plus-plus-s2ef-oc20-200k":{"fetched_at":"2026-07-20T20:09:37.960754+00:00","verified_at":"2026-07-20T20:09:58.444760+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"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":"2026-07-31T18:17:47.920527+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"equiformer-v2-31m-s2ef-oc20-all-md":{"fetched_at":"2026-07-20T20:21:06.384332+00:00","verified_at":"2026-07-31T18:18:36.117169+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"equiformer-v2-83m-s2ef-oc20-2m":{"fetched_at":"2026-07-20T20:21:33.323971+00:00","verified_at":"2026-07-31T00:45:31.892606+00:00","verified_device":"cuda","last_error":null,"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-07-31T00:54:25.936111+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"escn-l6-m3-lay20-s2ef-oc20-all-md":{"fetched_at":"2026-07-20T20:34:01.649114+00:00","verified_at":"2026-07-31T00:54:55.498297+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"escn-l6-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-07-20T20:34:48.217991+00:00","verified_at":"2026-07-31T00:54:00.833674+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"escn-l4-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-07-20T20:35:22.280265+00:00","verified_at":"2026-07-31T00:53:35.130422+00:00","verified_device":"cuda","last_error":null,"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":"2026-07-31T18:27:55.026157+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"esen-sm-conserving-all-omol":{"fetched_at":"2026-07-18T01:54:34.092323+00:00","verified_at":"2026-07-31T01:10:13.868629+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"esen-sm-direct-all-omol":{"fetched_at":"2026-07-18T01:55:21.266324+00:00","verified_at":"2026-07-31T01:10:37.576783+00:00","verified_device":"cuda","last_error":null,"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-07-18T03:53:42.849071+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-08-03T19:07:30.354807+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-mp-0-medium":{"fetched_at":"2026-07-18T02:10:38.115353+00:00","verified_at":"2026-08-03T19:07:15.593289+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-mp-0-large":{"fetched_at":"2026-07-18T02:10:53.149921+00:00","verified_at":"2026-08-03T19:07:00.555256+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-off23-small":{"fetched_at":"2026-07-18T02:11:07.975353+00:00","verified_at":"2026-08-03T19:08:30.073186+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-off23-medium":{"fetched_at":"2026-07-18T02:11:22.830181+00:00","verified_at":"2026-08-03T19:08:15.557545+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-off23-large":{"fetched_at":"2026-07-18T02:11:37.728101+00:00","verified_at":"2026-08-03T19:08:00.774512+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"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-08-03T19:06:45.296441+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-mpa-0-medium":{"fetched_at":"2026-07-30T23:57:26.498782+00:00","verified_at":"2026-08-03T19:07:45.626030+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"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-08-04T01:29:42.152619+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEmatpesr2scanomatftmodel","size":79470738}],"weights_recorded_at":"2026-08-04T01:29:42.152637+00:00"},"mace-omol-0-extra-large":{"fetched_at":"2026-08-03T18:03:27.581411+00:00","verified_at":"2026-08-03T19:08:46.605348+00:00","verified_device":"cuda","last_error":null,"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-07-31T19:00:10.703190+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"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":"2026-07-31T19:12:57.608753+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-conservative-20-omat":{"fetched_at":"2026-07-18T02:33:01.402092+00:00","verified_at":"2026-07-31T19:11:57.708203+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-direct-inf-omat":{"fetched_at":"2026-07-18T02:33:48.238058+00:00","verified_at":"2026-07-31T19:14:12.022093+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-direct-20-omat":{"fetched_at":"2026-07-18T02:34:23.077709+00:00","verified_at":"2026-07-31T19:13:34.998446+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-conservative-inf-mpa":{"fetched_at":"2026-07-18T02:34:57.953047+00:00","verified_at":"2026-07-31T19:12:27.830558+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-conservative-20-mpa":{"fetched_at":"2026-07-18T02:35:37.786464+00:00","verified_at":null,"verified_device":null,"last_error":"verify: 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-v3-direct-inf-mpa":{"fetched_at":"2026-07-18T02:36:11.635141+00:00","verified_at":"2026-07-31T19:13:53.656342+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"orb-v3-direct-20-mpa":{"fetched_at":"2026-07-18T02:36:34.279807+00:00","verified_at":"2026-07-31T19:13:16.466665+00:00","verified_device":"cuda","last_error":null,"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":"2026-07-20T20:44:30.598484+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"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":"2026-07-20T20:53:10.303886+00:00","verified_device":"cuda","last_error":null,"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-07-20T20:53:36.958316+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"scn-s2ef-oc20-2m":{"fetched_at":"2026-07-20T20:53:50.713253+00:00","verified_at":"2026-07-20T20:54:01.359550+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null}}},"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":"2026-07-31T19:23:38.769167+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"sevennet-l3i5":{"fetched_at":"2026-07-18T02:57:57.379501+00:00","verified_at":"2026-07-31T02:18:50.058385+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-07-31T02:20:19.776465+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-07-31T02:19:29.271957+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-07-31T02:20:39.416257+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":"2026-07-31T19:37:30.562614+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"uma-s-1p2":{"fetched_at":"2026-07-18T03:07:41.520380+00:00","verified_at":"2026-07-31T19:37:54.677022+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"uma-m-1p1":{"fetched_at":"2026-07-18T03:08:53.351374+00:00","verified_at":"2026-08-04T03:43:06.899686+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-08-04T03:43:06.899702+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}}},"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":"verify: 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-md-direct-all-omol":{"fetched_at":"2026-08-03T18:03:18.239815+00:00","verified_at":"2026-08-03T18:44:24.490143+00:00","verified_device":"cuda","last_error":null,"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":"2026-08-03T19:13:46.554287+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-polar-1-m":{"fetched_at":"2026-07-31T00:02:47.844007+00:00","verified_at":"2026-08-03T19:13:31.298741+00:00","verified_device":"cuda","last_error":null,"weight_files":null,"weights_recorded_at":null},"mace-polar-1-l":{"fetched_at":"2026-07-31T00:02:48.613868+00:00","verified_at":"2026-08-03T19:13:15.661436+00:00","verified_device":"cuda","last_error":null,"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":"verify: 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":"2026-07-31T01:43:56.123201+00:00","verified_device":"cuda","last_error":null,"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":"2026-08-04T03:27:02.056416+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-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":"verify: 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":"2026-08-04T01:04:42.096520+00:00","verified_device":"cuda","last_error":null,"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-08-05T00:14:50.398802+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-08-05T00:14:50.398821+00:00"},"pet-oam-xl":{"fetched_at":"2026-08-03T18:04:29.291262+00:00","verified_at":null,"verified_device":null,"last_error":"verify: 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}}},"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":"2026-08-05T00:26:28.197752+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"home/.cache/tace/TECE-OAM-RRA-1.0.pt","size":890523407}],"weights_recorded_at":"2026-08-05T00:26:28.197767+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":"verify: RuntimeError: Worker calculation failed:\nTraceback (most recent call last):\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/rootstock/worker.py\", line 275, in run\n energy, forces, virial = self._calculate(positions, cell)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/rootstock/worker.py\", line 182, in _calculate\n energy = atoms.get_potential_energy()\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/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 \"/work/hdd/data/rootstock/envs/deepmd/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 \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/ase/calculators/calculator.py\", line 517, in get_property\n self.calculate(atoms, [name], system_changes)\n File \"/work/hdd/data/rootstock/envs/deepmd/env_source.py\", line 240, in calculate\n super().calculate(atoms, list(properties), system_changes)\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/calculator.py\", line 143, in calculate\n e, f, v = self.dp.eval(\n ^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/infer/deep_pot.py\", line 206, in eval\n results = self.deep_eval.eval(\n ^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/pt/infer/deep_eval.py\", line 556, in eval\n out = self._eval_func(self._eval_model, numb_test, natoms)(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/pt/infer/deep_eval.py\", line 637, in eval_func\n return self.auto_batch_size.execute_all(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/utils/batch_size.py\", line 228, in execute_all\n n_batch, result = self.execute(execute_with_batch_size, index, natoms)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/utils/batch_size.py\", line 129, in execute\n raise e\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/utils/batch_size.py\", line 126, in execute\n n_batch, result = callable(max(batch_nframes, 1), start_index)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/utils/batch_size.py\", line 199, in execute_with_batch_size\n return (end_index - start_index), callable(\n ^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/pt/infer/deep_eval.py\", line 713, in _eval_model\n charge_spin_input = to_torch_tensor(charge_spin.reshape(nframes, 2))\n ^^^^^^^^^^^^^^^^^^^\nAttributeError: 'list' object has no attribute 'reshape'\n","weight_files":[{"path":"cache/deepmd/pretrained/models/DPA-3.3-1M.pt","size":41612295}],"weights_recorded_at":"2026-09-04T21:13:48.725421+00:00"},"dpa-3.1-3m":{"fetched_at":"2026-09-04T21:10:30.947721+00:00","verified_at":"2026-09-04T21:12:28.107017+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/deepmd/pretrained/models/DPA-3.1-3M.pt","size":47176032}],"weights_recorded_at":"2026-09-04T21:12:28.107034+00:00"},"dpa-3.2-5m":{"fetched_at":"2026-09-04T21:10:32.190658+00:00","verified_at":"2026-09-04T21:13:03.241982+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/deepmd/pretrained/models/DPA-3.2-5M.pt","size":64800782}],"weights_recorded_at":"2026-09-04T21:13:03.241996+00:00"},"dpa4-air-omat24":{"fetched_at":"2026-09-04T21:11:03.916056+00:00","verified_at":"2026-09-04T21:14:29.707694+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/deepmd/pretrained/models/DPA4-Air-OMat24-v20260805.pt","size":20713035}],"weights_recorded_at":"2026-09-04T21:14:29.707709+00:00"},"dpa3-omol-large":{"fetched_at":"2026-09-04T21:11:02.892992+00:00","verified_at":null,"verified_device":null,"last_error":"verify: RuntimeError: Worker calculation failed:\nTraceback (most recent call last):\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/rootstock/worker.py\", line 275, in run\n energy, forces, virial = self._calculate(positions, cell)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/rootstock/worker.py\", line 182, in _calculate\n energy = atoms.get_potential_energy()\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/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 \"/work/hdd/data/rootstock/envs/deepmd/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 \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/ase/calculators/calculator.py\", line 517, in get_property\n self.calculate(atoms, [name], system_changes)\n File \"/work/hdd/data/rootstock/envs/deepmd/env_source.py\", line 240, in calculate\n super().calculate(atoms, list(properties), system_changes)\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/calculator.py\", line 143, in calculate\n e, f, v = self.dp.eval(\n ^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/infer/deep_pot.py\", line 206, in eval\n results = self.deep_eval.eval(\n ^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/pt/infer/deep_eval.py\", line 556, in eval\n out = self._eval_func(self._eval_model, numb_test, natoms)(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/pt/infer/deep_eval.py\", line 637, in eval_func\n return self.auto_batch_size.execute_all(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/utils/batch_size.py\", line 228, in execute_all\n n_batch, result = self.execute(execute_with_batch_size, index, natoms)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/utils/batch_size.py\", line 129, in execute\n raise e\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/utils/batch_size.py\", line 126, in execute\n n_batch, result = callable(max(batch_nframes, 1), start_index)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/utils/batch_size.py\", line 199, in execute_with_batch_size\n return (end_index - start_index), callable(\n ^^^^^^^^^\n File \"/work/hdd/data/rootstock/envs/deepmd/lib/python3.11/site-packages/deepmd/pt/infer/deep_eval.py\", line 713, in _eval_model\n charge_spin_input = to_torch_tensor(charge_spin.reshape(nframes, 2))\n ^^^^^^^^^^^^^^^^^^^\nAttributeError: 'list' object has no attribute 'reshape'\n","weight_files":[{"path":"cache/deepmd/pretrained/models/DPA3-Omol-Large.pt","size":224057430}],"weights_recorded_at":"2026-09-04T21:14:07.763731+00:00"},"dpa4-mini-omat24":{"fetched_at":"2026-09-04T21:11:00.583238+00:00","verified_at":"2026-09-04T21:14:49.532955+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/deepmd/pretrained/models/DPA4-Mini-OMat24-v20260805.pt","size":2701347}],"weights_recorded_at":"2026-09-04T21:14:49.532969+00:00"},"dpa4-nano-omat24":{"fetched_at":"2026-09-04T21:11:00.030849+00:00","verified_at":"2026-09-04T21:15:09.084311+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/deepmd/pretrained/models/DPA4-Nano-OMat24-v20260805.pt","size":1997147}],"weights_recorded_at":"2026-09-04T21:15:09.084325+00:00"},"dpa4-neo-omat24":{"fetched_at":"2026-09-04T21:11:19.407213+00:00","verified_at":"2026-09-04T21:15:29.376559+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/deepmd/pretrained/models/DPA4-Neo-OMat24-v20260805.pt","size":4583819}],"weights_recorded_at":"2026-09-04T21:15:29.376574+00:00"},"dpa4-plus-omat24":{"fetched_at":"2026-09-04T21:11:20.501886+00:00","verified_at":"2026-09-04T21:15:49.095332+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/deepmd/pretrained/models/DPA4-Plus-OMat24-v20260805.pt","size":35550605}],"weights_recorded_at":"2026-09-04T21:15:49.095346+00:00"},"dpa-2.4-7m":{"fetched_at":"2026-09-04T21:10:30.286150+00:00","verified_at":"2026-09-04T21:11:45.005029+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/deepmd/pretrained/models/dpa-2.4-7M.pt","size":169224765}],"weights_recorded_at":"2026-09-04T21:11:45.005043+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.2","python_version":"3.10","last_updated":"2026-08-24T06:43:39.192433+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-08-24T06:43:35.647838+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231203mace128L1_epoch199model","size":44422970}],"weights_recorded_at":"2026-08-24T06:43:35.647844+00:00"},"mace-mp-0-small":{"fetched_at":"2026-07-20T19:56:00.318387+00:00","verified_at":"2026-08-24T06:43:35.647814+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231210mace128L0_energy_epoch249model","size":32581838}],"weights_recorded_at":"2026-08-24T06:43:35.647828+00:00"},"mace-mp-0-large":{"fetched_at":"2026-07-20T19:59:39.486721+00:00","verified_at":"2026-08-24T06:43:35.647849+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE_MPtrj_20229model","size":133803220}],"weights_recorded_at":"2026-08-24T06:43:35.647854+00:00"},"mace-off23-small":{"fetched_at":"2026-07-20T20:01:57.226028+00:00","verified_at":"2026-08-24T06:43:35.647858+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_small.model","size":7347350}],"weights_recorded_at":"2026-08-24T06:43:35.647862+00:00"},"mace-off23-medium":{"fetched_at":"2026-07-20T20:04:28.688189+00:00","verified_at":"2026-08-24T06:43:35.647866+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_medium.model","size":18350596}],"weights_recorded_at":"2026-08-24T06:43:35.647870+00:00"},"mace-off23-large":{"fetched_at":"2026-07-20T20:07:12.519950+00:00","verified_at":"2026-08-24T06:43:35.647874+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_large.model","size":55492786}],"weights_recorded_at":"2026-08-24T06:43:35.647877+00:00"},"mace-matpes-r2scan-0":{"fetched_at":"2026-08-03T17:54:44.773719+00:00","verified_at":"2026-08-24T06:43:35.647887+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEmatpesr2scanomatftmodel","size":79470738}],"weights_recorded_at":"2026-08-24T06:43:35.647893+00:00"},"mace-mh-1":{"fetched_at":"2026-08-03T17:54:45.010039+00:00","verified_at":"2026-08-24T06:43:35.647897+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macemh1model","size":59208139}],"weights_recorded_at":"2026-08-24T06:43:35.647901+00:00"},"mace-mpa-0-medium":{"fetched_at":"2026-08-03T17:55:21.444729+00:00","verified_at":"2026-08-24T06:43:35.647881+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macempa0mediummodel","size":79462305}],"weights_recorded_at":"2026-08-24T06:43:35.647884+00:00"},"mace-omol-0-extra-large":{"fetched_at":"2026-08-03T17:55:23.154847+00:00","verified_at":"2026-08-24T06:43:35.647904+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-08-24T06:43:35.647909+00:00"},"mace:custom":{"fetched_at":null,"verified_at":"2026-08-24T06:43:35.648016+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-08-24T06:43:35.647912+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-08-24T06:43:35.647916+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-08-24T06:43:35.647920+00:00"},"mattersim-v1-0-0-1m":{"fetched_at":"2026-07-20T20:13:50.936758+00:00","verified_at":"2026-08-24T06:43:35.647923+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-08-24T06:43:35.647926+00:00"},"mattersim:custom":{"fetched_at":null,"verified_at":"2026-08-24T06:43:35.648019+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-08-24T06:43:35.647929+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-d3-v2-20241011.ckpt","size":100967743}],"weights_recorded_at":"2026-08-24T06:43:35.647933+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-08-24T06:43:35.647936+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-08-24T06:43:35.647942+00:00"},"uma-s-1p2":{"fetched_at":"2026-07-20T20:30:39.154652+00:00","verified_at":"2026-08-24T06:43:35.647949+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-08-24T06:43:35.647955+00:00"},"uma-m-1p1":{"fetched_at":"2026-07-20T20:34:07.816107+00:00","verified_at":"2026-08-24T06:43:35.647958+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-08-24T06:43:35.647963+00:00"},"uma:custom":{"fetched_at":null,"verified_at":"2026-08-24T06:43:35.648022+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-08-24T06:43:35.647974+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-08-24T06:43:35.647979+00:00"},"allscaip-md-conserving-all-omol":{"fetched_at":"2026-08-03T17:51:44.672552+00:00","verified_at":"2026-08-24T06:43:35.647966+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-08-24T06:43:35.647971+00:00"},"allscaip:custom":{"fetched_at":null,"verified_at":"2026-08-24T06:43:35.648025+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-08-24T06:43:35.647989+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Mmodel","size":68133235}],"weights_recorded_at":"2026-08-24T06:43:35.647993+00:00"},"mace-polar-1-l":{"fetched_at":"2026-08-03T17:58:22.706652+00:00","verified_at":"2026-08-24T06:43:35.647996+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Lmodel","size":129808715}],"weights_recorded_at":"2026-08-24T06:43:35.647999+00:00"},"mace-polar-1-s":{"fetched_at":"2026-08-03T17:58:58.732967+00:00","verified_at":"2026-08-24T06:43:35.647982+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Smodel","size":33375439}],"weights_recorded_at":"2026-08-24T06:43:35.647986+00:00"},"mace-polar:custom":{"fetched_at":null,"verified_at":"2026-08-24T06:43:35.648027+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:5f080a65f9f830c598706724af721348c0c51890ccb524b9859e12abe54f4d7e","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","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-08-24T06:43:35.648012+00:00"},"orbmol-v1-conservative":{"fetched_at":"2026-08-12T20:24:41.375905+00:00","verified_at":"2026-08-24T06:43:35.648002+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-08-24T06:43:35.648006+00:00"},"orbmol:custom":{"fetched_at":null,"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: 46553 files, 16542 MB in 22.3s\nWarp CUDA error: Failed to get function cuGetProcAddress\n/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/orb_models/common/training/util.py:20: UserWarning: Setting global torch default dtype to torch.float32.\n warnings.warn(f\"Setting global torch default dtype to {torch_dtype}.\")\nTraceback (most recent call last):\n File \"/tmp/rootstock_spawn_lm50d84x/wrapper.py\", line 46, in \n run_worker(\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/rootstock/worker.py\", line 378, in run_worker\n calculator = setup_fn(checkpoint, device, **setup_kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/env_source.py\", line 178, in setup_from_path\n orbff, atoms_adapter = load_fn(\n ^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/orb_models/forcefield/pretrained.py\", line 447, in orbmol_v2\n model = load_model(\n ^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/orb_models/forcefield/pretrained.py\", line 115, in load_model\n model.load_state_dict(state_dict, strict=True)\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/orb_models/forcefield/models/conservative_regressor.py\", line 412, in load_state_dict\n load_regressor_state_dict(\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/orb_models/common/models/load.py\", line 87, in load_regressor_state_dict\n bad_keys = super(type(model), model).load_state_dict(state_dict, strict=False, assign=assign)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/sw/frontier/ums/ums047/rootstock/envs/orbmol/lib/python3.12/site-packages/torch/nn/modules/module.py\", line 2629, in load_state_dict\n raise RuntimeError(\nRuntimeError: Error(s) in loading state_dict for ConservativeForcefieldRegressor:\n\tsize mismatch for heads.energy.mlp.NN-0.weight: copying a param with shape torch.Size([256, 256]) from checkpoint, the shape in current model is torch.Size([256, 257]).","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.0","python_version":"unknown","last_updated":"2026-08-08T03:48:17.894494+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":{},"lock_hash":"sha256:c570c03b9916ce0e672c3eb8408b9aa605abc4fb2cf1aae1c9e846e42471cc9c","checkpoints":{"mace-mh-1":{"fetched_at":"2026-08-07T18:38:54.862793+00:00","verified_at":"2026-08-08T00:05:26.406992+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macemh1model","size":59208139}],"weights_recorded_at":"2026-08-08T00:05:26.407006+00:00"},"mace-matpes-r2scan-0":{"fetched_at":"2026-08-07T18:38:55.030933+00:00","verified_at":"2026-08-08T00:02:25.811554+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEmatpesr2scanomatftmodel","size":79470738}],"weights_recorded_at":"2026-08-08T00:02:25.811568+00:00"},"mace-mp-0-medium":{"fetched_at":"2026-08-07T18:39:11.227407+00:00","verified_at":"2026-08-08T00:10:32.529472+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231203mace128L1_epoch199model","size":44422970}],"weights_recorded_at":"2026-08-08T00:10:32.529486+00:00"},"mace-mp-0-large":{"fetched_at":"2026-08-07T18:39:12.628241+00:00","verified_at":"2026-08-08T00:08:14.520477+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE_MPtrj_20229model","size":133803220}],"weights_recorded_at":"2026-08-08T00:08:14.520491+00:00"},"mace-mp-0-small":{"fetched_at":"2026-08-07T18:39:27.560268+00:00","verified_at":"2026-08-08T00:13:22.937776+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231210mace128L0_energy_epoch249model","size":32581838}],"weights_recorded_at":"2026-08-08T00:13:22.937791+00:00"},"mace-mpa-0-medium":{"fetched_at":"2026-08-07T18:39:27.578362+00:00","verified_at":"2026-08-08T00:16:40.970485+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macempa0mediummodel","size":79462305}],"weights_recorded_at":"2026-08-08T00:16:40.970499+00:00"},"mace-off23-large":{"fetched_at":"2026-08-07T18:39:34.646763+00:00","verified_at":"2026-08-08T00:19:34.275862+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_large.model","size":55492786}],"weights_recorded_at":"2026-08-08T00:19:34.275875+00:00"},"mace-off23-medium":{"fetched_at":"2026-08-07T18:39:39.457787+00:00","verified_at":"2026-08-08T00:21:54.987973+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_medium.model","size":18350596}],"weights_recorded_at":"2026-08-08T00:21:54.988025+00:00"},"mace-off23-small":{"fetched_at":"2026-08-07T18:39:41.383497+00:00","verified_at":"2026-08-08T00:24:54.580919+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_small.model","size":7347350}],"weights_recorded_at":"2026-08-08T00:24:54.580933+00:00"},"mace-omol-0-extra-large":{"fetched_at":"2026-08-07T18:39:43.466419+00:00","verified_at":"2026-08-08T00:27:01.780105+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-08-08T00:27:01.780119+00:00"}}},"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":{},"lock_hash":"sha256:37e7cd799c8736978027d1cb9918aa06410e3bac86a82dfba59cbf31c5fd7eab","checkpoints":{"allscaip-md-conserving-all-omol":{"fetched_at":"2026-08-07T19:00:42.320790+00:00","verified_at":"2026-08-07T21:14:12.300704+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-08-07T21:14:12.300722+00:00"},"allscaip-md-direct-all-omol":{"fetched_at":"2026-08-07T19:00:39.080181+00:00","verified_at":"2026-08-07T21:22:12.786098+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-08-07T21:22:12.786117+00:00"}}},"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":{},"lock_hash":"sha256:387dd9038769210056cbc5601717edbb8d2f9b55a8f0b184d513fb9bd10b0bd5","checkpoints":{"ani-1x":{"fetched_at":"2026-08-07T18:16:55.906080+00:00","verified_at":"2026-08-07T21:30:07.294177+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-08-07T21:30:07.294194+00:00"},"ani-1ccx":{"fetched_at":"2026-08-07T18:16:56.177074+00:00","verified_at":"2026-08-07T21:27:28.273465+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-08-07T21:27:28.273483+00:00"},"ani-2x":{"fetched_at":"2026-08-07T18:17:09.535866+00:00","verified_at":"2026-08-07T21:33:30.821225+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-08-07T21:33:30.821242+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":{},"lock_hash":"sha256:aa0b72f3991bfa8c4d8465982e75df9635fa1b626ece8e541900f0ed8637abb6","checkpoints":{"chgnet-default":{"fetched_at":"2026-08-07T18:18:20.972193+00:00","verified_at":"2026-08-07T21:40:46.470626+00:00","verified_device":"cuda","last_error":null,"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":{},"lock_hash":null,"checkpoints":{"dimenet-plus-plus-s2ef-oc20-20m":{"fetched_at":"2026-08-07T18:19:32.658646+00:00","verified_at":"2026-08-07T21:57:31.908744+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/dimenetpp_20M.pt","size":21841849}],"weights_recorded_at":"2026-08-07T21:57:31.908763+00:00"},"dimenet-plus-plus-s2ef-oc20-2m":{"fetched_at":"2026-08-07T18:33:36.719332+00:00","verified_at":"2026-08-07T22:04:50.509140+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/dimenetpp_2M.pt","size":21841593}],"weights_recorded_at":"2026-08-07T22:04:50.509154+00:00"},"dimenet-plus-plus-s2ef-oc20-all":{"fetched_at":"2026-08-07T18:33:37.904281+00:00","verified_at":"2026-08-07T22:11:42.374667+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/dimenetpp_all.pt","size":21841849}],"weights_recorded_at":"2026-08-07T22:11:42.374681+00:00"},"dimenet-plus-plus-s2ef-oc20-200k":{"fetched_at":"2026-08-07T18:33:38.187393+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 stderr (tail) ---\n[Worker] Prewarmed page cache: 44477 files, 6345 MB in 498.0s","weight_files":[{"path":"cache/dimenetpp_200k.pt","size":21841593}],"weights_recorded_at":"2026-08-07T18:33:38.187412+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":{},"lock_hash":null,"checkpoints":{"equiformer-v2-31m-s2ef-oc20-all-md":{"fetched_at":"2026-08-07T18:34:48.975897+00:00","verified_at":"2026-08-07T22:26:36.550751+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_31M_ec4_allmd.pt","size":125685325}],"weights_recorded_at":"2026-08-07T22:26:36.550766+00:00"},"equiformer-v2-83m-s2ef-oc20-2m":{"fetched_at":"2026-08-07T18:34:56.875435+00:00","verified_at":"2026-08-07T22:35:04.536668+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_83M_2M.pt","size":336843381}],"weights_recorded_at":"2026-08-07T22:35:04.536682+00:00"},"equiformer-v2-153m-s2ef-oc20-all-md":{"fetched_at":"2026-08-07T18:35:08.566741+00:00","verified_at":"2026-08-07T22:21:08.521894+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_153M_ec4_allmd.pt","size":619855697}],"weights_recorded_at":"2026-08-07T22:21:08.521909+00:00"}}},"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":{},"lock_hash":null,"checkpoints":{"escn-l6-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-08-07T18:35:27.546087+00:00","verified_at":"2026-08-07T22:52:07.413866+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-08-07T22:52:07.413879+00:00"},"escn-l6-m2-lay12-s2ef-oc20-all-md":{"fetched_at":"2026-08-07T18:35:29.989827+00:00","verified_at":"2026-08-07T23:00:46.878120+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-08-07T23:00:46.878135+00:00"},"escn-l6-m3-lay20-s2ef-oc20-all-md":{"fetched_at":"2026-08-07T18:35:33.656833+00:00","verified_at":"2026-08-07T23:07:16.270553+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-08-07T23:07:16.270568+00:00"},"escn-l4-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-08-07T18:35:53.977072+00:00","verified_at":"2026-08-07T22:45:00.609635+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-08-07T22:45:00.609649+00:00"}}},"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":{},"lock_hash":"sha256:37e7cd799c8736978027d1cb9918aa06410e3bac86a82dfba59cbf31c5fd7eab","checkpoints":{"esen-sm-conserving-all-omol":{"fetched_at":"2026-08-07T19:00:30.424489+00:00","verified_at":"2026-08-07T23:26:47.788886+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-08-07T23:26:47.788902+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":"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 stderr (tail) ---\n[Worker] Prewarmed page cache: 44622 files, 7929 MB in 588.3s","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-07T19:00:36.064541+00:00"},"esen-sm-direct-all-omol":{"fetched_at":"2026-08-07T19:00:56.382873+00:00","verified_at":"2026-08-07T23:35:25.025226+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-08-07T23:35:25.025242+00:00"}}},"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":{},"lock_hash":"sha256:1e7fb7c8dc56e7a942e52d5e1d50ff69579f7e581ce3e64c9cc62266f2300a20","checkpoints":{"m3gnet-mp-2021-2-8-pes":{"fetched_at":"2026-08-07T18:39:18.387057+00:00","verified_at":"2026-08-07T23:58:28.951646+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":{},"lock_hash":"sha256:82da6f7c8c18c84cf9c0bf34830b86029a965670658e380aeeb38d50d38511fd","checkpoints":{"mace-polar-1-s":{"fetched_at":"2026-08-07T18:40:24.843571+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 stderr (tail) ---\n[Worker] Prewarmed page cache: 33116 files, 7585 MB in 498.1s\n/global/cfs/cdirs/m5268/rootstock/envs/mace_polar/lib/python3.11/site-packages/e3nn/o3/_wigner.py:10: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.\n _Jd, _W3j_flat, _W3j_indices = torch.load(os.path.join(os.path.dirname(__file__), 'constants.pt'))","weight_files":[{"path":"cache/mace/MACEPOLAR1Smodel","size":33375439}],"weights_recorded_at":"2026-08-07T18:40:24.843595+00:00"},"mace-polar-1-m":{"fetched_at":"2026-08-07T18:40:25.467111+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 stderr (tail) ---\n[Worker] Prewarmed page cache: 33116 files, 7585 MB in 549.4s","weight_files":[{"path":"cache/mace/MACEPOLAR1Mmodel","size":68133235}],"weights_recorded_at":"2026-08-07T18:40:25.467130+00:00"},"mace-polar-1-l":{"fetched_at":"2026-08-07T18:40:26.911489+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 stderr (tail) ---\n[Worker] Prewarmed page cache: 33116 files, 7585 MB in 426.3s\n/global/cfs/cdirs/m5268/rootstock/envs/mace_polar/lib/python3.11/site-packages/e3nn/o3/_wigner.py:10: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.\n _Jd, _W3j_flat, _W3j_indices = torch.load(os.path.join(os.path.dirname(__file__), 'constants.pt'))","weight_files":[{"path":"cache/mace/MACEPOLAR1Lmodel","size":129808715}],"weights_recorded_at":"2026-08-07T18:40:26.911514+00:00"}}},"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":{},"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":"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 stderr (tail) ---\n[Worker] Prewarmed page cache: 53988 files, 8248 MB in 580.4s","weight_files":[{"path":"home/.local/mattersim/pretrained_models/mattersim-v1.0.0-1M.pth","size":17932943}],"weights_recorded_at":"2026-08-07T18:41:34.575502+00:00"},"mattersim-v1-0-0-5m":{"fetched_at":"2026-08-07T18:41:35.132251+00:00","verified_at":"2026-08-08T01:11:07.695357+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-08-08T01:11:07.695371+00:00"}}},"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":{},"lock_hash":"sha256:8282748753e7201545484a51c53d2e7c9e1d5c38c66511efd2ee579ec62b71c6","checkpoints":{"orb-d3-v2":{"fetched_at":"2026-08-07T18:41:18.788701+00:00","verified_at":"2026-08-08T01:16:16.895824+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-d3-v2-20241011.ckpt","size":100967743}],"weights_recorded_at":"2026-08-08T01:16:16.895839+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":{},"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":"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":[{"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":"2026-08-07T23:52:43.213986+00:00","verified_device":"cuda","last_error":null,"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":{},"lock_hash":"sha256:1b14ab95f8183b2845ff04c3d145511c15b378f2a2a6f660d83545fa10594eeb","checkpoints":{"orb-v3-conservative-inf-omat":{"fetched_at":"2026-08-07T18:42:05.251153+00:00","verified_at":"2026-08-08T01:46:01.721813+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-08-08T01:46:01.721828+00:00"},"orb-v3-conservative-inf-mpa":{"fetched_at":"2026-08-07T18:42:05.487867+00:00","verified_at":"2026-08-08T01:40:49.154621+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-08-08T01:40:49.154636+00:00"},"orb-v3-conservative-20-mpa":{"fetched_at":"2026-08-07T18:42:05.833271+00:00","verified_at":"2026-08-08T01:27:03.136948+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-08-08T01:27:03.136963+00:00"},"orb-v3-conservative-20-omat":{"fetched_at":"2026-08-07T18:42:06.048676+00:00","verified_at":"2026-08-08T01:33:02.660555+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-08-08T01:33:02.660569+00:00"},"orb-v3-direct-20-mpa":{"fetched_at":"2026-08-07T18:42:26.108950+00:00","verified_at":"2026-08-08T01:49:21.227184+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-08-08T01:49:21.227199+00:00"},"orb-v3-direct-20-omat":{"fetched_at":"2026-08-07T18:42:26.198770+00:00","verified_at":"2026-08-08T01:56:00.372680+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-08-08T01:56:00.372695+00:00"},"orb-v3-direct-inf-mpa":{"fetched_at":"2026-08-07T18:42:26.341487+00:00","verified_at":"2026-08-08T01:59:32.249259+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-08-08T01:59:32.249272+00:00"},"orb-v3-direct-inf-omat":{"fetched_at":"2026-08-07T18:42:26.487865+00:00","verified_at":"2026-08-08T02:05:02.802276+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-08-08T02:05:02.802290+00:00"}}},"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":{},"lock_hash":null,"checkpoints":{"painn-s2ef-oc20-all":{"fetched_at":"2026-08-07T18:44:12.251376+00:00","verified_at":"2026-08-08T02:12:24.854964+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/painn_h512_s2ef_all.pt","size":80340525}],"weights_recorded_at":"2026-08-08T02:12:24.854979+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":{},"lock_hash":"sha256:a70651437572b0f9850932bae01a42e6cd89feef1661ccbcc687386448ea65e9","checkpoints":{"pet-omatpes-l":{"fetched_at":"2026-08-07T18:44:13.271159+00:00","verified_at":"2026-08-08T02:20:11.855478+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-08-08T02:20:11.855494+00:00"},"pet-oam-xl":{"fetched_at":"2026-08-07T18:44:27.538890+00:00","verified_at":"2026-08-08T02:17:14.296359+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-08-08T02:17:14.296376+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":{},"lock_hash":null,"checkpoints":{"scn-s2ef-oc20-2m":{"fetched_at":"2026-08-07T18:44:24.117756+00:00","verified_at":"2026-08-08T02:25:37.425905+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/scn_t1_b1_s2ef_2M.pt","size":494398909}],"weights_recorded_at":"2026-08-08T02:25:37.425919+00:00"},"scn-t4-b2-s2ef-oc20-2m":{"fetched_at":"2026-08-07T18:44:48.572433+00:00","verified_at":"2026-08-08T02:31:22.018488+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/scn_t4_b2_s2ef_2M.pt","size":506994109}],"weights_recorded_at":"2026-08-08T02:31:22.018502+00:00"},"scn-s2ef-oc20-all-md":{"fetched_at":"2026-08-07T18:44:53.943614+00:00","verified_at":"2026-08-08T02:28:40.105763+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/scn_all_md_s2ef.pt","size":675885597}],"weights_recorded_at":"2026-08-08T02:28:40.105778+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":{},"lock_hash":"sha256:f0907ac3ba945602b4ecef938a6f6b838750486f828f87e55248231beda8c649","checkpoints":{"sevennet-0":{"fetched_at":"2026-08-07T18:45:55.176461+00:00","verified_at":"2026-08-08T02:41:14.931314+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":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 stderr (tail) ---\n[Worker] Prewarmed page cache: 37327 files, 7883 MB in 531.5s","weight_files":null,"weights_recorded_at":null},"sevennet-omat":{"fetched_at":"2026-08-07T18:45:59.080797+00:00","verified_at":"2026-08-08T03:03:52.811676+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-08-08T02:59:10.426533+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-08-08T03:08:43.865253+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":{},"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":"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 stderr (tail) ---\n[Worker] Prewarmed page cache: 38503 files, 5420 MB in 572.4s","weight_files":[{"path":"cache/huggingface/.agent_harnesses.json","size":5698},{"path":"home/.cache/tace/CACHEDIR.TAG","size":191},{"path":"home/.cache/tace/TECE-OAM-RRA-1.0.pt","size":890523407},{"path":"home/.cache/tace/models--xvzemin--tace-foundations/blobs/9f36562582d931347c3904f763e820edcaf5f27c3f13beb6776e49fcc7de38bb","size":890523407}],"weights_recorded_at":"2026-08-07T18:47:47.347426+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":{},"lock_hash":null,"checkpoints":{"tensornet-matpes-pbe-2025-2":{"fetched_at":"2026-08-07T18:47:11.391749+00:00","verified_at":"2026-08-08T03:27:21.743331+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-08-08T03:27:21.743347+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":{},"lock_hash":"sha256:37e7cd799c8736978027d1cb9918aa06410e3bac86a82dfba59cbf31c5fd7eab","checkpoints":{"uma-m-1p1":{"fetched_at":"2026-08-07T19:02:45.891498+00:00","verified_at":"2026-08-08T03:35:42.038240+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-08-08T03:35:42.038257+00:00"},"uma-s-1p1":{"fetched_at":"2026-08-07T19:01:33.536471+00:00","verified_at":"2026-08-08T03:42:13.018227+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-08-08T03:42:13.018243+00:00"},"uma-s-1p2":{"fetched_at":"2026-08-07T19:01:40.358988+00:00","verified_at":"2026-08-08T03:47:57.096445+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-08-08T03:47:57.096462+00:00"}}}}},{"schema_version":5,"cluster":"polaris","root":"/eagle/projects/Rootstock/rootstock","maintainer":{"name":"Owen Price Skelly","email":"OwenPriceSkelly@uchicago.edu"},"rootstock_version":"1.6.3","python_version":"unknown","last_updated":"2026-08-29T07:50:20.886523+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-08-21T11:07:24.268537+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-08-29T07:50:18.788431+00:00"},"allscaip-md-conserving-all-omol":{"fetched_at":"2026-08-03T17:03:57.183635+00:00","verified_at":"2026-08-21T11:07:24.268510+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-08-29T07:50:18.788405+00:00"},"allscaip:custom":{"fetched_at":null,"verified_at":"2026-08-21T11:07:24.268819+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-08-21T11:07:24.268572+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-08-21T11:07:24.268581+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_31M_ec4_allmd.pt","size":125685325}],"weights_recorded_at":"2026-08-29T07:50:18.788486+00:00"},"equiformer-v2-83m-s2ef-oc20-2m":{"fetched_at":"2026-08-03T17:07:10.622992+00:00","verified_at":"2026-08-21T11:07:24.268587+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_83M_2M.pt","size":336843381}],"weights_recorded_at":"2026-08-29T07:50:18.788495+00:00"},"equiformer-v2-153m-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:07:12.126551+00:00","verified_at":"2026-08-21T11:07:24.268575+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_153M_ec4_allmd.pt","size":619855697}],"weights_recorded_at":"2026-08-29T07:50:18.788478+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 2.396e-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-08-21T11:07:24.268609+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-08-29T07:50:18.788525+00:00"},"escn-l6-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-08-03T17:08:03.373277+00:00","verified_at":"2026-08-21T11:07:24.268604+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-08-29T07:50:18.788517+00:00"},"escn-l6-m2-lay12-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:08:03.392125+00:00","verified_at":"2026-08-21T11:07:24.268593+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-08-29T07:50:18.788502+00:00"},"escn-l6-m3-lay20-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:08:14.400434+00:00","verified_at":"2026-08-21T11:07:24.268598+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-08-29T07:50:18.788510+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 5.914e-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-08-21T11:07:24.268628+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-08-29T07:50:18.788552+00:00"},"esen-sm-conserving-all-omol":{"fetched_at":"2026-08-03T17:11:00.922460+00:00","verified_at":"2026-08-21T11:07:24.268621+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-08-29T07:50:18.788544+00:00"},"esen-md-direct-all-omol":{"fetched_at":"2026-08-03T17:11:00.992499+00:00","verified_at":"2026-08-21T11:07:24.268615+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-08-29T07:50:18.788534+00:00"},"esen:custom":{"fetched_at":null,"verified_at":"2026-08-21T11:07:24.268825+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-08-21T11:07:24.268634+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-08-21T11:07:24.268637+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-08-21T11:07:24.268646+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231203mace128L1_epoch199model","size":44422970}],"weights_recorded_at":"2026-08-29T07:50:18.788577+00:00"},"mace-mh-1":{"fetched_at":"2026-08-03T17:12:25.771339+00:00","verified_at":"2026-08-21T11:07:24.268683+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macemh1model","size":59208139}],"weights_recorded_at":"2026-08-29T07:50:18.788634+00:00"},"mace-matpes-r2scan-0":{"fetched_at":"2026-08-03T17:12:26.162740+00:00","verified_at":"2026-08-21T11:07:24.268678+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEmatpesr2scanomatftmodel","size":79470738}],"weights_recorded_at":"2026-08-29T07:50:18.788626+00:00"},"mace-mp-0-large":{"fetched_at":"2026-08-03T17:12:26.906848+00:00","verified_at":"2026-08-21T11:07:24.268651+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE_MPtrj_20229model","size":133803220}],"weights_recorded_at":"2026-08-29T07:50:18.788585+00:00"},"mace-off23-medium":{"fetched_at":"2026-08-03T17:12:44.272336+00:00","verified_at":"2026-08-21T11:07:24.268662+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_medium.model","size":18350596}],"weights_recorded_at":"2026-08-29T07:50:18.788602+00:00"},"mace-mp-0-small":{"fetched_at":"2026-08-03T17:12:44.380401+00:00","verified_at":"2026-08-21T11:07:24.268640+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231210mace128L0_energy_epoch249model","size":32581838}],"weights_recorded_at":"2026-08-29T07:50:18.788569+00:00"},"mace-mpa-0-medium":{"fetched_at":"2026-08-03T17:12:44.712260+00:00","verified_at":"2026-08-21T11:07:24.268673+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macempa0mediummodel","size":79462305}],"weights_recorded_at":"2026-08-29T07:50:18.788619+00:00"},"mace-off23-large":{"fetched_at":"2026-08-03T17:12:44.790639+00:00","verified_at":"2026-08-21T11:07:24.268668+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_large.model","size":55492786}],"weights_recorded_at":"2026-08-29T07:50:18.788610+00:00"},"mace-off23-small":{"fetched_at":"2026-08-03T17:13:00.903240+00:00","verified_at":"2026-08-21T11:07:24.268656+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_small.model","size":7347350}],"weights_recorded_at":"2026-08-29T07:50:18.788593+00:00"},"mace-omol-0-extra-large":{"fetched_at":"2026-08-03T17:13:02.181806+00:00","verified_at":"2026-08-21T11:07:24.268687+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-08-29T07:50:18.788641+00:00"},"mace:custom":{"fetched_at":null,"verified_at":"2026-08-21T11:07:24.268829+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-08-21T11:07:24.268693+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Smodel","size":33375439}],"weights_recorded_at":"2026-08-29T07:50:18.788649+00:00"},"mace-polar-1-m":{"fetched_at":"2026-08-03T17:15:38.434404+00:00","verified_at":"2026-08-21T11:07:24.268698+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Mmodel","size":68133235}],"weights_recorded_at":"2026-08-29T07:50:18.788657+00:00"},"mace-polar-1-l":{"fetched_at":"2026-08-03T17:15:38.454464+00:00","verified_at":"2026-08-21T11:07:24.268703+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Lmodel","size":129808715}],"weights_recorded_at":"2026-08-29T07:50:18.788664+00:00"},"mace-polar:custom":{"fetched_at":null,"verified_at":"2026-08-21T11:07:24.268831+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-08-21T11:07:24.268733+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-08-29T07:50:18.788710+00:00"},"orb-v3-conservative-20-omat":{"fetched_at":"2026-08-03T17:18:18.543840+00:00","verified_at":"2026-08-21T11:07:24.268713+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-08-29T07:50:18.788679+00:00"},"orb-v3-direct-20-mpa":{"fetched_at":"2026-08-03T17:18:51.755916+00:00","verified_at":"2026-08-21T11:07:24.268743+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-08-29T07:50:18.788724+00:00"},"orb-v3-conservative-inf-omat":{"fetched_at":"2026-08-03T17:18:52.017384+00:00","verified_at":"2026-08-21T11:07:24.268708+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-08-29T07:50:18.788671+00:00"},"orb-v3-conservative-inf-mpa":{"fetched_at":"2026-08-03T17:18:52.075825+00:00","verified_at":"2026-08-21T11:07:24.268728+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-08-29T07:50:18.788702+00:00"},"orb-v3-direct-20-omat":{"fetched_at":"2026-08-03T17:18:53.566555+00:00","verified_at":"2026-08-21T11:07:24.268723+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-08-29T07:50:18.788695+00:00"},"orb-v3-direct-inf-omat":{"fetched_at":"2026-08-03T17:19:14.412381+00:00","verified_at":"2026-08-21T11:07:24.268718+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-08-29T07:50:18.788686+00:00"},"orb-v3-direct-inf-mpa":{"fetched_at":"2026-08-03T17:19:16.914781+00:00","verified_at":"2026-08-21T11:07:24.268738+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-08-29T07:50:18.788717+00:00"},"orb-v3:custom":{"fetched_at":null,"verified_at":"2026-08-21T11:07:24.268833+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-08-21T11:07:24.268748+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-d3-v2-20241011.ckpt","size":100967743}],"weights_recorded_at":"2026-08-29T07:50:18.788731+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-08-21T11:07:24.268776+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-08-21T11:07:24.268779+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-08-21T11:07:24.268782+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-08-21T11:07:24.268785+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-08-21T11:07:24.268787+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-08-21T11:07:24.268790+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-08-29T07:50:18.792958+00:00"},"uma-m-1p1":{"fetched_at":"2026-08-03T17:25:21.596525+00:00","verified_at":"2026-08-21T11:07:24.268804+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-08-29T07:50:18.792976+00:00"},"uma-s-1p2":{"fetched_at":"2026-08-03T17:26:22.685148+00:00","verified_at":"2026-08-21T11:07:24.268797+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-08-29T07:50:18.792968+00:00"},"uma:custom":{"fetched_at":null,"verified_at":"2026-08-21T11:07:24.268835+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-08-21T11:07:24.268811+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-08-29T07:50:18.792987+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-08-21T11:07:24.268565+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-08-21T11:07:24.268569+00:00"},"ani-1ccx":{"fetched_at":"2026-08-06T18:03:01.300850+00:00","verified_at":"2026-08-21T11:07:24.268558+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-08-21T11:07:24.268562+00:00"},"ani-2x":{"fetched_at":"2026-08-06T18:03:01.427309+00:00","verified_at":"2026-08-21T11:07:24.268550+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-08-21T11:07:24.268554+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-08-21T11:07:24.268770+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-08-21T11:07:24.268774+00:00"},"pet-oam-xl":{"fetched_at":"2026-08-06T18:03:24.011446+00:00","verified_at":"2026-08-21T11:07:24.268753+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-08-21T11:07:24.268757+00:00"}}}}},{"schema_version":5,"cluster":"sophia","root":"/eagle/projects/Rootstock/rootstock","maintainer":{"name":"Owen Price Skelly","email":"OwenPriceSkelly@uchicago.edu"},"rootstock_version":"1.6.3","python_version":"unknown","last_updated":"2026-08-29T07:50:20.886523+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-08-29T07:50:18.788422+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-08-29T07:50:18.788431+00:00"},"allscaip-md-conserving-all-omol":{"fetched_at":"2026-08-03T17:03:57.183635+00:00","verified_at":"2026-08-29T07:50:18.787612+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-08-29T07:50:18.788405+00:00"},"allscaip:custom":{"fetched_at":null,"verified_at":"2026-08-29T07:50:18.792993+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-08-29T07:50:18.788460+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-08-29T07:50:18.788464+00:00"},"ani-1ccx":{"fetched_at":"2026-08-03T17:03:11.776896+00:00","verified_at":"2026-08-29T07:50:18.788451+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-08-29T07:50:18.788456+00:00"},"ani-2x":{"fetched_at":"2026-08-03T17:04:29.856289+00:00","verified_at":"2026-08-29T07:50:18.788441+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-08-29T07:50:18.788446+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-08-29T07:50:18.788468+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-08-29T07:50:18.788482+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_31M_ec4_allmd.pt","size":125685325}],"weights_recorded_at":"2026-08-29T07:50:18.788486+00:00"},"equiformer-v2-83m-s2ef-oc20-2m":{"fetched_at":"2026-08-03T17:07:10.622992+00:00","verified_at":"2026-08-29T07:50:18.788490+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_83M_2M.pt","size":336843381}],"weights_recorded_at":"2026-08-29T07:50:18.788495+00:00"},"equiformer-v2-153m-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:07:12.126551+00:00","verified_at":"2026-08-29T07:50:18.788473+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/eq2_153M_ec4_allmd.pt","size":619855697}],"weights_recorded_at":"2026-08-29T07:50:18.788478+00:00"},"equiformer:custom":{"fetched_at":null,"verified_at":"2026-08-29T07:50:18.792999+00:00","verified_device":"cuda","last_error":null,"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-08-29T07:50:18.788521+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-08-29T07:50:18.788525+00:00"},"escn-l6-m2-lay12-s2ef-oc20-2m":{"fetched_at":"2026-08-03T17:08:03.373277+00:00","verified_at":"2026-08-29T07:50:18.788514+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-08-29T07:50:18.788517+00:00"},"escn-l6-m2-lay12-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:08:03.392125+00:00","verified_at":"2026-08-29T07:50:18.788498+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-08-29T07:50:18.788502+00:00"},"escn-l6-m3-lay20-s2ef-oc20-all-md":{"fetched_at":"2026-08-03T17:08:14.400434+00:00","verified_at":"2026-08-29T07:50:18.788506+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-08-29T07:50:18.788510+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 3.978e-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-08-29T07:50:18.788548+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-08-29T07:50:18.788552+00:00"},"esen-sm-conserving-all-omol":{"fetched_at":"2026-08-03T17:11:00.922460+00:00","verified_at":"2026-08-29T07:50:18.788539+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-08-29T07:50:18.788544+00:00"},"esen-md-direct-all-omol":{"fetched_at":"2026-08-03T17:11:00.992499+00:00","verified_at":"2026-08-29T07:50:18.788529+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-08-29T07:50:18.788534+00:00"},"esen:custom":{"fetched_at":null,"verified_at":"2026-08-29T07:50:18.793005+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-08-29T07:50:18.788556+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-08-29T07:50:18.788560+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-08-29T07:50:18.788573+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231203mace128L1_epoch199model","size":44422970}],"weights_recorded_at":"2026-08-29T07:50:18.788577+00:00"},"mace-mh-1":{"fetched_at":"2026-08-03T17:12:25.771339+00:00","verified_at":"2026-08-29T07:50:18.788630+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macemh1model","size":59208139}],"weights_recorded_at":"2026-08-29T07:50:18.788634+00:00"},"mace-matpes-r2scan-0":{"fetched_at":"2026-08-03T17:12:26.162740+00:00","verified_at":"2026-08-29T07:50:18.788623+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEmatpesr2scanomatftmodel","size":79470738}],"weights_recorded_at":"2026-08-29T07:50:18.788626+00:00"},"mace-mp-0-large":{"fetched_at":"2026-08-03T17:12:26.906848+00:00","verified_at":"2026-08-29T07:50:18.788581+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE_MPtrj_20229model","size":133803220}],"weights_recorded_at":"2026-08-29T07:50:18.788585+00:00"},"mace-off23-medium":{"fetched_at":"2026-08-03T17:12:44.272336+00:00","verified_at":"2026-08-29T07:50:18.788598+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_medium.model","size":18350596}],"weights_recorded_at":"2026-08-29T07:50:18.788602+00:00"},"mace-mp-0-small":{"fetched_at":"2026-08-03T17:12:44.380401+00:00","verified_at":"2026-08-29T07:50:18.788565+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/20231210mace128L0_energy_epoch249model","size":32581838}],"weights_recorded_at":"2026-08-29T07:50:18.788569+00:00"},"mace-mpa-0-medium":{"fetched_at":"2026-08-03T17:12:44.712260+00:00","verified_at":"2026-08-29T07:50:18.788615+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/macempa0mediummodel","size":79462305}],"weights_recorded_at":"2026-08-29T07:50:18.788619+00:00"},"mace-off23-large":{"fetched_at":"2026-08-03T17:12:44.790639+00:00","verified_at":"2026-08-29T07:50:18.788606+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_large.model","size":55492786}],"weights_recorded_at":"2026-08-29T07:50:18.788610+00:00"},"mace-off23-small":{"fetched_at":"2026-08-03T17:13:00.903240+00:00","verified_at":"2026-08-29T07:50:18.788589+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACE-OFF23_small.model","size":7347350}],"weights_recorded_at":"2026-08-29T07:50:18.788593+00:00"},"mace-omol-0-extra-large":{"fetched_at":"2026-08-03T17:13:02.181806+00:00","verified_at":"2026-08-29T07:50:18.788637+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-08-29T07:50:18.788641+00:00"},"mace:custom":{"fetched_at":null,"verified_at":"2026-08-29T07:50:18.793009+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-08-29T07:50:18.788645+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Smodel","size":33375439}],"weights_recorded_at":"2026-08-29T07:50:18.788649+00:00"},"mace-polar-1-m":{"fetched_at":"2026-08-03T17:15:38.434404+00:00","verified_at":"2026-08-29T07:50:18.788653+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Mmodel","size":68133235}],"weights_recorded_at":"2026-08-29T07:50:18.788657+00:00"},"mace-polar-1-l":{"fetched_at":"2026-08-03T17:15:38.454464+00:00","verified_at":"2026-08-29T07:50:18.788660+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/mace/MACEPOLAR1Lmodel","size":129808715}],"weights_recorded_at":"2026-08-29T07:50:18.788664+00:00"},"mace-polar:custom":{"fetched_at":null,"verified_at":"2026-08-29T07:50:18.793012+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-08-29T07:50:18.788706+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-08-29T07:50:18.788710+00:00"},"orb-v3-conservative-20-omat":{"fetched_at":"2026-08-03T17:18:18.543840+00:00","verified_at":"2026-08-29T07:50:18.788675+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-08-29T07:50:18.788679+00:00"},"orb-v3-direct-20-mpa":{"fetched_at":"2026-08-03T17:18:51.755916+00:00","verified_at":"2026-08-29T07:50:18.788720+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-08-29T07:50:18.788724+00:00"},"orb-v3-conservative-inf-omat":{"fetched_at":"2026-08-03T17:18:52.017384+00:00","verified_at":"2026-08-29T07:50:18.788667+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-08-29T07:50:18.788671+00:00"},"orb-v3-conservative-inf-mpa":{"fetched_at":"2026-08-03T17:18:52.075825+00:00","verified_at":"2026-08-29T07:50:18.788699+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-08-29T07:50:18.788702+00:00"},"orb-v3-direct-20-omat":{"fetched_at":"2026-08-03T17:18:53.566555+00:00","verified_at":"2026-08-29T07:50:18.788691+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-08-29T07:50:18.788695+00:00"},"orb-v3-direct-inf-omat":{"fetched_at":"2026-08-03T17:19:14.412381+00:00","verified_at":"2026-08-29T07:50:18.788682+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-08-29T07:50:18.788686+00:00"},"orb-v3-direct-inf-mpa":{"fetched_at":"2026-08-03T17:19:16.914781+00:00","verified_at":"2026-08-29T07:50:18.788713+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-08-29T07:50:18.788717+00:00"},"orb-v3:custom":{"fetched_at":null,"verified_at":"2026-08-29T07:50:18.793016+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-08-29T07:50:18.788727+00:00","verified_device":"cuda","last_error":null,"weight_files":[{"path":"cache/orb/orb-d3-v2-20241011.ckpt","size":100967743}],"weights_recorded_at":"2026-08-29T07:50:18.788731+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-08-29T07:50:18.792920+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-08-29T07:50:18.792927+00:00"},"pet-oam-xl":{"fetched_at":"2026-08-03T17:21:45.473434+00:00","verified_at":"2026-08-29T07:50:18.788735+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-08-29T07:50:18.788740+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-08-29T07:50:18.792933+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-08-29T07:50:18.792938+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-08-29T07:50:18.792942+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-08-29T07:50:18.792946+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-08-29T07:50:18.792949+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-08-29T07:50:18.792953+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-08-29T07:50:18.792958+00:00"},"uma-m-1p1":{"fetched_at":"2026-08-03T17:25:21.596525+00:00","verified_at":"2026-08-29T07:50:18.792971+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-08-29T07:50:18.792976+00:00"},"uma-s-1p2":{"fetched_at":"2026-08-03T17:26:22.685148+00:00","verified_at":"2026-08-29T07:50:18.792963+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-08-29T07:50:18.792968+00:00"},"uma:custom":{"fetched_at":null,"verified_at":"2026-08-29T07:50:18.793019+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-08-29T07:50:18.792980+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-08-29T07:50:18.792987+00:00"}}}}},{"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."}}}}}]}