Skip to content

Commit

Permalink
Add CLI option for geomopt structure path (#88)
Browse files Browse the repository at this point in the history 
* Add CLI option for geomopt structure and tidy code
* Tidy command help
* Clarify CLI kwargs inputs

---------

Co-authored-by: ElliottKasoar <ElliottKasoar@users.noreply.github.com>
Co-authored-by: Alin Marin Elena <alin@elena.re>
  • Loading branch information
@ElliottKasoar @alinelena
3 people authored Mar 22, 2024
1 parent 15d8fc0 commit b9d05bb
Show file tree
Hide file tree
Showing 3 changed files with 114 additions and 65 deletions.
118 changes: 80 additions & 38 deletions janus_core/cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,17 +90,22 @@ def parse_typer_dicts(typer_dicts: list[TyperDict]) -> list[dict]:

# Shared type aliases
StructPath = Annotated[
Path, typer.Option("--struct", help="Path of structure to simulate")
Path, typer.Option("--struct", help="Path of structure to simulate.")
]
Architecture = Annotated[
str, typer.Option("--arch", help="MLIP architecture to use for calculations")
str, typer.Option("--arch", help="MLIP architecture to use for calculations.")
]
Device = Annotated[str, typer.Option(help="Device to run calculations on")]
Device = Annotated[str, typer.Option(help="Device to run calculations on.")]
ReadKwargs = Annotated[
TyperDict,
typer.Option(
parser=parse_dict_class,
help="Keyword arguments to pass to ase.io.read [default: {}]",
help=(
"""
Keyword arguments to pass to ase.io.read. Must be passed as a dictionary
wrapped in quotes, e.g. "{'key' : value}". [default: "{}"]
"""
),
metavar="DICT",
),
]
Expand All @@ -109,9 +114,11 @@ def parse_typer_dicts(typer_dicts: list[TyperDict]) -> list[dict]:
typer.Option(
parser=parse_dict_class,
help=(
"Keyword arguments to pass to selected calculator. For the default "
"architecture ('mace_mp'), {'model':'small'} is set by default "
"[default: {}]"
"""
Keyword arguments to pass to selected calculator. Must be passed as a
dictionary wrapped in quotes, e.g. "{'key' : value}". For the default
architecture ('mace_mp'), "{'model':'small'}" is set unless overwritten.
"""
),
metavar="DICT",
),
Expand All @@ -121,16 +128,19 @@ def parse_typer_dicts(typer_dicts: list[TyperDict]) -> list[dict]:
typer.Option(
parser=parse_dict_class,
help=(
"Keyword arguments to pass to ase.io.write when saving "
"results [default: {}]"
"""
Keyword arguments to pass to ase.io.write when saving results. Must be
passed as a dictionary wrapped in quotes, e.g. "{'key' : value}".
[default: "{}"]
"""
),
metavar="DICT",
),
]
LogFile = Annotated[Path, typer.Option("--log", help="Path to save logs to")]
LogFile = Annotated[Path, typer.Option("--log", help="Path to save logs to.")]


@app.command()
@app.command(help="Perform single point calculations and save to file.")
def singlepoint(
struct_path: StructPath,
architecture: Architecture = "mace_mp",
Expand Down Expand Up @@ -188,48 +198,65 @@ def singlepoint(
s_point.run(properties=properties, write_results=True, write_kwargs=write_kwargs)


@app.command()
@app.command(
help="Perform geometry optimization and save optimized structure to file.",
)
def geomopt( # pylint: disable=too-many-arguments,too-many-locals
struct_path: StructPath,
fmax: Annotated[
float, typer.Option("--max-force", help="Maximum force for convergence")
float, typer.Option("--max-force", help="Maximum force for convergence.")
] = 0.1,
steps: Annotated[
int, typer.Option("--steps", help="Maximum number of optimization steps")
int, typer.Option("--steps", help="Maximum number of optimization steps.")
] = 1000,
architecture: Architecture = "mace_mp",
device: Device = "cpu",
vectors_only: Annotated[
bool,
typer.Option(
"--vectors-only",
help=("Optimize cell vectors, as well as atomic positions."),
),
] = False,
fully_opt: Annotated[
bool,
typer.Option(
"--fully-opt",
help="Fully optimize the cell vectors, angles, and atomic positions",
help="Fully optimize the cell vectors, angles, and atomic positions.",
),
] = False,
vectors_only: Annotated[
bool,
opt_file: Annotated[
Path,
typer.Option(
"--vectors-only",
"--opt",
help=(
"Disable optimizing cell angles, so only cell vectors and atomic "
"positions are optimized. Requires --fully-opt to be set"
"Path to save optimized structure. Default is inferred from name "
"of structure file."
),
),
] = False,
] = None,
traj_file: Annotated[
str,
typer.Option(
"--traj", help="Path if saving optimization frames. [default: None]"
),
] = None,
read_kwargs: ReadKwargs = None,
calc_kwargs: CalcKwargs = None,
opt_kwargs: Annotated[
TyperDict,
typer.Option(
parser=parse_dict_class,
help=("Keyword arguments to pass to optimizer [default: {}]"),
help=(
"""
Keyword arguments to pass to optimizer. Must be passed as a dictionary
wrapped in quotes, e.g. "{'key' : value}". [default: "{}"]
"""
),
metavar="DICT",
),
] = None,
write_kwargs: WriteKwargs = None,
traj_file: Annotated[
str, typer.Option("--traj", help="Path to save optimization frames to")
] = None,
log_file: LogFile = "geomopt.log",
):
"""
Expand All @@ -249,10 +276,17 @@ def geomopt( # pylint: disable=too-many-arguments,too-many-locals
Default is "mace_mp".
device : Optional[str]
Device to run model on. Default is "cpu".
fully_opt : bool
Whether to optimize the cell as well as atomic positions. Default is False.
vectors_only : bool
Whether to allow only hydrostatic deformations. Default is False.
Whether to optimize cell vectors, as well as atomic positions, by setting
`hydrostatic_strain` in the filter function. Default is False.
fully_opt : bool
Whether to fully optimize the cell vectors, angles, and atomic positions.
Default is False.
opt_file : Optional[Path]
Path to save optimized structure, or last structure if optimization did not
converge. Default is inferred from name of structure file.
traj_file : Optional[str]
Path if saving optimization frames. Default is None.
read_kwargs : Optional[dict[str, Any]]
Keyword arguments to pass to ase.io.read. Default is {}.
calc_kwargs : Optional[dict[str, Any]]
Expand All @@ -262,18 +296,13 @@ def geomopt( # pylint: disable=too-many-arguments,too-many-locals
write_kwargs : Optional[ASEWriteArgs]
Keyword arguments to pass to ase.io.write when saving optimized structure.
Default is {}.
traj_file : Optional[str]
Path to save optimization trajectory to. Default is None.
log_file : Optional[Path]
Path to write logs to. Default is "geomopt.log".
"""
[read_kwargs, calc_kwargs, opt_kwargs, write_kwargs] = parse_typer_dicts(
[read_kwargs, calc_kwargs, opt_kwargs, write_kwargs]
)

if not fully_opt and vectors_only:
raise ValueError("--vectors-only requires --fully-opt to be set")

# Set up single point calculator
s_point = SinglePoint(
struct_path=struct_path,
Expand All @@ -284,18 +313,31 @@ def geomopt( # pylint: disable=too-many-arguments,too-many-locals
log_kwargs={"filename": log_file, "filemode": "w"},
)

# Check trajectory not duplicated
# Check optimized structure path not duplicated
if "filename" in write_kwargs:
raise ValueError("'filename' must be passed through the --opt option")

# Check trajectory path not duplicated
if "trajectory" in opt_kwargs:
raise ValueError("'trajectory' must be passed through the --traj option")

# Set same name to overwrite saved binary with xyz
# Set default filname for writing optimized structure if not specified
if opt_file:
write_kwargs["filename"] = opt_file
else:
write_kwargs["filename"] = f"{s_point.struct_name}-opt.xyz"

# Set same trajectory filenames to overwrite saved binary with xyz
opt_kwargs["trajectory"] = traj_file if traj_file else None
traj_kwargs = {"filename": traj_file} if traj_file else None

filter_kwargs = {"hydrostatic_strain": vectors_only} if fully_opt else None
# Set hydrostatic_strain
# If not passed --fully-opt or --vectors-only, will be unused
filter_kwargs = {"hydrostatic_strain": vectors_only}

# Use default filter if passed --fully-opt, otherwise override with None
fully_opt_dict = {} if fully_opt else {"filter_func": None}
# Use default filter if passed --fully-opt or --vectors-only
# Otherwise override with None
fully_opt_dict = {} if (fully_opt or vectors_only) else {"filter_func": None}

# Run geometry optimization and save output structure
optimize(
Expand Down
4 changes: 4 additions & 0 deletions janus_core/geom_opt.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,10 @@ def optimize( # pylint: disable=too-many-arguments,too-many-locals,too-many-bra
if logger:
logger.info("Using filter %s", filter_func.__name__)
logger.info("Using optimizer %s", optimizer.__name__)
if "hydrostatic_strain" in filter_kwargs:
logger.info(
"hydrostatic_strain: %s", filter_kwargs["hydrostatic_strain"]
)

else:
dyn = optimizer(atoms, **opt_kwargs)
Expand Down
57 changes: 30 additions & 27 deletions tests/test_geomopt_cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ def test_geomopt_help():

def test_geomopt():
"""Test geomopt calculation."""
results_path = Path("./Cl4Na4-opt.xyz").absolute()
results_path = Path("./NaCl-opt.xyz").absolute()
assert not results_path.exists()

result = runner.invoke(
Expand Down Expand Up @@ -53,8 +53,8 @@ def test_geomopt_log(tmp_path, caplog):
"geomopt",
"--struct",
DATA_PATH / "NaCl.cif",
"--write-kwargs",
f"{{'filename': '{str(results_path)}'}}",
"--opt",
results_path,
"--log",
f"{tmp_path}/test.log",
],
Expand All @@ -75,8 +75,8 @@ def test_geomopt_traj(tmp_path):
"geomopt",
"--struct",
DATA_PATH / "NaCl.cif",
"--write-kwargs",
f"{{'filename': '{str(results_path)}'}}",
"--opt",
results_path,
"--traj",
traj_path,
],
Expand All @@ -97,13 +97,14 @@ def test_fully_opt(tmp_path, caplog):
"geomopt",
"--struct",
DATA_PATH / "NaCl-deformed.cif",
"--write-kwargs",
f"{{'filename': '{str(results_path)}'}}",
"--opt",
results_path,
"--fully-opt",
],
)
assert result.exit_code == 0
assert "Using filter" in caplog.text
assert "hydrostatic_strain: False" in caplog.text

atoms = read(results_path)
expected = [5.68834069, 5.68893345, 5.68932555, 89.75938298, 90.0, 90.0]
Expand All @@ -122,36 +123,38 @@ def test_fully_opt_and_vectors(tmp_path, caplog):
DATA_PATH / "NaCl-deformed.cif",
"--fully-opt",
"--vectors-only",
"--write-kwargs",
f"{{'filename': '{str(results_path)}'}}",
"--opt",
results_path,
"--log",
f"{tmp_path}/test.log",
],
)
assert result.exit_code == 0
assert "Using filter" in caplog.text
assert "hydrostatic_strain: True" in caplog.text

atoms = read(results_path)
expected = [5.69139709, 5.69139709, 5.69139709, 89.0, 90.0, 90.0]
assert atoms.cell.cellpar() == pytest.approx(expected)


def test_vectors_not_fully_opt(tmp_path):
def test_vectors_not_fully_opt(tmp_path, caplog):
"""Test passing --vectors-only without --fully-opt."""
results_path = tmp_path / "NaCl-opt.xyz"
result = runner.invoke(
app,
[
"geomopt",
"--struct",
DATA_PATH / "NaCl.cif",
"--write-kwargs",
f"{{'filename': '{str(results_path)}'}}",
"--vectors-only",
],
)
assert result.exit_code == 1
assert isinstance(result.exception, ValueError)
with caplog.at_level("INFO", logger="janus_core.geom_opt"):
result = runner.invoke(
app,
[
"geomopt",
"--struct",
DATA_PATH / "NaCl.cif",
"--opt",
results_path,
"--vectors-only",
],
)
assert result.exit_code == 0
assert "hydrostatic_strain: True" in caplog.text


def duplicate_traj(tmp_path):
Expand Down Expand Up @@ -183,8 +186,8 @@ def test_restart(tmp_path):
"geomopt",
"--struct",
data_path,
"--write-kwargs",
f"{{'filename': '{str(results_path)}'}}",
"--opt",
results_path,
"--opt-kwargs",
f"{{'restart': '{str(restart_path)}'}}",
"--steps",
Expand All @@ -201,8 +204,8 @@ def test_restart(tmp_path):
"geomopt",
"--struct",
DATA_PATH / "NaCl.cif",
"--write-kwargs",
f"{{'filename': '{str(results_path)}'}}",
"--opt",
results_path,
"--opt-kwargs",
f"{{'restart': '{str(restart_path)}'}}",
"--steps",
Expand Down

0 comments on commit b9d05bb

Please sign in to comment.