Specify traj_kwargs through minimize_kwargs

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 --write-traj --traj-kwargs "{'filename':'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/cli/geomopt.py

@@ -18,7 +18,6 @@
     ReadKwargsLast,
     StructPath,
     Summary,
-    TrajKwargs,
     WriteKwargs,
 )
 from janus_core.cli.utils import yaml_converter_callback
@@ -48,6 +47,9 @@ def _set_minimize_kwargs(
     if "opt_kwargs" not in minimize_kwargs:
         minimize_kwargs["opt_kwargs"] = {}
 
+    if "traj_kwargs" not in minimize_kwargs:
+        minimize_kwargs["traj_kwargs"] = {}
+
     # Check hydrostatic_strain and scalar pressure not duplicated
     if "filter_kwargs" in minimize_kwargs:
         if "hydrostatic_strain" in minimize_kwargs["filter_kwargs"]:
@@ -130,9 +132,15 @@ def geomopt(
         ),
     ] = None,
     write_traj: Annotated[
-        bool, Option(help="Whether to save a trajectory file of optimization frames.")
+        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,
-    traj_kwargs: TrajKwargs = None,
     read_kwargs: ReadKwargsLast = None,
     calc_kwargs: CalcKwargs = None,
     minimize_kwargs: MinimizeKwargs = None,
@@ -191,9 +199,7 @@ def geomopt(
         converge. Default is inferred from name of structure file.
     write_traj
         Whether to save a trajectory file of optimization frames.
-    traj_kwargs
-        Keyword arguments to pass to ase.io.write when saving trajectory.
-        If "filename" is not included, it is inferred from `file_prefix`.
+        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.
@@ -229,10 +235,8 @@ def geomopt(
     # Check options from configuration file are all valid
     check_config(ctx)
 
-    [read_kwargs, calc_kwargs, minimize_kwargs, write_kwargs, traj_kwargs] = (
-        parse_typer_dicts(
-            [read_kwargs, calc_kwargs, minimize_kwargs, write_kwargs, traj_kwargs]
-        )
+    [read_kwargs, calc_kwargs, minimize_kwargs, write_kwargs] = parse_typer_dicts(
+        [read_kwargs, calc_kwargs, minimize_kwargs, write_kwargs]
     )
 
     # Read only first structure by default and ensure only one image is read
@@ -284,7 +288,6 @@ def geomopt(
         "write_results": True,
         "write_kwargs": write_kwargs,
         "write_traj": write_traj,
-        "traj_kwargs": traj_kwargs,
     }
 
     # Set up geometry optimization

janus_core/cli/types.py

@@ -271,21 +271,6 @@ def __str__(self) -> str:
     ),
 ]
 
-TrajKwargs = Annotated[
-    TyperDict | None,
-    Option(
-        parser=parse_dict_class,
-        help=(
-            """
-            Keyword arguments to pass to ase.io.write when saving trajectory. Must be
-            passed as a dictionary wrapped in quotes, e.g. "{'key': value}".
-            If "filename" is not included, it is inferred from `file_prefix`.
-            """
-        ),
-        metavar="DICT",
-    ),
-]
-
 PostProcessKwargs = Annotated[
     TyperDict | None,
     Option(

tests/test_geomopt_cli.py

@@ -116,8 +116,8 @@ def test_traj(tmp_path):
             "--out",
             results_path,
             "--write-traj",
-            "--traj-kwargs",
-            f"{{'filename':'{traj_path}'}}",
+            "--minimize-kwargs",
+            f"{{'traj_kwargs':{{'filename':'{traj_path}'}}}}",
             "--log",
             log_path,
             "--summary",
@@ -808,8 +808,8 @@ def test_traj_kwargs_no_write(tmp_path):
             DATA_PATH / "NaCl.cif",
             "--file-prefix",
             tmp_path / "NaCl",
-            "--traj-kwargs",
-            f"{{'filename':'{tmp_path / 'traj.extxyz'}'}}",
+            "--minimize-kwargs",
+            f"{{'traj_kwargs':{{'filename':'{tmp_path / 'traj.extxyz'}'}}}}",
         ],
     )
     assert result.exit_code == 1