fix MPMetaGGARelaxGenerator + MPGGARelaxGenerator missing @DataClass

materialsproject · Jul 24, 2023 · 8dd8a92 · 8dd8a92
1 parent 6e070db
commit 8dd8a92
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/src/atomate2/vasp/flows/mp.py b/src/atomate2/vasp/flows/mp.py
@@ -80,13 +80,14 @@ def make(self, structure: Structure, prev_vasp_dir: str | Path | None = None):
                 structure, prev_vasp_dir=prev_vasp_dir
             )
             jobs += [initial_relax]
+            structure = initial_relax.output.structure
+            prev_vasp_dir = initial_relax.output.dir_name
 
         self.final_relax_maker.input_set_generator.additional_vasp_files = (
             self.copy_vasp_files
         )
         final_relax = self.final_relax_maker.make(
-            structure=initial_relax.output.structure,
-            prev_vasp_dir=initial_relax.output.dir_name,
+            structure=structure, prev_vasp_dir=prev_vasp_dir
         )
         output = final_relax.output
         jobs += [final_relax]
@@ -157,13 +158,14 @@ def make(self, structure: Structure, prev_vasp_dir: str | Path | None = None):
                 structure, prev_vasp_dir=prev_vasp_dir
             )
             jobs += [initial_relax]
+            structure = initial_relax.output.structure
+            prev_vasp_dir = initial_relax.output.dir_name
 
         self.final_relax_maker.input_set_generator.additional_vasp_files = (
             self.copy_vasp_files
         )
         final_relax = self.final_relax_maker.make(
-            structure=initial_relax.output.structure,
-            prev_vasp_dir=initial_relax.output.dir_name,
+            structure=structure, prev_vasp_dir=prev_vasp_dir
         )
         output = final_relax.output
         jobs += [final_relax]

diff --git a/src/atomate2/vasp/sets/mp.py b/src/atomate2/vasp/sets/mp.py
@@ -6,7 +6,7 @@
 
 from __future__ import annotations
 
-from dataclasses import field
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -30,6 +30,7 @@
 )
 
 
+@dataclass
 class MPGGARelaxGenerator(VaspInputGenerator):
     """Class to generate MP-compatible VASP GGA relaxation input sets."""
 
@@ -67,6 +68,7 @@ def get_incar_updates(
         return {}
 
 
+@dataclass
 class MPMetaGGARelaxGenerator(VaspInputGenerator):
     """Class to generate MP-compatible VASP metaGGA relaxation input sets."""