mirror of https://github.com/abinit/abipy.git
Playing with fireworks
This commit is contained in:
parent
ee54952f89
commit
bb4ca846b2
|
@ -0,0 +1,95 @@
|
|||
#!/usr/bin/env python
|
||||
from fireworks import FireTaskBase, FWAction, FireWork, LaunchPad, ScriptTask
|
||||
from fireworks.utilities.fw_serializers import FWSerializable
|
||||
from fireworks.core.rocket_launcher import launch_rocket
|
||||
|
||||
from abipy.fworks.tasks import *
|
||||
|
||||
def main():
|
||||
# set up the LaunchPad and reset it
|
||||
launchpad = LaunchPad()
|
||||
launchpad.reset('', require_password=False)
|
||||
|
||||
# Build the flow
|
||||
flow = build_flow()
|
||||
flow.build_and_pickle_dump()
|
||||
|
||||
# create the FireWork consisting of a single task
|
||||
firework = FireWork(FireTaskWithFlow(flow=flow))
|
||||
|
||||
# store workflow
|
||||
launchpad.add_wf(firework)
|
||||
|
||||
#launch it locally
|
||||
#launch_rocket(launchpad)
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
import sys
|
||||
import os
|
||||
import abipy.data as abidata
|
||||
import abipy.abilab as abilab
|
||||
from abipy.tools.devtools import profile
|
||||
|
||||
|
||||
def make_scf_nscf_inputs():
|
||||
"""Returns two input files: GS run and NSCF on a high symmetry k-mesh."""
|
||||
pseudos = abidata.pseudos("14si.pspnc")
|
||||
#pseudos = data.pseudos("Si.GGA_PBE-JTH-paw.xml")
|
||||
|
||||
inp = abilab.AbiInput(pseudos=pseudos, ndtset=2)
|
||||
print(inp.pseudos)
|
||||
structure = inp.set_structure_from_file(abidata.cif_file("si.cif"))
|
||||
|
||||
# Global variables
|
||||
ecut = 6
|
||||
global_vars = dict(ecut=ecut,
|
||||
nband=8,
|
||||
timopt=-1,
|
||||
istwfk="*1",
|
||||
nstep=15,
|
||||
paral_kgb=0,
|
||||
#accesswff=3,
|
||||
)
|
||||
|
||||
if inp.ispaw:
|
||||
global_vars.update(pawecutdg=2*ecut)
|
||||
|
||||
inp.set_variables(**global_vars)
|
||||
|
||||
# Dataset 1 (GS run)
|
||||
inp[1].set_kmesh(ngkpt=[8,8,8], shiftk=[0,0,0])
|
||||
inp[1].set_variables(tolvrs=1e-6)
|
||||
|
||||
# Dataset 2 (NSCF run)
|
||||
kptbounds = [
|
||||
[0.5, 0.0, 0.0], # L point
|
||||
[0.0, 0.0, 0.0], # Gamma point
|
||||
[0.0, 0.5, 0.5], # X point
|
||||
]
|
||||
|
||||
inp[2].set_kpath(ndivsm=6, kptbounds=kptbounds)
|
||||
inp[2].set_variables(tolwfr=1e-12)
|
||||
|
||||
# Generate two input files for the GS and the NSCF run
|
||||
scf_input, nscf_input = inp.split_datasets()
|
||||
return scf_input, nscf_input
|
||||
|
||||
def build_flow():
|
||||
# Working directory (default is the name of the script with '.py' removed and "run_" replaced by "flow_")
|
||||
workdir = "flow_test"
|
||||
|
||||
# Instantiate the TaskManager.
|
||||
manager = abilab.TaskManager.from_user_config()
|
||||
|
||||
# Get the SCF and the NSCF input.
|
||||
scf_input, nscf_input = make_scf_nscf_inputs()
|
||||
|
||||
# Build the flow.
|
||||
return abilab.bandstructure_flow(workdir, manager, scf_input, nscf_input)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
#profile("main()", local_vars=locals(), global_vars=globals())
|
||||
main()
|
|
@ -0,0 +1,41 @@
|
|||
# coding: utf-8
|
||||
from __future__ import print_function, division, unicode_literals
|
||||
|
||||
from fireworks import FireTaskBase, FWAction, FireWork, explicit_serialize
|
||||
from abipy.abilab import AbinitFlow
|
||||
|
||||
@explicit_serialize
|
||||
class FireTaskWithFlow(FireTaskBase):
|
||||
|
||||
required_params = ["flow"]
|
||||
|
||||
#def __init__(self, *args, **kwargs):
|
||||
# #print("args", args, "kwargs", kwargs)
|
||||
# #self._flow = kwargs["flow"]
|
||||
# super(FireTaskWithFlow, self).__init__(*args, **kwargs)
|
||||
|
||||
def run_task(self, fw_spec):
|
||||
# print("entering run_task: %s " % str(self))
|
||||
|
||||
# Initialize the AbinitFlow from the filepath.
|
||||
# Most of the serious stuff is delegated to flow.
|
||||
# Note that FireTask converts all the objects to strings
|
||||
# hence self["flow"] is the workdir of the flow from which
|
||||
# we can reconstruct our object via pickle.
|
||||
flow = AbinitFlow.pickle_load(self["flow"])
|
||||
#print("In MyTask with flow %s" % flow)
|
||||
|
||||
flow.build_and_pickle_dump()
|
||||
try:
|
||||
all_ok = fw_spec["all_ok"]
|
||||
except KeyError:
|
||||
all_ok = False
|
||||
|
||||
if all_ok:
|
||||
print("all_ok will stop")
|
||||
return FWAction()
|
||||
else:
|
||||
print("launched %s" % flow.rapidfire(check_status=True))
|
||||
all_ok = flow.all_ok
|
||||
new_fw = FireWork(FireTaskWithFlow(flow=flow), {'all_ok': all_ok})
|
||||
return FWAction(stored_data={'all_ok': all_ok}, additions=new_fw)
|
|
@ -0,0 +1,11 @@
|
|||
dependencies:
|
||||
override:
|
||||
- pip install -q numpy
|
||||
- pip install -r requirements.txt
|
||||
#- pip install -r requirements-optional.txt
|
||||
- python setup.py develop
|
||||
#test:
|
||||
# override:
|
||||
# - nosetests:
|
||||
# environment:
|
||||
# PYTHONPATH: $PYTHONPATH:/usr/lib/python2.7/dist-packages/
|
|
@ -4,7 +4,6 @@ pydispatcher>=2.0.3
|
|||
pyyaml>=3.11
|
||||
scipy>=0.10
|
||||
netCDF4
|
||||
scipy>=0.10
|
||||
pymatgen>=3.0.0
|
||||
wxmplot
|
||||
#matplotlib>=1.1
|
||||
|
|
Loading…
Reference in New Issue