Add options for cell optimization via CLI

janus_core/cli.py

@@ -168,13 +168,21 @@ def singlepoint(
 
 
 @app.command()
-def geomopt(
+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")
     ] = 0.1,
     architecture: Architecture = "mace_mp",
     device: Device = "cpu",
+    fully_opt: Annotated[
+        bool,
+        typer.Option("--fully-opt", help="Optimise cell and atomic positions"),
+    ] = False,
+    vectors_only: Annotated[
+        bool,
+        typer.Option("--vectors-only", help="Allow only hydrostatic deformations"),
+    ] = False,
     read_kwargs: ReadKwargs = None,
     calc_kwargs: CalcKwargs = None,
     write_kwargs: WriteKwargs = None,
@@ -198,6 +206,10 @@ def geomopt(
         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.
     read_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to ase.io.read. Default is {}.
     calc_kwargs : Optional[dict[str, Any]]
@@ -220,6 +232,10 @@ def geomopt(
     if not isinstance(write_kwargs, dict):
         raise ValueError("write_kwargs must be a dictionary")
 
+    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,
         architecture=architecture,
@@ -231,11 +247,20 @@ def geomopt(
 
     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
+
+    # Use default filter if passed --fully-opt, otherwise override with None
+    fully_opt_dict = {} if fully_opt else {"filter_func": None}
 
+    # Run geometry optimization and save output structure
     optimize(
         s_point.struct,
         fmax=fmax,
+        filter_kwargs=filter_kwargs,
+        **fully_opt_dict,
         opt_kwargs=opt_kwargs,
+        write_results=True,
+        write_kwargs=write_kwargs,
         traj_kwargs=traj_kwargs,
         log_kwargs={"filename": log_file, "filemode": "a"},
     )

tests/test_geomopt_cli.py

@@ -2,6 +2,7 @@
 
 from pathlib import Path
 
+from ase import Atoms
 from ase.io import read
 from typer.testing import CliRunner
 
@@ -20,55 +21,124 @@ def test_geomopt_help():
     assert "Usage: root geomopt [OPTIONS]" in result.stdout
 
 
-def test_geomopt(tmp_path):
+def test_geomopt():
     """Test geomopt calculation."""
-    results_path = tmp_path / "NaCl-results.xyz"
+    results_path = Path(".").absolute() / "Cl4Na4-opt.xyz"
+    assert not results_path.exists()
 
     result = runner.invoke(
         app,
         [
             "geomopt",
             "--struct",
             DATA_PATH / "NaCl.cif",
-            "--write-kwargs",
-            f"{{'filename': '{str(results_path)}'}}",
             "--max-force",
             "0.2",
         ],
     )
     assert result.exit_code == 0
+    atoms = read(results_path)
+    assert isinstance(atoms, Atoms)
+    results_path.unlink()
 
 
 def test_geomopt_log(tmp_path, caplog):
     """Test log correctly written for geomopt."""
+    results_path = tmp_path / "NaCl-opt.xyz"
+
     with caplog.at_level("INFO", logger="janus_core.geom_opt"):
         result = runner.invoke(
             app,
             [
                 "geomopt",
                 "--struct",
                 DATA_PATH / "NaCl.cif",
+                "--write-kwargs",
+                f"{{'filename': '{str(results_path)}'}}",
                 "--log",
                 f"{tmp_path}/test.log",
             ],
         )
         assert result.exit_code == 0
         assert "Starting geometry optimization" in caplog.text
+        assert "Using filter" not in caplog.text
 
 
 def test_geomopt_traj(tmp_path):
     """Test log correctly written for geomopt."""
+    results_path = tmp_path / "NaCl-opt.xyz"
     traj_path = f"{tmp_path}/test.xyz"
+
     result = runner.invoke(
         app,
         [
             "geomopt",
             "--struct",
             DATA_PATH / "NaCl.cif",
+            "--write-kwargs",
+            f"{{'filename': '{str(results_path)}'}}",
             "--traj",
             traj_path,
         ],
     )
     assert result.exit_code == 0
     atoms = read(traj_path)
     assert "forces" in atoms.arrays
+
+
+def test_fully_opt(tmp_path):
+    """Test passing --fully-opt without --vectors-only"""
+    results_path = tmp_path / "NaCl-opt.xyz"
+
+    result = runner.invoke(
+        app,
+        [
+            "geomopt",
+            "--struct",
+            DATA_PATH / "NaCl.cif",
+            "--write-kwargs",
+            f"{{'filename': '{str(results_path)}'}}",
+            "--fully-opt",
+        ],
+    )
+    assert result.exit_code == 0
+
+
+def test_fully_opt_and_vectors(tmp_path, caplog):
+    """Test passing --fully-opt with --vectors-only."""
+    results_path = tmp_path / "NaCl-opt.xyz"
+    with caplog.at_level("INFO", logger="janus_core.geom_opt"):
+        result = runner.invoke(
+            app,
+            [
+                "geomopt",
+                "--struct",
+                DATA_PATH / "NaCl.cif",
+                "--fully-opt",
+                "--vectors-only",
+                "--write-kwargs",
+                f"{{'filename': '{str(results_path)}'}}",
+                "--log",
+                f"{tmp_path}/test.log",
+            ],
+        )
+        assert result.exit_code == 0
+        assert "Using filter" in caplog.text
+
+
+def test_vectors_not_fully_opt(tmp_path):
+    """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)

tests/test_single_point.py

@@ -97,7 +97,7 @@ def test_single_point_write():
     """Test writing singlepoint results."""
     data_path = DATA_PATH / "NaCl.cif"
     results_path = Path(".").absolute() / "NaCl-results.xyz"
-    assert not Path(results_path).exists()
+    assert not results_path.exists()
 
     single_point = SinglePoint(
         struct_path=data_path,