Add minimize kwargs for geomopt (#173)

* Add minimize-kwargs for geomopt * Update README * Update README.md Co-authored-by: Alin Marin Elena <alin@elena.re> --------- Co-authored-by: Alin Marin Elena <alin@elena.re>
stfc · May 31, 2024 · 30e3f46 · 30e3f46
1 parent ee0d76d
commit 30e3f46
Show file tree

Hide file tree

Showing 4 changed files with 87 additions and 28 deletions.
diff --git a/README.md b/README.md
@@ -110,7 +110,15 @@ Additional options may be specified. This shares most options with `singlepoint`
 janus geomopt --struct tests/data/NaCl.cif --arch mace_mp --calc-kwargs "{'model' : 'small'}" --vectors-only --traj 'NaCl-traj.xyz'
 ```
 
-This allows the cell to be optimised, allowing only hydrostatic deformation, and saves the optimization trajector in addition to the final structure and log.
+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.
+
+Further options for the optimizer and filter can be specified using the `--minimize-kwargs` option. For example:
+
+```shell
+janus geomopt --struct tests/data/NaCl.cif --arch mace_mp --calc-kwargs "{'model' : 'small'}" --fully-opt --minimize-kwargs "{'filter_kwargs': {'constant_volume' : True}, 'opt_kwargs': {'alpha': 100}}"
+```
+
+This allows the cell vectors and angles to be optimized, as well as the atomic positions, at constant volume, and sets the `alpha`, the initial guess for the Hessian, to 100 for the optimizer function.
 
 For all options, run `janus geomopt --help`.
 

diff --git a/janus_core/cli/geomopt.py b/janus_core/cli/geomopt.py
@@ -14,7 +14,7 @@
     CalcKwargs,
     Device,
     LogPath,
-    OptKwargs,
+    MinimizeKwargs,
     ReadKwargs,
     StructPath,
     Summary,
@@ -76,7 +76,7 @@ def geomopt(
     ] = None,
     read_kwargs: ReadKwargs = None,
     calc_kwargs: CalcKwargs = None,
-    opt_kwargs: OptKwargs = None,
+    minimize_kwargs: MinimizeKwargs = None,
     write_kwargs: WriteKwargs = None,
     log: LogPath = "geomopt.log",
     summary: Summary = "geomopt_summary.yml",
@@ -117,8 +117,8 @@ def geomopt(
         Keyword arguments to pass to ase.io.read. Default is {}.
     calc_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to the selected calculator. Default is {}.
-    opt_kwargs : Optional[ASEOptArgs]
-        Keyword arguments to pass to optimizer. Default is {}.
+    minimize_kwargs : Optional[dict[str, Any]]
+        Other keyword arguments to pass to geometry optimizer. Default is {}.
     write_kwargs : Optional[ASEWriteArgs]
         Keyword arguments to pass to ase.io.write when saving optimized structure.
         Default is {}.
@@ -133,8 +133,8 @@ def geomopt(
     # Check options from configuration file are all valid
     check_config(ctx)
 
-    [read_kwargs, calc_kwargs, opt_kwargs, write_kwargs] = parse_typer_dicts(
-        [read_kwargs, calc_kwargs, opt_kwargs, write_kwargs]
+    [read_kwargs, calc_kwargs, minimize_kwargs, write_kwargs] = parse_typer_dicts(
+        [read_kwargs, calc_kwargs, minimize_kwargs, write_kwargs]
     )
 
     # Set up single point calculator
@@ -151,24 +151,43 @@ def geomopt(
     if "filename" in write_kwargs:
         raise ValueError("'filename' must be passed through the --out option")
 
-    # Check trajectory path not duplicated
-    if "trajectory" in opt_kwargs:
-        raise ValueError("'trajectory' must be passed through the --traj option")
-
     # Set default filname for writing optimized structure if not specified
     if out:
         write_kwargs["filename"] = out
     else:
         write_kwargs["filename"] = f"{s_point.struct_name}-opt.xyz"
 
+    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
-    opt_kwargs["trajectory"] = traj if traj else None
-    traj_kwargs = {"filename": traj} if traj else None
+    if traj:
+        minimize_kwargs["opt_kwargs"]["trajectory"] = traj
+        minimize_kwargs["traj_kwargs"]["filename"] = traj
+
+    # Check hydrostatic_strain and scalar pressure not duplicated
+    if "filter_kwargs" in minimize_kwargs:
+        if "hydrostatic_strain" in minimize_kwargs["filter_kwargs"]:
+            raise ValueError(
+                "'hydrostatic_strain' must be passed through the --vectors-only option"
+            )
+        if "scalar_pressure" in minimize_kwargs["filter_kwargs"]:
+            raise ValueError(
+                "'scalar_pressure' must be passed through the --pressure option"
+            )
+    else:
+        minimize_kwargs["filter_kwargs"] = {}
 
-    # Set hydrostatic_strain
-    # If not passed --fully-opt or --vectors-only, will be unused
-    filter_kwargs = {"hydrostatic_strain": vectors_only}
-    filter_kwargs["scalar_pressure"] = pressure * units.bar
+    # Set hydrostatic_strain and scalar pressure
+    minimize_kwargs["filter_kwargs"]["hydrostatic_strain"] = vectors_only
+    minimize_kwargs["filter_kwargs"]["scalar_pressure"] = pressure * units.bar
 
     # Use default filter if passed --fully-opt or --vectors-only
     # Otherwise override with None
@@ -179,12 +198,10 @@ def geomopt(
         "struct": s_point.struct,
         "fmax": fmax,
         "steps": steps,
-        "filter_kwargs": filter_kwargs,
         **fully_opt_dict,
-        "opt_kwargs": opt_kwargs,
+        **minimize_kwargs,
         "write_results": True,
         "write_kwargs": write_kwargs,
-        "traj_kwargs": traj_kwargs,
         "log_kwargs": {"filename": log, "filemode": "a"},
     }
 

diff --git a/tests/data/geomopt_config.yml b/tests/data/geomopt_config.yml
@@ -1,4 +1,5 @@
 struct: "NaCl.cif"
 out: "NaCl-results.xyz"
-opt-kwargs:
-  alpha: 100
+minimize-kwargs:
+  opt-kwargs:
+    alpha: 100
diff --git a/tests/test_geomopt_cli.py b/tests/test_geomopt_cli.py
@@ -234,14 +234,16 @@ def test_duplicate_traj(tmp_path):
     log_path = tmp_path / "test.log"
     summary_path = tmp_path / "summary.yml"
 
+    minimize_kwargs = f"{{'opt_kwargs': {{'trajectory' : '{str(traj_path)}'}}}}"
+
     result = runner.invoke(
         app,
         [
             "geomopt",
             "--struct",
             DATA_PATH / "NaCl.cif",
-            "--opt-kwargs",
-            f"{{'trajectory': '{str(traj_path)}'}}",
+            "--minimize-kwargs",
+            minimize_kwargs,
             "--log",
             log_path,
             "--summary",
@@ -260,6 +262,8 @@ def test_restart(tmp_path):
     log_path = tmp_path / "test.log"
     summary_path = tmp_path / "summary.yml"
 
+    minimize_kwargs = f"{{'opt_kwargs': {{'restart': '{str(restart_path)}'}}}}"
+
     result = runner.invoke(
         app,
         [
@@ -268,8 +272,8 @@ def test_restart(tmp_path):
             data_path,
             "--out",
             results_path,
-            "--opt-kwargs",
-            f"{{'restart': '{str(restart_path)}'}}",
+            "--minimize-kwargs",
+            minimize_kwargs,
             "--steps",
             2,
             "--log",
@@ -290,8 +294,8 @@ def test_restart(tmp_path):
             DATA_PATH / "NaCl.cif",
             "--out",
             results_path,
-            "--opt-kwargs",
-            f"{{'restart': '{str(restart_path)}'}}",
+            "--minimize-kwargs",
+            minimize_kwargs,
             "--steps",
             2,
             "--log",
@@ -389,3 +393,32 @@ def test_invalid_config():
     )
     assert result.exit_code == 1
     assert isinstance(result.exception, ValueError)
+
+
+def test_const_volume(tmp_path):
+    """Test setting constant volume with --fully-opt."""
+    results_path = tmp_path / "NaCl-opt.xyz"
+    log_path = tmp_path / "test.log"
+    summary_path = tmp_path / "summary.yml"
+
+    minimize_kwargs = "{'filter_kwargs': {'constant_volume' : True}}"
+
+    result = runner.invoke(
+        app,
+        [
+            "geomopt",
+            "--struct",
+            DATA_PATH / "NaCl-deformed.cif",
+            "--out",
+            results_path,
+            "--fully-opt",
+            "--minimize-kwargs",
+            minimize_kwargs,
+            "--log",
+            log_path,
+            "--summary",
+            summary_path,
+        ],
+    )
+    assert result.exit_code == 0
+    assert_log_contains(log_path, includes=["constant_volume: True"])