Add the create_magnetic_allotrope function

In order to pass on the magnetic configuration between calculations, the
structure used for the follow-up calculation needs to have the right
magnetic kinds to be able to properly assign the magnetisation to each
site.

Here we add a calculation function that, based on the structure and
magnetic moments, returns a new `StructureData` with the required
magnetic kinds, as well as a `Dict` with the corresponding magnetic
moments for each kind.

pyproject.toml

@@ -67,6 +67,7 @@ aiida-quantumespresso = 'aiida_quantumespresso.cli:cmd_root'
 [project.entry-points.'aiida.calculations']
 'quantumespresso.cp' = 'aiida_quantumespresso.calculations.cp:CpCalculation'
 'quantumespresso.create_kpoints_from_distance' = 'aiida_quantumespresso.calculations.functions.create_kpoints_from_distance:create_kpoints_from_distance'
+'quantumespresso.create_magnetic_allotrope' = 'aiida_quantumespresso.calculations.functions.create_magnetic_allotrope:create_magnetic_allotrope'
 'quantumespresso.merge_ph_outputs' = 'aiida_quantumespresso.calculations.functions.merge_ph_outputs:merge_ph_outputs'
 'quantumespresso.dos' = 'aiida_quantumespresso.calculations.dos:DosCalculation'
 'quantumespresso.epw' = 'aiida_quantumespresso.calculations.epw:EpwCalculation'

src/aiida_quantumespresso/calculations/functions/create_magnetic_allotrope.py

@@ -0,0 +1,121 @@
+# -*- coding: utf-8 -*-
+"""Create a new magnetic allotrope from the given structure based on the desired magnetic moments."""
+from aiida.engine import calcfunction
+from aiida.orm import Float
+import numpy
+
+
+@calcfunction
+def create_magnetic_allotrope(structure, magnetic_moment_per_site, atol=lambda: Float(5E-1), ztol=lambda: Float(5E-2)):
+    """Create a new magnetic allotrope from the given structure based on a list of magnetic moments per site.
+
+    To create the new list of kinds, the algorithm loops over all the elements in the structure and makes a list of the
+    sites with that element and their corresponding magnetic moment. Next, it splits this list in three lists:
+
+    * Zero magnetic moments: Any site that has an absolute magnetic moment lower than ``ztol``
+    * Positive magnetic moments
+    * Negative magnetic moments
+
+    The algorithm then sorts the positive and negative lists from large to small absolute value, and loops over each of
+    list. New magnetic kinds will be created when the absolute difference between the magnetic moment of the current
+    kind and the site exceeds ``atol``.
+
+    The positive and negative magnetic moments are handled separately to avoid assigning two sites with opposite signs
+    in their magnetic moment to the same kind and make sure that each kind has the correct magnetic moment, i.e. the
+    largest magnetic moment in absolute value of the sites corresponding to that kind.
+
+    .. important:: the function currently does not support alloys.
+
+    :param structure: a `StructureData` instance.
+    :param magnetic_moment_per_site: list of magnetic moments for each site in the structure.
+    :param atol: the absolute tolerance on determining if two sites have the same magnetic moment.
+    :param ztol: threshold for considering a kind to have non-zero magnetic moment.
+    """
+    # pylint: disable=too-many-locals,too-many-branches,too-many-statements
+    import string
+
+    from aiida.orm import Dict, StructureData
+
+    if structure.is_alloy:
+        raise ValueError('Alloys are currently not supported.')
+
+    atol = atol.value
+    rtol = 0  # Relative tolerance used in the ``numpy.is_close()`` calls.
+    ztol = ztol.value
+
+    allotrope = StructureData(cell=structure.cell, pbc=structure.pbc)
+    allotrope_magnetic_moments = {}
+
+    for element in structure.get_symbols_set():
+
+        # Filter the sites and magnetic moments on the site element
+        element_sites, element_magnetic_moments = zip(
+            *[(site, magnetic_moment)
+              for site, magnetic_moment in zip(structure.sites, magnetic_moment_per_site)
+              if site.kind_name.rstrip(string.digits) == element]
+        )
+
+        # Split the sites and their magnetic moments by sign to filter out the sites with magnetic moment lower than
+        # `ztol`and deal with the positive and negative magnetic moment separately. This is important to avoid assigning
+        # two sites with opposite signs to the same kind and make sure that each kind has the correct magnetic moment,
+        # i.e. the largest magnetic moment in absolute value of the sites corresponding to that kind.
+        zero_sites = []
+        pos_sites = []
+        neg_sites = []
+
+        for site, magnetic_moment in zip(element_sites, element_magnetic_moments):
+
+            if abs(magnetic_moment) <= ztol:
+                zero_sites.append((site, 0))
+            elif magnetic_moment > 0:
+                pos_sites.append((site, magnetic_moment))
+            else:
+                neg_sites.append((site, magnetic_moment))
+
+        kind_index = -1
+        kind_names = []
+        kind_sites = []
+        kind_magnetic_moments = {}
+
+        for site_list in (zero_sites, pos_sites, neg_sites):
+
+            if not site_list:
+                continue
+
+            # Sort the site list in order to build the kind lists from large to small absolute magnetic moment.
+            site_list = sorted(site_list, key=lambda x: abs(x[1]), reverse=True)
+
+            sites, magnetic_moments = zip(*site_list)
+
+            kind_index += 1
+            current_kind_name = f'{element}{kind_index}'
+            kind_sites.append(sites[0])
+            kind_names.append(current_kind_name)
+            kind_magnetic_moments[current_kind_name] = magnetic_moments[0]
+
+            for site, magnetic_moment in zip(sites[1:], magnetic_moments[1:]):
+
+                if not numpy.isclose(magnetic_moment, kind_magnetic_moments[current_kind_name], rtol, atol):
+                    kind_index += 1
+                    current_kind_name = f'{element}{kind_index}'
+                    kind_magnetic_moments[current_kind_name] = magnetic_moment
+
+                kind_sites.append(site)
+                kind_names.append(current_kind_name)
+
+        # In case there is only a single kind for the element, remove the 0 kind index
+        if current_kind_name == f'{element}0':
+            kind_names = len(element_magnetic_moments) * [element]
+            kind_magnetic_moments = {element: kind_magnetic_moments[current_kind_name]}
+
+        allotrope_magnetic_moments.update(kind_magnetic_moments)
+
+        for name, site in zip(kind_names, kind_sites):
+            allotrope.append_atom(
+                name=name,
+                symbols=(element,),
+                weights=(1.0,),
+                position=site.position,
+            )
+
+    return {'allotrope': allotrope, 'magnetic_moments': Dict(dict=allotrope_magnetic_moments)}

tests/calculations/functions/test_create_magnetic_allotrope.py

@@ -0,0 +1,197 @@
+# -*- coding: utf-8 -*-
+"""Tests for the `create_magnetic_allotrope` calculation function."""
+from aiida.orm import Float, List
+from aiida.plugins import CalculationFactory
+import pytest
+
+create_magnetic_allotrope = CalculationFactory('quantumespresso.create_magnetic_allotrope')
+
+
+@pytest.mark.usefixtures('aiida_profile')
+def test_configuration_00(generate_structure_from_kinds):
+    """Test `create_magnetic_allotrope` calculation function.
+
+    Case: one kind but with equal magnetic moments.
+    Expected result: no new kind names should be introduced.
+    """
+    kind_names = ['Fe', 'Fe']
+    magnetic_moments = List(list=[0.2, 0.2])
+
+    structure = generate_structure_from_kinds(kind_names)
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments).values()
+
+    assert set(allotrope.get_kind_names()) == {'Fe', 'Fe'}
+    assert allotrope_magnetic_moments.get_dict() == {'Fe': 0.2}
+
+
+@pytest.mark.usefixtures('aiida_profile')
+def test_configuration_01(generate_structure_from_kinds):
+    """Test `create_magnetic_allotrope` calculation function.
+
+    Case: two kinds all with equal magnetic moments.
+    Expected result: no new kind names should be introduced.
+    """
+    kind_names = ['Fe', 'Fe', 'Ni', 'Ni', 'Ni']
+    magnetic_moments = List(list=[0.2, 0.2, 0.5, 0.5, 0.5])
+
+    structure = generate_structure_from_kinds(kind_names)
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments).values()
+
+    assert set(allotrope.get_kind_names()) == {'Fe', 'Ni'}
+    assert allotrope_magnetic_moments.get_dict() == {'Fe': 0.2, 'Ni': 0.5}
+
+
+@pytest.mark.usefixtures('aiida_profile')
+def test_configuration_02(generate_structure_from_kinds):
+    """Test `create_magnetic_allotrope` calculation function.
+
+    Case: only one kind but with unequal magnetic moments.
+    Expected result: two new kinds introduced one for each magnetic moment.
+    """
+    kind_names = ['Fe', 'Fe']
+    magnetic_moments = List(list=[0.2, 1.0])
+
+    structure = generate_structure_from_kinds(kind_names)
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments).values()
+
+    assert set(allotrope.get_kind_names()) == {'Fe0', 'Fe1'}
+    assert allotrope_magnetic_moments.get_dict() == {'Fe0': 1.0, 'Fe1': 0.2}
+
+
+@pytest.mark.usefixtures('aiida_profile')
+def test_configuration_03(generate_structure_from_kinds):
+    """Test `create_magnetic_allotrope` calculation function.
+
+    Case: only one kind but with three types of magnetic moments that are not grouped together.
+    Expected result: two new kinds introduced one for each magnetic moment.
+    """
+    kind_names = ['Fe', 'Fe', 'Fe', 'Fe']
+    magnetic_moments = List(list=[0.2, 0.8, 1.5, 0.8])
+
+    structure = generate_structure_from_kinds(kind_names)
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments).values()
+
+    assert set(allotrope.get_kind_names()) == {'Fe0', 'Fe1', 'Fe2'}
+    assert allotrope_magnetic_moments.get_dict() == {'Fe0': 1.5, 'Fe1': 0.8, 'Fe2': 0.2}
+
+
+@pytest.mark.usefixtures('aiida_profile')
+def test_configuration_04(generate_structure_from_kinds):
+    """Test `create_magnetic_allotrope` calculation function.
+
+    Case: only one kind but with four different values of magnetic moments but middle two are within tolerance.
+    Expected result: two new kinds introduced one for each magnetic moment.
+    """
+    kind_names = ['Fe', 'Fe', 'Fe', 'Fe']
+    magnetic_moments = List(list=[0.0, 0.50, 0.45, 0.40])
+
+    structure = generate_structure_from_kinds(kind_names)
+
+    # Default tolerances: just two different kinds
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments).values()
+    assert set(allotrope.get_kind_names()) == {'Fe0', 'Fe1'}
+    assert [site.kind_name for site in allotrope.sites] == ['Fe0', 'Fe1', 'Fe1', 'Fe1']
+    assert allotrope_magnetic_moments.get_dict() == {'Fe0': 0.0, 'Fe1': 0.5}
+
+    # Lower atol to 0.05: 0.5 & 0.45 now one kind, 0.4 new kind -> three different kinds
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments,
+                                                                      atol=Float(0.05)).values()
+    assert set(allotrope.get_kind_names()) == {'Fe0', 'Fe1', 'Fe2'}
+    assert [site.kind_name for site in allotrope.sites] == ['Fe0', 'Fe1', 'Fe1', 'Fe2']
+    assert allotrope_magnetic_moments.get_dict() == {'Fe0': 0.0, 'Fe1': 0.5, 'Fe2': 0.4}
+
+    # Increase atol to 0.1, again only two different kinds
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments,
+                                                                      atol=Float(0.1)).values()
+    assert set(allotrope.get_kind_names()) == {'Fe0', 'Fe1'}
+    assert [site.kind_name for site in allotrope.sites] == ['Fe0', 'Fe1', 'Fe1', 'Fe1']
+    assert allotrope_magnetic_moments.get_dict() == {'Fe0': 0.0, 'Fe1': 0.5}
+
+    # Really strict tolerance or atol = 0.01: All sites get different kinds
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments,
+                                                                      atol=Float(1E-2)).values()
+    assert set(allotrope.get_kind_names()) == {'Fe0', 'Fe1', 'Fe2', 'Fe3'}
+    assert [site.kind_name for site in allotrope.sites] == ['Fe0', 'Fe1', 'Fe2', 'Fe3']
+    assert allotrope_magnetic_moments.get_dict() == {'Fe0': 0.0, 'Fe1': 0.5, 'Fe2': 0.45, 'Fe3': 0.4}
+
+
+@pytest.mark.usefixtures('aiida_profile')
+def test_configuration_05(generate_structure_from_kinds):
+    """Test `create_magnetic_allotrope` calculation function.
+
+    Case: One kind, only negative magnetic moments with one close to zero
+    Expected result: Depends on tolerance, see below
+    """
+    kind_names = ['Fe', 'Fe', 'Fe', 'Fe']
+    magnetic_moments = List(list=[-0.5, -0.6, -1.5, -0.01])
+
+    structure = generate_structure_from_kinds(kind_names)
+
+    # Default tolerance values, one zero site and two magnetic
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments).values()
+    assert set(allotrope.get_kind_names()) == {'Fe0', 'Fe1', 'Fe2'}
+    assert [site.kind_name for site in allotrope.sites] == ['Fe0', 'Fe1', 'Fe2', 'Fe2']
+    assert allotrope_magnetic_moments.get_dict() == {'Fe0': 0.0, 'Fe1': -1.5, 'Fe2': -0.6}
+
+    # Strict absolute tolerance, one zero site and three magnetic
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments,
+                                                                      atol=Float(0.05)).values()
+    assert set(allotrope.get_kind_names()) == {'Fe0', 'Fe1', 'Fe2', 'Fe3'}
+    assert [site.kind_name for site in allotrope.sites] == ['Fe0', 'Fe1', 'Fe2', 'Fe3']
+    assert allotrope_magnetic_moments.get_dict() == {'Fe0': 0.0, 'Fe1': -1.5, 'Fe2': -0.6, 'Fe3': -0.5}
+
+    # Strict absolute and zero tolerance, four magnetic sites
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(
+        structure, magnetic_moments, atol=Float(0.05), ztol=Float(1E-3)
+    ).values()
+    assert set(allotrope.get_kind_names()) == {'Fe0', 'Fe1', 'Fe2', 'Fe3'}
+    assert [site.kind_name for site in allotrope.sites] == ['Fe0', 'Fe1', 'Fe2', 'Fe3']
+    assert allotrope_magnetic_moments.get_dict() == {'Fe0': -1.5, 'Fe1': -0.6, 'Fe2': -0.5, 'Fe3': -0.01}
+
+
+@pytest.mark.usefixtures('aiida_profile')
+def test_configuration_06(generate_structure_from_kinds):
+    """Test `create_magnetic_allotrope` calculation function.
+
+    Case: Two kinds, magnetic moments with different signs for the first (Fe)
+    Expected result: Depends on tolerance, see below
+    """
+    kind_names = ['Fe', 'Fe', 'Fe', 'Fe', 'Ni', 'Ni']
+    magnetic_moments = List(list=[-0.1, 0.1, -0.2, 0.01, 0.2, 0.25])
+
+    structure = generate_structure_from_kinds(kind_names)
+
+    # Default tolerance values, one zero and two magnetic sites for Fe, one magnetic site for Ni
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments).values()
+    assert set(allotrope.get_kind_names()) == {'Fe0', 'Fe1', 'Fe2', 'Ni'}
+    assert allotrope_magnetic_moments.get_dict() == {'Fe0': 0.0, 'Fe1': 0.1, 'Fe2': -0.2, 'Ni': 0.25}
+
+    # Very strict absolute tolerance, all different sites
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments,
+                                                                      atol=Float(0.02)).values()
+    assert set(allotrope.get_kind_names()) == {'Fe0', 'Fe1', 'Fe2', 'Fe3', 'Ni0', 'Ni1'}
+    assert allotrope_magnetic_moments.get_dict() == {
+        'Fe0': 0.0,
+        'Fe1': 0.1,
+        'Fe2': -0.2,
+        'Fe3': -0.1,
+        'Ni0': 0.25,
+        'Ni1': 0.2
+    }
+
+
+@pytest.mark.usefixtures('aiida_profile')
+def test_configuration_07(generate_structure_from_kinds):
+    """Test `create_magnetic_allotrope` calculation function.
+
+    Case: Two different symbols but the same magnetic moment.
+    Expected result: One kind with name equal to the element symbol
+    """
+    kind_names = ['Fe0', 'Fe1']
+    magnetic_moments = List(list=[0.1, 0.1])
+
+    structure = generate_structure_from_kinds(kind_names)
+
+    allotrope, allotrope_magnetic_moments = create_magnetic_allotrope(structure, magnetic_moments).values()
+    assert set(allotrope.get_kind_names()) == {'Fe'}
+    assert allotrope_magnetic_moments.get_dict() == {'Fe': 0.1}

tests/conftest.py

@@ -393,6 +393,29 @@ def _generate_structure(structure_id='silicon'):
     return _generate_structure
 
 
+@pytest.fixture
+def generate_structure_from_kinds():
+    """Return a dummy `StructureData` instance with the specified kind names."""
+
+    def _generate_structure_from_kinds(site_kind_names):
+        """Return a dummy `StructureData` instance with the specified kind names."""
+        import re
+
+        from aiida import orm
+
+        if not isinstance(site_kind_names, (list, tuple)):
+            site_kind_names = (site_kind_names,)
+
+        structure = orm.StructureData(cell=[[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+
+        for kind_name in site_kind_names:
+            structure.append_atom(name=kind_name, symbols=re.sub('[0-9]', '', kind_name), position=(0., 0., 0.))
+
+        return structure
+
+    return _generate_structure_from_kinds
+
+
 @pytest.fixture
 def generate_kpoints_mesh():
     """Return a `KpointsData` node."""