Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add file prefix options to singlepoint and geomopt #452

Merged
merged 11 commits into from
Feb 28, 2025
2 changes: 1 addition & 1 deletion docs/source/user_guide/command_line.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
17 changes: 6 additions & 11 deletions janus_core/calculations/geom_opt.py
Original file line number Diff line number Diff line change
Expand Up @@ -246,17 +246,12 @@ def __init__(
if "trajectory" not in self.opt_kwargs:
# Overwrite .traj binary with .extxyz by default.
self.opt_kwargs["trajectory"] = str(self.traj_kwargs["filename"])
else:
if self.traj_kwargs:
raise ValueError(
"traj_kwargs given, but trajectory writing not enabled."
)
if "trajectory" in self.opt_kwargs:
raise ValueError(
'"trajectory" given in opt_kwargs,'
"but trajectory writing not enabled."
)

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()

Expand Down
25 changes: 14 additions & 11 deletions janus_core/cli/geomopt.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
ReadKwargsLast,
StructPath,
Summary,
TrajKwargs,
WriteKwargs,
)
from janus_core.cli.utils import yaml_converter_callback
Expand Down Expand Up @@ -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"]:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
15 changes: 0 additions & 15 deletions janus_core/cli/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
22 changes: 12 additions & 10 deletions tests/test_geomopt_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -788,12 +788,14 @@ def test_file_prefix(tmp_path):
],
)
assert result.exit_code == 0
tmp_path_contents = list(tmp_path.glob("*"))
assert len(tmp_path_contents) == 1
assert tmp_path_contents[0].stem == "test"
subdir_contents = list((tmp_path / "test/").glob("*"))
assert len(subdir_contents) == 4
assert all(file.stem[:4] for file in subdir_contents)
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):
Expand All @@ -806,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
Expand Down
Loading