Skip to content

Adding relax workflow for ABACUS #349

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zhubonan wants to merge 11 commits into
base: master
Choose a base branch
from
Open

Adding relax workflow for ABACUS #349

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
8aeabf2
Initial implementation for Abacus
zhubonan Mar 31, 2025
7f8050f
Update protocol and conversion
zhubonan Mar 31, 2025
12aac5e
Update the extractor the workcchain definition
zhubonan Mar 31, 2025
404ac16
Updated abacus protocols
zhubonan May 17, 2025
ecd5a58
Added with additional protocols
zhubonan Aug 19, 2025
166b8bc
Merge branch 'master' into feature/abacus
zhubonan Nov 29, 2025
e2b5064
Fix dependency and DOI to readme
zhubonan Dec 17, 2025
8a861b9
Avoids top-level import
zhubonan Dec 17, 2025
d556c5d
Merge branch 'feature/abacus' of github.com:MCresearch/aiida-common-w…
zhubonan Dec 17, 2025
e00351e
Make minimum python version to be 3.10
zhubonan Dec 17, 2025
f8eabc5
Fix pre-commit
zhubonan Dec 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .github/workflows/ci.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ jobs:

strategy:
matrix:
python-version: ['3.9', '3.10', '3.11', '3.12']
python-version: ['3.10', '3.11', '3.12']

services:
rabbitmq:
Expand Down Expand Up @@ -70,7 +70,7 @@ jobs:

strategy:
matrix:
python-version: ['3.9']
python-version: ['3.10']

steps:
- uses: actions/checkout@v2
Expand Down
2 changes: 1 addition & 1 deletion .readthedocs.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ version: 2
build:
os: ubuntu-22.04
tools:
python: '3.9'
python: '3.10'
apt_packages:
- gfortran # This is necessary for the `sisl` dependency of `aiida-siesta`

Expand Down
1 change: 1 addition & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ In addition, if you run the common workflows, please also cite:

Engine | DOIs or URLs to be cited
---------------- | ----------------------------
ABACUS | [10.1063/5.0297563](https://doi.org/10.1063/5.0297563)
ABINIT | [10.1016/j.cpc.2016.04.003](https://doi.org/10.1016/j.cpc.2016.04.003) [10.1016/j.cpc.2019.107042](https://doi.org/10.1016/j.cpc.2019.107042) [10.1063/1.5144261](https://doi.org/10.1063/1.5144261)
BigDFT | [10.1063/5.0004792](https://doi.org/10.1063/5.0004792)
CASTEP | [10.1524/zkri.220.5.567.65075](https://doi.org/10.1524/zkri.220.5.567.65075)
Expand Down
60 changes: 60 additions & 0 deletions abacus.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
from aiida import orm
from aiida.engine import run_get_node
from aiida.plugins import WorkflowFactory
from aiida_common_workflows.common import ElectronicType, RelaxType, SpinType
from aiida_common_workflows.plugins import get_entry_point_name_from_class, load_workflow_entry_point
from ase.build import bulk

PLUGIN_NAME = 'abacus'
CODE_LABEL = 'abacus@localhost'


structure = orm.StructureData(ase=bulk('Si', 'diamond', 5.4))
print(f'Structure PK: {structure.pk}')

sub_process_cls = load_workflow_entry_point('relax', 'abacus')
sub_process_cls_name = get_entry_point_name_from_class(sub_process_cls).name
generator = sub_process_cls.get_input_generator()

engine_types = generator.spec().inputs['engines']
engines = {}
# There should be only one
for engine in engine_types:
engines[engine] = {
'code': CODE_LABEL,
'options': {
'resources': {'num_machines': 1, 'tot_num_mpiprocs': 1},
'max_wallclock_seconds': 1700, # A bit less than 30 minutes (so we fit in the debug queue=partition)
},
}

inputs = {
'structure': structure,
'generator_inputs': { # code-agnostic inputs for the relaxation
'engines': engines,
'protocol': 'fast',
'relax_type': RelaxType.NONE,
'electronic_type': ElectronicType.METAL,
'spin_type': SpinType.NONE,
},
'sub_process_class': sub_process_cls_name,
# 'sub_process' : { # optional code-dependent overrides
# 'base': {
# ## In order to make this work, you have perform the change discussed
# ## at the bottom of the file in the eos.py file.
# ## Otherwise the whole namespace is replaced.
# 'abacus': {
# 'parameters': orm.Dict(dict={
# 'input': {
# 'cal_stress': True
# }
# })
# }
# }
# }
}

cls = WorkflowFactory('common_workflows.eos')
results, node = run_get_node(cls, **inputs)
print(f'Submitted workflow with PK = {node.pk} for {structure.get_formula()}')
print(f'Results: {results}')
6 changes: 5 additions & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@ classifiers = [
'Operating System :: POSIX :: Linux',
'Operating System :: MacOS :: MacOS X',
'Programming Language :: Python',
'Programming Language :: Python :: 3.9',
'Programming Language :: Python :: 3.10',
'Programming Language :: Python :: 3.11',
'Programming Language :: Python :: 3.12'
Expand All @@ -35,6 +34,7 @@ requires-python = '>=3.9'
'common_workflows.dissociation_curve' = 'aiida_common_workflows.workflows.dissociation:DissociationCurveWorkChain'
'common_workflows.em' = 'aiida_common_workflows.workflows.em:EnergyMagnetizationWorkChain'
'common_workflows.eos' = 'aiida_common_workflows.workflows.eos:EquationOfStateWorkChain'
'common_workflows.relax.abacus' = 'aiida_common_workflows.workflows.relax.abacus.workchain:AbacusCommonRelaxWorkChain'
'common_workflows.relax.abinit' = 'aiida_common_workflows.workflows.relax.abinit.workchain:AbinitCommonRelaxWorkChain'
'common_workflows.relax.bigdft' = 'aiida_common_workflows.workflows.relax.bigdft.workchain:BigDftCommonRelaxWorkChain'
'common_workflows.relax.castep' = 'aiida_common_workflows.workflows.relax.castep.workchain:CastepCommonRelaxWorkChain'
Expand All @@ -50,12 +50,16 @@ requires-python = '>=3.9'
'common_workflows.relax.wien2k' = 'aiida_common_workflows.workflows.relax.wien2k.workchain:Wien2kCommonRelaxWorkChain'

[project.optional-dependencies]
abacus = [
'aiida-abacus~=0.3.1'
]
abinit = [
'abipy==0.9.6',
'aiida-abinit~=0.5.0'
]
all_plugins = [
'abipy==0.9.6',
'aiida-abacus~=0.3.1',
'aiida-abinit~=0.5.0',
'aiida-ase~=3.0',
'aiida-bigdft~=0.3.0',
Expand Down
5 changes: 5 additions & 0 deletions src/aiida_common_workflows/workflows/relax/abacus/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
"""Module with the implementations of the common structure relaxation workchain for Abacus"""
from .generator import *
from .workchain import *

__all__ = generator.__all__ + workchain.__all__
29 changes: 29 additions & 0 deletions src/aiida_common_workflows/workflows/relax/abacus/extractors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
"""
Post-processing extractor functions for the ``QuantumEspressoCommonRelaxWorkChain``.
"""

from aiida.common import LinkType
from aiida.orm import WorkChainNode

from .workchain import AbacusCommonRelaxWorkChain


def get_ts_energy(common_relax_workchain: AbacusCommonRelaxWorkChain) -> float:
"""Return the T * S value of a concluded ``QuantumEspressoCommonRelaxWorkChain``.

Here, T is the fictitious temperature due to the use of smearing and S is the entropy. This "smearing contribution"
to the free energy in Abacus is expressed as -T * S:

E_KS(sigma->0)

This energy is printed at every electronic cycle.
:param common_relax_workchain: ``AbacusCommonRelaxWorkChain`` for which to extract the smearing energy.
:returns: The T*S value in eV.
"""
if not isinstance(common_relax_workchain, WorkChainNode):
return ValueError('The input is not a workchain (instance of `WorkChainNode`)')
if common_relax_workchain.process_class != AbacusCommonRelaxWorkChain:
return ValueError('The input workchain is not a `AbacusCommonRelaxWorkChain`')

abacus_relax_wc = common_relax_workchain.base.links.get_outgoing(link_type=LinkType.CALL_WORK).one().node
return abacus_relax_wc.outputs.misc['ts_contribution']
168 changes: 168 additions & 0 deletions src/aiida_common_workflows/workflows/relax/abacus/generator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,168 @@
"""Implementation of `aiida_common_workflows.common.relax.generator.CommonRelaxInputGenerator` for Abacus."""

from importlib import resources

import yaml
from aiida import engine, orm, plugins

from aiida_common_workflows.common import ElectronicType, RelaxType, SpinType
from aiida_common_workflows.generators import ChoiceType, CodeType

from ..generator import CommonRelaxInputGenerator

__all__ = ('AbacusCommonRelaxInputGenerator',)

StructureData = plugins.DataFactory('core.structure')


class AbacusCommonRelaxInputGenerator(CommonRelaxInputGenerator):
"""Input generator for the common relax workflow implementation of Abacus."""

def __init__(self, *args, **kwargs):
"""Construct an instance of the input generator, validating the class attributes."""
process_class = kwargs.get('process_class', None)

if process_class is not None:
self._default_protocol = process_class._process_class.get_default_protocol()
self._protocols = process_class._process_class.get_available_protocols()
self._protocols.update({key: value['description'] for key, value in self._load_local_protocols().items()})

super().__init__(*args, **kwargs)

@staticmethod
def _load_local_protocols():
"""Load the protocols defined in the ``aiida-common-workflows`` package."""
from .. import abacus

with resources.open_text(abacus, 'protocol.yml') as handle:
protocol_dict = yaml.safe_load(handle)
return protocol_dict

@classmethod
def define(cls, spec):
"""Define the specification of the input generator.

The ports defined on the specification are the inputs that will be accepted by the ``get_builder`` method.
"""
super().define(spec)
spec.inputs['protocol'].valid_type = ChoiceType(('fast', 'moderate', 'precise', 'verification-PBE-v1'))
spec.inputs['protocol'].valid_type = ChoiceType(
(
'fast',
'moderate',
'precise',
'verification-PBE-v1',
'verification-PBE-v1-lcao-dzp',
'verification-PBE-v1-lcao-tzdp',
'verification-PBE-v1-lcao-dzp-sg15',
'verification-PBE-v1-lcao-tzdp-sg15',
'verification-PBE-v1-lcao-apns-efficiency',
)
)
# spec.inputs['protocol'].valid_types = None
spec.inputs['spin_type'].valid_type = ChoiceType((SpinType.NONE, SpinType.COLLINEAR))
spec.inputs['relax_type'].valid_type = ChoiceType(tuple(RelaxType))
spec.inputs['electronic_type'].valid_type = ChoiceType((ElectronicType.METAL, ElectronicType.INSULATOR))
spec.inputs['engines']['relax']['code'].valid_type = CodeType('abacus.abacus')

def _construct_builder(self, **kwargs) -> engine.ProcessBuilder: # noqa: PLR0915
"""Construct a process builder based on the provided keyword arguments.

The keyword arguments will have been validated against the input generator specification.
"""
from aiida_abacus.common import CONSTANTS, ElectronicType, RelaxType, SpinType, recursive_merge

structure = kwargs['structure']
engines = kwargs['engines']
protocol = kwargs['protocol']
spin_type = kwargs['spin_type']
relax_type = kwargs['relax_type']
electronic_type = kwargs['electronic_type']
# magnetization_per_site = kwargs.get('magnetization_per_site', None)
threshold_forces = kwargs.get('threshold_forces', None)
threshold_stress = kwargs.get('threshold_stress', None)
reference_workchain = kwargs.get('reference_workchain', None)

if isinstance(electronic_type, str):
electronic_type = ElectronicType(electronic_type)
else:
electronic_type = ElectronicType(electronic_type.value)

if isinstance(relax_type, str):
relax_type = RelaxType(relax_type)
else:
relax_type = RelaxType(relax_type.value)

if isinstance(spin_type, str):
spin_type = SpinType(spin_type)
else:
spin_type = SpinType(spin_type.value)

# if magnetization_per_site:
# kind_to_magnetization = set(zip([site.kind_name for site in structure.sites], magnetization_per_site))

# if len(structure.kinds) != len(kind_to_magnetization):
# structure, initial_magnetic_moments = create_magnetic_allotrope(structure, magnetization_per_site)
# else:
# initial_magnetic_moments = dict(kind_to_magnetization)
# else:
initial_magnetic_moments = None

# Currently, the `aiida-abcus` workflows will expect one of the basic protocols to be passed to the
# `get_builder_from_protocol()` method. Here, we switch to using the default protocol for the
# `aiida-abacus` plugin and pass the local protocols as `overrides`.
if protocol not in self.process_class._process_class.get_available_protocols():
overrides = self._load_local_protocols()[protocol]
protocol = self._default_protocol
else:
overrides = {}

options_overrides = {
'base': {'abacus': {'metadata': {'options': engines['relax']['options']}}},
'base_final_scf': {'abacus': {'metadata': {'options': engines['relax']['options']}}},
}
overrides = recursive_merge(overrides, options_overrides)

builder = self.process_class._process_class.get_builder_from_protocol(
engines['relax']['code'],
structure,
protocol=protocol,
overrides=overrides,
relax_type=relax_type,
electronic_type=electronic_type,
spin_type=spin_type,
initial_magnetic_moments=initial_magnetic_moments,
)

if threshold_forces is not None:
threshold = threshold_forces * CONSTANTS.bohr_to_ang.value / CONSTANTS.ry_to_ev.value
parameters = builder.base['abacus']['parameters'].get_dict()
parameters.setdefault('input', {})['force_thr'] = threshold
builder.base['abacus']['parameters'] = orm.Dict(dict=parameters)

if threshold_stress is not None:
threshold = threshold_stress * CONSTANTS.ev_ang3_to_kbar.value # Abacus uses kBar for stress threshold
parameters = builder.base['abacus']['parameters'].get_dict()
parameters.setdefault('input', {})['stress_thr'] = threshold
builder.base['abacus']['parameters'] = orm.Dict(dict=parameters)

if reference_workchain:
relax = reference_workchain.base.links.get_outgoing(node_class=orm.WorkChainNode).one().node
base = sorted(relax.called, key=lambda x: x.ctime)[-1]
calc = sorted(base.called, key=lambda x: x.ctime)[-1]
kpoints = calc.inputs.kpoints

builder.base.pop('kpoints_distance', None)
builder.base.pop('kpoints_force_parity', None)
builder.base_final_scf.pop('kpoints_distance', None)
builder.base_final_scf.pop('kpoints_force_parity', None)

builder.base['kpoints'] = kpoints
builder.base_final_scf['kpoints'] = kpoints

# Currently the builder is set for the `AbacusRelaxWorkChain`, but we should return one for the wrapper
# workchain
# `AbacusCommonRelaxWorkChain` for which this input generator is built
builder._process_class = self.process_class

return builder
Loading