Add scale_factor_increment as parameter to md flow

src/mpmorph/atomate1/firetasks/dbtasks.py

@@ -2,7 +2,6 @@
 import os
 import re
 import zlib
-from collections import defaultdict
 from typing import List
 
 import gridfs

src/mpmorph/atomate1/firetasks/mdtasks.py

@@ -3,7 +3,6 @@
 
 import numpy as np
 from fireworks import FireTaskBase, FWAction, Workflow, explicit_serialize
-from pymatgen.core import Structure
 from pymatgen.io.vasp import Poscar
 from pymatgen.io.vasp.outputs import Vasprun
 from scipy import stats

src/mpmorph/atomate1/fireworks/core.py

@@ -7,8 +7,8 @@
 from fireworks import Firework
 from pymatgen.io.vasp.sets import MPMDSet, MPRelaxSet, MPStaticSet
 
-from mpmorph.firetasks.dbtasks import VaspMDToDb
-from mpmorph.firetasks.glue_tasks import PreviousStructureTask, SaveStructureTask
+from ..firetasks.dbtasks import VaspMDToDb
+from ..firetasks.glue_tasks import PreviousStructureTask, SaveStructureTask
 
 """
 These fireworks were adapted from atomate.vasp.fireworks.core specifically for this 

src/mpmorph/atomate1/fireworks/powerups.py

@@ -1,10 +1,10 @@
-from mpmorph.firetasks.dbtasks import TrajectoryDBTask, VaspMDToDb
-from mpmorph.firetasks.glue_tasks import (
+from ..firetasks.dbtasks import TrajectoryDBTask, VaspMDToDb
+from ..firetasks.glue_tasks import (
     PassPVTask,
     PreviousStructureTask,
     SaveStructureTask,
 )
-from mpmorph.firetasks.mdtasks import (
+from ..firetasks.mdtasks import (
     ConvergeTask,
     DiffusionTask,
     PVRescaleTask,

src/mpmorph/atomate1/workflows/converge.py

@@ -3,8 +3,8 @@
 
 from fireworks import Workflow
 
-from mpmorph.fireworks import powerups
-from mpmorph.fireworks.core import MDFW
+from ..fireworks import powerups
+from ..fireworks.core import MDFW
 from mpmorph.utils import recursive_update
 
 

src/mpmorph/atomate1/workflows/quench.py

@@ -1,8 +1,7 @@
 import numpy as np
 from fireworks import Workflow
 
-from mpmorph.fireworks import powerups
-from mpmorph.fireworks.core import MDFW, OptimizeFW, StaticFW
+from ..fireworks.core import MDFW, OptimizeFW, StaticFW
 from mpmorph.utils import recursive_update
 
 

src/mpmorph/atomate2/flows/amorphous_limit_flow.py

@@ -9,7 +9,7 @@
 from .vasp import get_md_flow_vasp
 from .fast_quench import get_fast_quench_flow, get_original_relax_flow
 from .utils import get_frames_from_trajectory
-from ..runners.amorphous_maker import get_random_packed
+from ...runners.amorphous_maker import get_random_packed
 from monty.json import MontyDecoder
 
 from typing import Union

src/mpmorph/atomate2/flows/chgnet_md_flow.py

@@ -1,11 +1,11 @@
 from jobflow import Flow
-from mpmorph.jobs.core import CHGNetMDMaker
+from ..jobs.core import CHGNetMDMaker
 
-from mpmorph.jobs.equilibrate_volume import EquilibriumVolumeSearchMaker
+from ..jobs.equilibrate_volume import EquilibriumVolumeSearchMaker
 from pymatgen.core.structure import Structure
 
-from mpmorph.jobs.pv_from_calc import PVFromCHGNet
-from mpmorph.jobs.tasks.chgnet_input import CHGNetMDInputs
+from ..jobs.pv_from_calc import PVFromCHGNet
+from ..jobs.tasks.chgnet_input import CHGNetMDInputs
 from .utils import get_md_flow
 
 EQUILIBRATE_VOLUME_FLOW = "EQUILIBRATE_VOLUME_FLOW"

src/mpmorph/atomate2/flows/lammps.py

@@ -1,6 +1,6 @@
 from jobflow import Flow
 
-from mpmorph.jobs.lammps.lammps_basic_const_temp import BasicLammpsConstantTempMaker
+from ..jobs.lammps.lammps_basic_const_temp import BasicLammpsConstantTempMaker
 from pymatgen.core.structure import Structure
 
 from .utils import collect_vt_results

src/mpmorph/atomate2/flows/m3gnet.py

@@ -1,11 +1,11 @@
 from jobflow import Flow
-from mpmorph.jobs.core import M3GNetMDMaker
+from ..jobs.core import M3GNetMDMaker
 
-from mpmorph.jobs.equilibrate_volume import EquilibriumVolumeSearchMaker
+from ..jobs.equilibrate_volume import EquilibriumVolumeSearchMaker
 from pymatgen.core.structure import Structure
 
-from mpmorph.jobs.pv_from_calc import PVFromM3GNet
-from mpmorph.jobs.tasks.m3gnet_input import M3GNetMDInputs
+from ..jobs.pv_from_calc import PVFromM3GNet
+from ..jobs.tasks.m3gnet_input import M3GNetMDInputs
 from .utils import get_md_flow, collect_vt_results, VOLUME_TEMPERATURE_SWEEP
 
 EQUILIBRATE_VOLUME_FLOW = "EQUILIBRATE_VOLUME_FLOW"

src/mpmorph/atomate2/flows/utils.py

@@ -4,11 +4,11 @@
 
 from jobflow import Flow, Maker, job
 
-from mpmorph.jobs.equilibrate_volume import EquilibriumVolumeSearchMaker
+from ..jobs.equilibrate_volume import EquilibriumVolumeSearchMaker
 from pymatgen.core.structure import Structure
 from pymatgen.core.trajectory import Trajectory
 
-from mpmorph.jobs.pv_from_calc import PVExtractor
+from ..jobs.pv_from_calc import PVExtractor
 from ..jobs.equilibrate_volume import PVFromMDFlowMaker, GetPVDocFromMDMaker
 
 
@@ -40,6 +40,7 @@ def get_md_flow(
     structure,
     converge_first,
     initial_vol_scale,
+    scale_factor_increment: float = 0.2,
     flow_name: str = "MD_FLOW",
 ):
     struct = structure.copy()
@@ -52,7 +53,8 @@ def get_md_flow(
             pv_md_maker = pv_md_maker,
             pv_extractor = pv_extractor,
             production_run_maker = production_md_maker,
-            flow_name=flow_name
+            flow_name=flow_name,
+            scale_factor_increment=scale_factor_increment
         )
     else:
         return Flow([production_md_maker.make(struct)], name=flow_name)
@@ -63,6 +65,7 @@ def get_converge_flow(
     pv_md_maker: Maker,
     pv_extractor: PVExtractor,
     production_run_maker: Maker,
+    scale_factor_increment: float = 0.2,
     flow_name: str = M3GNET_MD_CONVERGED_VOL_FLOW
 ):
     eq_vol_maker = EquilibriumVolumeSearchMaker(
@@ -71,7 +74,8 @@ def get_converge_flow(
             extract_maker=GetPVDocFromMDMaker(
                 pv_extractor=pv_extractor
             )
-        )
+        ),
+        scale_factor_increment=scale_factor_increment
     )
 
     equil_vol_job = eq_vol_maker.make(structure)

src/mpmorph/atomate2/flows/vasp.py

@@ -1,7 +1,7 @@
 from atomate2.vasp.jobs.core import MDMaker
 from atomate2.vasp.sets.core import MDSetGenerator
 from pymatgen.core.structure import Structure
-from mpmorph.jobs.pv_from_calc import PVFromVasp
+from ..jobs.pv_from_calc import PVFromVasp
 
 
 from .utils import get_md_flow
@@ -16,6 +16,7 @@ def get_md_flow_vasp(
     steps_prod: int = None,
     steps_pv: int = None,
     initial_vol_scale: int = 1,
+    scale_factor_increment: float = 0.2,
     production_md_set_generator: MDSetGenerator = None,
     pv_md_set_generator: MDSetGenerator = None
 ):
@@ -73,5 +74,6 @@ def get_md_flow_vasp(
         structure=structure,
         converge_first=converge_first,
         initial_vol_scale=initial_vol_scale,
+        scale_factor_increment=scale_factor_increment,
         flow_name=flow_name
     )

src/mpmorph/atomate2/jobs/equilibrate_volume.py

@@ -1,18 +1,17 @@
 from dataclasses import dataclass
 from typing import List, Tuple
 import numpy as np
-from emmet.core.tasks import TaskDoc
 
 from pymatgen.core.structure import Structure
 from jobflow import Maker, job, Flow, Response
 
 from .pv_from_calc import PVExtractor
 from ..schemas.pv_data_doc import MDPVDataDoc
-from ..runners import rescale_volume
+from ...runners import rescale_volume
 
 
-MAX_MD_JOBS = 9  # if you can't converge with five additional calcs you're doing something wrong...
-OFFSET = 0.1  # gives it enough room to slosh back
+MAX_MD_JOBS = 9  # if you can't converge with 6 additional calcs you're doing something wrong...
+BUFFER = 0.1  # gives it enough room to slosh back
 
 def get_scaled_structure(struct: Structure, scale_factor: float):
     copy: Structure = struct.copy()
@@ -55,7 +54,7 @@ class EquilibriumVolumeSearchMaker(Maker):
 
     name: str = "EQUIL_VOL_SEARCH"
     pv_from_md_maker: PVFromMDFlowMaker = None
-    initial_scale_factors: Tuple[float] = (0.8, 1, 1.2)
+    scale_factor_increment: float = 0.2
 
     @job
     def make(
@@ -69,9 +68,11 @@ def make(
             )
 
         if pv_data_docs is None:
+            initial_scale_factors = [1 - self.scale_factor_increment, 1, 1 + self.scale_factor_increment]
+
             scaled_structs = [
                 get_scaled_structure(original_structure, factor)
-                for factor in self.initial_scale_factors
+                for factor in initial_scale_factors
             ]
 
             new_jobs = [
@@ -129,13 +130,13 @@ def make(
 
 
 def get_new_max_volume(equil_guess, original_structure):
-    return equil_guess / original_structure.volume + OFFSET
+    return equil_guess / original_structure.volume + BUFFER
 
 def expand_upper_bound(old_max_vol, original_structure):
     return old_max_vol / original_structure.volume + 0.2
 
 def get_new_min_volume(equil_guess, original_structure):
-    return equil_guess / original_structure.volume - OFFSET
+    return equil_guess / original_structure.volume - BUFFER
 
 def expand_lower_bound(old_max_vol, original_structure):
     return old_max_vol / original_structure.volume - 0.2

src/mpmorph/database.py

@@ -1,10 +1,8 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import json
-import zlib
 from collections import defaultdict
 
-import gridfs
 import numpy as np
 from atomate.utils.utils import get_logger
 from atomate.vasp.database import VaspCalcDb