From 80882e84286547bc3bc380c804c18116419fc47e Mon Sep 17 00:00:00 2001
From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
Date: Wed, 26 Feb 2025 18:03:04 +0000
Subject: [PATCH] Fix writing xyz files (#449)

* Test writing xyz file

* Fix writing xyz files
---
 janus_core/helpers/struct_io.py |  2 +-
 tests/test_geom_opt.py          | 16 ++++++++++++++--
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/janus_core/helpers/struct_io.py b/janus_core/helpers/struct_io.py
index 2c4eba03..8eea7ba5 100644
--- a/janus_core/helpers/struct_io.py
+++ b/janus_core/helpers/struct_io.py
@@ -280,7 +280,7 @@ def output_structs(
         )
 
         # write_results is only a valid kwarg for extxyz
-        if write_format in ("extxyz"):
+        if write_format == "extxyz":
             write_kwargs.setdefault("write_results", not invalidate_calc)
         else:
             write_kwargs.pop("write_results", None)
diff --git a/tests/test_geom_opt.py b/tests/test_geom_opt.py
index cd73d862..b4a57dec 100644
--- a/tests/test_geom_opt.py
+++ b/tests/test_geom_opt.py
@@ -100,17 +100,18 @@ def test_traj_reformat(tmp_path):
     )
 
     traj_path_binary = tmp_path / "NaCl.traj"
-    traj_path_xyz = tmp_path / "NaCl-traj.extxyz"
+    traj_path_xyz = tmp_path / "NaCl-traj.xyz"
 
     optimizer = GeomOpt(
         single_point.struct,
         opt_kwargs={"trajectory": str(traj_path_binary)},
-        traj_kwargs={"filename": traj_path_xyz},
+        traj_kwargs={"filename": traj_path_xyz, "format": "xyz"},
     )
     optimizer.run()
     traj = read(traj_path_xyz, index=":")
 
     assert len(traj) == 3
+    assert traj[0].info == {}
 
 
 def test_missing_traj_kwarg(tmp_path):
@@ -357,3 +358,14 @@ def test_logging(tmp_path):
     assert log_file.exists()
     assert "emissions" in single_point.struct.info
     assert single_point.struct.info["emissions"] > 0
+
+
+def test_write_xyz(tmp_path):
+    """Test writing a non-extended xyz file."""
+    optimizer = GeomOpt(
+        struct=DATA_PATH / "NaCl-deformed.cif",
+        arch="mace_mp",
+        calc_kwargs={"model": MODEL_PATH},
+        fmax=0.1,
+    )
+    optimizer.run()