feature: Support writing predictions to extxyz files

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),
 }

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

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)
 

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):