python3Packages.mace-torch: init at 0.3.16

2026-06-12 23:25:35 +00:00
parent 8b2f221d44
commit f03ca2e04c
2 changed files with 140 additions and 0 deletions
@@ -0,0 +1,138 @@
+{
+  lib,
+  stdenv,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  setuptools,
+
+  # dependencies
+  ase,
+  configargparse,
+  e3nn,
+  gitpython,
+  h5py,
+  lmdb,
+  matplotlib,
+  matscipy,
+  numpy,
+  opt-einsum,
+  orjson,
+  pandas,
+  prettytable,
+  python-hostlist,
+  pyyaml,
+  torch,
+  torch-ema,
+  torchmetrics,
+  tqdm,
+
+  # tests
+  pytestCheckHook,
+  writableTmpDirAsHomeHook,
+}:
+
+buildPythonPackage (finalAttrs: {
+  pname = "mace-torch";
+  version = "0.3.16";
+  pyproject = true;
+  __structuredAttrs = true;
+
+  src = fetchFromGitHub {
+    owner = "acesuit";
+    repo = "mace";
+    tag = "v${finalAttrs.version}";
+    hash = "sha256-sJ/3c7kPe99vkliixUMqqQi2OiL3CCUdlcgpDZ/PUHA=";
+  };
+
+  build-system = [
+    setuptools
+  ];
+
+  pythonRelaxDeps = [
+    "e3nn"
+  ];
+  dependencies = [
+    ase
+    configargparse
+    e3nn
+    gitpython
+    h5py
+    lmdb
+    matplotlib
+    matscipy
+    numpy
+    opt-einsum
+    orjson
+    pandas
+    prettytable
+    python-hostlist
+    pyyaml
+    torch
+    torch-ema
+    torchmetrics
+    tqdm
+  ];
+
+  pythonImportsCheck = [ "mace" ];
+
+  nativeCheckInputs = [
+    pytestCheckHook
+    writableTmpDirAsHomeHook
+  ];
+
+  disabledTests = [
+    # _pickle.PickleError: ScriptFunction cannot be pickled
+    "test_run_eval_fail_with_wrong_model"
+
+    # ValueError: too many values to unpack (expected 2)
+    "test_mace_mp"
+
+    # AssertionError (tensors not close)
+    "test_run_train_dipole"
+    "test_run_train_dipole_polar"
+
+    # RuntimeError: Model download failed and no local model found
+    "test_calculator_descriptor"
+    "test_compile_foundation"
+    "test_extract_config"
+    "test_finite_difference_hessian"
+    "test_foundations"
+    "test_initial_metrics_replay_head_mh0"
+    "test_initial_metrics_replay_head_mh1"
+    "test_initial_metrics_replay_head_omol"
+    "test_mace_mh_1_elements_subset_reproduces_energy_forces"
+    "test_mace_mp_energies"
+    "test_mace_mp_stresses"
+    "test_mace_off"
+    "test_mace_omol_elements_subset_reproduces_energy_forces"
+    "test_multi_reference"
+    "test_multihead_finetuning_different_formats"
+    "test_multihead_finetuning_does_not_modify_default_keyspec"
+    "test_potential_energy_and_hessian"
+    "test_remove_pt_head_omol_multihead"
+    "test_run_train_foundation"
+    "test_run_train_foundation_elements"
+    "test_run_train_foundation_elements_multihead"
+    "test_run_train_foundation_multihead"
+    "test_run_train_foundation_multihead_json"
+    "test_run_train_foundation_multihead_pseudolabeling"
+    "test_run_train_freeze"
+    "test_run_train_mh_foundation"
+    "test_run_train_soft_freeze"
+  ]
+  ++ lib.optionals stdenv.hostPlatform.isDarwin [
+    # symbol not found in flat namespace '___kmpc_barrier'
+    "test_mace"
+    "test_mace_compile_stress"
+  ];
+
+  meta = {
+    description = "Fast and accurate machine learning interatomic potentials with higher order equivariant message passing";
+    homepage = "https://github.com/acesuit/mace";
+    changelog = "https://github.com/acesuit/mace/releases/tag/${finalAttrs.src.tag}";
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ GaetanLepage ];
+  };
+})
@@ -9578,6 +9578,8 @@ self: super: with self; {

  macaddress = callPackage ../development/python-modules/macaddress { };

+  mace-torch = callPackage ../development/python-modules/mace-torch { };
+
  macfsevents = callPackage ../development/python-modules/macfsevents { };

  macholib = callPackage ../development/python-modules/macholib { };