PyMatGen to AiiDA

pymatgen_to_aiida.py

def pymatgen_to_aiida(pymatgen_structure, StructureData):

    # analyze magnetism
    magnetic_structure = CollinearMagneticStructureAnalyzer(pymatgen_structure, make_primitive=False)

    has_spin = 0
    if magnetic_structure.is_magnetic:
        if magnetic_structure.is_collinear:
            has_spin=2
        else:
            has_spin=4
    else:
        has_spin=1
        return (StructureData(pymatgen=pymatgen_structure), has_spin, {})

    # generate collinear spin structure
    if (has_spin == 2):
        structure_with_spin = magnetic_structure.get_structure_with_spin()
    else:
        structure_with_spin = pymatgen_structure
    
    # initial definition of cell (is PBC specified by default?)
    cell = structure_with_spin.lattice.matrix.tolist()
    aiida_structure = StructureData(cell=cell)

    # collect magnetic elements by name. For each of them create kinds
    magnetic_elements_kinds = {}
    for site in structure_with_spin:

        # check spin and element name. Information is in slightly different places
        spin = site.specie.spin              if has_spin == 2 else site.magmom.moment
        element = site.specie.element.symbol if has_spin == 2 else site.specie.symbol
        
        kind_name = None
        if not np.allclose(np.abs(spin), 0.0):
            # checks if element was already found to be magnetic in a previous site
            # otherwise return an empty dictionary to be filled with the information
            # of this site
            kinds_for_element = magnetic_elements_kinds.get(element, {})
            
            # If the spin of this site is for this element is the same we found
            # previously, just give it the same kind, otherwise add it as 
            # a new kind for this element type.
            for kind, kind_spin in kinds_for_element.items():
                if np.allclose (spin, kind_spin):
                    kind_name = kind
                    break
            else:
                kind_name = '{}{}'.format(element, len(kinds_for_element)+1)
                kinds_for_element[kind_name] = spin
            
            # store the updated list of kinds for this element in the full dictionary.
            magnetic_elements_kinds[element] = kinds_for_element

        # prepare to add site to AiiDA structure...
        inputs = {
            'symbols': [x.symbol for x in site.species_and_occu.keys()],
            'weights': [x for x in site.species_and_occu.values()],
            'position': site.coords.tolist()
        }
        # ...and also its kind, in case we chose one...
        if kind_name is not None:
            inputs['name'] = kind_name
        # ...and it's done for this one!
        aiida_structure.append_atom(**inputs)
    
    # Finally return structre, type of magnetic order (1 or 2 for the time being) 
    # and a dictionary for magnetic elements, where all the kinds are reported together
    # with the value of the spin.
    return (aiida_structure, has_spin, magnetic_elements_kinds )