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