From 54286e16f56edd3b51594be742d1ae6c58e2c60c Mon Sep 17 00:00:00 2001 From: k-harris27 <120191386+k-harris27@users.noreply.github.com> Date: Fri, 28 Feb 2025 16:26:46 +0000 Subject: [PATCH] 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 | 2 +- janus_core/calculations/geom_opt.py | 41 +++++++++++++------- janus_core/calculations/single_point.py | 6 +++ janus_core/cli/geomopt.py | 48 ++++++++++++------------ janus_core/cli/singlepoint.py | 9 +++++ tests/test_geom_opt.py | 5 ++- tests/test_geomopt_cli.py | 50 ++++++++++++++++++++++++- 7 files changed, 120 insertions(+), 41 deletions(-) diff --git a/docs/source/user_guide/command_line.rst b/docs/source/user_guide/command_line.rst index 01e28f21..90157bf3 100644 --- a/docs/source/user_guide/command_line.rst +++ b/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. diff --git a/janus_core/calculations/geom_opt.py b/janus_core/calculations/geom_opt.py index a8828a4d..6db16e80 100644 --- a/janus_core/calculations/geom_opt.py +++ b/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, diff --git a/janus_core/calculations/single_point.py b/janus_core/calculations/single_point.py index 45c35836..0474d731 100644 --- a/janus_core/calculations/single_point.py +++ b/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 diff --git a/janus_core/cli/geomopt.py b/janus_core/cli/geomopt.py index fa7352b8..4adb89bd 100644 --- a/janus_core/cli/geomopt.py +++ b/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,8 +37,6 @@ 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. @@ -47,20 +44,8 @@ def _set_minimize_kwargs( 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 diff --git a/janus_core/cli/singlepoint.py b/janus_core/cli/singlepoint.py index 87659c6d..4ee457e5 100644 --- a/janus_core/cli/singlepoint.py +++ b/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, diff --git a/tests/test_geom_opt.py b/tests/test_geom_opt.py index b4a57dec..53753232 100644 --- a/tests/test_geom_opt.py +++ b/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() diff --git a/tests/test_geomopt_cli.py b/tests/test_geomopt_cli.py index 69d9af27..4f17885b 100644 --- a/tests/test_geomopt_cli.py +++ b/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]