From bb4ca846b22a7a2bbdef39d98f790b01695a20a6 Mon Sep 17 00:00:00 2001
From: gmatteo <gmatteo@gmail.com>
Date: Sat, 20 Sep 2014 02:53:00 +0200
Subject: [PATCH] Playing with fireworks

---
 abipy/fworks/__init__.py |  0
 abipy/fworks/abifw.py    | 95 ++++++++++++++++++++++++++++++++++++++++
 abipy/fworks/tasks.py    | 41 +++++++++++++++++
 circle.yml               | 11 +++++
 requirements.txt         |  1 -
 5 files changed, 147 insertions(+), 1 deletion(-)
 create mode 100644 abipy/fworks/__init__.py
 create mode 100755 abipy/fworks/abifw.py
 create mode 100644 abipy/fworks/tasks.py
 create mode 100644 circle.yml

diff --git a/abipy/fworks/__init__.py b/abipy/fworks/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/abipy/fworks/abifw.py b/abipy/fworks/abifw.py
new file mode 100755
index 00000000..bb7de446
--- /dev/null
+++ b/abipy/fworks/abifw.py
@@ -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()
diff --git a/abipy/fworks/tasks.py b/abipy/fworks/tasks.py
new file mode 100644
index 00000000..c08af76c
--- /dev/null
+++ b/abipy/fworks/tasks.py
@@ -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)
diff --git a/circle.yml b/circle.yml
new file mode 100644
index 00000000..d0bd9d5c
--- /dev/null
+++ b/circle.yml
@@ -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/
diff --git a/requirements.txt b/requirements.txt
index eaf6c4f2..547b2bca 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -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