Add file prefix options to singlepoint and geomopt (#452)

* Add file_prefix option to singlepoint

* Add file_prefix option to geomopt

* Rename option to write_traj

* Update user guide

---------

Co-authored-by: Jacob Wilkins <46597752+oerc0122@users.noreply.github.com>

docs/source/user_guide/command_line.rst

@@ -197,7 +197,7 @@ Additional options may be specified. This shares most options with ``singlepoint
 
 .. code-block:: bash
 
-    janus geomopt --struct tests/data/NaCl.cif --arch mace_mp --model-path small --opt-cell-lengths --traj 'NaCl-traj.extxyz'
+    janus geomopt --struct tests/data/NaCl.cif --arch mace_mp --model-path small --opt-cell-lengths --write-traj --minimize-kwargs "{'traj_kwargs':{'filename':'NaCl-traj.extxyz'}}"
 
 
 This allows the cell vectors to be optimised, allowing only hydrostatic deformation, and saves the optimization trajectory in addition to the final structure and log.

janus_core/calculations/geom_opt.py

@@ -58,6 +58,8 @@ class GeomOpt(BaseCalculation):
         Default is True if attach_logger is True, else False.
     tracker_kwargs
         Keyword arguments to pass to `config_tracker`. Default is {}.
+    file_prefix
+        Prefix for output filenames. Default is inferred from structure.
     fmax
         Set force convergence criteria for optimizer in units eV/Å. Default is 0.1.
     steps
@@ -85,9 +87,11 @@ class GeomOpt(BaseCalculation):
     write_kwargs
         Keyword arguments to pass to ase.io.write to save optimized structure.
         Default is {}.
+    write_traj
+        Whether to save a trajectory file of optimization frames.
     traj_kwargs
         Keyword arguments to pass to ase.io.write to save optimization trajectory.
-        Must include "filename" keyword. Default is {}.
+        "filename" keyword is inferred from `file_prefix` if not given. Default is {}.
     """
 
     def __init__(
@@ -103,6 +107,7 @@ def __init__(
         log_kwargs: dict[str, Any] | None = None,
         track_carbon: bool | None = None,
         tracker_kwargs: dict[str, Any] | None = None,
+        file_prefix: PathLike | None = None,
         fmax: float = 0.1,
         steps: int = 1000,
         symmetrize: bool = False,
@@ -114,6 +119,7 @@ def __init__(
         opt_kwargs: ASEOptArgs | None = None,
         write_results: bool = False,
         write_kwargs: OutputKwargs | None = None,
+        write_traj: bool = False,
         traj_kwargs: OutputKwargs | None = None,
     ) -> None:
         """
@@ -146,6 +152,8 @@ def __init__(
             Default is True if attach_logger is True, else False.
         tracker_kwargs
             Keyword arguments to pass to `config_tracker`. Default is {}.
+        file_prefix
+            Prefix for output filenames. Default is inferred from structure.
         fmax
             Set force convergence criteria for optimizer in units eV/Å. Default is 0.1.
         steps
@@ -173,9 +181,12 @@ def __init__(
         write_kwargs
             Keyword arguments to pass to ase.io.write to save optimized structure.
             Default is {}.
+        write_traj
+            Whether to save a trajectory file of optimization frames.
         traj_kwargs
             Keyword arguments to pass to ase.io.write to save optimization trajectory.
-            Must include "filename" keyword. Default is {}.
+            "filename" keyword is inferred from `file_prefix` if not given.
+            Default is {}.
         """
         read_kwargs, filter_kwargs, opt_kwargs, write_kwargs, traj_kwargs = (
             none_to_dict(
@@ -194,17 +205,9 @@ def __init__(
         self.opt_kwargs = opt_kwargs
         self.write_results = write_results
         self.write_kwargs = write_kwargs
+        self.write_traj = write_traj
         self.traj_kwargs = traj_kwargs
 
-        # Validate parameters
-        if self.traj_kwargs and "filename" not in self.traj_kwargs:
-            raise ValueError("'filename' must be included in `traj_kwargs`")
-
-        if self.traj_kwargs and "trajectory" not in self.opt_kwargs:
-            raise ValueError(
-                "'trajectory' must be a key in `opt_kwargs` to save the trajectory."
-            )
-
         # Read last image by default
         read_kwargs.setdefault("index", -1)
 
@@ -223,17 +226,29 @@ def __init__(
             log_kwargs=log_kwargs,
             track_carbon=track_carbon,
             tracker_kwargs=tracker_kwargs,
+            file_prefix=file_prefix,
         )
 
         if not self.struct.calc:
             raise ValueError("Please attach a calculator to `struct`.")
 
-        # Set output file
+        # Set output files
         self.write_kwargs.setdefault("filename", None)
         self.write_kwargs["filename"] = self._build_filename(
             "opt.extxyz", filename=self.write_kwargs["filename"]
         ).absolute()
 
+        if self.write_traj:
+            self.traj_kwargs.setdefault(
+                "filename", self._build_filename("traj.extxyz").absolute()
+            )
+            self.opt_kwargs.setdefault("trajectory", str(self.traj_kwargs["filename"]))
+        elif self.traj_kwargs:
+            raise ValueError("traj_kwargs given, but trajectory writing not enabled.")
+        elif "trajectory" in self.opt_kwargs:
+            raise ValueError(
+                '"trajectory" given in opt_kwargs,but trajectory writing not enabled.'
+            )
         # Configure optimizer dynamics
         self.set_optimizer()
 
@@ -342,7 +357,7 @@ def run(self) -> None:
         )
 
         # Reformat trajectory file from binary
-        if self.traj_kwargs:
+        if self.write_traj:
             traj = read(self.opt_kwargs["trajectory"], index=":")
             output_structs(
                 traj,

janus_core/calculations/single_point.py

@@ -40,6 +40,8 @@ class SinglePoint(BaseCalculation):
         Device to run model on. Default is "cpu".
     model_path
         Path to MLIP model. Default is `None`.
+    file_prefix
+        Prefix for output filenames. Default is inferred from structure.
     read_kwargs
         Keyword arguments to pass to ase.io.read. By default,
         read_kwargs["index"] is ":".
@@ -82,6 +84,7 @@ def __init__(
         arch: Architectures = "mace_mp",
         device: Devices = "cpu",
         model_path: PathLike | None = None,
+        file_prefix: PathLike | None = None,
         read_kwargs: ASEReadArgs | None = None,
         calc_kwargs: dict[str, Any] | None = None,
         set_calc: bool | None = None,
@@ -108,6 +111,8 @@ def __init__(
             Device to run MLIP model on. Default is "cpu".
         model_path
             Path to MLIP model. Default is `None`.
+        file_prefix
+            Prefix for output filenames. Default is inferred from structure.
         read_kwargs
             Keyword arguments to pass to ase.io.read. By default,
             read_kwargs["index"] is ":".
@@ -162,6 +167,7 @@ def __init__(
             log_kwargs=log_kwargs,
             track_carbon=track_carbon,
             tracker_kwargs=tracker_kwargs,
+            file_prefix=file_prefix,
         )
 
         # Properties validated using calculator

janus_core/cli/geomopt.py

@@ -27,7 +27,6 @@
 
 def _set_minimize_kwargs(
     minimize_kwargs: dict[str, Any],
-    traj: str | None,
     opt_cell_lengths: bool,
     pressure: float,
 ) -> None:
@@ -38,29 +37,15 @@ def _set_minimize_kwargs(
     ----------
     minimize_kwargs
         Other keyword arguments to pass to geometry optimizer.
-    traj
-        Path if saving optimization frames.
     opt_cell_lengths
         Whether to optimize cell vectors, as well as atomic positions, by setting
         `hydrostatic_strain` in the filter function.
     pressure
         Scalar pressure when optimizing cell geometry, in GPa. Passed to the filter
         function if either `opt_cell_lengths` or `opt_cell_fully` is True.
     """
-    if "opt_kwargs" in minimize_kwargs:
-        # Check trajectory path not duplicated
-        if "trajectory" in minimize_kwargs["opt_kwargs"]:
-            raise ValueError("'trajectory' must be passed through the --traj option")
-    else:
-        minimize_kwargs["opt_kwargs"] = {}
-
-    if "traj_kwargs" not in minimize_kwargs:
-        minimize_kwargs["traj_kwargs"] = {}
-
-    # Set same trajectory filenames to overwrite saved binary with xyz
-    if traj:
-        minimize_kwargs["opt_kwargs"]["trajectory"] = traj
-        minimize_kwargs["traj_kwargs"]["filename"] = traj
+    minimize_kwargs.setdefault("opt_kwargs", {})
+    minimize_kwargs.setdefault("traj_kwargs", {})
 
     # Check hydrostatic_strain and scalar pressure not duplicated
     if "filter_kwargs" in minimize_kwargs:
@@ -130,6 +115,10 @@ def geomopt(
             help="Atom displacement tolerance for spglib symmetry determination, in Å."
         ),
     ] = 0.001,
+    file_prefix: Annotated[
+        Path | None,
+        Option(help="Prefix for output filenames. Default is inferred from structure."),
+    ] = None,
     out: Annotated[
         Path | None,
         Option(
@@ -139,10 +128,16 @@ def geomopt(
             ),
         ),
     ] = None,
-    traj: Annotated[
-        str,
-        Option(help="Path if saving optimization frames."),
-    ] = None,
+    write_traj: Annotated[
+        bool,
+        Option(
+            help=(
+                "Whether to save a trajectory file of optimization frames. "
+                'If traj_kwargs["filename"] is not specified, it is inferred '
+                "from `file_prefix`."
+            ),
+        ),
+    ] = False,
     read_kwargs: ReadKwargsLast = None,
     calc_kwargs: CalcKwargs = None,
     minimize_kwargs: MinimizeKwargs = None,
@@ -194,11 +189,14 @@ def geomopt(
     symmetry_tolerance
         Atom displacement tolerance for spglib symmetry determination, in Å.
         Default is 0.001.
+    file_prefix
+        Prefix for output filenames. Default is inferred from structure.
     out
         Path to save optimized structure, or last structure if optimization did not
         converge. Default is inferred from name of structure file.
-    traj
-        Path if saving optimization frames. Default is None.
+    write_traj
+        Whether to save a trajectory file of optimization frames.
+        If traj_kwargs["filename"] is not specified, it is inferred from `file_prefix`.
     read_kwargs
         Keyword arguments to pass to ase.io.read. By default,
             read_kwargs["index"] is -1.
@@ -247,7 +245,7 @@ def geomopt(
     if out:
         write_kwargs["filename"] = out
 
-    _set_minimize_kwargs(minimize_kwargs, traj, opt_cell_lengths, pressure)
+    _set_minimize_kwargs(minimize_kwargs, opt_cell_lengths, pressure)
 
     if opt_cell_fully or opt_cell_lengths:
         # Use default filter unless filter function explicitly passed
@@ -281,10 +279,12 @@ def geomopt(
         "steps": steps,
         "symmetrize": symmetrize,
         "symmetry_tolerance": symmetry_tolerance,
+        "file_prefix": file_prefix,
         **opt_cell_fully_dict,
         **minimize_kwargs,
         "write_results": True,
         "write_kwargs": write_kwargs,
+        "write_traj": write_traj,
     }
 
     # Set up geometry optimization

janus_core/cli/singlepoint.py

@@ -42,6 +42,12 @@ def singlepoint(
             ),
         ),
     ] = None,
+    file_prefix: Annotated[
+        Path | None,
+        Option(
+            help=("Prefix for output filenames. Default is inferred from structure.")
+        ),
+    ] = None,
     out: Annotated[
         Path | None,
         Option(
@@ -78,6 +84,8 @@ def singlepoint(
         Path to MLIP model. Default is `None`.
     properties
         Physical properties to calculate. Default is ("energy", "forces", "stress").
+    file_prefix
+        Prefix for output filenames. Default is inferred from structure.
     out
         Path to save structure with calculated results. Default is inferred from name
         of the structure file.
@@ -133,6 +141,7 @@ def singlepoint(
         "write_results": True,
         "arch": arch,
         "device": device,
+        "file_prefix": file_prefix,
         "model_path": model_path,
         "read_kwargs": read_kwargs,
         "calc_kwargs": calc_kwargs,

tests/test_geom_opt.py

@@ -84,7 +84,9 @@ def test_saving_traj(tmp_path):
         calc_kwargs={"model": MODEL_PATH},
     )
     optimizer = GeomOpt(
-        single_point.struct, opt_kwargs={"trajectory": str(tmp_path / "NaCl.traj")}
+        single_point.struct,
+        write_traj=True,
+        opt_kwargs={"trajectory": str(tmp_path / "NaCl.traj")},
     )
     optimizer.run()
     traj = read(tmp_path / "NaCl.traj", index=":")
@@ -105,6 +107,7 @@ def test_traj_reformat(tmp_path):
     optimizer = GeomOpt(
         single_point.struct,
         opt_kwargs={"trajectory": str(traj_path_binary)},
+        write_traj=True,
         traj_kwargs={"filename": traj_path_xyz, "format": "xyz"},
     )
     optimizer.run()

tests/test_geomopt_cli.py

@@ -60,6 +60,7 @@ def test_geomopt():
     finally:
         # Ensure files deleted if command fails
         read_atoms(results_path)
+        results_path.unlink(missing_ok=True)
         log_path.unlink(missing_ok=True)
         summary_path.unlink(missing_ok=True)
         clear_log_handlers()
@@ -114,8 +115,9 @@ def test_traj(tmp_path):
             DATA_PATH / "NaCl.cif",
             "--out",
             results_path,
-            "--traj",
-            traj_path,
+            "--write-traj",
+            "--minimize-kwargs",
+            f"{{'traj_kwargs':{{'filename':'{traj_path}'}}}}",
             "--log",
             log_path,
             "--summary",
@@ -769,3 +771,47 @@ def test_units(tmp_path):
     assert "units" in atoms.info
     for prop, units in expected_units.items():
         assert atoms.info["units"][prop] == units
+
+
+def test_file_prefix(tmp_path):
+    """Test file prefix creates directories and affects all files."""
+    file_prefix = tmp_path / "test/test"
+    result = runner.invoke(
+        app,
+        [
+            "geomopt",
+            "--struct",
+            DATA_PATH / "NaCl.cif",
+            "--file-prefix",
+            file_prefix,
+            "--write-traj",
+        ],
+    )
+    assert result.exit_code == 0
+    test_path = tmp_path / "test"
+    assert list(tmp_path.iterdir()) == [test_path]
+    assert set(test_path.iterdir()) == {
+        test_path / "test-opt.extxyz",
+        test_path / "test-traj.extxyz",
+        test_path / "test-geomopt-summary.yml",
+        test_path / "test-geomopt-log.yml",
+    }
+
+
+def test_traj_kwargs_no_write(tmp_path):
+    """Test traj_kwargs without flag to write traj raises error."""
+    result = runner.invoke(
+        app,
+        [
+            "geomopt",
+            "--struct",
+            DATA_PATH / "NaCl.cif",
+            "--file-prefix",
+            tmp_path / "NaCl",
+            "--minimize-kwargs",
+            f"{{'traj_kwargs':{{'filename':'{tmp_path / 'traj.extxyz'}'}}}}",
+        ],
+    )
+    assert result.exit_code == 1
+    assert isinstance(result.exception, ValueError)
+    assert "trajectory writing not enabled." in result.exception.args[0]