Enable more refined supercell settings and adapt DFT settings in phonon workflow

.gitignore

@@ -52,6 +52,7 @@ coverage.xml
 .pytest_cache/
 cover/
 job*/
+submit*
 
 # Translations
 *.mo

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ default_language_version:
   python: python3
 repos:
 - repo: https://github.com/charliermarsh/ruff-pre-commit
-  rev: v0.1.1
+  rev: v0.6.9
   hooks:
   - id: ruff
     args: [--fix]
@@ -22,11 +22,11 @@ repos:
   - id: end-of-file-fixer
   - id: trailing-whitespace
 - repo: https://github.com/psf/black
-  rev: 23.10.0
+  rev: 24.10.0
   hooks:
   - id: black
 - repo: https://github.com/asottile/blacken-docs
-  rev: 1.16.0
+  rev: 1.19.0
   hooks:
   - id: blacken-docs
     additional_dependencies: [black]
@@ -39,7 +39,7 @@ repos:
   - id: rst-directive-colons
   - id: rst-inline-touching-normal
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v1.6.1
+  rev: v1.11.2
   hooks:
   - id: mypy
     files: ^autoplex/.*\.py$
@@ -48,7 +48,7 @@ repos:
     - types-pkg_resources==0.1.2
     - types-paramiko
 - repo: https://github.com/codespell-project/codespell
-  rev: v2.2.6
+  rev: v2.3.0
   hooks:
   - id: codespell
     stages: [commit, commit-msg]

autoplex/auto/phonons/flows.py

@@ -5,10 +5,11 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
+from autoplex.data.phonons.flows import TightDFTStaticMaker
+
 if TYPE_CHECKING:
     from atomate2.common.schemas.phonons import PhononBSDOSDoc
     from atomate2.vasp.jobs.base import BaseVaspMaker
-    from emmet.core.math import Matrix3D
     from pymatgen.core.structure import Structure
 
 from jobflow import Flow, Maker
@@ -17,12 +18,14 @@
     complete_benchmark,
     dft_phonopy_gen_data,
     dft_random_gen_data,
+    generate_supercells,
     get_iso_atom,
+    run_supercells,
 )
 from autoplex.benchmark.phonons.jobs import write_benchmark_metrics
 from autoplex.fitting.common.flows import MLIPFitMaker
 
-__all__ = ["CompleteDFTvsMLBenchmarkWorkflow"]
+__all__ = ["CompleteDFTvsMLBenchmarkWorkflow", "DFTSupercellSettingsMaker"]
 
 
 # Volker's idea: provide several default flows with different setting/setups
@@ -61,8 +64,6 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker):
         Default=10.
     displacements: list[float]
         displacement distance for phonons
-    min_length: float
-        min length of the supercell that will be built
     symprec: float
         Symmetry precision to use in the
         reduction of symmetry to find the primitive/conventional cell
@@ -71,8 +72,6 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker):
     uc: bool.
         If True, will generate randomly distorted structures (unitcells)
         and add static computation jobs to the flow.
-    supercell_matrix: Matrix3D or None
-        The matrix to construct the supercell.
     distort_type : int.
         0- volume distortion, 1- angle distortion, 2- volume and angle distortion. Default=0.
     volume_scale_factor_range : list[float]
@@ -117,8 +116,8 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker):
         List of SOAP delta values that are checked.
     n_sparse_list: list
         List of GAP n_sparse values that are checked.
-    adaptive_supercell_settings: bool
-        prevent too big rattled supercells or too tight phonopy supercell settings.
+    supercell_settings: dict
+        settings for supercell generation
     benchmark_kwargs: dict
         kwargs for the benchmark flows
     """
@@ -130,13 +129,11 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker):
     phonon_displacement_maker: BaseVaspMaker = None
     n_structures: int = 10
     displacements: list[float] = field(default_factory=lambda: [0.01])
-    min_length: int = 20
     symprec: float = 1e-4
     uc: bool = False
     volume_custom_scale_factors: list[float] | None = None
     volume_scale_factor_range: list[float] | None = None
     rattle_std: float = 0.01
-    supercell_matrix: Matrix3D | None = None
     distort_type: int = 0
     min_distance: float = 1.5
     angle_percentage_scale: float = 10
@@ -150,7 +147,7 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker):
     atomwise_regularization_list: list | None = None
     soap_delta_list: list | None = None
     n_sparse_list: list | None = None
-    adaptive_supercell_settings: bool = True
+    supercell_settings: dict = field(default_factory=lambda: {"min_length": 15})
     benchmark_kwargs: dict = field(default_factory=dict)
 
     def make(
@@ -225,7 +222,6 @@ def make(
                     volume_custom_scale_factors=self.volume_custom_scale_factors,
                     volume_scale_factor_range=self.volume_scale_factor_range,
                     rattle_std=self.rattle_std,
-                    supercell_matrix=self.supercell_matrix,
                     distort_type=self.distort_type,
                     min_distance=self.min_distance,
                     rattle_type=self.rattle_type,
@@ -234,23 +230,30 @@ def make(
                     angle_max_attempts=self.angle_max_attempts,
                     angle_percentage_scale=self.angle_percentage_scale,
                     w_angle=self.w_angle,
-                    adaptive_rattled_supercell_settings=self.adaptive_supercell_settings,
+                    supercell_settings=self.supercell_settings,
                 )
                 flows.append(addDFTrand)
                 fit_input.update({mp_id: addDFTrand.output})
             if self.add_dft_phonon_struct:
                 addDFTphon = self.add_dft_phonons(
-                    structure,
-                    self.displacements,
-                    self.symprec,
-                    self.phonon_displacement_maker,
-                    self.min_length,
-                    self.adaptive_supercell_settings,
+                    structure=structure,
+                    displacements=self.displacements,
+                    symprec=self.symprec,
+                    phonon_displacement_maker=self.phonon_displacement_maker,
+                    supercell_settings=self.supercell_settings,
                 )
                 flows.append(addDFTphon)
                 fit_input.update({mp_id: addDFTphon.output})
             if self.add_dft_random_struct and self.add_dft_phonon_struct:
-                fit_input.update({mp_id: {**addDFTrand.output, **addDFTphon.output}})
+                fit_input.update(
+                    {
+                        mp_id: {
+                            "rand_struc_dir": addDFTrand.output["rand_struc_dir"],
+                            "phonon_dir": addDFTphon.output["phonon_dir"],
+                            "phonon_data": addDFTphon.output["phonon_data"],
+                        }
+                    }
+                )
             if self.add_rss_struct:
                 raise NotImplementedError
 
@@ -281,7 +284,7 @@ def make(
                 **fit_kwargs,
             )
             flows.append(add_data_fit)
-            if "regularization" in fit_kwargs and fit_kwargs["regularization"]:
+            if fit_kwargs.get("regularization"):
                 hyper_list.append(
                     {
                         "f="
@@ -291,7 +294,7 @@ def make(
             hyper_list.append(
                 {"f=" + str(atomwise_regularization_parameter): "default"}
             )
-            if "separated" in fit_kwargs and fit_kwargs["separated"]:
+            if fit_kwargs.get("separated"):
                 hyper_list.append(
                     {"f=" + str(atomwise_regularization_parameter): "default phonon"}
                 )
@@ -306,7 +309,6 @@ def make(
                     complete_bm = complete_benchmark(
                         ibenchmark_structure=ibenchmark_structure,
                         benchmark_structure=benchmark_structure,
-                        min_length=self.min_length,
                         ml_model=ml_model,
                         ml_path=add_data_fit.output["mlip_path"],
                         mp_ids=mp_ids,
@@ -316,6 +318,7 @@ def make(
                         symprec=self.symprec,
                         phonon_displacement_maker=self.phonon_displacement_maker,
                         dft_references=dft_references,
+                        supercell_settings=self.supercell_settings,
                         **self.benchmark_kwargs,
                     )
                     flows.append(complete_bm)
@@ -373,7 +376,6 @@ def make(
                                     complete_bm = complete_benchmark(
                                         ibenchmark_structure=ibenchmark_structure,
                                         benchmark_structure=benchmark_structure,
-                                        min_length=self.min_length,
                                         ml_model=ml_model,
                                         ml_path=loop_data_fit.output["mlip_path"],
                                         mp_ids=mp_ids,
@@ -383,6 +385,7 @@ def make(
                                         symprec=self.symprec,
                                         phonon_displacement_maker=self.phonon_displacement_maker,
                                         dft_references=dft_references,
+                                        supercell_settings=self.supercell_settings,
                                         **self.benchmark_kwargs,
                                     )
                                     flows.append(complete_bm)
@@ -399,14 +402,13 @@ def make(
 
         return Flow(jobs=flows, output=collect_bm, name=self.name)
 
+    @staticmethod
     def add_dft_phonons(
-        self,
         structure: Structure,
         displacements: list[float],
         symprec: float,
         phonon_displacement_maker: BaseVaspMaker,
-        min_length: float,
-        adaptive_phonopy_supercell_settings: bool = True,
+        supercell_settings: dict,
     ):
         """Add DFT phonon runs for reference structures.
 
@@ -423,39 +425,30 @@ def add_dft_phonons(
             and to handle all symmetry-related tasks in phonopy
         phonon_displacement_maker: BaseVaspMaker
             Maker used to compute the forces for a supercell.
-        min_length: float
-             min length of the supercell that will be built
-        adaptive_phonopy_supercell_settings: bool
-            prevent too tight phonopy supercell settings
-        """
-        additonal_dft_phonon = dft_phonopy_gen_data(
-            structure,
-            displacements,
-            symprec,
-            phonon_displacement_maker,
-            min_length,
-            adaptive_phonopy_supercell_settings,
-        )
+        supercell_settings: dict
+            supercell settings
 
-        return Flow(
-            jobs=additonal_dft_phonon,  # flows
-            output={
-                "phonon_dir": additonal_dft_phonon.output["dirs"],
-                "phonon_data": additonal_dft_phonon.output["data"],
-            },
-            name=self.name,
+        """
+        dft_phonons = dft_phonopy_gen_data(
+            structure=structure,
+            displacements=displacements,
+            symprec=symprec,
+            phonon_displacement_maker=phonon_displacement_maker,
+            supercell_settings=supercell_settings,
         )
+        # let's append a name
+        dft_phonons.name = "single-atom displaced supercells"
+        return dft_phonons
 
+    @staticmethod
     def add_dft_random(
-        self,
         structure: Structure,
         mp_id: str,
         phonon_displacement_maker: BaseVaspMaker,
         uc: bool = False,
         volume_custom_scale_factors: list[float] | None = None,
         volume_scale_factor_range: list[float] | None = None,
         rattle_std: float = 0.01,
-        supercell_matrix: Matrix3D | None = None,
         distort_type: int = 0,
         n_structures: int = 10,
         min_distance: float = 1.5,
@@ -465,7 +458,7 @@ def add_dft_random(
         rattle_seed: int = 42,
         rattle_mc_n_iter: int = 10,
         w_angle: list[float] | None = None,
-        adaptive_rattled_supercell_settings: bool = True,
+        supercell_settings: dict | None = None,
     ):
         """Add DFT phonon runs for randomly displaced structures.
 
@@ -480,8 +473,6 @@ def add_dft_random(
         uc: bool.
             If True, will generate randomly distorted structures (unitcells)
             and add static computation jobs to the flow.
-        supercell_matrix: Matrix3D or None
-            The matrix to construct the supercell.
         distort_type : int.
             0- volume distortion, 1- angle distortion, 2- volume and angle distortion. Default=0.
         n_structures : int.
@@ -520,8 +511,8 @@ def add_dft_random(
             Number of Monte Carlo iterations.
             Larger number of iterations will generate larger displacements.
             Default=10.
-        adaptive_rattled_supercell_settings: bool
-            prevent too big rattled supercells
+        supercell_settings: dict
+            settings for supercells
         """
         additonal_dft_random = dft_random_gen_data(
             structure=structure,
@@ -532,7 +523,6 @@ def add_dft_random(
             volume_custom_scale_factors=volume_custom_scale_factors,
             volume_scale_factor_range=volume_scale_factor_range,
             rattle_std=rattle_std,
-            supercell_matrix=supercell_matrix,
             distort_type=distort_type,
             rattle_seed=rattle_seed,
             rattle_mc_n_iter=rattle_mc_n_iter,
@@ -541,10 +531,54 @@ def add_dft_random(
             angle_percentage_scale=angle_percentage_scale,
             w_angle=w_angle,
             min_distance=min_distance,
-            adaptive_rattled_supercell_settings=adaptive_rattled_supercell_settings,
+            supercell_settings=supercell_settings,
         )
-        return Flow(
-            jobs=additonal_dft_random,
-            output={"rand_struc_dir": additonal_dft_random.output},
-            name=self.name,
+        additonal_dft_random.name = "rattled supercells"
+        return additonal_dft_random
+
+
+@dataclass
+class DFTSupercellSettingsMaker(Maker):
+    """
+    Maker to test the DFT and supercell settings.
+
+    This maker is used to test your queue settings for the rattled and phonon supercells.
+    Although the cells are not displaced, it provides an impression of the required memory
+    and other resources as the process runs without symmetry considerations.
+
+    Parameters
+    ----------
+    name (str): The name of the maker. Default is "test dft and supercell settings".
+    supercell_settings (dict): Settings for the supercells. Default is {"min_length": 15}.
+    DFT_Maker (BaseVaspMaker): The DFT maker to be used. Default is TightDFTStaticMaker.
+
+    """
+
+    name: str = "test dft and supercell settings"
+    supercell_settings: dict = field(default_factory=lambda: {"min_length": 15})
+    DFT_Maker: BaseVaspMaker = field(default_factory=TightDFTStaticMaker)
+
+    def make(self, structure_list: list[Structure], mp_ids: list[str]):
+        """
+        Generate and runs supercell jobs for the given list of structures.
+
+        Args:
+            structure_list (list[Structure]): List of structures to process.
+            mp_ids (list[str]): List of MP IDs.
+
+        Returns
+        -------
+            Flow: A Flow object containing the jobs and their output.
+        """
+        job_list = []
+
+        # Modify to run for more than one cell
+        supercell_job = generate_supercells(structure_list, self.supercell_settings)
+        job_list.append(supercell_job)
+
+        supercell_job = run_supercells(
+            structure_list, supercell_job.output, mp_ids, self.DFT_Maker
         )
+        job_list.append(supercell_job)
+
+        return Flow(jobs=job_list, output=supercell_job.output, name=self.name)