Add mesh CLI option and fix thermal mesh (#307)

* Add mesh CLI option and fix thermal mesh * Convert supercell CLI option to tuple --------- Co-authored-by: Alin Marin Elena <[email protected]> Co-authored-by: Jacob Wilkins <[email protected]>
stfc · Oct 3, 2024 · b6d8430 · b6d8430
1 parent 11ad3b6
commit b6d8430
Show file tree

Hide file tree

Showing 5 changed files with 80 additions and 59 deletions.
diff --git a/docs/source/user_guide/command_line.rst b/docs/source/user_guide/command_line.rst
@@ -314,7 +314,7 @@ Calculate phonons with a 2x2x2 supercell, after geometry optimization (using the
 
 .. code-block:: bash
 
-    janus phonons --struct tests/data/NaCl.cif --supercell 2x2x2 --minimize --arch mace_mp --model-path small
+    janus phonons --struct tests/data/NaCl.cif --supercell 2 2 2 --minimize --arch mace_mp --model-path small
 
 
 This will save the Phonopy parameters, including displacements and force constants, to ``NaCl-phonopy.yml`` and ``NaCl-force_constants.hdf5``,
@@ -324,7 +324,7 @@ Additionally, the ``--bands`` option can be added to calculate the band structur
 
 .. code-block:: bash
 
-    janus phonons --struct tests/data/NaCl.cif --supercell 2x2x2 --minimize --arch mace_mp --model-path small --bands
+    janus phonons --struct tests/data/NaCl.cif --supercell 2 2 2 --minimize --arch mace_mp --model-path small --bands
 
 
 If you need eigenvectors and group velocities written, add the ``--write-full`` option. This will generate a much larger file, but can be used to visualise phonon modes.
@@ -333,7 +333,7 @@ Further calculations, including thermal properties, DOS, and PDOS, can also be c
 
 .. code-block:: bash
 
-    janus phonons --struct tests/data/NaCl.cif --supercell 2x3x4 --dos --pdos --thermal --temp-start 0 --temp-end 300 --temp-step 50
+    janus phonons --struct tests/data/NaCl.cif --supercell 2 3 4 --dos --pdos --thermal --temp-start 0 --temp-end 300 --temp-step 50
 
 
 This will create additional output files: ``NaCl-thermal.dat`` for the thermal properties (heat capacity, entropy, and free energy)

diff --git a/janus_core/calculations/phonons.py b/janus_core/calculations/phonons.py
@@ -60,6 +60,8 @@ class Phonons(BaseCalculation):
         Size of supercell for calculation. Default is 2.
     displacement : float
         Displacement for force constants calculation, in A. Default is 0.01.
+    mesh : tuple[int, int, int]
+        Mesh for sampling. Default is (10, 10, 10).
     symmetrize : bool
         Whether to symmetrize force constants after calculation.
         Default is False.
@@ -106,7 +108,7 @@ class Phonons(BaseCalculation):
         Calculate band structure and optionally write and plot results.
     write_bands(bands_file, save_plots, plot_file)
         Write results of band structure calculations.
-    calc_thermal_props(write_thermal)
+    calc_thermal_props(mesh, write_thermal)
         Calculate thermal properties and optionally write results.
     write_thermal_props(thermal_file)
         Write results of thermal properties calculations.
@@ -138,6 +140,7 @@ def __init__(
         calcs: MaybeSequence[PhononCalcs] = (),
         supercell: MaybeList[int] = 2,
         displacement: float = 0.01,
+        mesh: tuple[int, int, int] = (10, 10, 10),
         symmetrize: bool = False,
         minimize: bool = False,
         minimize_kwargs: Optional[dict[str, Any]] = None,
@@ -186,6 +189,8 @@ def __init__(
             Size of supercell for calculation. Default is 2.
         displacement : float
             Displacement for force constants calculation, in A. Default is 0.01.
+        mesh : tuple[int, int, int]
+            Mesh for sampling. Default is (10, 10, 10).
         symmetrize : bool
             Whether to symmetrize force constants after calculations.
             Default is False.
@@ -219,6 +224,7 @@ def __init__(
 
         self.calcs = calcs
         self.displacement = displacement
+        self.mesh = mesh
         self.symmetrize = symmetrize
         self.minimize = minimize
         self.minimize_kwargs = minimize_kwargs
@@ -490,13 +496,18 @@ def write_bands(
             bplt.savefig(plot_file)
 
     def calc_thermal_props(
-        self, write_thermal: Optional[bool] = None, **kwargs
+        self,
+        mesh: Optional[tuple[int, int, int]] = None,
+        write_thermal: Optional[bool] = None,
+        **kwargs,
     ) -> None:
         """
         Calculate thermal properties and optionally write results.
 
         Parameters
         ----------
+        mesh : Optional[tuple[int, int, int]]
+            Mesh for sampling. Default is self.mesh.
         write_thermal : Optional[bool]
             Whether to write out thermal properties to file. Default is
             self.write_results.
@@ -506,6 +517,9 @@ def calc_thermal_props(
         if write_thermal is None:
             write_thermal = self.write_results
 
+        if mesh is None:
+            mesh = self.mesh
+
         # Calculate phonons if not already in results
         if "phonon" not in self.results:
             # Use general (self.write_results) setting for writing force constants
@@ -515,7 +529,7 @@ def calc_thermal_props(
             self.logger.info("Starting thermal properties calculation")
             self.tracker.start_task("Thermal calculation")
 
-        self.results["phonon"].run_mesh()
+        self.results["phonon"].run_mesh(mesh)
         self.results["phonon"].run_thermal_properties(
             t_step=self.temp_step, t_max=self.temp_max, t_min=self.temp_min
         )
@@ -563,7 +577,7 @@ def write_thermal_props(self, thermal_file: Optional[PathLike] = None) -> None:
     def calc_dos(
         self,
         *,
-        mesh: MaybeList[float] = (10, 10, 10),
+        mesh: Optional[tuple[int, int, int]] = None,
         write_dos: Optional[bool] = None,
         **kwargs,
     ) -> None:
@@ -572,8 +586,8 @@ def calc_dos(
 
         Parameters
         ----------
-        mesh : MaybeList[float]
-            Mesh for sampling. Default is (10, 10, 10).
+        mesh : Optional[tuple[int, int, int]]
+            Mesh for sampling. Default is self.mesh.
         write_dos : Optional[bool]
             Whether to write out results to file. Default is True.
         **kwargs
@@ -582,6 +596,9 @@ def calc_dos(
         if write_dos is None:
             write_dos = self.write_results
 
+        if mesh is None:
+            mesh = self.mesh
+
         # Calculate phonons if not already in results
         if "phonon" not in self.results:
             # Use general (self.write_results) setting for writing force constants
@@ -665,7 +682,7 @@ def write_dos(
     def calc_pdos(
         self,
         *,
-        mesh: MaybeList[float] = (10, 10, 10),
+        mesh: Optional[tuple[int, int, int]] = None,
         write_pdos: Optional[bool] = None,
         **kwargs,
     ) -> None:
@@ -674,8 +691,8 @@ def calc_pdos(
 
         Parameters
         ----------
-        mesh : MaybeList[float]
-            Mesh for sampling. Default is (10, 10, 10).
+        mesh : Optional[tuple[int, int, int]]
+            Mesh for sampling. Default is self.mesh.
         write_pdos : Optional[bool]
             Whether to write out results to file. Default is self.write_results.
         **kwargs
@@ -684,6 +701,9 @@ def calc_pdos(
         if write_pdos is None:
             write_pdos = self.write_results
 
+        if mesh is None:
+            mesh = self.mesh
+
         # Calculate phonons if not already in results
         if "phonon" not in self.results:
             # Use general (self.write_results) setting for writing force constants

diff --git a/janus_core/cli/phonons.py b/janus_core/cli/phonons.py
@@ -21,6 +21,7 @@
 from janus_core.cli.utils import (
     carbon_summary,
     check_config,
+    dict_tuples_to_lists,
     end_summary,
     parse_typer_dicts,
     save_struct_calc,
@@ -39,30 +40,22 @@ def phonons(
     ctx: Context,
     struct: StructPath,
     supercell: Annotated[
-        str,
-        Option(help="Supercell lattice vectors in the form '1x2x3'."),
-    ] = "2x2x2",
+        tuple[int, int, int], Option(help="Supercell lattice vectors.")
+    ] = (2, 2, 2),
     displacement: Annotated[
-        float,
-        Option(help="Displacement for force constants calculation, in A."),
+        float, Option(help="Displacement for force constants calculation, in A.")
     ] = 0.01,
+    mesh: Annotated[
+        tuple[int, int, int], Option(help="Mesh numbers along a, b, c axes.")
+    ] = (10, 10, 10),
     bands: Annotated[
         bool,
         Option(help="Whether to compute band structure."),
     ] = False,
-    dos: Annotated[
-        bool,
-        Option(help="Whether to calculate the DOS."),
-    ] = False,
-    pdos: Annotated[
-        bool,
-        Option(
-            help="Whether to calculate the PDOS.",
-        ),
-    ] = False,
+    dos: Annotated[bool, Option(help="Whether to calculate the DOS.")] = False,
+    pdos: Annotated[bool, Option(help="Whether to calculate the PDOS.")] = False,
     thermal: Annotated[
-        bool,
-        Option(help="Whether to calculate thermal properties."),
+        bool, Option(help="Whether to calculate thermal properties.")
     ] = False,
     temp_min: Annotated[
         float,
@@ -80,18 +73,14 @@ def phonons(
         bool, Option(help="Whether to symmetrize force constants.")
     ] = False,
     minimize: Annotated[
-        bool,
-        Option(
-            help="Whether to minimize structure before calculations.",
-        ),
+        bool, Option(help="Whether to minimize structure before calculations.")
     ] = False,
     fmax: Annotated[
         float, Option(help="Maximum force for optimization convergence.")
     ] = 0.1,
     minimize_kwargs: MinimizeKwargs = None,
     hdf5: Annotated[
-        bool,
-        Option(help="Whether to save force constants in hdf5."),
+        bool, Option(help="Whether to save force constants in hdf5.")
     ] = True,
     plot_to_file: Annotated[
         bool,
@@ -133,11 +122,12 @@ def phonons(
         Typer (Click) Context. Automatically set.
     struct : Path
         Path of structure to simulate.
-    supercell : str
-        Supercell lattice vectors. Must be passed in the form '1x2x3'. Default is
-        2x2x2.
+    supercell : tuple[int, int, int]
+        Supercell lattice vectors. Default is (2, 2, 2).
     displacement : float
         Displacement for force constants calculation, in A. Default is 0.01.
+    mesh : tuple[int, int, int]
+        Mesh for sampling. Default is (10, 10, 10).
     bands : bool
         Whether to calculate and save the band structure. Default is False.
     dos : bool
@@ -209,17 +199,6 @@ def phonons(
         raise ValueError("'fmax' must be passed through the --fmax option")
     minimize_kwargs["fmax"] = fmax
 
-    try:
-        supercell = [int(x) for x in supercell.split("x")]
-    except ValueError as exc:
-        raise ValueError(
-            "Please pass lattice vectors as integers in the form 1x2x3"
-        ) from exc
-
-    # Validate supercell list
-    if len(supercell) != 3:
-        raise ValueError("Please pass three lattice vectors in the form 1x2x3")
-
     calcs = []
     if bands:
         calcs.append("bands")
@@ -247,6 +226,7 @@ def phonons(
         "calcs": calcs,
         "supercell": supercell,
         "displacement": displacement,
+        "mesh": mesh,
         "symmetrize": symmetrize,
         "minimize": minimize,
         "minimize_kwargs": minimize_kwargs,
@@ -283,6 +263,9 @@ def phonons(
         log=log,
     )
 
+    # Convert all tuples to list in inputs nested dictionary
+    dict_tuples_to_lists(inputs)
+
     # Save summary information before calculations begin
     start_summary(command="phonons", summary=summary, inputs=inputs)
 

diff --git a/janus_core/cli/utils.py b/janus_core/cli/utils.py
@@ -36,6 +36,22 @@ def dict_paths_to_strs(dictionary: dict) -> None:
             dictionary[key] = str(value)
 
 
+def dict_tuples_to_lists(dictionary: dict) -> None:
+    """
+    Recursively iterate over dictionary, converting tuple values to lists.
+
+    Parameters
+    ----------
+    dictionary : dict
+        Dictionary to be converted.
+    """
+    for key, value in dictionary.items():
+        if isinstance(value, dict):
+            dict_paths_to_strs(value)
+        elif isinstance(value, tuple):
+            dictionary[key] = list(value)
+
+
 def dict_remove_hyphens(dictionary: dict) -> dict:
     """
     Recursively iterate over dictionary, replacing hyphens with underscores in keys.

diff --git a/tests/test_phonons_cli.py b/tests/test_phonons_cli.py
@@ -228,7 +228,9 @@ def test_plot(tmp_path):
             "--struct",
             DATA_PATH / "NaCl.cif",
             "--supercell",
-            "1x1x1",
+            1,
+            1,
+            1,
             "--pdos",
             "--dos",
             "--bands",
@@ -268,7 +270,9 @@ def test_supercell(tmp_path):
             "--struct",
             DATA_PATH / "NaCl.cif",
             "--supercell",
-            "1x2x3",
+            1,
+            2,
+            3,
             "--no-hdf5",
             "--file-prefix",
             file_prefix,
@@ -285,10 +289,7 @@ def test_supercell(tmp_path):
     assert params["supercell_matrix"] == [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
 
 
-test_data = ["2", "2.1x2.1x2.1", "2x2xa"]
-
-
-@pytest.mark.parametrize("supercell", test_data)
+@pytest.mark.parametrize("supercell", [(2,), (2, 2), (2, 2, "a"), ("2x2x2",)])
 def test_invalid_supercell(supercell, tmp_path):
     """Test errors are raise for invalid supercells."""
     file_prefix = tmp_path / "test"
@@ -300,13 +301,12 @@ def test_invalid_supercell(supercell, tmp_path):
             "--struct",
             DATA_PATH / "NaCl.cif",
             "--supercell",
-            supercell,
+            *supercell,
             "--file-prefix",
             file_prefix,
         ],
     )
-    assert result.exit_code == 1
-    assert isinstance(result.exception, ValueError)
+    assert result.exit_code == 1 or result.exit_code == 2
 
 
 def test_minimize_kwargs(tmp_path):
@@ -379,7 +379,9 @@ def test_valid_traj_input(read_kwargs, tmp_path):
             "--struct",
             DATA_PATH / "NaCl-traj.xyz",
             "--supercell",
-            "1x1x1",
+            1,
+            1,
+            1,
             "--read-kwargs",
             read_kwargs,
             "--no-hdf5",