Adding support for aiida-atomistic StructureData

src/aiida_quantumespresso/calculations/__init__.py

@@ -18,13 +18,23 @@
 from aiida_quantumespresso.data.hubbard_structure import HubbardStructureData
 from aiida_quantumespresso.utils.convert import convert_input_to_namelist_entry
 from aiida_quantumespresso.utils.hubbard import HubbardUtils
+from aiida_quantumespresso.utils.magnetic import MagneticUtils
 
 from .base import CalcJob
 from .helpers import QEInputValidationError
 
 LegacyUpfData = DataFactory('core.upf')
 UpfData = DataFactory('pseudo.upf')
 
+LegacyStructureData = DataFactory('core.structure')  # pylint: disable=invalid-name
+
+try:
+    StructureData = DataFactory('atomistic.structure')
+except exceptions.MissingEntryPointError:
+    structures_classes = (LegacyStructureData,)
+else:
+    structures_classes = (LegacyStructureData, StructureData)
+
 
 class BasePwCpInputGenerator(CalcJob):
     """Base `CalcJob` for implementations for pw.x and cp.x of Quantum ESPRESSO."""
@@ -94,6 +104,8 @@ class BasePwCpInputGenerator(CalcJob):
 
     _use_kpoints = False
 
+    supported_properties = ['magmoms', 'hubbard']
+
     @classproperty
     def xml_filenames(cls):
         """Return a list of XML output filenames that can be written by a calculation.
@@ -116,7 +128,7 @@ def define(cls, spec):
         spec.input('metadata.options.input_filename', valid_type=str, default=cls._DEFAULT_INPUT_FILE)
         spec.input('metadata.options.output_filename', valid_type=str, default=cls._DEFAULT_OUTPUT_FILE)
         spec.input('metadata.options.withmpi', valid_type=bool, default=True)  # Override default withmpi=False
-        spec.input('structure', valid_type=orm.StructureData,
+        spec.input('structure', valid_type=structures_classes,
             help='The input structure.')
         spec.input('parameters', valid_type=orm.Dict,
             help='The input parameters that are to be used to construct the input file.')
@@ -168,6 +180,21 @@ def validate_inputs(cls, value, port_namespace):
         if any(key not in port_namespace for key in ('pseudos', 'structure')):
             return
 
+        if not isinstance(value['structure'], LegacyStructureData):
+            # we have the atomistic StructureData, so we need to check if all the defined properties are supported
+            plugin_check = value['structure'].check_plugin_support(cls.supported_properties)
+            if len(plugin_check) > 0:
+                raise NotImplementedError(
+                    f'The input structure contains one or more unsupported properties \
+                    for this process: {plugin_check}'
+                )
+
+        if value['structure'].is_alloy or value['structure'].has_vacancies:
+            raise exceptions.InputValidationError(
+                'The structure is an alloy or has vacancies. This is not allowed for \
+                aiida-quantumespresso input structures.'
+            )
+
         # At this point, both ports are part of the namespace, and both are required so return an error message if any
         # of the two is missing.
         for key in ('pseudos', 'structure'):
@@ -702,9 +729,17 @@ def _generate_PWCPinputdata(cls, parameters, settings, pseudos, structure, kpoin
             kpoints_card = ''.join(kpoints_card_list)
             del kpoints_card_list
 
-        # HUBBARD CARD
-        hubbard_card = HubbardUtils(structure).get_hubbard_card() if isinstance(structure, HubbardStructureData) \
-            else None
+        # HUBBARD CARD and MAGNETIC NAMELIST
+        hubbard_card = None
+        magnetic_namelist = None
+        if isinstance(structure, HubbardStructureData):
+            hubbard_card = HubbardUtils(structure).get_hubbard_card()
+        elif len(structures_classes) == 2 and not isinstance(structure, LegacyStructureData):
+            # this means that we have the atomistic StructureData.
+            hubbard_card = HubbardUtils(structure).get_hubbard_card() if 'hubbard' \
+                in structure.get_defined_properties() else None
+            magnetic_namelist = MagneticUtils(structure).generate_magnetic_namelist(input_params) if 'magmoms' in \
+                structure.get_defined_properties()  else None
 
         # =================== NAMELISTS AND CARDS ========================
         try:
@@ -734,6 +769,14 @@ def _generate_PWCPinputdata(cls, parameters, settings, pseudos, structure, kpoin
                     'namelists using the NAMELISTS inside the `settings` input node'
                 ) from exception
 
+        if magnetic_namelist is not None:
+            if input_params['SYSTEM'].get('nspin', 1) == 1 and not input_params['SYSTEM'].get('noncolin', False):
+                raise exceptions.InputValidationError(
+                    'The structure has magnetic moments but the inputs are not set for \
+                    a magnetic calculation (`nspin`, `noncolin`)'
+                )
+            input_params['SYSTEM'].update(magnetic_namelist)
+
         inputfile = ''
         for namelist_name in namelists_toprint:
             inputfile += f'&{namelist_name}\n'

src/aiida_quantumespresso/calculations/functions/seekpath_structure_analysis.py

@@ -1,10 +1,18 @@
 # -*- coding: utf-8 -*-
 """Calcfunction to primitivize a structure and return high symmetry k-point path through its Brillouin zone."""
+from aiida.common import exceptions
 from aiida.engine import calcfunction
 from aiida.orm import Data
+from aiida.plugins import DataFactory
 
 from aiida_quantumespresso.data.hubbard_structure import HubbardStructureData
 
+try:
+    StructureData = DataFactory('atomistic.structure')
+    HAS_ATOMISTIC = True
+except exceptions.MissingEntryPointError:
+    HAS_ATOMISTIC = False
+
 
 @calcfunction
 def seekpath_structure_analysis(structure, **kwargs):
@@ -32,7 +40,10 @@ def seekpath_structure_analysis(structure, **kwargs):
 
     result = get_explicit_kpoints_path(structure, **unwrapped_kwargs)
 
-    if isinstance(structure, HubbardStructureData):
+    if HAS_ATOMISTIC:
+        if isinstance(structure, StructureData):
+            raise NotImplementedError('This function does not yet support the conversion into atomistic instances.')
+    elif isinstance(structure, HubbardStructureData):
         result['primitive_structure'] = update_structure_with_hubbard(result['primitive_structure'], structure)
         result['conv_structure'] = update_structure_with_hubbard(result['conv_structure'], structure)
 

src/aiida_quantumespresso/calculations/pw.py

@@ -4,11 +4,20 @@
 import warnings
 
 from aiida import orm
+from aiida.common import exceptions
 from aiida.common.lang import classproperty
+from aiida.orm import StructureData as LegacyStructureData
 from aiida.plugins import factories
 
 from aiida_quantumespresso.calculations import BasePwCpInputGenerator
 
+try:
+    StructureData = factories.DataFactory('atomistic.structure')
+except exceptions.MissingEntryPointError:
+    structures_classes = (LegacyStructureData,)
+else:
+    structures_classes = (LegacyStructureData, StructureData)
+
 
 class PwCalculation(BasePwCpInputGenerator):
     """`CalcJob` implementation for the pw.x code of Quantum ESPRESSO."""
@@ -69,13 +78,13 @@ def define(cls, spec):
             'will not fail if the XML file is missing in the retrieved folder.')
         spec.input('kpoints', valid_type=orm.KpointsData,
             help='kpoint mesh or kpoint path')
-        spec.input('hubbard_file', valid_type=orm.SinglefileData, required=False,
+        spec.input('hubbard_file', valid_type=structures_classes, required=False,
             help='SinglefileData node containing the output Hubbard parameters from a HpCalculation')
         spec.inputs.validator = cls.validate_inputs
 
         spec.output('output_parameters', valid_type=orm.Dict,
             help='The `output_parameters` output node of the successful calculation.')
-        spec.output('output_structure', valid_type=orm.StructureData, required=False,
+        spec.output('output_structure', valid_type=structures_classes, required=False,
             help='The `output_structure` output node of the successful calculation if present.')
         spec.output('output_trajectory', valid_type=orm.TrajectoryData, required=False)
         spec.output('output_band', valid_type=orm.BandsData, required=False,

src/aiida_quantumespresso/parsers/parse_raw/base.py

@@ -2,7 +2,15 @@
 """A basic parser for the common format of QE."""
 import re
 
-from aiida.orm import StructureData
+from aiida.orm import StructureData as LegacyStructureData
+from aiida.plugins import DataFactory
+
+try:
+    StructureData = DataFactory('atomistic.structure')
+except exceptions.MissingEntryPointError:
+    structures_classes = (LegacyStructureData,)
+else:
+    structures_classes = (LegacyStructureData, StructureData)
 
 __all__ = ('convert_qe_time_to_sec', 'convert_qe_to_aiida_structure', 'convert_qe_to_kpoints')
 
@@ -40,25 +48,46 @@ def convert_qe_time_to_sec(timestr):
 
 
 def convert_qe_to_aiida_structure(output_dict, input_structure=None):
-    """Convert the dictionary parsed from the Quantum ESPRESSO output into ``StructureData``."""
+    """Convert the dictionary parsed from the Quantum ESPRESSO output into ``StructureData``.
+    If we have an ``orm.StructureData`` as input, we return an ``orm.StructureData`` instance,
+    otherwise we always return an aiida-atomistic ``StructureData``.
+    """
 
     cell_dict = output_dict['cell']
 
     # Without an input structure, try to recreate the structure from the output
     if not input_structure:
 
-        structure = StructureData(cell=cell_dict['lattice_vectors'])
+        if isinstance(input_structure, LegacyStructureData):
+            structure = LegacyStructureData()
+            structure.set_cell=cell_dict['lattice_vectors']
+
+            for kind_name, position in output_dict['atoms']:
+                symbol = re.sub(r'\d+', '', kind_name)
+                structure.append_atom(position=position, symbols=symbol, name=kind_name)
+
+        else:
+            structure = StructureDataMutable()
+            structure.set_cell=cell_dict['lattice_vectors']
+
+            for kind_name, position in output_dict['atoms']:
+                symbol = re.sub(r'\d+', '', kind_name)
+                structure.append_atom(positions=position, symbols=symbol, kinds=kind_name)
 
-        for kind_name, position in output_dict['atoms']:
-            symbol = re.sub(r'\d+', '', kind_name)
-            structure.append_atom(position=position, symbols=symbol, name=kind_name)
+            structure = StructureData.from_mutable(structure)
 
     else:
 
-        structure = input_structure.clone()
-        structure.reset_cell(cell_dict['lattice_vectors'])
-        new_pos = [i[1] for i in cell_dict['atoms']]
-        structure.reset_sites_positions(new_pos)
+        if isinstance(input_structure, LegacyStructureData):
+            structure = input_structure.clone()
+            structure.reset_cell(cell_dict['lattice_vectors'])
+            new_pos = [i[1] for i in cell_dict['atoms']]
+            structure.reset_sites_positions(new_pos)
+        elif isinstance(input_structure, StructureData):
+            structure = input_structure.get_value()
+            structure.set_cell(cell_dict['lattice_vectors'])
+            for site,position in zip(structure.sites,[i[1] for i in cell_dict['atoms']]):
+                site.position = position
 
     return structure
 

src/aiida_quantumespresso/utils/magnetic.py

@@ -0,0 +1,96 @@
+# -*- coding: utf-8 -*-
+"""Utility class for handling the :class:`aiida_quantumespresso.data.hubbard_structure.HubbardStructureData`."""
+# pylint: disable=no-name-in-module
+
+from aiida import orm
+from aiida.common.exceptions import MissingEntryPointError
+from aiida.engine import calcfunction
+from aiida.orm import StructureData as LegacyStructureData
+from aiida.plugins import DataFactory
+
+try:
+    StructureData = DataFactory('atomistic.structure')
+except MissingEntryPointError:
+    structures_classes = (LegacyStructureData,)
+else:
+    structures_classes = (LegacyStructureData, StructureData)
+
+
+class MagneticUtils:  # pylint: disable=too-few-public-methods
+    """Class to manage the magnetic structure of the atomistic `LegacyStructureData`.
+
+    It contains methods to manipulate the magne tic structure in such a way to produce
+    the correct input for QuantumESPRESSO calculations.
+    """
+
+    def __init__(
+        self,
+        structure: structures_classes,
+    ):
+        """Set a the `StructureData` to manipulate."""
+        if isinstance(structure, StructureData):
+            if 'magmoms' not in structure.get_defined_properties():
+                raise ValueError('The input structure does not contain magnetic moments.')
+            self.structure = structure
+        else:
+            raise ValueError('input is not of type atomistic `StructureData')
+
+    def generate_magnetic_namelist(self, parameters):
+        """Generate the magnetic namelist for Quantum ESPRESSO.
+
+        :param parameters: dictionary of inputs for the Quantum ESPRESSO calculation.
+        """
+        if 'nspin' not in parameters['SYSTEM'] and 'noncolin' not in parameters['SYSTEM']:
+            raise ValueError("The input parameters must contain the 'nspin' or the 'noncolin' key.")
+
+        namelist = {'starting_magnetization': {}, 'angle1': {}, 'angle2': {}}
+
+        if parameters['SYSTEM'].get('nspin', None) == 2:
+            namelist.pop('angle1')
+            namelist.pop('angle2')
+            if self.structure.is_collinear:
+                for kind, magmom in zip(self.structure.kinds, self.structure.magmoms):
+                    # this should be fixed, now only magmom_z is considered...
+                    if magmom[2] != 0:
+                        namelist['starting_magnetization'][kind] = magmom[2]
+            else:
+                raise NotImplementedError(
+                    'The input structure is not collinear, but you choose collinear calculations.'
+                )
+        elif parameters['SYSTEM']['noncolin']:
+            for site in self.structure.sites:
+                for variable, value in namelist.items():
+                    value[site.kinds] = site.get_magmom_coord(coord='spherical')[variable]
+
+        return namelist
+
+
+@calcfunction
+def generate_structure_with_magmoms(input_structure: structures_classes, input_magnetic_moments: orm.List):
+    """Generate a new structure with the magnetic moments for each site.
+
+    :param input_structure: the input structure to add the magnetic moments.
+    :param input_magnetic_moments: the magnetic moments for each site, represented as a float (see below).
+
+    For now, only supports collinear magnetic moments, i.e. atomic magnetizations (along z axis).
+    """
+    magmoms = input_magnetic_moments.get_list()
+    if len(input_structure.sites) != len(magmoms):
+        raise ValueError('The input structure and the magnetic moments must have the same length.')
+
+    mutable_structure = input_structure.get_value()
+    mutable_structure.clear_sites()
+    for site, magmom in zip(input_structure.sites, magmoms):
+        mutable_structure.add_atom(
+            **{
+                'positions': site.positions,
+                'symbols': site.symbols,
+                'kinds': site.kinds,
+                'weights': site.weights,
+                'magmoms': [0, 0, magmom] if isinstance(magmom, float) else magmom  # 3D vector
+            }
+        )
+
+    output_structure = StructureData.from_mutable(mutable_structure, detect_kinds=True)
+
+    return output_structure