GPAW

GPAW#

The gpaw density-functional theory (DFT) simulation code provides a Python interface supporting the message passing interface (MPI) for Python - mpi4py for parallelization. So executorlib is used to orchestrate multiple gpaw simulation each using multiple CPU cores for parallelization. These kind of hierarchical workflows are one of the core strength of executorlib. While the same could be achieved by writing the whole simulation workflow using mpi4py this would drastically increase the complexity. With executorlib the user can quickly up-scale their simulation workflow without the need to address the parallel execution explicitly, rather the parallelization is introduced on a per-function level, by submitting the functions to the FluxClusterExecutor or FluxJobExecutor.

import subprocess
from ase.build import bulk
from atomistics.workflows import (
    analyse_results_for_energy_volume_curve,
    get_tasks_for_energy_volume_curve,
)
import matplotlib.pyplot as plt
import pprint
from tqdm import tqdm
from time import sleep

/srv/conda/envs/notebook/lib/python3.10/site-packages/atomistics/workflows/__init__.py:77: UserWarning: PhonopyWorkflow(), QuasiHarmonicWorkflow(), get_band_structure(), get_dynamical_matrix(), get_hesse_matrix(), get_thermal_properties_for_harmonic_approximation(), get_tasks_for_harmonic_approximation(), analyse_results_for_harmonic_approximation(), get_tasks_for_quasi_harmonic_approximation(), analyse_results_for_quasi_harmonic_approximation(), get_thermal_properties_for_quasi_harmonic_approximation(), plot_band_structure() and plot_dos() are not available as the import of the module named 'structuretoolkit' failed.
  raise_warning(module_list=phonopy_workflows, import_error=e)

The only function which is executed using mpi4py is the evaluate_with_gpaw() function:

def evaluate_with_gpaw(task_dict, kpts, encut):
    import os
    os.environ["OMP_NUM_THREADS"] = "1"
    os.environ["GPAW_MPI4PY"] = "1"

    from mpi4py import MPI
    from gpaw import GPAW, PW
    from gpaw.mpi4pywrapper import MPI4PYWrapper

    structure = task_dict["calc_energy"].copy()
    structure.calc = GPAW(
        xc="PBE",
        mode=PW(encut),
        kpts=kpts,
        communicator=MPI4PYWrapper(MPI.COMM_WORLD),
    )
    return structure.get_potential_energy()

As a first step of the workflow the equilibrium structure of four Aluminium atoms is strained by 5%.

task_dict = get_tasks_for_energy_volume_curve(
    structure=bulk("Al", a=4.03),
    num_points=7,
    vol_range=0.05,
    axes=("x", "y", "z"),
)

The resulting dictionary of structures task_dict is transformed to simplify the parallel execution:

task_loop_dict = {k: {"calc_energy": v} for k, v in task_dict["calc_energy"].items()}

The status of the flux cluster is validated using the flux resource list command and the flux jobs -a command, just to highlight flux was initialized correctly and has access to the available resources.

pprint.pp(subprocess.check_output(["flux", "resource", "list"], universal_newlines=True).split("\n"))

['     STATE NNODES   NCORES    NGPUS NODELIST',
 '      free      1       24        0 '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 ' allocated      0        0        0 ',
 '      down      0        0        0 ',
 '']

pprint.pp(subprocess.check_output(["flux", "jobs", "-a"], universal_newlines=True).split("\n"))

['       JOBID USER     NAME       ST NTASKS NNODES     TIME INFO', '']

FluxClusterExecutor#

The FluxClusterExecutor is used in this demonstration primarily because flux can be installed on any workstation for testing. The SlurmClusterExecutor could be used analogously.

from executorlib import FluxClusterExecutor

The for each strained structure a calculation task is submitted to the FluxClusterExecutor. After the successful submission the current status of the flux queue is printed using flux jobs -a. Finally, the results are collected by gathering the concurrent.futures.Future objects.

future_dict = {}
with FluxClusterExecutor() as exe:
    for k, v in task_loop_dict.items():
        future_dict[k] = exe.submit(
            evaluate_with_gpaw, 
            task_dict=v, 
            kpts=(3, 3, 3), 
            encut=300,
            resource_dict={"cores": 2},
        )
    sleep(1)
    pprint.pp(subprocess.check_output(["flux", "jobs", "-a"], universal_newlines=True).split("\n"))
    result_dict = {
        k: f.result()
        for k, f in tqdm(future_dict.items())
    }
    sleep(1)
    pprint.pp(subprocess.check_output(["flux", "jobs", "-a"], universal_newlines=True).split("\n"))

['       JOBID USER     NAME       ST NTASKS NNODES     TIME INFO',
 '    ƒ3exNRKu jovyan   executorl+  R      2      1   0.104s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ3BjTcCf jovyan   executorl+  R      2      1   1.191s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ31P2d51 jovyan   executorl+  R      2      1   1.646s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '']

100%|██████████| 7/7 [01:08<00:00,  9.86s/it]

['       JOBID USER     NAME       ST NTASKS NNODES     TIME INFO',
 '    ƒ6RrzUJb jovyan   executorl+  R      2      1   1.108m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ5WD8Tno jovyan   executorl+  R      2      1   1.145m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ4YJoYUf jovyan   executorl+  R      2      1   1.181m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ4ESWCVd jovyan   executorl+  R      2      1   1.193m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ3exNRKu jovyan   executorl+  R      2      1   1.215m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ3BjTcCf jovyan   executorl+  R      2      1   1.233m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ31P2d51 jovyan   executorl+  R      2      1   1.241m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '']

The resulting energies for the different volumes are fitted using a 3rd order polynomial to derive the bulk modulus as second derivative multiplied by the equilibrium volume.

fit_dict = analyse_results_for_energy_volume_curve(
    output_dict={"energy": result_dict},
    task_dict=task_dict,
    fit_type="polynomial",
    fit_order=3,
)

The final energy volume curve plot summarizes the results of this calculation.

plt.plot(fit_dict["volume"], fit_dict["energy"], label="$B_0$= %0.2f GPa" % fit_dict["bulkmodul_eq"])
plt.xlabel("Volume [$\AA^3$]")
plt.ylabel("Energy [eV]")
plt.legend()

<matplotlib.legend.Legend at 0x7945475ffd90>
_images/205ff1966944e78b13101a49f9465926d3731ead8edfec8906e93d443a71baf6.png

FluxJobExecutor#

In analogy to the FluxClusterExecutor the FluxJobExecutor can be applied to distribute simulation within a given queuing system allocation. The calculation of the bulk modulus with gpaw is implemented in the same way.

from executorlib import FluxJobExecutor

The for each strained structure a calculation task is submitted to the FluxJobExecutor. After the successful submission the current status of the flux queue is printed using flux jobs -a. Finally, the results are collected by gathering the concurrent.futures.Future objects.

future_dict = {}
with FluxJobExecutor() as exe:
    for k, v in task_loop_dict.items():
        future_dict[k] = exe.submit(
            evaluate_with_gpaw, 
            task_dict=v, 
            kpts=(3, 3, 3), 
            encut=300,
            resource_dict={"cores": 2},
        )
    sleep(1)
    pprint.pp(subprocess.check_output(["flux", "jobs", "-a"], universal_newlines=True).split("\n"))
    result_dict = {
        k: f.result()
        for k, f in tqdm(future_dict.items())
    }
    sleep(1)
    pprint.pp(subprocess.check_output(["flux", "jobs", "-a"], universal_newlines=True).split("\n"))

['       JOBID USER     NAME       ST NTASKS NNODES     TIME INFO',
 '    ƒeqRVxAU jovyan   python      R      2      1   4.138s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒeqRVxAT jovyan   python      R      2      1   4.138s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒeqQ1xt8 jovyan   python      R      2      1   4.139s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒeqQ1xt7 jovyan   python      R      2      1   4.139s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒeqNXybn jovyan   python      R      2      1   4.139s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒeqNXybm jovyan   python      R      2      1   4.139s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒeqM3zKR jovyan   python      R      2      1   4.141s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ6RrzUJb jovyan   executorl+  F      2      1   1.192m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ4ESWCVd jovyan   executorl+  F      2      1   1.268m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ5WD8Tno jovyan   executorl+  F      2      1   1.218m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ3exNRKu jovyan   executorl+  F      2      1   1.264m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ3BjTcCf jovyan   executorl+  F      2      1   1.271m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ4YJoYUf jovyan   executorl+  F      2      1   1.212m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ31P2d51 jovyan   executorl+  F      2      1   1.251m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '']

100%|██████████| 7/7 [00:25<00:00,  3.60s/it]

['       JOBID USER     NAME       ST NTASKS NNODES     TIME INFO',
 '    ƒeqRVxAU jovyan   python      R      2      1   30.78s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒeqRVxAT jovyan   python      R      2      1   30.78s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒeqQ1xt8 jovyan   python      R      2      1   30.78s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒeqQ1xt7 jovyan   python      R      2      1   30.78s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒeqNXybn jovyan   python      R      2      1   30.78s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒeqNXybm jovyan   python      R      2      1   30.78s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒeqM3zKR jovyan   python      R      2      1   30.78s '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ6RrzUJb jovyan   executorl+  F      2      1   1.192m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ4ESWCVd jovyan   executorl+  F      2      1   1.268m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ5WD8Tno jovyan   executorl+  F      2      1   1.218m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ3exNRKu jovyan   executorl+  F      2      1   1.264m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ3BjTcCf jovyan   executorl+  F      2      1   1.271m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ4YJoYUf jovyan   executorl+  F      2      1   1.212m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '    ƒ31P2d51 jovyan   executorl+  F      2      1   1.251m '
 'jupyter-jan-janssen-exe-ntegration-test-jc5mdph0',
 '']

The resulting energies for the different volumes are fitted using a 3rd order polynomial to derive the bulk modulus as second derivative multiplied by the equilibrium volume.

fit_dict = analyse_results_for_energy_volume_curve(
    output_dict={"energy": result_dict},
    task_dict=task_dict,
    fit_type="polynomial",
    fit_order=3,
)

The final energy volume curve plot summarizes the results of this calculation.

plt.plot(fit_dict["volume"], fit_dict["energy"], label="$B_0$= %0.2f GPa" % fit_dict["bulkmodul_eq"])
plt.xlabel("Volume [$\AA^3$]")
plt.ylabel("Energy [eV]")
plt.legend()

<matplotlib.legend.Legend at 0x7945345f1690>
_images/205ff1966944e78b13101a49f9465926d3731ead8edfec8906e93d443a71baf6.png