MACE fixes & tests

Author: kazeevn

.gitignore

@@ -40,4 +40,7 @@ generated.tar.gz
 generated_public.tar.gz
 uv.toml
 src/wyckoff_transformer/wyckoffs_enumerated_by_ss.json
-dist/
+dist/
+relax/
+mp_20_test.json.gz
+*.log

README.md

@@ -122,6 +122,53 @@ wyformer-generate <output-file> --hf-model SymmetryAdvantage/<model-name> --sg-d
 There are two ways to generate 3D structures from Wyckoff representations: DiffCSP++ and CHGNet. They later can be relaxed with CHGNet and/or DFT.
 #### DiffCSP++
 Wyckoffs can be relaxed with modified [DiffCSP++ code](https://github.com/kazeevn/DiffCSPNew/tree/master)
+#### CrySPR + MACE
+[CrySPR](https://chemrxiv.org/engage/chemrxiv/article-details/66b308a501103d79c5fd9b91) scheme using [`pyxtal`](https://pyxtal.readthedocs.io/en/latest/index.html) and a [MACE](https://github.com/ACEsuit/mace) ML force field is integrated directly into the package. Install the optional extra first:
+```bash
+pip install "wyckoff-transformer[relax]"
+```
+Then run from the command line:
+```bash
+wyformer-cryspr WyckoffTransformer_mp_20.json \
+    --model https://github.com/ACEsuit/mace-foundations/releases/download/mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model \
+    --output-dir results/ --start 0 --end 1000
+# model_name defaults to the model file stem
+head results/2023-12-10-mace-128-L0_energy_epoch-249_results.csv
+model,id,formula,energy,energy_per_atom
+...
+2023-12-10-mace-128-L0_energy_epoch-249,35,H6O8Si2,-97.98,...
+```
+URL-based models are downloaded once and cached in `~/.cache/wyckoff_transformer/mace_models/`. A local path is accepted too: `--model /path/to/model.model`.
+
+Output layout is identical to the CHGNet variant below. Key options:
+- `--n-trials N` — number of random PyXtal trials per structure (default 6)
+- `--fmax F` — force convergence criterion in eV/Å (default 0.01)
+- `--model-name NAME` — label for the results CSV (default: model file stem)
+- `--device auto|cpu|cuda` — PyTorch device selection (default: auto)
+
+To use as a library:
+```python
+from wyckoff_transformer.cryspr.calculator import build_mace_calculator
+from wyckoff_transformer.cryspr.generator import single_pyxtal, func_run
+import json
+
+calculator = build_mace_calculator(
+    "https://github.com/ACEsuit/mace-foundations/releases/download/"
+    "mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model"
+)
+
+with open("WyckoffTransformer_mp_20.json") as f:
+    data = json.load(f)
+
+atoms, formula, energy, energy_per_atom = func_run(
+    id_gene=99,
+    wyckoffgene=data[99],
+    calculator=calculator,
+    output_dir="results/",
+    n_trials=6,
+)
+```
+
 #### CrySPR + CHGNet
 [CrySPR](https://chemrxiv.org/engage/chemrxiv/article-details/66b308a501103d79c5fd9b91) scheme that combines [`pyxtal`](https://pyxtal.readthedocs.io/en/latest/index.html) with CHGNet (or any other machine-learning interatomic potentials)
 ```bash

pyproject.toml

@@ -49,6 +49,7 @@ cdvae = ["cdvae-property-models"]
 publish = ["build", "twine", "packaging >=25"]
 research = ["ase", "matbench-discovery", "ipykernel", "ipywidgets", "jupyter", "ipympl"]
 relax = ["mace-torch >=0.3", "ase >=3.23"]
+cuequivariance-cu13 = ["cuequivariance-torch >=0.9", "triton >=3", "cuequivariance-ops-torch-cu13 >=0.9"]
 
 [build-system]
 requires = ["hatchling", "pyxtal>=1.1.3", "numpy>=1.26.0,<2", "scipy", "scikit-learn>=1.5.0", "pydantic", "pandas"]

src/wyckoff_transformer/cli/relax.py

@@ -1,6 +1,7 @@
 """CLI entry point for CrySPR: crystal structure prediction via PyXtal + MACE."""
 import argparse
 import json
+import gzip
 import logging
 from pathlib import Path
 
@@ -93,7 +94,11 @@ def main() -> None:
         format="%(asctime)s %(levelname)s %(name)s: %(message)s",
     )
 
-    with open(args.input) as f:
+    if args.input.suffix == ".gz":
+        opener = gzip.open
+    else:
+        opener = open
+    with opener(args.input, mode="rt", encoding="utf-8") as f:
         data = json.load(f)
 
     end = args.end if args.end is not None else len(data)
@@ -108,7 +113,7 @@ def main() -> None:
 
     results = []
     for i, wyckoffgene in enumerate(selected, start=args.start):
-        atoms, formula, energy, energy_per_atom = func_run(
+        atoms, formula, energy, energy_per_atom, cif = func_run(
             id_gene=i,
             wyckoffgene=wyckoffgene,
             calculator=calculator,
@@ -123,6 +128,7 @@ def main() -> None:
             "formula": formula,
             "energy": energy,
             "energy_per_atom": energy_per_atom,
+            "cif": cif,
         })
 
     results_csv = args.output_dir / f"{model_name}_results.csv"

src/wyckoff_transformer/cli/tests/test_relax.py

@@ -37,7 +37,7 @@ def test_missing_model_exits(self):
     @patch("wyckoff_transformer.cli.relax.build_mace_calculator")
     @patch("wyckoff_transformer.cli.relax.func_run")
     def test_basic_run_writes_csv(self, mock_func_run, mock_build_calc):
-        mock_func_run.return_value = (MagicMock(), "NaCl", -10.0, -1.25)
+        mock_func_run.return_value = (MagicMock(), "NaCl", -10.0, -1.25, "data_\nCIF content")
         mock_build_calc.return_value = MagicMock()
 
         out_dir = self.tmp_path / "output"
@@ -54,12 +54,13 @@ def test_basic_run_writes_csv(self, mock_func_run, mock_build_calc):
         self.assertEqual(len(csv_files), 1)
         df = pd.read_csv(csv_files[0])
         self.assertIn("energy", df.columns)
+        self.assertIn("cif", df.columns)
         self.assertEqual(len(df), 1)
 
     @patch("wyckoff_transformer.cli.relax.build_mace_calculator")
     @patch("wyckoff_transformer.cli.relax.func_run")
     def test_model_name_used_as_csv_stem(self, mock_func_run, mock_build_calc):
-        mock_func_run.return_value = (None, None, None, None)
+        mock_func_run.return_value = (None, None, None, None, None)
         mock_build_calc.return_value = MagicMock()
 
         out_dir = self.tmp_path / "out2"
@@ -80,7 +81,7 @@ def test_start_end_slices_input(self, mock_func_run, mock_build_calc):
         # Write a 3-element input file
         input3 = self.tmp_path / "genes3.json"
         input3.write_text(json.dumps([_NACL_GENE] * 3))
-        mock_func_run.return_value = (None, None, None, None)
+        mock_func_run.return_value = (None, None, None, None, None)
         mock_build_calc.return_value = MagicMock()
 
         out_dir = self.tmp_path / "out3"
@@ -102,7 +103,7 @@ def test_start_end_slices_input(self, mock_func_run, mock_build_calc):
     @patch("wyckoff_transformer.cli.relax.build_mace_calculator")
     @patch("wyckoff_transformer.cli.relax.func_run")
     def test_url_model_name_derived_from_stem(self, mock_func_run, mock_build_calc):
-        mock_func_run.return_value = (None, None, None, None)
+        mock_func_run.return_value = (None, None, None, None, None)
         mock_build_calc.return_value = MagicMock()
 
         out_dir = self.tmp_path / "out4"

src/wyckoff_transformer/cryspr/generator.py

@@ -73,7 +73,7 @@ def func_run(
         fix_symmetry: bool = True,
         fmax: float = 0.01,
         optimizer: type[Optimizer] = BFGS,
-) -> tuple[Optional[Atoms], Optional[str], Optional[float], Optional[float]]:
+) -> tuple[Optional[Atoms], Optional[str], Optional[float], Optional[float], Optional[str]]:
     """Generate and relax crystal structures for one Wyckoff gene.
 
     Runs *n_trials* independent PyXtal generation + MACE relaxation cycles.
@@ -92,9 +92,10 @@ def func_run(
         optimizer: ASE local optimisation algorithm class.
 
     Returns:
-        Tuple ``(atoms, formula, energy, energy_per_atom)`` for the
-        lowest-energy successful trial, or ``(None, None, None, None)``
-        when all trials fail.
+        Tuple ``(atoms, formula, energy, energy_per_atom, cif)`` for the
+        lowest-energy successful trial, where *cif* is the text of the
+        symmetrized ``*_2_cell+pos_symmetrized.cif`` file.
+        Returns ``(None, None, None, None, None)`` when all trials fail.
     """
     output_dir = Path(output_dir)
     gene_dir = output_dir / str(id_gene)
@@ -141,7 +142,7 @@ def func_run(
             "[%s-%s] All %d trials failed or produced no structure.",
             model_name, id_gene, n_trials,
         )
-        return None, None, None, None
+        return None, None, None, None, None
 
     lowest_key = min(energy_by_trial, key=energy_by_trial.__getitem__)
 
@@ -162,4 +163,11 @@ def func_run(
     energy = energy_by_trial[lowest_key]
     energy_per_atom = energy / len(atoms)
 
-    return atoms, formula, energy, energy_per_atom
+    # Read the symmetrized final CIF; fall back to the raw cell+pos CIF if absent.
+    lowest_dir = gene_dir / lowest_key
+    cif_candidates = sorted(lowest_dir.glob("*_2_cell+pos_symmetrized.cif"))
+    if not cif_candidates:
+        cif_candidates = sorted(lowest_dir.glob("*_cell+pos.cif"))
+    cif_content: Optional[str] = cif_candidates[0].read_text() if cif_candidates else None
+
+    return atoms, formula, energy, energy_per_atom, cif_content

src/wyckoff_transformer/cryspr/relaxer.py

@@ -1,21 +1,46 @@
 """ASE-based structure relaxation with optional symmetry and cell constraints."""
 import logging
+import os
 from pathlib import Path
 from typing import Optional
 
+os.environ.setdefault("SPGLIB_OLD_ERROR_HANDLING", "0")
+
 from ase import Atoms
 from ase.calculators.calculator import Calculator
 from ase.constraints import FixAtoms, FixSymmetry
 from ase.filters import FrechetCellFilter as CellFilter
 from ase.io import write
 from ase.optimize import BFGS
 from ase.optimize.optimize import Optimizer
-from ase.spacegroup import get_spacegroup
-from pymatgen.io.ase import AseAtomsAdaptor
+import spglib
 
 logger = logging.getLogger(__name__)
 
 
+def _get_spacegroup_info(atoms: Atoms, symprec: float) -> tuple[str, int]:
+    """Return the international symbol and number from spglib."""
+    try:
+        dataset = spglib.get_symmetry_dataset(
+            (atoms.cell.array, atoms.get_scaled_positions(), atoms.numbers),
+            symprec=symprec,
+        )
+    except spglib.SpglibError as exc:
+        logger.warning("Failed to determine symmetry via spglib: %s", exc)
+        return "unknown", 0
+
+    if dataset is None:
+        return "unknown", 0
+
+    symbol = getattr(dataset, "international", None)
+    number = getattr(dataset, "number", None)
+    if symbol is None:
+        symbol = dataset["international"]
+    if number is None:
+        number = dataset["number"]
+    return str(symbol), int(number)
+
+
 def run_ase_relaxer(
         atoms_in: Atoms,
         calculator: Calculator,
@@ -41,8 +66,8 @@ def run_ase_relaxer(
         fix_symmetry: Apply a :class:`~ase.constraints.FixSymmetry` constraint.
         fix_fractional: Fix all atomic positions (ions immobile).
         hydrostatic_strain: Restrict cell filter to isotropic strain only.
-        symprec: Symmetry tolerance in Å used by :func:`~ase.spacegroup.get_spacegroup`
-            and :class:`~ase.constraints.FixSymmetry`.
+        symprec: Symmetry tolerance in Å used by :mod:`spglib` and
+            :class:`~ase.constraints.FixSymmetry`.
         fmax: Force convergence criterion in eV/Å.
         steps_limit: Maximum number of optimisation steps.
         wdir: Directory for the output CIF file.
@@ -58,15 +83,15 @@ def run_ase_relaxer(
 
     if fix_fractional:
         atoms.set_constraint([FixAtoms(indices=list(range(len(atoms))))])
-    spg0 = get_spacegroup(atoms, symprec=symprec)
+    spg0_symbol, spg0_number = _get_spacegroup_info(atoms, symprec=symprec)
     if fix_symmetry:
         atoms.set_constraint([FixSymmetry(atoms, symprec=symprec)])
     target = cell_filter(atoms, hydrostatic_strain=hydrostatic_strain) if cell_filter is not None else atoms
 
     E0 = atoms.get_potential_energy()
     logger.info(
         "Start relaxation: E₀ = %.5f eV, symmetry = %s (%d), fix_sym = %s, relax_cell = %s",
-        E0, spg0.symbol, spg0.no, fix_symmetry, cell_filter is not None,
+        E0, spg0_symbol, spg0_number, fix_symmetry, cell_filter is not None,
     )
 
     log_arg = str(logfile) if logfile is not None else "-"
@@ -81,11 +106,11 @@ def run_ase_relaxer(
     write(filename=str(cif_path), images=atoms, format="cif")
 
     E1 = atoms.get_potential_energy()
-    spg1 = get_spacegroup(atoms, symprec=symprec)
+    spg1_symbol, spg1_number = _get_spacegroup_info(atoms, symprec=symprec)
     cell_diff = (atoms.cell.cellpar() / atoms_in.cell.cellpar() - 1.0) * 100
     logger.info(
         "End relaxation: E₁ = %.5f eV, symmetry = %s (%d), max|F| = %.4f eV/Å",
-        E1, spg1.symbol, spg1.no, abs(atoms.get_forces()).max(),
+        E1, spg1_symbol, spg1_number, abs(atoms.get_forces()).max(),
     )
     logger.debug("Cell diff (%%): %s", cell_diff)
 
@@ -130,10 +155,10 @@ def stepwise_relax(
     full_formula = atoms.get_chemical_formula(mode="metal")
     reduced_formula = atoms.get_chemical_formula(mode="metal", empirical=True)
 
-    structure0 = AseAtomsAdaptor.get_structure(atoms)
-    structure0.to(
+    write(
         filename=str(wdir / f"{reduced_formula}_{full_formula}_0_initial_symmetrized.cif"),
-        symprec=symprec,
+        images=atoms,
+        format="cif",
     )
 
     # Stage 1: fix cell, relax atomic positions
@@ -152,10 +177,10 @@ def stepwise_relax(
         wdir=wdir,
         logfile=logfile1,
     )
-    structure1 = AseAtomsAdaptor.get_structure(atoms1)
-    structure1.to(
+    write(
         filename=str(wdir / f"{reduced_formula}_{full_formula}_1_fix-cell_symmetrized.cif"),
-        symprec=symprec,
+        images=atoms1,
+        format="cif",
     )
 
     # Stage 2: relax both cell and atomic positions
@@ -174,10 +199,10 @@ def stepwise_relax(
         wdir=wdir,
         logfile=logfile2,
     )
-    structure2 = AseAtomsAdaptor.get_structure(atoms2)
-    structure2.to(
+    write(
         filename=str(wdir / f"{reduced_formula}_{full_formula}_2_cell+pos_symmetrized.cif"),
-        symprec=symprec,
+        images=atoms2,
+        format="cif",
     )
 
     return atoms2

src/wyckoff_transformer/cryspr/tests/test_cryspr.py

@@ -155,7 +155,7 @@ def test_returns_none_tuple_when_all_trials_fail(self):
                     output_dir=Path(tmp),
                     n_trials=2,
                 )
-        self.assertEqual(result, (None, None, None, None))
+        self.assertEqual(result, (None, None, None, None, None))
 
 
 # ---------------------------------------------------------------------------
@@ -174,7 +174,7 @@ def setUpClass(cls):
     def test_nacl_relaxation_produces_negative_energy(self):
         from wyckoff_transformer.cryspr.generator import func_run
         with tempfile.TemporaryDirectory() as tmp:
-            atoms, formula, energy, energy_per_atom = func_run(
+            atoms, formula, energy, energy_per_atom, cif = func_run(
                 id_gene=0,
                 wyckoffgene=NACL_GENE,
                 calculator=self.calculator,
@@ -185,3 +185,5 @@ def test_nacl_relaxation_produces_negative_energy(self):
         self.assertIsNotNone(formula)
         self.assertLess(energy, 0.0, "Relaxed NaCl energy should be negative")
         self.assertLess(energy_per_atom, 0.0)
+        self.assertIsNotNone(cif, "CIF content should be returned for a successful relaxation")
+        self.assertIn("_cell_length_a", cif)