Skip to content

Commit

Permalink
update workgraph
Browse files Browse the repository at this point in the history
  • Loading branch information
@superstar54
superstar54 committed Jun 24, 2024
1 parent 5c97ec6 commit d7f4e3c
Show file tree
Hide file tree
Showing 3 changed files with 98 additions and 75 deletions.
53 changes: 44 additions & 9 deletions aiida_bader/workgraph/qe_bader.py
View file
Original file line number Diff line number Diff line change
@@ -1,26 +1,61 @@
# -*- coding: utf-8 -*-
"""QeBaderWorkGraph of the AiiDA bader plugin"""
from aiida import orm
from aiida_workgraph import WorkGraph, task

from aiida_workgraph import WorkGraph

@task.graph_builder(outputs=[["bader.charge", "charge"]])
def bader_workgraph(
structure: orm.StructureData = None,
pw_code: orm.Code = None,
pp_code: orm.Code = None,
bader_code: orm.Code = None,
inputs: dict = None,
):
"""Workgraph for Bader charge analysis.
1. Run the SCF calculation.
2. Run the PP calculation for valence charge density.
3. Run the PP calculation for all-electron charge density.
4. Run the Bader charge analysis.
"""

def QeBaderWorkGraph():
from aiida_quantumespresso.workflows.pw.base import PwBaseWorkChain
from aiida_quantumespresso.calculations.pp import PpCalculation
from aiida_bader.calculations import BaderCalculation

wg = WorkGraph("charge-density")
pw_node = wg.nodes.new(PwBaseWorkChain, name="scf")
pp_valence = wg.nodes.new(
PpCalculation, name="pp_valence", parent_folder=pw_node.outputs["remote_folder"]
inputs = {} if inputs is None else inputs
wg = WorkGraph("BaderCharge")
# -------- scf -----------
scf_task = wg.tasks.new(PwBaseWorkChain, name="scf")
scf_inputs = inputs.get("scf", {})
scf_inputs.update({"pw.structure": structure, "pw.code": pw_code})
scf_task.set(scf_inputs)
# -------- pp valence -----------
pp_valence = wg.tasks.new(
PpCalculation,
name="pp_valence",
code=pp_code,
parent_folder=scf_task.outputs["remote_folder"],
)
pp_all = wg.nodes.new(
PpCalculation, name="pp_all", parent_folder=pw_node.outputs["remote_folder"]
pp_valence_inputs = inputs.get("pp_valence", {})
pp_valence.set(pp_valence_inputs)
# -------- pp all -----------
pp_all = wg.tasks.new(
PpCalculation,
name="pp_all",
code=pp_code,
parent_folder=scf_task.outputs["remote_folder"],
)
wg.nodes.new(
pp_all_inputs = inputs.get("pp_all", {})
pp_all.set(pp_all_inputs)
# -------- bader -----------
bader_task = wg.tasks.new(
BaderCalculation,
name="bader",
code=bader_code,
charge_density_folder=pp_valence.outputs["remote_folder"],
reference_charge_density_folder=pp_all.outputs["remote_folder"],
)
bader_inputs = inputs.get("bader", {})
bader_task.set(bader_inputs)
return wg
118 changes: 53 additions & 65 deletions docs/source/quick_start.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
"\n",
" export OMP_NUM_THREADS=1\n",
"append_text: ' '\n",
"with_mpi: false\n",
"```\n",
"Then run the following command to set up the `bader` code:\n",
"\n",
Expand Down Expand Up @@ -105,8 +106,11 @@
],
"source": [
"from aiida import load_profile\n",
"from aiida.orm import Dict, KpointsData, StructureData, load_code, load_group, load_node\n",
"from aiida.orm import Dict, KpointsData, StructureData, load_code\n",
"from aiida_pseudo.data.pseudo import UpfData\n",
"from ase.build import molecule\n",
"from workgraph_collections.qe.bader import bader_workgraph\n",
"\n",
"\n",
"load_profile()\n",
"#===============================================================================\n",
Expand All @@ -115,13 +119,13 @@
"pp_code = load_code(\"qe-7.2-pp@localhost\")\n",
"bader_code = load_code(\"bader@localhost\")\n",
"# ===============================================================================\n",
"# create input structure node\n",
"h2o = molecule(\"H2O\")\n",
"h2o.center(vacuum=3.0)\n",
"h2o.pbc = True\n",
"structure = StructureData(ase=h2o)\n",
"# create input structure\n",
"atoms = molecule(\"H2O\")\n",
"atoms.center(vacuum=3.0)\n",
"atoms.pbc = True\n",
"structure = StructureData(ase=atoms)\n",
"# create input parameters node\n",
"pw_paras = Dict(\n",
"scf_paras = Dict(\n",
" {\n",
" \"CONTROL\": {\n",
" \"calculation\": \"scf\",\n",
Expand All @@ -137,74 +141,58 @@
")\n",
"kpoints = KpointsData()\n",
"kpoints.set_kpoints_mesh([1, 1, 1])\n",
"# Load the pseudopotential family.\n",
"# check if group exists\n",
"pseudo_group = load_group(\"psl_1.0.0_pbe_kjpaw\")\n",
"pseudo_keys = pseudo_group.base.extras.get(\"pseudos\")\n",
"pseudos = {kind.name: load_node(pseudo_keys[kind.name]) for kind in structure.kinds}\n",
"# Load the pseudopotential kjpaw.\n",
"pseudos = {\"H\": UpfData(\"/home/xing/data/ase/espresso_pseudo/kjpaw/H.pbe-kjpaw_psl.1.0.0.UPF\"),\n",
" \"O\": UpfData(\"/home/xing/data/ase/espresso_pseudo/kjpaw/O.pbe-n-kjpaw_psl.1.0.0.UPF\")}\n",
"#\n",
"#\n",
"metadata = {\n",
" \"options\": {\n",
" \"resources\": {\n",
" \"num_machines\": 1,\n",
" \"num_mpiprocs_per_machine\": 2,\n",
" \"num_mpiprocs_per_machine\": 1,\n",
" },\n",
" }\n",
"}\n",
"# pp parameters\n",
"pp_paras = Dict({\n",
" \"INPUTPP\": {\"plot_num\": 21},\n",
" \"PLOT\": {\"iflag\": 3},\n",
" })\n",
"# ===============================================================================\n",
"from aiida_bader.worktrees.qe_bader import QeBaderWorkTree\n",
"# prepare inputs and submit\n",
"wt = QeBaderWorkTree()\n",
"wt.nodes[\"pw_base\"].set({\"pw\": {\n",
" \"code\": pw_code,\n",
" \"structure\": structure,\n",
" \"parameters\": pw_paras,\n",
" \"pseudos\": pseudos,\n",
" \"metadata\": metadata,\n",
"inputs = {\"scf\":{\n",
" \"pw\": {\n",
" \"parameters\": scf_paras,\n",
" \"pseudos\": pseudos,\n",
" \"metadata\": metadata\n",
" },\n",
" \"kpoints\": kpoints,\n",
" },\n",
"\n",
" \"pp_valence\": {\n",
" \"parameters\": Dict({\n",
" \"INPUTPP\": {\"plot_num\": 0},\n",
" \"PLOT\": {\"iflag\": 3},\n",
" }),\n",
" \"metadata\": metadata,\n",
" },\n",
" \"kpoints\": kpoints,\n",
" })\n",
"wt.nodes[\"pp\"].set({\"code\": pp_code,\n",
" \"parameters\": pp_paras,\n",
" \"metadata\": metadata,\n",
" })\n",
"wt.nodes[\"bader\"].set({\"code\": bader_code,})\n",
"# submit\n",
"wt.submit(wait=True, timeout=300)\n"
]
},
{
"cell_type": "markdown",
"id": "d25beb02-ee82-4a27-ae48-edc5c147904c",
"metadata": {},
"source": [
"### Check status and results\n"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "9ebf35aa",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"State of WorkTree: FINISHED\n",
"Result of bader : [10.033972 0.426889 0.426893]\n"
]
}
],
"source": [
"wt.update()\n",
"print(\"State of WorkTree: {}\".format(wt.state))\n",
"print('Result of bader : {}'.format(wt.nodes[\"bader\"].node.outputs.bader_charge.get_array(\"charge\")))"
" \"pp_all\": {\n",
" \"parameters\": Dict({\n",
" \"INPUTPP\": {\"plot_num\": 21},\n",
" \"PLOT\": {\"iflag\": 3},\n",
" }),\n",
" \"metadata\": metadata,\n",
" },\n",
" \"bader\": {\n",
" \"metadata\": metadata,\n",
" }\n",
"}\n",
"# prepare inputs and submit\n",
"wg = bader_workgraph(structure=structure, pw_code=pw_code,\n",
" pp_code=pp_code, bader_code=bader_code,\n",
" inputs=inputs)\n",
"wg.submit(wait=True, timeout=300)\n",
"#------------------------- Print the output -------------------------\n",
"charges = wg.tasks[\"bader\"].node.outputs.bader_charge.get_array(\"charge\")\n",
"print(\"Bader charges:\")\n",
"print(\"Index Symbol Charge\")\n",
"for i, charge in enumerate(charges):\n",
" print(f\"{i:5d} {atoms.get_chemical_symbols()[i]:5s} {charge:5.3f}\")\n"
]
},
{
Expand Down
2 changes: 1 addition & 1 deletion examples/qe_bader_workchain.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from aiida import load_profile
from aiida.orm import Dict, KpointsData, StructureData, load_code, load_group, load_node
from aiida.orm import Dict, StructureData, load_code
from ase.build import molecule
from aiida.plugins import WorkflowFactory
from aiida.engine import submit
Expand Down

0 comments on commit d7f4e3c

Please sign in to comment.