Add implementation for PySCF

PySCF is a Python-based Simulations of Chemistry Framework and the
implementation is provided through the `aiida-pyscf` plugin.

Author: Sebastiaan Huber

pyproject.toml

@@ -43,6 +43,7 @@ requires-python = '>=3.9'
 'common_workflows.relax.gpaw' = 'aiida_common_workflows.workflows.relax.gpaw.workchain:GpawCommonRelaxWorkChain'
 'common_workflows.relax.nwchem' = 'aiida_common_workflows.workflows.relax.nwchem.workchain:NwchemCommonRelaxWorkChain'
 'common_workflows.relax.orca' = 'aiida_common_workflows.workflows.relax.orca.workchain:OrcaCommonRelaxWorkChain'
+'common_workflows.relax.pyscf' = 'aiida_common_workflows.workflows.relax.pyscf.workchain:PyscfCommonRelaxWorkChain'
 'common_workflows.relax.quantum_espresso' = 'aiida_common_workflows.workflows.relax.quantum_espresso.workchain:QuantumEspressoCommonRelaxWorkChain'
 'common_workflows.relax.siesta' = 'aiida_common_workflows.workflows.relax.siesta.workchain:SiestaCommonRelaxWorkChain'
 'common_workflows.relax.vasp' = 'aiida_common_workflows.workflows.relax.vasp.workchain:VaspCommonRelaxWorkChain'
@@ -64,6 +65,7 @@ all_plugins = [
   'aiida-gaussian~=2.0',
   'aiida-nwchem~=3.0',
   'aiida-orca~=0.6.0',
+  'aiida-pyscf~=0.5.1',
   'aiida-quantumespresso~=4.4',
   'aiida-siesta~=2.0',
   'aiida-vasp~=3.1',
@@ -105,6 +107,9 @@ orca = [
 pre-commit = [
   'pre-commit~=3.6'
 ]
+pyscf = [
+  'aiida-pyscf~=0.5.1'
+]
 quantum_espresso = [
   'aiida-quantumespresso~=4.4'
 ]

src/aiida_common_workflows/workflows/relax/generator.py

@@ -59,7 +59,7 @@ def define(cls, spec):
         )
         spec.input(
             'magnetization_per_site',
-            valid_type=list,
+            valid_type=(list, tuple),
             required=False,
             help='The initial magnetization of the system. Should be a list of floats, where each float represents the '
             'spin polarization in units of electrons, meaning the difference between spin up and spin down '

src/aiida_common_workflows/workflows/relax/pyscf/__init__.py

@@ -0,0 +1,5 @@
+"""Module with the implementations of the common structure relaxation workchain for pyscf."""
+from .generator import *
+from .workchain import *
+
+__all__ = generator.__all__ + workchain.__all__

src/aiida_common_workflows/workflows/relax/pyscf/generator.py

@@ -0,0 +1,105 @@
+"""Implementation of `aiida_common_workflows.common.relax.generator.CommonRelaxInputGenerator` for pyscf."""
+import pathlib
+import warnings
+
+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__ = ('PyscfCommonRelaxInputGenerator',)
+
+StructureData = plugins.DataFactory('structure')
+
+
+class PyscfCommonRelaxInputGenerator(CommonRelaxInputGenerator):
+    """Input generator for the common relax workflow implementation of pyscf."""
+
+    def __init__(self, *args, **kwargs):
+        """Construct an instance of the input generator, validating the class attributes."""
+        process_class = kwargs.get('process_class', None)
+        super().__init__(*args, **kwargs)
+        self._initialize_protocols()
+
+    def _initialize_protocols(self):
+        """Initialize the protocols class attribute by parsing them from the configuration file."""
+        with (pathlib.Path(__file__).parent / 'protocol.yml').open() as handle:
+            self._protocols = yaml.safe_load(handle)
+            self._default_protocol = 'moderate'
+
+    @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['spin_type'].valid_type = ChoiceType((SpinType.NONE, SpinType.COLLINEAR))
+        spec.inputs['relax_type'].valid_type = ChoiceType((RelaxType.NONE, RelaxType.POSITIONS))
+        spec.inputs['electronic_type'].valid_type = ChoiceType((ElectronicType.METAL, ElectronicType.INSULATOR))
+        spec.inputs['engines']['relax']['code'].valid_type = CodeType('pyscf.base')
+
+    def _construct_builder(
+        self,
+        structure,
+        engines,
+        protocol,
+        spin_type,
+        relax_type,
+        electronic_type,
+        magnetization_per_site=None,
+        **kwargs,
+    ) -> engine.ProcessBuilder:
+        """Construct a process builder based on the provided keyword arguments.
+
+        The keyword arguments will have been validated against the input generator specification.
+        """
+        if not self.is_valid_protocol(protocol):
+            raise ValueError(
+                f'selected protocol {protocol} is not valid, please choose from: {", ".join(self.get_protocol_names())}'
+            )
+
+        protocol_inputs = self.get_protocol(protocol)
+        parameters = protocol_inputs.pop('parameters')
+
+        if relax_type == RelaxType.NONE:
+            parameters.pop('optimizer')
+
+        if spin_type == SpinType.COLLINEAR:
+            parameters['mean_field']['method'] = 'DKS'
+            parameters['mean_field']['collinear'] = 'mcol'
+
+        num_electrons = structure.get_pymatgen_molecule().nelectrons
+
+        if spin_type == SpinType.NONE and num_electrons % 2 == 1:
+            raise ValueError('structure has odd number of electrons, please select `spin_type = SpinType.COLLINEAR`')
+
+        if spin_type == SpinType.COLLINEAR:
+            if magnetization_per_site is None:
+                multiplicity = 1
+            else:
+                warnings.warn('magnetization_per_site site-resolved info is disregarded, only total spin is processed.')
+                # ``magnetization_per_site`` is in units of Bohr magnetons, multiple by 0.5 to get atomic units
+                total_spin = 0.5 * abs(sum(magnetization_per_site))
+                multiplicity = 2 * total_spin + 1
+
+            # In case of even/odd electrons, find closest odd/even multiplicity
+            if num_electrons % 2 == 0:
+                # round guess to nearest odd integer
+                spin_multiplicity = int(round((multiplicity - 1) / 2) * 2 + 1)
+            else:
+                # round guess to nearest even integer; 0 goes to 2
+                spin_multiplicity = max([int(round(multiplicity / 2) * 2), 2])
+
+            parameters['structure']['spin'] = int((spin_multiplicity - 1) / 2)
+
+        builder = self.process_class.get_builder()
+        builder.pyscf.code = engines['relax']['code']
+        builder.pyscf.structure = structure
+        builder.pyscf.parameters = orm.Dict(parameters)
+        builder.pyscf.metadata.options = engines['relax']['options']
+
+        return builder

src/aiida_common_workflows/workflows/relax/pyscf/protocol.yml

@@ -0,0 +1,35 @@
+fast:
+  description: Protocol to relax a structure with low precision at minimal computational cost for testing purposes.
+  parameters:
+    mean_field:
+      method: UKS
+    structure:
+      basis: def2-svp
+    optimizer:
+      solver: geomeTRIC
+      convergence_parameters:
+        convergence_energy: 1E-6
+moderate:
+  description: Protocol to relax a structure with normal precision at moderate computational cost.
+  parameters:
+    mean_field:
+      method: UKS
+      xc: pbe
+    structure:
+      basis: def2-tzvp
+    optimizer:
+      solver: geomeTRIC
+      # convergence_parameters:
+      #     convergence_energy: 1E-6
+precise:
+  description: Protocol to relax a structure with high precision at higher computational cost.
+  parameters:
+    mean_field:
+      method: UHF
+      # xc: 'pbe'
+    structure:
+      basis: def2-qzvp
+    optimizer:
+      solver: geomeTRIC
+      # convergence_parameters:
+      #     convergence_energy: 1E-6

src/aiida_common_workflows/workflows/relax/pyscf/workchain.py

@@ -0,0 +1,44 @@
+"""Implementation of `aiida_common_workflows.common.relax.workchain.CommonRelaxWorkChain` for pyscf."""
+import numpy
+from aiida import orm
+from aiida.engine import calcfunction
+from aiida.plugins import WorkflowFactory
+
+from ..workchain import CommonRelaxWorkChain
+from .generator import PyscfCommonRelaxInputGenerator
+
+__all__ = ('PyscfCommonRelaxWorkChain',)
+
+
+@calcfunction
+def extract_energy_from_parameters(parameters):
+    """Return the total energy from the given parameters node."""
+    total_energy = parameters.get_attribute('total_energy')
+    return {'total_energy': orm.Float(total_energy)}
+
+
+@calcfunction
+def extract_forces_from_parameters(parameters):
+    """Return the forces from the given parameters node."""
+    forces = orm.ArrayData()
+    forces.set_array('forces', numpy.array(parameters.get_attribute('forces')))
+    return {'forces': forces}
+
+
+class PyscfCommonRelaxWorkChain(CommonRelaxWorkChain):
+    """Implementation of `aiida_common_workflows.common.relax.workchain.CommonRelaxWorkChain` for pyscf."""
+
+    _process_class = WorkflowFactory('pyscf.base')
+    _generator_class = PyscfCommonRelaxInputGenerator
+
+    def convert_outputs(self):
+        """Convert the outputs of the sub workchain to the common output specification."""
+        outputs = self.ctx.workchain.outputs
+        total_energy = extract_energy_from_parameters(outputs.parameters)['total_energy']
+        forces = extract_forces_from_parameters(outputs.parameters)['forces']
+
+        if 'structure' in outputs:
+            self.out('relaxed_structure', outputs.structure)
+
+        self.out('total_energy', total_energy)
+        self.out('forces', forces)