Add PDoS kwargs option

Author: ElliottKasoar

janus_core/calculations/phonons.py

@@ -85,6 +85,8 @@ class Phonons(BaseCalculation):
         Default is None.
     dos_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to run_total_dos. Default is {}.
+    pdos_kwargs : Optional[dict[str, Any]]
+        Keyword arguments to pass to run_projected_dos. Default is {}.
     temp_min : float
         Start temperature for thermal properties calculations, in K. Default is 0.0.
     temp_max : float
@@ -164,6 +166,7 @@ def __init__(
         n_qpoints: int = 51,
         paths: Optional[PathLike] = None,
         dos_kwargs: Optional[dict[str, Any]] = None,
+        pdos_kwargs: Optional[dict[str, Any]] = None,
         temp_min: float = 0.0,
         temp_max: float = 1000.0,
         temp_step: float = 50.0,
@@ -230,6 +233,8 @@ def __init__(
            structure. Default is None.
         dos_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to run_total_dos. Default is {}.
+        pdos_kwargs : Optional[dict[str, Any]]
+            Keyword arguments to pass to run_projected_dos. Default is {}.
         temp_min : float
             Start temperature for thermal calculations, in K. Default is 0.0.
         temp_max : float
@@ -252,8 +257,16 @@ def __init__(
         enable_progress_bar : bool
             Whether to show a progress bar during phonon calculations. Default is False.
         """
-        (read_kwargs, displacement_kwargs, minimize_kwargs, dos_kwargs) = none_to_dict(
-            (read_kwargs, displacement_kwargs, minimize_kwargs, dos_kwargs)
+        (read_kwargs, displacement_kwargs, minimize_kwargs, dos_kwargs, pdos_kwargs) = (
+            none_to_dict(
+                (
+                    read_kwargs,
+                    displacement_kwargs,
+                    minimize_kwargs,
+                    dos_kwargs,
+                    pdos_kwargs,
+                )
+            )
         )
 
         self.calcs = calcs
@@ -266,6 +279,7 @@ def __init__(
         self.n_qpoints = n_qpoints
         self.paths = paths
         self.dos_kwargs = dos_kwargs
+        self.pdos_kwargs = pdos_kwargs
         self.temp_min = temp_min
         self.temp_max = temp_max
         self.temp_step = temp_step
@@ -778,7 +792,7 @@ def calc_pdos(
         self.results["phonon"].run_mesh(
             mesh, with_eigenvectors=True, is_mesh_symmetry=False
         )
-        self.results["phonon"].run_projected_dos()
+        self.results["phonon"].run_projected_dos(**self.pdos_kwargs)
 
         if self.logger:
             emissions = self.tracker.stop_task().emissions

janus_core/cli/phonons.py

@@ -15,6 +15,7 @@
     LogPath,
     MinimizeKwargs,
     ModelPath,
+    PDoSKwargs,
     ReadKwargsLast,
     StructPath,
     Summary,
@@ -63,7 +64,9 @@ def phonons(
         ),
     ] = None,
     dos: Annotated[bool, Option(help="Whether to calculate the DOS.")] = False,
+    dos_kwargs: DoSKwargs = None,
     pdos: Annotated[bool, Option(help="Whether to calculate the PDOS.")] = False,
+    pdos_kwargs: PDoSKwargs = None,
     thermal: Annotated[
         bool, Option(help="Whether to calculate thermal properties.")
     ] = False,
@@ -89,7 +92,6 @@ def phonons(
         float, Option(help="Maximum force for optimization convergence.")
     ] = 0.1,
     minimize_kwargs: MinimizeKwargs = None,
-    dos_kwargs: DoSKwargs = None,
     hdf5: Annotated[
         bool, Option(help="Whether to save force constants in hdf5.")
     ] = True,
@@ -151,8 +153,12 @@ def phonons(
         Default is None.
     dos : bool
         Whether to calculate and save the DOS. Default is False.
+    dos_kwargs : Optional[dict[str, Any]]
+        Other keyword arguments to pass to run_total_dos. Default is {}.
     pdos : bool
         Whether to calculate and save the PDOS. Default is False.
+    pdos_kwargs : Optional[dict[str, Any]]
+        Other keyword arguments to pass to run_projected_dos. Default is {}.
     thermal : bool
         Whether to calculate thermal properties. Default is False.
     temp_min : float
@@ -173,8 +179,6 @@ def phonons(
         Default is 0.1.
     minimize_kwargs : Optional[dict[str, Any]]
         Other keyword arguments to pass to geometry optimizer. Default is {}.
-    dos_kwargs : Optional[dict[str, Any]]
-        Other keyword arguments to pass to run_total_dos. Default is {}.
     hdf5 : bool
         Whether to save force constants in hdf5 format. Default is True.
     plot_to_file : bool
@@ -226,8 +230,16 @@ def phonons(
         calc_kwargs,
         minimize_kwargs,
         dos_kwargs,
+        pdos_kwargs,
     ) = parse_typer_dicts(
-        [displacement_kwargs, read_kwargs, calc_kwargs, minimize_kwargs, dos_kwargs]
+        [
+            displacement_kwargs,
+            read_kwargs,
+            calc_kwargs,
+            minimize_kwargs,
+            dos_kwargs,
+            pdos_kwargs,
+        ]
     )
 
     # Read only first structure by default and ensure only one image is read
@@ -273,6 +285,7 @@ def phonons(
         "n_qpoints": n_qpoints,
         "paths": paths,
         "dos_kwargs": dos_kwargs,
+        "pdos_kwargs": pdos_kwargs,
         "temp_min": temp_min,
         "temp_max": temp_max,
         "temp_step": temp_step,

janus_core/cli/types.py

@@ -173,6 +173,20 @@ def __str__(self):
     ),
 ]
 
+PDoSKwargs = Annotated[
+    Optional[TyperDict],
+    Option(
+        parser=parse_dict_class,
+        help=(
+            """
+            Keyword arguments to pass to run_projected_dos. Must be passed as a
+            dictionary wrapped in quotes, e.g. "{'key' : value}".
+            """
+        ),
+        metavar="DICT",
+    ),
+]
+
 EnsembleKwargs = Annotated[
     Optional[TyperDict],
     Option(

tests/test_phonons_cli.py

@@ -208,13 +208,19 @@ def test_pdos(tmp_path):
             "--struct",
             DATA_PATH / "NaCl.cif",
             "--pdos",
+            "--pdos-kwargs",
+            "{'freq_min': -1, 'freq_max': 0, 'xyz_projection': True}",
             "--no-hdf5",
             "--file-prefix",
             file_prefix,
         ],
     )
     assert result.exit_code == 0
     assert pdos_results.exists()
+    with open(pdos_results, encoding="utf8") as file:
+        lines = file.readlines()
+    assert lines[1].split()[0] == "-1.0000000000"
+    assert lines[-1].split()[0] == "0.0000000000"
 
 
 def test_plot(tmp_path):