From 912fd7583a2dab2fa177d0accd1ff77fc1bf0c85 Mon Sep 17 00:00:00 2001 From: Anthony Onwuli Date: Thu, 27 Aug 2026 18:20:14 +0100 Subject: [PATCH 1/3] feat: Support extxyz as a prediction writer --- src/metatrain/utils/data/writers/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/metatrain/utils/data/writers/__init__.py b/src/metatrain/utils/data/writers/__init__.py index 5dedf6cfc1..36bfc065c2 100644 --- a/src/metatrain/utils/data/writers/__init__.py +++ b/src/metatrain/utils/data/writers/__init__.py @@ -40,6 +40,7 @@ def factory( PREDICTIONS_WRITERS: Dict[str, WriterFactory] = { ".xyz": _make_factory(ASEWriter), + ".extxyz": _make_factory(ASEWriter), ".mts": _make_factory(MetatensorWriter), ".zip": _make_factory(DiskDatasetWriter), } From 17c57aa4113274eb313c8e8729b00930f696b04e Mon Sep 17 00:00:00 2001 From: Anthony Onwuli Date: Thu, 27 Aug 2026 18:20:42 +0100 Subject: [PATCH 2/3] chore: Add tests for writing extxyz as well as for evals --- tests/cli/test_eval_model.py | 15 +++++++++++++++ tests/utils/data/test_writers.py | 10 ++++++++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/tests/cli/test_eval_model.py b/tests/cli/test_eval_model.py index af69122ea6..b52192b0c3 100644 --- a/tests/cli/test_eval_model.py +++ b/tests/cli/test_eval_model.py @@ -223,6 +223,21 @@ def test_eval_no_targets(monkeypatch, tmp_path, model, options): assert Path("output.xyz").is_file() +def test_eval_extxyz_input_and_output(monkeypatch, tmp_path, model, options): + """Test that eval can read systems from and write predictions to extxyz files.""" + monkeypatch.chdir(tmp_path) + + shutil.copy(RESOURCES_PATH / "qm9_reduced_100.xyz", "qm9_reduced_100.extxyz") + options.pop("targets") + options["systems"] = "qm9_reduced_100.extxyz" + + eval_model(model=model, options=options, output="predictions.extxyz") + + frames = read("predictions.extxyz", ":") + assert len(frames) == 100 + assert "energy" in frames[0].info + + @pytest.mark.parametrize("suffix", [".zip", "/"]) def test_eval_no_targets_disallowed_for_dataset_writers( monkeypatch, tmp_path, model, options, suffix diff --git a/tests/utils/data/test_writers.py b/tests/utils/data/test_writers.py index 6a24daa318..e79d2eb552 100644 --- a/tests/utils/data/test_writers.py +++ b/tests/utils/data/test_writers.py @@ -229,7 +229,13 @@ def test_write_xyz_cell(monkeypatch, tmp_path): @pytest.mark.parametrize( "filename", - ("test_output.xyz", "test_output.mts", "test_output.zip", "test_output/"), + ( + "test_output.xyz", + "test_output.extxyz", + "test_output.mts", + "test_output.zip", + "test_output/", + ), ) @pytest.mark.parametrize("fileformat", (None, "same_as_filename")) @pytest.mark.parametrize("cell", (None, torch.eye(3))) @@ -253,7 +259,7 @@ def test_write_predictions(filename, fileformat, cell, monkeypatch, tmp_path): writer.write(systems, predictions) writer.finish() - if filename.endswith(".xyz"): + if filename.endswith((".xyz", ".extxyz")): frames = read(filename, index=":") assert len(frames) == len(systems) for i, frame in enumerate(frames): From 6985afe034da920e218a069ba67d05b43c42675c Mon Sep 17 00:00:00 2001 From: Anthony Onwuli Date: Thu, 27 Aug 2026 18:20:53 +0100 Subject: [PATCH 3/3] chore: Add a test for reading extxyz files --- tests/utils/data/test_readers.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/utils/data/test_readers.py b/tests/utils/data/test_readers.py index db4a71c86c..5b3e76732b 100644 --- a/tests/utils/data/test_readers.py +++ b/tests/utils/data/test_readers.py @@ -16,10 +16,11 @@ @pytest.mark.parametrize("reader", (None, "ase")) -def test_read_systems(reader, monkeypatch, tmp_path): +@pytest.mark.parametrize("suffix", (".xyz", ".extxyz")) +def test_read_systems(reader, suffix, monkeypatch, tmp_path): monkeypatch.chdir(tmp_path) - filename = "systems.xyz" + filename = f"systems.{suffix}" systems = ase_systems() ase.io.write(filename, systems)