Workchain_AmarantaTemplate_18dec_debugged.py

import sys, math
import numpy as np
import subprocess
from aiida.orm import StructureData, Code, Str, Float, Int, Dict, List, KpointsData
from aiida.common.extendeddicts import AttributeDict
from aiida.plugins import WorkflowFactory, DataFactory
from aiida.engine import calcfunction, WorkChain, ToContext, append_, submit, while_, if_
from aiida_vasp.utils.workchains import prepare_process_inputs
from ase.build import make_supercell, sort
from ase.io import read, write

SingleFileData = DataFactory('core.singlefile')
FolderData = DataFactory('core.folder')




class AmarantaWorkChain(WorkChain):
    _next_workchain_string = 'vasp.vasp'
    _next_workchain = WorkflowFactory(_next_workchain_string)
   
    @classmethod
    def define(cls, spec):
        super(AmarantaWorkChain, cls).define(spec)
       
        spec.expose_inputs(cls._next_workchain, exclude=['parameters','kpoints','settings'])
        spec.expose_inputs(SubWorkChain,        exclude=['structure_sub','kpoints_sub','magmom_sub','magatom_sub','spin_sub','which_incar','which_label'])
       
        spec.output('supercell_1NN', valid_type=StructureData, help='Minimal supercell 1NN')
        spec.output('supercell_2NN', valid_type=StructureData, help='Minimal supercell 2NN')
        spec.output('supercell_3NN', valid_type=StructureData, help='Minimal supercell 3NN')
        spec.output('couples', valid_type=Dict, help='couples')
        spec.output('kmesh_1NN', valid_type=KpointsData, help='K-mesh 1NN')
        spec.output('kmesh_2NN', valid_type=KpointsData, help='K-mesh 2NN')
        spec.output('kmesh_3NN', valid_type=KpointsData, help='K-mesh 3NN')   
        spec.output('spin', valid_type=Float, help='Spin modulus')
        spec.output('J_dict', valid_type=Dict, help='J_dict')
        spec.output('M_dict', valid_type=Dict, help='M_dict')
        spec.output('units', valid_type=Dict, help='units')

        spec.outline(
            cls.print_welcome,
            cls.run_ASE,
            if_(cls.run_full_wc)(
                cls.print_run_for_spin,
                cls.init_aiida_for_spin, 
                cls.run_aiida_for_spin,
                cls.elaborate_results_aiida_for_spin,
                cls.print_run_main,
                cls.init_main,
                while_(cls.index_is_less_than_number_of_J)(
                    if_(cls.index_is_within_selected)(
                        cls.init_run_subwc,
                        cls.retrieve_J
                    ),
                    cls.increment_index_by_one,
                ),
                cls.finalize_dict,
                cls.print_goodbye
            )
        )
    
    

    
               
    def print_welcome(self):
        str_log = '\n'
        str_log += ('\n************************************************************************************')
        str_log += ('\n*                                                                                  *')
        str_log += ('\n*                                    AMaRaNTA                                      *')
        str_log += ('\n*              Automating MAgnetic paRAmeters iN a Tensorial Approach              *')
        str_log += ('\n*                                                                                  *')
        str_log += ('\n************************************************************************************')
        str_log += '\n'
        self.report(str_log)
 
 
 
 
    def run_ASE(self):
        #### executes "run_ase_calculation" calcfunction in order to
        #### build supercells, k-meshes and magmom sets
        #### safety_distance is the parameter governing supercell size
        input_structure = self.inputs.structure
        safety_distance = 10 
        self.ctx.ase_results = run_ase_calculation(input_structure, safety_distance)
        self.out("supercell_1NN", self.ctx.ase_results["final_structure_1"])
        self.out("supercell_2NN", self.ctx.ase_results["final_structure_2"])
        self.out("supercell_3NN", self.ctx.ase_results["final_structure_3"])
        self.out("couples", self.ctx.ase_results["couples"])
        self.out("kmesh_1NN", self.ctx.ase_results["kmesh_1"])
        self.out("kmesh_2NN", self.ctx.ase_results["kmesh_2"])
        self.out("kmesh_3NN", self.ctx.ase_results["kmesh_3"])
        self.ctx.magmom_unitcell = self.ctx.ase_results["magmom_uc"]
        self.ctx.n_magatoms = self.ctx.ase_results["n_magatoms"]
        self.ctx.magatom = self.ctx.ase_results["magatom"]
        self.report(self.ctx.ase_results["str_out"].value)
        
        
        
        
    def run_full_wc(self):
        #### if True, whole workchain executed
        #### if False, only executes ASE step...
        #### ...good for debugging, to see if workchain starts in pathological cases
        return False
       
       
       
       
       
    def print_run_for_spin(self):
        str_log = '\n'
        str_log += ('\n==============================================================')
        str_log += ('\n           PRELIMINARY RUN --- for spin modulus only          ')
        str_log += ('\n==============================================================')
        str_log += ('\n')
        self.report(str_log)
       
       
       
       
       
    def init_aiida_for_spin(self):
        self.ctx.inputs = AttributeDict()
        self.ctx.inputs.update(self.exposed_inputs(self._next_workchain))      
       
        #### ENMAX and RWIGS values read from PARAMS dict; ENCUT = 1.3*max(ENMAX)
        encut_list = []
        rwigs_list = []
        ldaul_list = [2]
        ldauu_list = [2.80]
        ldauj_list = [0.80]

        pseudolist = dict(self.inputs.potential_mapping)

        self.ctx.inputs.potential_mapping = Dict(pseudolist)
        pseudolist = dict(sorted(pseudolist.items()))
        
        #### here I have to reorder, forcing magnetic species to appear as 1st
        magatom = self.ctx.magatom.value
        
        ordered_dict = {magatom: pseudolist.get(magatom)}  
        for key, value in pseudolist.items():
            if key != magatom:
                ordered_dict[key] = value

        for elem in ordered_dict.values():
            RWIGS, ENMAX = PARAMS_dict[elem]
            encut_list.append(ENMAX)
            rwigs_list.append(RWIGS)   
            
        elist = np.array(encut_list)
        rwigslist = np.array(rwigs_list)
        encut = 1.3*np.max(elist)
        rwigs = np.array2string(rwigslist)
        rwigs = rwigs.replace("[", "").replace("]", "")
       
        magmom_unitcell = self.ctx.magmom_unitcell
        magmom_values_str = ' '.join(map(str, self.ctx.magmom_unitcell.get_dict()['magmom_uc']))
        magmom_values_str = f"'{magmom_values_str}'"
        magmom_values_str = magmom_values_str.replace("'", "")
        
        #### U flags
        st = self.inputs.structure.get_ase()
        species = st.get_chemical_symbols()
        species = set(species)
        n_species = len(species)
        n_eff = n_species-1

        for k in range(0,n_eff):
             ldaul_list.append(-1)
             ldauu_list.append(0.00)
             ldauj_list.append(0.00)

        ldaul = np.array(ldaul_list)
        ldaul = np.array2string(ldaul)
        ldaul = ldaul.replace("[", "").replace("]", "")

        ldauu = np.array(ldauu_list)
        ldauu = np.array2string(ldauu)
        ldauu = ldauu.replace("[", "").replace("]", "")

        ldauj = np.array(ldauj_list)
        ldauj = np.array2string(ldauj)
        ldauj = ldauj.replace("[", "").replace("]", "")

        incar = {'incar': {'NPAR': 4, 'KPAR': 4, 'ENCUT': encut, 'ISMEAR': 0, 'SIGMA': 0.05, 'ISYM': -1, 'EDIFF': 1E-6, 'NELM': 200, 'PREC': 'ACCURATE', 'ISTART': 0, 'ICHARG': 2, 'LWAVE': 'FALSE', 'LREAL': 'AUTO', 'IBRION': -1, 'AMIX': 0.2, 'BMIX': 0.00001, 'AMIX_MAG': 0.8, 'BMIX_MAG': 0.00001, 'ISPIN': 2, 'LNONCOLLINEAR': '.FALSE.', 'LSORBIT':'.FALSE.', 'LORBIT': 11, 'MAGMOM' : magmom_values_str, 'RWIGS': rwigs, 'LDAU' : '.TRUE.', 'LDAUTYPE': 1, 'LDAUL' : ldaul, 'LDAUU' : ldauu, 'LDAUJ' : ldauj, 'LDAUPRINT' : '0', 'LMAXMIX' : '4'}}

        self.ctx.inputs.parameters = Dict(incar)
       
        kpoints = KpointsData()
        kpoints.set_kpoints_mesh([12, 12, 1])
        self.ctx.inputs.kpoints = kpoints
       
        settings = AttributeDict()
        settings.parser_settings = {'output_params': ['total_energies', 'maximum_force']}
        settings.parser_settings = {'add_site_magnetization':True, 'add_kpoints':True, 'add_structure':True}
        self.ctx.inputs.settings = Dict(settings)
       
        self.ctx.inputs = prepare_process_inputs(self.ctx.inputs, namespaces=['dynamics','verify'])  
      
      
      
      
           
    def run_aiida_for_spin(self):
        inputs = self.ctx.inputs
        running = self.submit(self._next_workchain, **inputs)
        self.report('launching {}<{}> '.format(self._next_workchain.__name__, running.pk))
        return self.to_context(workchains=append_(running))
      
      
      
      
             
    def elaborate_results_aiida_for_spin(self): 
        workchain = self.ctx.workchains[-1]       
        site_magnetization = workchain.outputs.site_magnetization.get_dict()
        spin = Float(round(site_magnetization['site_magnetization']['full_cell'][0]))
        spin = spin/self.ctx.n_magatoms
        spin.store()
        self.report("SPIN = " + str(spin.value) + "\n")
        self.out("spin", spin)
        self.ctx.spin = spin
      
      
      
      
       
    def print_run_main(self):
        str_log = '\n'
        str_log += ('\n==============================================================')
        str_log += ('\n               MAIN RUN --- compute J parameters              ')
        str_log += ('\n==============================================================')
        str_log += '\n'
        self.report(str_log)
      
      
      
      
        
    def init_main(self):
        #### self.ctx.selected_indices controls how many elements (and which) are calculated
        #### self.ctx.index not to be touched
        self.ctx.index = 0
        
        #### run all
        self.ctx.selected_indices = [i for i in range(12)]
        #### run only selected
        #self.ctx.selected_indices = [9]
        
        #### order:
        #### 0 = "J_xx"
        #### 1 = "J_yy"
        #### 2 = "J_zz"
        #### 3 = "J_xy"
        #### 4 = "J_yx"
        #### 5 = "J_yz"
        #### 6 = "J_zy"
        #### 7 = "J_xz"
        #### 8 = "J_zx"
        #### 9 = "J_scalar_2NN"
        #### 10 = "J_scalar_3NN"
        #### 11 = "A_scalar" 
       
       	self.ctx.J_dict = {}
       	self.ctx.M_dict = {}
       
       
        
    def index_is_less_than_number_of_J(self):
        return self.ctx.index < 12
        
    
    
    

    def index_is_within_selected(self):
    	return self.ctx.index in self.ctx.selected_indices
	 
     
     
     
             
    def increment_index_by_one(self):
        self.ctx.index += 1      
     
     
     
     
        
    def init_run_subwc(self):
        self.ctx.inputs = AttributeDict()       
        self.ctx.inputs.update(self.exposed_inputs(SubWorkChain))
        
        #### depending of which calculation is run (J1, J2, J3, SIA)
        #### provide correct POSCAR, KPOINTS, INCAR (collinear or not)       
        if self.ctx.index <= 8:
            self.ctx.inputs.structure_sub = self.ctx.ase_results["final_structure_1"]
            self.ctx.inputs.kpoints_sub   = self.ctx.ase_results["kmesh_1"]
            self.ctx.inputs.which_incar   = Int(2)
        elif self.ctx.index == 9:
            self.ctx.inputs.structure_sub = self.ctx.ase_results["final_structure_2"]
            self.ctx.inputs.kpoints_sub   = self.ctx.ase_results["kmesh_2"]
            self.ctx.inputs.which_incar   = Int(1)
        elif self.ctx.index == 10:
            self.ctx.inputs.structure_sub = self.ctx.ase_results["final_structure_3"]
            self.ctx.inputs.kpoints_sub   = self.ctx.ase_results["kmesh_3"]
            self.ctx.inputs.which_incar   = Int(1)
        elif self.ctx.index > 10:
            self.ctx.inputs.structure_sub = self.ctx.ase_results["final_structure_1"]
            self.ctx.inputs.kpoints_sub   = self.ctx.ase_results["kmesh_1"]
            self.ctx.inputs.which_incar   = Int(2)
        
        #### this specifies in which order parameters are computed
        #### do not touch!!! use method "init_main" if you want to
        #### compute specific parameters only
        keys = ["J_xx", "J_yy", "J_zz", "J_xy", "J_yx", "J_yz", "J_zy", "J_xz", "J_zx", "J_scalar_2NN", "J_scalar_3NN", "A_scalar"]     
        selected_key = keys[self.ctx.index]
        self.ctx.selected_key = selected_key               
        self.ctx.inputs.magmom_sub    = List(self.ctx.ase_results["magmom_all"][selected_key])
        self.ctx.inputs.magatom_sub = self.ctx.magatom
        self.ctx.inputs.which_label = Str(selected_key)
        
        self.ctx.inputs.spin_sub    = self.ctx.spin      
        #### if you want to speed things up and skip the computation of spin
        #### (for instance because you already know its value...)
        #### use this line below instead, but do remember you are using it!!!!
        #self.ctx.inputs.spin_sub    = Float(2.0)
        
        self.report('\n \n==============================================================' + "\nindex: " + str(self.ctx.index) + "                   parameter to compute: " + str(selected_key) + '\n==============================================================\n')
            
        inputs = self.ctx.inputs
        
        running = self.submit(SubWorkChain, **inputs)    
        self.to_context(**{'test': running})
        return ToContext(child = running)





    def retrieve_J(self):
        J = self.ctx.child.outputs.J_parameter
        self.ctx.J_dict[self.ctx.selected_key] = J
        
        M = self.ctx.child.outputs.magnetizations
        self.ctx.M_dict[self.ctx.selected_key] = M
        
        
       


    def finalize_dict(self):
        J_dict_final = Dict(self.ctx.J_dict)
        J_dict_final.store()
        self.out("J_dict", J_dict_final)
        
        M_dict_final = Dict(self.ctx.M_dict)
        M_dict_final.store()
        self.out("M_dict", M_dict_final)
        
        units_dict = Dict(dict(distances = "Angstrom", J_parameters = "meV (not divided by S^2)", magnetizations = "Bohr magneton", coordinates = "cartesian, Angstrom"))
        units_dict.store()
        self.out("units", units_dict) 
        
        
        
    def print_goodbye(self):
        str_log = '\n'
        str_log += ('\n==============================================================')
        str_log += ('\n                           Job done                           ')
        str_log += ('\n==============================================================')
        str_log += '\n'
        self.report(str_log)





















class SubWorkChain(WorkChain):
    _next_workchain_string = 'vasp.vasp'
    _next_workchain = WorkflowFactory(_next_workchain_string)
   
    @classmethod
    def define(cls, spec):
        super(SubWorkChain, cls).define(spec)
       
        spec.expose_inputs(cls._next_workchain, exclude=['structure','parameters','kpoints','settings'])
                     
        spec.input('structure_sub',valid_type=StructureData, help='Minimal supercell')
        spec.input('kpoints_sub',valid_type=KpointsData, help='K points')
        spec.input('magmom_sub',valid_type=List, help='magmom')
        spec.input('magatom_sub',valid_type=Str, help='magatom')
        spec.input('spin_sub',valid_type=Float, help='spin')
        spec.input('which_incar',valid_type=Int, help='which incar')
        spec.input('which_label',valid_type=Str, help='which label')
        
        spec.output('J_parameter', valid_type=Float, help='J_parameter')
        spec.output('magnetizations', valid_type=List, help='magnetizations')
        spec.output('txtfile',valid_type=SingleFileData, help='txtfile')

        spec.outline(
            cls.initialize,
            while_(cls.n_is_less_than_four)(
                cls.init_run_aiida_main,
                cls.elaborate_results_main,
                cls.increment_n_by_one
            ),
            cls.finalize
        )
    
    
    
    
    
    def initialize(self):
        self.ctx.n = 0
        self.ctx.total_energies = []
        self.ctx.magnetizations = []
     
     
     
     
        
    def n_is_less_than_four(self):
        return self.ctx.n < 4    
    
    
    
    
        
    def increment_n_by_one(self):
        self.ctx.n += 1    
        
        
        
        
         
    def init_run_aiida_main(self): 
                
        self.ctx.inputs = AttributeDict()
        self.ctx.inputs.update(self.exposed_inputs(self._next_workchain))
       
        self.ctx.inputs.structure = self.inputs.structure_sub
        self.ctx.inputs.kpoints = self.inputs.kpoints_sub
        incar_flag = int(self.inputs.which_incar)
        magmom_sub = self.inputs.magmom_sub
        S = float(self.inputs.spin_sub)
       
        encut_list = []
        rwigs_list = []
        ldaul_list = [2]
        ldauu_list = [2.80]
        ldauj_list = [0.80]

        pseudolist = dict(self.inputs.potential_mapping)
        
        self.ctx.inputs.potential_mapping = Dict(pseudolist)
        pseudolist = dict(sorted(pseudolist.items()))
        #### here I have to reorder, forcing magnetic species to appear as 1st
        magatom = self.inputs.magatom_sub.value 
        
        ordered_dict = {magatom: pseudolist.get(magatom)}
        for key, value in pseudolist.items():
            if key != magatom:
                ordered_dict[key] = value
        
        for elem in ordered_dict.values():
            RWIGS, ENMAX = PARAMS_dict[elem]
            encut_list.append(ENMAX)
            rwigs_list.append(RWIGS)       
                
        elist = np.array(encut_list)
        rwigslist = np.array(rwigs_list)
        encut = 1.3*np.max(elist)
        rwigs = np.array2string(rwigslist)
        rwigs = rwigs.replace("[", "").replace("]", "")
               
        settings = AttributeDict()
        settings.parser_settings = {'output_params': ['total_energies', 'maximum_force']}
        settings.parser_settings = {'add_site_magnetization':True, 'add_kpoints':True, 'add_structure':True}
        self.ctx.inputs.settings = Dict(settings)
        
        i = self.ctx.n
        
        rescaled_magmom = [[value * S for value in sub_array] for sub_array in magmom_sub]
        magmom_values_str = str(rescaled_magmom[i]).replace("[", "").replace("]", "").replace(",", "")
        
        #### U flags
        st = self.ctx.inputs.structure.get_ase()
        species = st.get_chemical_symbols()
        species = set(species)
        n_species = len(species)
        n_eff = n_species-1
        
        for k in range(0,n_eff):
             ldaul_list.append(-1)
             ldauu_list.append(0.00)
             ldauj_list.append(0.00)

        ldaul = np.array(ldaul_list)
        ldaul = np.array2string(ldaul)
        ldaul = ldaul.replace("[", "").replace("]", "")

        ldauu = np.array(ldauu_list)
        ldauu = np.array2string(ldauu)
        ldauu = ldauu.replace("[", "").replace("]", "")
        
        ldauj = np.array(ldauj_list)
        ldauj = np.array2string(ldauj)
        ldauj = ldauj.replace("[", "").replace("]", "")

        self.report("\nMAGMOM " +str(magmom_values_str))
        
        INCAR_COLL = {'incar': {'NPAR': 8, 'KPAR': 4, 'ENCUT': encut, 'ISMEAR': 0, 'SIGMA': 0.05, 'ISYM': -1, 'EDIFF': 1E-6, 'NELM': 200, 'PREC': 'ACCURATE', 'ISTART': 0, 'ICHARG': 2, 'LWAVE': 'FALSE',  'LCHARG' : 'FALSE', 'LREAL': 'AUTO', 'AMIX': 0.1, 'BMIX': 0.00001, 'ISPIN': 2, 'LNONCOLLINEAR': '.FALSE.', 'LSORBIT':'.FALSE.', 'LORBIT': 11, 'MAGMOM' : magmom_values_str, 'RWIGS': rwigs, 'LDAU' : '.TRUE.', 'LDAUTYPE': 1, 'LDAUL' : ldaul, 'LDAUU' : ldauu, 'LDAUJ' : ldauj, 'LDAUPRINT' : '0', 'LMAXMIX' : '4' }}

        INCAR_NONCOLL = {'incar': {'NPAR': 8, 'KPAR': 4, 'ENCUT': encut, 'ISMEAR': 0, 'SIGMA': 0.05, 'ISYM': -1, 'EDIFF': 1E-6, 'NELM': 200, 'PREC': 'ACCURATE', 'ISTART': 0, 'ICHARG': 2, 'LWAVE': 'FALSE',  'LCHARG' : 'FALSE', 'LREAL': 'AUTO', 'AMIX': 0.1, 'BMIX': 0.00001, 'ISPIN': 2, 'LNONCOLLINEAR': '.TRUE.', 'LSORBIT':'.TRUE.', 'LORBIT': 11, 'MAGMOM' : magmom_values_str, 'M_CONSTR' : magmom_values_str, 'I_CONSTRAINED_M' : 1, 'LAMBDA': 10, 'RWIGS': rwigs,  'SAXIS' : '0.0000 0.0000 1.0000', 'LDAU' : '.TRUE.', 'LDAUTYPE': 1, 'LDAUL' : ldaul, 'LDAUU' : ldauu, 'LDAUJ' : ldauj, 'LDAUPRINT' : '0', 'LMAXMIX' : '4' }}
        
        if incar_flag == 1:
            self.ctx.inputs.parameters = Dict(INCAR_COLL)
        else:
            self.ctx.inputs.parameters = Dict(INCAR_NONCOLL)
            
        self.ctx.inputs = prepare_process_inputs(self.ctx.inputs, namespaces=['dynamics','verify'])  
        inputs = self.ctx.inputs
       
        running = self.submit(self._next_workchain, **inputs)

        self.to_context(workchains=append_(running))
        workchain = self.ctx.workchains[-1]
       
     
     
     
       
    def elaborate_results_main(self): 
        workchain = self.ctx.workchains[-1]  
           
        misc = workchain.outputs.misc.get_dict()
        total_energy = Float(misc['total_energies']['energy_extrapolated'])
        total_energy.store()
        self.ctx.total_energies.append(total_energy)
        self.report("ENERGY:   " + str(total_energy.value) + " eV \n \n")
        
        site_mag = workchain.outputs.site_magnetization.get_dict()
        mags = site_mag['site_magnetization']['sphere']
        ase_struct = self.ctx.inputs.structure.get_ase()
        atom_species_dict = {(atom.index+1): atom.symbol for atom in ase_struct}
        for key in ['x', 'y', 'z']:
             for atom_number, atom_data in mags[key]['site_moment'].items():
                  atom_data['species'] = atom_species_dict[int(atom_number)]

        self.ctx.magnetizations.append(mags)
       


     
     
        
    def finalize(self):
        total_energies = self.ctx.total_energies
        J_parameter = find_energy_four_state(total_energies)
        
        label = self.inputs.which_label
        self.report(str(label.value) + ":               " + str(J_parameter.value) + " meV \n \n")
        J_parameter.store()
        self.out('J_parameter', J_parameter)
        
        magnetizations = self.ctx.magnetizations
        magnetizations = List(magnetizations)
        magnetizations.store()
        self.out('magnetizations',magnetizations)
        
        incar_flag = int(self.inputs.which_incar)
        if incar_flag == 1:
            coll = True
        else:
            coll = False
        
        filename = Str('/home/federico/Amaranta_txt_files/' + label.value + '.txt')
        save_info(total_energies, J_parameter, magnetizations, filename, coll)
        file = create_singlefile(filename)
        self.out('txtfile', file)

        this_struct = self.inputs.structure_sub.get_ase()
        chem_symb = this_struct.get_chemical_symbols()
        N_magatoms = Int(sum(1 for symbol in chem_symb if symbol == self.inputs.magatom_sub.value))
        printed_info = print_info(N_magatoms, total_energies, J_parameter, magnetizations, coll)
        self.report(printed_info.value)
      
      
        



         
         
         
         
         
         
         
         
         
         
         
         
         
         
         
         
         
         
         
         
         
         
@calcfunction
def find_energy_four_state(list_of_energies): 
     #### implements 4-state formula
     energy = 1000*(list_of_energies[0]-list_of_energies[1]-list_of_energies[2]+list_of_energies[3])/4
     return Float(energy)        

         
         
         
         
         
         


 
@calcfunction    
def save_info(total_energies,J_parameter,magnetization_results, filename, coll):
    filename = str(filename.value)

    with open(filename, 'w') as f:
        f.write(f"ENERGIES \n")
        for item in total_energies:
            f.write(str(item) + "\n")
        f.write("\n")
        J_parameter = float(J_parameter)
        f.write(f"J PARAMETER = {J_parameter} \n")
        f.write("\n \n \n \n")
            
        if bool(coll) == False:
            count_state = 1
            for direction_data in magnetization_results:			#loops over each of 4 dicts in list
                f.write(f"STATE {count_state} \n")
                f.write('========================================== \n')
                for direction, data in direction_data.items():		#loops over x,y,z 
                    f.write(f" magnetization ({direction})\n")
                    f.write("# of ion       s       p       d       tot\n")
                    f.write("------------------------------------------\n")
                
                    # Iterate over site-wise moment data
                    for site, moment_data in data["site_moment"].items():
                       s, p, d, tot = moment_data["s"], moment_data["p"], moment_data["d"], moment_data["tot"]
                       f.write(f"{site:<13}{s:<8}{p:<8}{d:<8}{tot:<8}\n")
                
                    f.write("------------------------------------------\n")
                f.write(f"\n \n \n")
                count_state += 1  
                
        else:
            count_state = 1
            for direction_data in magnetization_results:			#loops over each of 4 dicts in list
                f.write(f"STATE {count_state} \n")
                f.write('========================================== \n')
                f.write(f" magnetization (x)\n")
                f.write("# of ion       s       p       d       tot\n")
                f.write("------------------------------------------\n")
                
                # Iterate over site-wise moment data
                for site, moment_data in direction_data["x"]["site_moment"].items():
                    s, p, d, tot = moment_data["s"], moment_data["p"], moment_data["d"], moment_data["tot"]
                    f.write(f"{site:<13}{s:<8}{p:<8}{d:<8}{tot:<8}\n")
                
                f.write("------------------------------------------\n")
                f.write(f"\n \n \n")
                count_state += 1 
       
        
 

  





@calcfunction
def create_singlefile(filename):
     path = str(filename.value)
     singlefile = SingleFileData(path)
     return singlefile










@calcfunction    
def print_info(N_magatoms,total_energies,J_parameter,magnetization_results, coll):     
     str_out = '\n \n'
     
     if bool(coll) == False:     
          count_state = 1
          for data in magnetization_results:			#loops over each of 4 dicts in list
                 str_out += (f"STATE {count_state}" + "\n")
                 str_out += ('===========================================\n')
                 str_out += ("Atom         M        Mx/M    My/M    Mz/M\n")
                 for atom_num, atom_data in data['x']['site_moment'].items():
                        if int(atom_num) <= N_magatoms:
                               Mx = data['x']['site_moment'][atom_num]['tot']
                               My = data['y']['site_moment'][atom_num]['tot']
                               Mz = data['z']['site_moment'][atom_num]['tot']
                               M = math.sqrt(Mx**2+My**2+Mz**2)
                               str_out += (f"{atom_num:<13}{round(M,3):<9}{round(Mx/M,3):<8}{round(My/M,3):<8}{round(Mz/M,3):<8}\n")
                        else:
                               Mx = data['x']['site_moment'][atom_num]['tot']
                               My = data['y']['site_moment'][atom_num]['tot']
                               Mz = data['z']['site_moment'][atom_num]['tot']
                               M = math.sqrt(Mx**2+My**2+Mz**2)
                               if M>0.2:
                                      str_out += (f"{atom_num:<13}{round(M,3):<9}{round(Mx/M,3):<8}{round(My/M,3):<8}{round(Mz/M,3):<8}     !!!!!!!\n")
                 str_out += '\n'
                 count_state += 1 
                 
     else:     
          count_state = 1
          for data in magnetization_results:			#loops over each of 4 dicts in list
                 str_out += (f"STATE {count_state}" + "\n")
                 str_out += ('==============\n')
                 str_out += ("Atom         M\n")
                 for atom_num, atom_data in data['x']['site_moment'].items():
                        if int(atom_num) <= N_magatoms:
                               M = data['x']['site_moment'][atom_num]['tot']
                               str_out += (f"{atom_num:<13}{round(M,3):<9}\n")
                        else:
                               M = data['x']['site_moment'][atom_num]['tot']
                               if M>0.2:
                                      str_out += (f"{atom_num:<13}{round(M,3):<9}     !!!!!!!\n")
                 str_out += '\n'
                 count_state += 1 
     
     str_out = Str(str_out)
     return str_out 
     
     
     
     
@calcfunction
def run_ase_calculation(initial_structure, safety_distance):
    
    str_out = ''
    
    atoms = initial_structure.get_ase()
    structure = sort(atoms)   #in case input files is not ordered VASP-like (NB this sort orders alphabetically)

    ### loop over chemical symbols to identify magnetic species (from Va to Ni). Ansatz: one magnetic species only!
    counter = 0

    for atom in structure:
        if 22 < atom.number < 29:
                magatom = atom.symbol
                counter += 1
                break

    if counter==0:
        str_out += ('No magnetism found!' + '\n')
        sys.exit()

    n_magatoms = len(structure[[atom.index for atom in structure if atom.symbol==magatom]])
    n_tot_atoms = len(structure)

    str_out += '\n \n' + 'Material is ' + str(structure.symbols) + '\n'
   
    item = np.array([])
    item = np.append(item, np.ones(n_magatoms))
    item = np.append(item, np.zeros(n_tot_atoms-n_magatoms))
   
    magmom_dict_unitcell = {
    "magmom_uc": item,
    }
   
    magmom_uc = Dict(magmom_dict_unitcell)
   
    ######################################### SELECTION OF PAIRS

    ### trick to place magnetic atoms on top: call them Ac, create supercell, sort and then call them back the right way
    for atom in structure:              
        if atom.symbol == magatom:
                atom.symbol = 'Ac'

    superstructure = sort(structure.repeat([5,5,1])) #here: take this SC large enough (e.g. 4x4 not enough for small unit cells such as NiI2)
                                                 #MIC true in the following: to interpret more easily, choose odd mesh

    for atom in superstructure:
        if atom.symbol == 'Ac':
                atom.symbol = magatom
               
    ##################### OLD WAY!!!!!!!!!!!!!!!!!! atom 0 is selected automatically as first in the couple
    ### atom 0 is selected automatically as the closest magnetic atom to the centre of 5x5 supercell
    ### get distances between that and the other magnetic atoms, identify 1st, 2nd and 3rd nn's and select first of them as partner in the couple
    ### distances rounded because atoms belonging to the same shell of neighbours can actually have slightly different distances wrt to atom 0 for numerical reasons

    ### Compute the supercell center using fractional coordinates
    cell = superstructure.get_cell()  # Get lattice vectors (3x3 matrix)
    center_fractional = np.array([0.5, 0.5, 0.5])  # Center in fractional coordinates
    supercell_center = np.dot(center_fractional, cell)  # Convert to Cartesian coordinates

    # Reduce center to xy-plane
    supercell_center_xy = supercell_center[:2]  # Take only x and y components
    
    ### Find the magnetic atom closest to the supercell center in the xy-plane
    magnetic_atoms_indices = [atom.index for atom in superstructure if atom.symbol == magatom]
    distances_to_center_xy = [
        np.linalg.norm(superstructure[index].position[:2] - supercell_center_xy) for index in magnetic_atoms_indices
    ]
    s0 = magnetic_atoms_indices[np.argmin(distances_to_center_xy)]
    s0_position = superstructure[s0].position
    #####################s0 = 0

    #####################for atom in superstructure:
    distances_original = superstructure.get_distances(s0,magnetic_atoms_indices,mic=True)

    distances = distances_original

    for i in range(len(distances)): #at first I tried to just round to 2nd digit, but this is safer, remember pathological case of CrISe...
        for j in range(len(distances)):
            x = distances[i]
            y = distances[j]
            if (abs(x-y)/(y + 0.00001)) < 0.01:
                distances[j] = distances[i]

    nn1_dist = np.min(distances[distances != 0])
    nn1_indexes = np.flatnonzero(distances == nn1_dist) 
    nn1_number = len(nn1_indexes)
    s1 = nn1_indexes[0]
    s1_position = superstructure[s1].position

    nn2_dist = np.min(distances[distances > nn1_dist])
    nn2_indexes = np.flatnonzero(distances == nn2_dist)
    nn2_number = len(nn2_indexes)
    s2 = nn2_indexes[0]
    s2_position = superstructure[s2].position

    nn3_dist = np.min(distances[distances > nn2_dist])
    nn3_indexes = np.flatnonzero(distances == nn3_dist)
    nn3_number = len(nn3_indexes)
    s3 = nn3_indexes[0]
    s3_position = superstructure[s3].position
    
    str_out += '\n'
    str_out += ('--------------------------------------------------- s0 is: ' + str(np.around(s0_position,3)))
    str_out += ('\n' + 'The 1nn couple is: ' + str(s0) + ',' + str(s1) + ' with distance: ' + str(round(nn1_dist,4)) + ' ----  s1 is: ' + str(np.around(s1_position,3)))
    str_out += ('\n' + 'The 2nn couple is: ' + str(s0) + ',' + str(s2) + ' with distance: ' + str(round(nn2_dist,4)) + ' ----  s2 is: ' + str(np.around(s2_position,3)))
    str_out += ('\n' + 'The 3nn couple is: ' + str(s0) + ',' + str(s3) + ' with distance: ' + str(round(nn3_dist,4)) + ' ----  s3 is: ' + str(np.around(s3_position,3)))
    str_out += '\n'
    
    couples_dict = {
            "1NN couple": [np.around(s0_position,3), np.around(s1_position,3)],
            "2NN couple": [np.around(s0_position,3), np.around(s2_position,3)],
            "3NN couple": [np.around(s0_position,3), np.around(s3_position,3)],
            }
    couples = Dict(couples_dict)

    f = open("/home/federico/Amaranta_txt_files/NN_list.txt","w")
    f.write("### NN order, number, distance\n")
    n1 = (1,nn1_number,round(nn1_dist,4))
    n2 = (2,nn2_number,round(nn2_dist,4))
    n3 = (3,nn3_number,round(nn3_dist,4))
    for item in n1:
        f.write(str(item)+" ")
    f.write("\n")
    for item in n2:
        f.write(str(item)+" ")
    f.write("\n")
    for item in n3:
        f.write(str(item)+" ")
    f.write("\n")
    f.close()
   
    safe_dist = int(safety_distance)		# minimum distance accepted between original pair of neighbours and replicas
    str_out += ('\n'+'Safety distance chosen: ' + str(safe_dist) + ' Ang' + '\n \n')

    ######################################### OPTIMAL SC FOR 1st NN
    
    for j in range(2,10):
        for k in range (2,j+1):
                for atom in structure:              
                        if atom.symbol == magatom:
                                atom.symbol = 'Ac'

                candidate_SC_nn1 = sort(structure.repeat([j,k,1]))

                for atom in candidate_SC_nn1:
                        if atom.symbol == 'Ac':
                                atom.symbol = magatom
                
                s0_found = False
                s1_found = False

                #for atom in candidate_SC_nn1:
                        #if np.allclose(atom.position, s1_position, atol=1e-6):
                                #s1 = atom.index
                                #s1_found = True
                                #break
                for atom in candidate_SC_nn1:
                    if not s0_found and np.allclose(atom.position, s0_position, atol=1e-6):
                        s0 = atom.index
                        s0_found = True
                    elif not s1_found and np.allclose(atom.position, s1_position, atol=1e-6):
                        s1 = atom.index
                        s1_found = True
                    if s0_found and s1_found:
                        break                                
                
                #very important to define dist_min = 0 here AND = 100 just after the continue statement
                #if s1 not found, the continue will correctly skip rest of k-loop
                #but before incrementing j, "dist_min > safe_dist" is evaluated
                #and this may cause the j-loop to be aborted as well, undesirably
                #if s1 is found, then dist_min must be set large enough for the rest of code in k-loop
                dist_min = 0

                #if not s1_found:
                if not (s0_found and s1_found):
                        continue
                    
                super_superstructure = candidate_SC_nn1.repeat([2,2,1])
       
                dist_min = 100
                dist_min_vec = np.array([100,100,100,100])

                for m in range (1,4):
                        s0_img = s0+m*len(candidate_SC_nn1)
                        s1_img = s1+m*len(candidate_SC_nn1)

                        d_0_0img = float(super_superstructure.get_distances(s0,s0_img))
                        d_0_1img = float(super_superstructure.get_distances(s0,s1_img))
                        d_1_0img = float(super_superstructure.get_distances(s1,s0_img))
                        d_1_1img = float(super_superstructure.get_distances(s1,s1_img))

                        dist_couple_img = np.array([d_0_0img,d_0_1img,d_1_0img,d_1_1img])
                        if dist_min > dist_couple_img.min():
                                dist_min = dist_couple_img.min()
                                dist_min_vec = dist_couple_img

                if dist_min > safe_dist:
                        opt_j_nn1_direct = j
                        opt_k_nn1_direct = k
                        break

        if dist_min > safe_dist:
                break
    
    s0 = None
    s1 = None

    for k in range(2,10):
        for j in range (2,k+1):
                for atom in structure:              
                        if atom.symbol == magatom:
                                atom.symbol = 'Ac'

                candidate_SC_nn1 = sort(structure.repeat([j,k,1]))

                for atom in candidate_SC_nn1:
                        if atom.symbol == 'Ac':
                                atom.symbol = magatom
                
                s0_found = False       
                s1_found = False

                #for atom in candidate_SC_nn1:
                        #if np.allclose(atom.position, s1_position, atol=1e-6):
                                #s1 = atom.index
                                #s1_found = True
                                #break
                for atom in candidate_SC_nn1:
                    if not s0_found and np.allclose(atom.position, s0_position, atol=1e-6):
                        s0 = atom.index
                        s0_found = True
                    elif not s1_found and np.allclose(atom.position, s1_position, atol=1e-6):
                        s1 = atom.index
                        s1_found = True
                    if s0_found and s1_found:
                        break

                #very important to define dist_min = 0 here AND = 100 just after the continue statement
                #if s1 not found, the continue will correctly skip rest of k-loop
                #but before incrementing j, "dist_min > safe_dist" is evaluated
                #and this may cause the j-loop to be aborted as well, undesirably
                #if s1 is found, then dist_min must be set large enough for the rest of code in k-loop
                dist_min = 0

                #if not s1_found:
                if not (s0_found and s1_found):
                        continue

                super_superstructure = candidate_SC_nn1.repeat([2,2,1])
       
                dist_min = 100
                dist_min_vec = np.array([100,100,100,100])

                for m in range (1,4):
                        s0_img = s0+m*len(candidate_SC_nn1)
                        s1_img = s1+m*len(candidate_SC_nn1)

                        d_0_0img = float(super_superstructure.get_distances(s0,s0_img))
                        d_0_1img = float(super_superstructure.get_distances(s0,s1_img))
                        d_1_0img = float(super_superstructure.get_distances(s1,s0_img))
                        d_1_1img = float(super_superstructure.get_distances(s1,s1_img))

                        dist_couple_img = np.array([d_0_0img,d_0_1img,d_1_0img,d_1_1img])
                        if dist_min > dist_couple_img.min():
                                dist_min = dist_couple_img.min()
                                dist_min_vec = dist_couple_img

                if dist_min > safe_dist:
                        opt_j_nn1_reversed = j
                        opt_k_nn1_reversed = k
                        break
       
        if dist_min > safe_dist:
                break
                
    for atom in structure:              
        if atom.symbol == magatom:
                atom.symbol = 'Ac'
                
    str_out += ("opt_j_nn1_direct " + str(opt_j_nn1_direct) + "       opt_k_nn1_direct " + str(opt_k_nn1_direct)+"\n")
    str_out += ("opt_j_nn1_reversed " + str(opt_j_nn1_reversed) + "       opt_k_nn1_reversed " +str(opt_k_nn1_reversed)+"\n")

    if opt_j_nn1_direct*opt_k_nn1_direct < opt_j_nn1_reversed*opt_k_nn1_reversed:
        opt_j_nn1 = opt_j_nn1_direct
        opt_k_nn1 = opt_k_nn1_direct
        SC_nn1 = sort(structure.repeat([opt_j_nn1,opt_k_nn1,1]))
        str_out += "Chosen direct values \n"
    elif opt_j_nn1_direct*opt_k_nn1_direct == opt_j_nn1_reversed*opt_k_nn1_reversed:
        opt_j_nn1 = opt_j_nn1_direct
        opt_k_nn1 = opt_k_nn1_direct
        SC_nn1 = sort(structure.repeat([opt_j_nn1,opt_k_nn1,1]))
        str_out += "Direct = reversed values \n"
    else:
        opt_j_nn1 = opt_j_nn1_reversed
        opt_k_nn1 = opt_k_nn1_reversed
        SC_nn1 = sort(structure.repeat([opt_j_nn1,opt_k_nn1,1]))
        str_out += "Chosen reversed values \n"

    for atom in SC_nn1:
        if atom.symbol == 'Ac':
                atom.symbol = magatom
    
    #CRUCIAL! at this point index s1 refers to the last evaluated "reversed" supercell
    #if this is chosen as final supercell, then no problem
    #but if the direct is chosen, s1 may point to the wrong atom!!!!
    #hence the need to reassign its value
    for atom in SC_nn1:
        if np.allclose(atom.position, s0_position, atol=1e-6):
            s0 = atom.index
            break
    for atom in SC_nn1:
        if np.allclose(atom.position, s1_position, atol=1e-6):
            s1 = atom.index
            break  
   
    SC_nn1.positions[[0,s0]] = SC_nn1.positions[[s0,0]] #to have pair of atoms appear in beginning of POSCAR
    SC_nn1.positions[[1,s1]] = SC_nn1.positions[[s1,1]] #to have pair of atoms appear in beginning of POSCAR
    
    final_structure_1 = StructureData(ase=SC_nn1)

    str_out += ('Minimal supercell for 1nn: ---- ' +str(opt_j_nn1) + ' x ' + str(opt_k_nn1) + "\n")
    str_out += ('with distances: ' + str(np.around(dist_min_vec,4)) + "\n \n")
    write('/home/federico/Amaranta_txt_files/POSCAR_1NN',SC_nn1,format='vasp')
   
    ######################################### OPTIMAL SC FOR 2nd NN
    
    for j in range(2,10):
        for k in range (2,j+1):
                for atom in structure:              
                        if atom.symbol == magatom:
                                atom.symbol = 'Ac'

                candidate_SC_nn2 = sort(structure.repeat([j,k,1]))

                for atom in candidate_SC_nn2:
                        if atom.symbol == 'Ac':
                                atom.symbol = magatom
                
                s0_found = False  
                s2_found = False

                #for atom in candidate_SC_nn2:
                        #if np.allclose(atom.position, s2_position, atol=1e-6):
                                #s2 = atom.index
                                #s2_found = True
                                #break
                for atom in candidate_SC_nn2:
                    if not s0_found and np.allclose(atom.position, s0_position, atol=1e-6):
                        s0 = atom.index
                        s0_found = True
                    elif not s2_found and np.allclose(atom.position, s2_position, atol=1e-6):
                        s2 = atom.index
                        s2_found = True
                    if s0_found and s2_found:
                        break
                
                #very important to define dist_min = 0 here AND = 100 just after the continue statement
                #if s2 not found, the continue will correctly skip rest of k-loop
                #but before incrementing j, "dist_min > safe_dist" is evaluated
                #and this may cause the j-loop to be aborted as well, undesirably
                #if s2 is found, then dist_min must be set large enough for the rest of code in k-loop
                dist_min = 0

                #if not s2_found:
                if not (s0_found and s2_found):
                        continue

                super_superstructure = candidate_SC_nn2.repeat([2,2,1])

                dist_min = 100
                dist_min_vec = np.array([100,100,100,100])

                for m in range (1,4):
                        s0_img = s0+m*len(candidate_SC_nn2)
                        s2_img = s2+m*len(candidate_SC_nn2)

                        d_0_0img = float(super_superstructure.get_distances(s0,s0_img))
                        d_0_2img = float(super_superstructure.get_distances(s0,s2_img))
                        d_2_0img = float(super_superstructure.get_distances(s2,s0_img))
                        d_2_2img = float(super_superstructure.get_distances(s2,s2_img))
     
                        dist_couple_img = np.array([d_0_0img,d_0_2img,d_2_0img,d_2_2img])
                        if dist_min > dist_couple_img.min():
                                dist_min = dist_couple_img.min()
                                dist_min_vec = dist_couple_img
               
                if dist_min > safe_dist:
                        opt_j_nn2_direct = j
                        opt_k_nn2_direct = k
                        break
        
        if dist_min > safe_dist:
                break

    s0 = None
    s2 = None

    for k in range(2,10):
        for j in range (2,k+1):
                for atom in structure:              
                        if atom.symbol == magatom:
                                atom.symbol = 'Ac'

                candidate_SC_nn2 = sort(structure.repeat([j,k,1]))

                for atom in candidate_SC_nn2:
                        if atom.symbol == 'Ac':
                                atom.symbol = magatom
                
                s0_found = False
                s2_found = False

                #for atom in candidate_SC_nn2:
                        #if np.allclose(atom.position, s2_position, atol=1e-6):
                                #s2 = atom.index
                                #s2_found = True
                                #break
                for atom in candidate_SC_nn2:
                    if not s0_found and np.allclose(atom.position, s0_position, atol=1e-6):
                        s0 = atom.index
                        s0_found = True
                    elif not s2_found and np.allclose(atom.position, s2_position, atol=1e-6):
                        s2 = atom.index
                        s2_found = True
                    if s0_found and s2_found:
                        break

                #very important to define dist_min = 0 here AND = 100 just after the continue statement
                #if s2 not found, the continue will correctly skip rest of k-loop
                #but before incrementing j, "dist_min > safe_dist" is evaluated
                #and this may cause the j-loop to be aborted as well, undesirably
                #if s2 is found, then dist_min must be set large enough for the rest of code in k-loop
                dist_min = 0

                #if not s2_found:
                if not (s0_found and s2_found):
                        continue

                super_superstructure = candidate_SC_nn2.repeat([2,2,1])

                dist_min = 100
                dist_min_vec = np.array([100,100,100,100])

                for m in range (1,4):
                        s0_img = s0+m*len(candidate_SC_nn2)
                        s2_img = s2+m*len(candidate_SC_nn2)

                        d_0_0img = float(super_superstructure.get_distances(s0,s0_img))
                        d_0_2img = float(super_superstructure.get_distances(s0,s2_img))
                        d_2_0img = float(super_superstructure.get_distances(s2,s0_img))
                        d_2_2img = float(super_superstructure.get_distances(s2,s2_img))
     
                        dist_couple_img = np.array([d_0_0img,d_0_2img,d_2_0img,d_2_2img])
                        if dist_min > dist_couple_img.min():
                                dist_min = dist_couple_img.min()
                                dist_min_vec = dist_couple_img

                if dist_min > safe_dist:
                        opt_j_nn2_reversed = j
                        opt_k_nn2_reversed = k
                        break
                       
        if dist_min > safe_dist:
                break                
               
    for atom in structure:              
        if atom.symbol == magatom:
                atom.symbol = 'Ac'
            
    str_out += ("opt_j_nn2_direct " + str(opt_j_nn2_direct) + "       opt_k_nn2_direct " + str(opt_k_nn2_direct)+"\n")
    str_out += ("opt_j_nn2_reversed " + str(opt_j_nn2_reversed) + "       opt_k_nn2_reversed " +str(opt_k_nn2_reversed)+"\n")

    if opt_j_nn2_direct*opt_k_nn2_direct < opt_j_nn2_reversed*opt_k_nn2_reversed:
        opt_j_nn2 = opt_j_nn2_direct
        opt_k_nn2 = opt_k_nn2_direct
        SC_nn2 = sort(structure.repeat([opt_j_nn2,opt_k_nn2,1]))
        str_out += "Chosen direct values \n"
    elif opt_j_nn2_direct*opt_k_nn2_direct == opt_j_nn2_reversed*opt_k_nn2_reversed:
        opt_j_nn2 = opt_j_nn2_direct
        opt_k_nn2 = opt_k_nn2_direct
        SC_nn2 = sort(structure.repeat([opt_j_nn2,opt_k_nn2,1]))
        str_out += "Direct = reversed values \n"
    else:
        opt_j_nn2 = opt_j_nn2_reversed
        opt_k_nn2 = opt_k_nn2_reversed
        SC_nn2 = sort(structure.repeat([opt_j_nn2,opt_k_nn2,1]))
        str_out += "Chosen reversed values \n"
                         
    for atom in SC_nn2:
        if atom.symbol == 'Ac':
                atom.symbol = magatom
    
    #CRUCIAL! at this point index s2 refers to the last evaluated "reversed" supercell
    #if this is chosen as final supercell, then no problem
    #but if the direct is chosen, s2 may point to the wrong atom!!!!
    #hence the need to reassign its value
    for atom in SC_nn2:
        if np.allclose(atom.position, s0_position, atol=1e-6):
            s0 = atom.index
            break
    for atom in SC_nn2:
        if np.allclose(atom.position, s2_position, atol=1e-6):
            s2 = atom.index
            break
    
    SC_nn2.positions[[0,s0]] = SC_nn2.positions[[s0,0]] #to have pair of atoms appear in beginning of POSCAR
    SC_nn2.positions[[1,s2]] = SC_nn2.positions[[s2,1]] #to have pair of atoms appear in beginning of POSCAR
         
    final_structure_2 = StructureData(ase=SC_nn2)
               
    str_out += ('Minimal supercell for 2nn: ---- ' +str(opt_j_nn2) + ' x ' + str(opt_k_nn2) + "\n")
    str_out += ('with distances: ' + str(np.around(dist_min_vec,4)) + "\n \n")
    write('/home/federico/Amaranta_txt_files/POSCAR_2NN',SC_nn2,format='vasp')
   
    ######################################### OPTIMAL SC FOR 3rd NN

    for j in range(2,10):
        for k in range (2,j+1):
                for atom in structure:              
                        if atom.symbol == magatom:
                                atom.symbol = 'Ac'

                candidate_SC_nn3 = sort(structure.repeat([j,k,1]))

                for atom in candidate_SC_nn3:
                        if atom.symbol == 'Ac':
                                atom.symbol = magatom
                
                s0_found = False
                s3_found = False

                #for atom in candidate_SC_nn3:
                        #if np.allclose(atom.position, s3_position, atol=1e-6):
                                #s3 = atom.index
                                #s3_found = True
                                #break
                for atom in candidate_SC_nn3:
                    if not s0_found and np.allclose(atom.position, s0_position, atol=1e-6):
                        s0 = atom.index
                        s0_found = True
                    elif not s3_found and np.allclose(atom.position, s3_position, atol=1e-6):
                        s3 = atom.index
                        s3_found = True
                    if s0_found and s3_found:
                        break

                #very important to define dist_min = 0 here AND = 100 just after the continue statement
                #if s3 not found, the continue will correctly skip rest of k-loop
                #but before incrementing j, "dist_min > safe_dist" is evaluated
                #and this may cause the j-loop to be aborted as well, undesirably
                #if s3 is found, then dist_min must be set large enough for the rest of code in k-loop
                dist_min = 0

                #if not s3_found:
                if not (s0_found and s3_found):
                        continue

                super_superstructure = candidate_SC_nn3.repeat([2,2,1])
       
                dist_min = 100
                dist_min_vec = np.array([100,100,100,100])

                for m in range (1,4):
                        s0_img = s0+m*len(candidate_SC_nn3)
                        s3_img = s3+m*len(candidate_SC_nn3)

                        d_0_0img = float(super_superstructure.get_distances(s0,s0_img))
                        d_0_3img = float(super_superstructure.get_distances(s0,s3_img))
                        d_3_0img = float(super_superstructure.get_distances(s3,s0_img))
                        d_3_3img = float(super_superstructure.get_distances(s3,s3_img))

                        dist_couple_img = np.array([d_0_0img,d_0_3img,d_3_0img,d_3_3img])
                        if dist_min > dist_couple_img.min():
                                dist_min = dist_couple_img.min()
                                dist_min_vec = dist_couple_img

                if dist_min > safe_dist:
                        opt_j_nn3_direct = j
                        opt_k_nn3_direct = k
                        break

        if dist_min > safe_dist:
                break
    
    s0 = None    
    s3 = None
    
    for k in range(2,10):
        for j in range (2,k+1):
                for atom in structure:              
                        if atom.symbol == magatom:
                                atom.symbol = 'Ac'

                candidate_SC_nn3 = sort(structure.repeat([j,k,1]))
              
                for atom in candidate_SC_nn3:
                        if atom.symbol == 'Ac':
                                atom.symbol = magatom
                
                s0_found = False
                s3_found = False

                #for atom in candidate_SC_nn3:
                        #if np.allclose(atom.position, s3_position, atol=1e-6):
                                #s3 = atom.index
                                #s3_found = True
                                #break
                for atom in candidate_SC_nn3:
                    if not s0_found and np.allclose(atom.position, s0_position, atol=1e-6):
                        s0 = atom.index
                        s0_found = True
                    elif not s3_found and np.allclose(atom.position, s3_position, atol=1e-6):
                        s3 = atom.index
                        s3_found = True
                    if s0_found and s3_found:
                        break
                
                #very important to define dist_min = 0 here AND = 100 just after the continue statement
                #if s3 not found, the continue will correctly skip rest of k-loop
                #but before incrementing j, "dist_min > safe_dist" is evaluated
                #and this may cause the j-loop to be aborted as well, undesirably
                #if s3 is found, then dist_min must be set large enough for the rest of code in k-loop
                dist_min = 0

                #if not s3_found:
                if not (s0_found and s3_found):
                        continue

                super_superstructure = candidate_SC_nn3.repeat([2,2,1])
       
                dist_min = 100
                dist_min_vec = np.array([100,100,100,100])

                for m in range (1,4):
                        s0_img = s0+m*len(candidate_SC_nn3)
                        s3_img = s3+m*len(candidate_SC_nn3)

                        d_0_0img = float(super_superstructure.get_distances(s0,s0_img))
                        d_0_3img = float(super_superstructure.get_distances(s0,s3_img))
                        d_3_0img = float(super_superstructure.get_distances(s3,s0_img))
                        d_3_3img = float(super_superstructure.get_distances(s3,s3_img))

                        dist_couple_img = np.array([d_0_0img,d_0_3img,d_3_0img,d_3_3img])
                        if dist_min > dist_couple_img.min():
                                dist_min = dist_couple_img.min()
                                dist_min_vec = dist_couple_img

                if dist_min > safe_dist:
                        opt_j_nn3_reversed = j
                        opt_k_nn3_reversed = k
                        break

        if dist_min > safe_dist:
                break            
               
    for atom in structure:              
        if atom.symbol == magatom:
                atom.symbol = 'Ac'
    
    str_out += ("opt_j_nn3_direct " + str(opt_j_nn3_direct) + "       opt_k_nn3_direct " + str(opt_k_nn3_direct)+"\n")
    str_out += ("opt_j_nn3_reversed " + str(opt_j_nn3_reversed) + "       opt_k_nn3_reversed " +str(opt_k_nn3_reversed)+"\n")

    if opt_j_nn3_direct*opt_k_nn3_direct < opt_j_nn3_reversed*opt_k_nn3_reversed:
        opt_j_nn3 = opt_j_nn3_direct
        opt_k_nn3 = opt_k_nn3_direct
        SC_nn3 = sort(structure.repeat([opt_j_nn3,opt_k_nn3,1]))
        str_out += "Chosen direct values \n"
    elif opt_j_nn3_direct*opt_k_nn3_direct == opt_j_nn3_reversed*opt_k_nn3_reversed:
        opt_j_nn3 = opt_j_nn3_direct
        opt_k_nn3 = opt_k_nn3_direct
        SC_nn3 = sort(structure.repeat([opt_j_nn3,opt_k_nn3,1]))
        str_out += "Direct = reversed values \n"
    else:
        opt_j_nn3 = opt_j_nn3_reversed
        opt_k_nn3 = opt_k_nn3_reversed
        SC_nn3 = sort(structure.repeat([opt_j_nn3,opt_k_nn3,1]))
        str_out += "Chosen reversed values \n"

    for atom in SC_nn3:
        if atom.symbol == 'Ac':
                atom.symbol = magatom
    
    #CRUCIAL! at this point index s3 refers to the last evaluated "reversed" supercell
    #if this is chosen as final supercell, then no problem
    #but if the direct is chosen, s3 may point to the wrong atom!!!!
    #hence the need to reassign its value
    for atom in SC_nn3:
        if np.allclose(atom.position, s0_position, atol=1e-6):
            s0 = atom.index
            break
    for atom in SC_nn3:
        if np.allclose(atom.position, s3_position, atol=1e-6):
            s3 = atom.index
            break
    
    SC_nn3.positions[[0,s0]] = SC_nn3.positions[[s0,0]] #to have pair of atoms appear in beginning of POSCAR
    SC_nn3.positions[[1,s3]] = SC_nn3.positions[[s3,1]] #to have pair of atoms appear in beginning of POSCAR
    
    final_structure_3 = StructureData(ase=SC_nn3)
           
    str_out += ('Minimal supercell for 3nn ---- ' +str(opt_j_nn3) + ' x ' + str(opt_k_nn3) + "\n")
    str_out += ('with distances: ' + str(np.around(dist_min_vec,4)) + "\n \n")
    write('/home/federico/Amaranta_txt_files/POSCAR_3NN',SC_nn3,format='vasp')
       
    ######################################### prepare KPOINTS 
    
    kp_density = 50
    v1 = structure.cell.cellpar()[0]
    v2 = structure.cell.cellpar()[1]
    kp_uc_x = kp_density/v1
    kp_uc_y = kp_density/v2
    n_kp_axis_nn1_j = round(kp_uc_x/opt_j_nn1)
    n_kp_axis_nn1_k = round(kp_uc_y/opt_k_nn1)
    n_kp_axis_nn2_j = round(kp_uc_x/opt_j_nn2)
    n_kp_axis_nn2_k = round(kp_uc_y/opt_k_nn2)
    n_kp_axis_nn3_j = round(kp_uc_x/opt_j_nn3)
    n_kp_axis_nn3_k = round(kp_uc_y/opt_k_nn3)
    
    str_out += ('k-point linear density: ' +str(kp_density) +'\n')
    str_out += ('unit cell parameters: ' + str(round(v1,3)) + ' ' + str(round(v2,3)) + '\n')
    str_out += ('k-mesh for unit cell: ' + str(round(kp_uc_x,3)) + ' x ' + str(round(kp_uc_y,3)) +' x 1' + '\n')
    str_out += ('k-mesh for 1nn: ' + str(n_kp_axis_nn1_j) + ' x ' + str(n_kp_axis_nn1_k) + ' x 1' + '\n')
    str_out += ('k-mesh for 2nn: ' + str(n_kp_axis_nn2_j) + ' x ' + str(n_kp_axis_nn2_k) + ' x 1' + '\n')
    str_out += ('k-mesh for 3nn: ' + str(n_kp_axis_nn3_j) + ' x ' + str(n_kp_axis_nn3_k) + ' x 1' + '\n')
    
    mesh_1 = [n_kp_axis_nn1_j, n_kp_axis_nn1_k, 1]
    kpoints_data_1 = KpointsData()
    kpoints_data_1.set_kpoints_mesh(mesh_1)
    mesh_2 = [n_kp_axis_nn2_j, n_kp_axis_nn2_k, 1]
    kpoints_data_2 = KpointsData()
    kpoints_data_2.set_kpoints_mesh(mesh_2)
    mesh_3 = [n_kp_axis_nn3_j, n_kp_axis_nn3_k, 1]
    kpoints_data_3 = KpointsData()
    kpoints_data_3.set_kpoints_mesh(mesh_3)
   
    ######################################### prepare MAGMOMs for 4-state calcs, for all J's
    
    ############ Jxx 
       
    magmom1 = np.array([1.0, 0.0, 0.0, 1.0, 0.0, 0.0])
    magmom2 = np.array([1.0, 0.0, 0.0, -1.0, 0.0, 0.0])
    magmom3 = np.array([-1.0, 0.0, 0.0, 1.0, 0.0, 0.0])
    magmom4 = np.array([-1.0, 0.0, 0.0, -1.0, 0.0, 0.0])
    magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    new_magmomlist_xx = []
   
    for item in magmomlist:
        item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 2) * [0.0, 0.0, 1.0])
        item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        new_magmomlist_xx.append(item)
        
    ############ Jyy 
       
    magmom1 = np.array([0.0, 1.0, 0.0, 0.0, 1.0, 0.0])
    magmom2 = np.array([0.0, 1.0, 0.0, 0.0, -1.0, 0.0])
    magmom3 = np.array([0.0, -1.0, 0.0, 0.0, 1.0, 0.0])
    magmom4 = np.array([0.0, -1.0, 0.0, 0.0, -1.0, 0.0])
    magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    new_magmomlist_yy = []
   
    for item in magmomlist:
        item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 2) * [0.0, 0.0, 1.0])
        item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        new_magmomlist_yy.append(item)
        
    ############ Jzz 
       
    magmom1 = np.array([0.0, 0.0, 1.0, 0.0, 0.0, 1.0])
    magmom2 = np.array([0.0, 0.0, 1.0, 0.0, 0.0, -1.0])
    magmom3 = np.array([0.0, 0.0, -1.0, 0.0, 0.0, 1.0])
    magmom4 = np.array([0.0, 0.0, -1.0, 0.0, 0.0, -1.0])
    magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    new_magmomlist_zz = []
   
    for item in magmomlist:
        item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 2) * [1.0, 0.0, 0.0])
        item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        new_magmomlist_zz.append(item)
        
    ############ Jxy 
       
    magmom1 = np.array([1.0, 0.0, 0.0, 0.0, 1.0, 0.0])
    magmom2 = np.array([1.0, 0.0, 0.0, 0.0, -1.0, 0.0])
    magmom3 = np.array([-1.0, 0.0, 0.0, 0.0, 1.0, 0.0])
    magmom4 = np.array([-1.0, 0.0, 0.0, 0.0, -1.0, 0.0])
    magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    new_magmomlist_xy = []
   
    for item in magmomlist:
        item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 2) * [0.0, 0.0, 1.0])
        item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        new_magmomlist_xy.append(item)
        
    ############ Jyx 
       
    magmom1 = np.array([0.0, 1.0, 0.0, 1.0, 0.0, 0.0])
    magmom2 = np.array([0.0, 1.0, 0.0, -1.0, 0.0, 0.0])
    magmom3 = np.array([0.0, -1.0, 0.0, 1.0, 0.0, 0.0])
    magmom4 = np.array([0.0, -1.0, 0.0, -1.0, 0.0, 0.0])
    magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    new_magmomlist_yx = []
   
    for item in magmomlist:
        item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 2) * [0.0, 0.0, 1.0])
        item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        new_magmomlist_yx.append(item)
   
    ############ Jyz 
       
    magmom1 = np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0])
    magmom2 = np.array([0.0, 1.0, 0.0, 0.0, 0.0, -1.0])
    magmom3 = np.array([0.0, -1.0, 0.0, 0.0, 0.0, 1.0])
    magmom4 = np.array([0.0, -1.0, 0.0, 0.0, 0.0, -1.0])
    magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    new_magmomlist_yz = []
   
    for item in magmomlist:
        item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 2) * [1.0, 0.0, 0.0])
        item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        new_magmomlist_yz.append(item)
        
    ############ Jzy 
       
    magmom1 = np.array([0.0, 0.0, 1.0, 0.0, 1.0, 0.0])
    magmom2 = np.array([0.0, 0.0, 1.0, 0.0, -1.0, 0.0])
    magmom3 = np.array([0.0, 0.0, -1.0, 0.0, 1.0, 0.0])
    magmom4 = np.array([0.0, 0.0, -1.0, 0.0, -1.0, 0.0])
    magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    new_magmomlist_zy = []
   
    for item in magmomlist:
        item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 2) * [1.0, 0.0, 0.0])
        item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        new_magmomlist_zy.append(item)
        
    ############ Jxz 
       
    magmom1 = np.array([1.0, 0.0, 0.0, 0.0, 0.0, 1.0])
    magmom2 = np.array([1.0, 0.0, 0.0, 0.0, 0.0, -1.0])
    magmom3 = np.array([-1.0, 0.0, 0.0, 0.0, 0.0, 1.0])
    magmom4 = np.array([-1.0, 0.0, 0.0, 0.0, 0.0, -1.0])
    magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    new_magmomlist_xz = []
   
    for item in magmomlist:
        item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 2) * [0.0, 1.0, 0.0])
        item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        new_magmomlist_xz.append(item)        
            
    ############ Jzx 
       
    magmom1 = np.array([0.0, 0.0, 1.0, 1.0, 0.0, 0.0])
    magmom2 = np.array([0.0, 0.0, 1.0, -1.0, 0.0, 0.0])
    magmom3 = np.array([0.0, 0.0, -1.0, 1.0, 0.0, 0.0])
    magmom4 = np.array([0.0, 0.0, -1.0, -1.0, 0.0, 0.0])
    magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    new_magmomlist_zx = []
   
    for item in magmomlist:
        item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 2) * [0.0, 1.0, 0.0])
        item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        new_magmomlist_zx.append(item)
        
    ############ J_2NN
       
    magmom1 = np.array([1.0, 1.0])
    magmom2 = np.array([1.0, -1.0])
    magmom3 = np.array([-1.0, 1.0])
    magmom4 = np.array([-1.0, -1.0])
    magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    new_magmomlist_J2NN = []
   
    for item in magmomlist:
        item = np.append(item, (n_magatoms * opt_j_nn2 * opt_k_nn2 - 2) * [1.0])
        item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn2 * opt_k_nn2) * [0.0])
        new_magmomlist_J2NN.append(item) 
        
    ############ J_3NN
       
    magmom1 = np.array([1.0, 1.0])
    magmom2 = np.array([1.0, -1.0])
    magmom3 = np.array([-1.0, 1.0])
    magmom4 = np.array([-1.0, -1.0])
    magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    new_magmomlist_J3NN = []
   
    for item in magmomlist:
        item = np.append(item, (n_magatoms * opt_j_nn3 * opt_k_nn3 - 2) * [1.0])
        item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn3 * opt_k_nn3) * [0.0])
        new_magmomlist_J3NN.append(item)
        
    ############ A_scalar
       
    magmom1 = np.array([0.0, 0.0, 1.0])
    magmom2 = np.array([1.0, 0.0, 0.0])
    magmom3 = np.array([-1.0, 0.0, 0.0])
    magmom4 = np.array([0.0, 0.0, -1.0])
    magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    new_magmomlist_A_scalar = []
   
    for item in magmomlist:
        item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 1) * [0.0, 1.0, 0.0])
        item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        new_magmomlist_A_scalar.append(item)
    
    #### here is SIA tensorial, which we decided not to include
            
    ############ Axy
    
    #ratio = math.sqrt(2)/2
       
    #magmom1 = np.array([ratio, ratio, 0.0])
    #magmom2 = np.array([ratio, -ratio, 0.0])
    #magmom3 = np.array([-ratio, ratio, 0.0])
    #magmom4 = np.array([-ratio, -ratio, 0.0])
    #magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    #new_magmomlist_Axy = []
   
    #for item in magmomlist:
        #item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 1) * [0.0, 0.0, 1.0])
        #item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        #new_magmomlist_Axy.append(item) 
        
    ############ Axz
           
    #magmom1 = np.array([ratio, 0.0, ratio])
    #magmom2 = np.array([ratio, 0.0, -ratio])
    #magmom3 = np.array([-ratio, 0.0, ratio])
    #magmom4 = np.array([-ratio, 0.0, -ratio])
    #magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    #new_magmomlist_Axz = []
   
    #for item in magmomlist:
        #item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 1) * [0.0, 1.0, 0.0])
        #item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        #new_magmomlist_Axz.append(item)
        
    ############ Ayz
           
    #magmom1 = np.array([0.0, ratio, ratio])
    #magmom2 = np.array([0.0, ratio, -ratio])
    #magmom3 = np.array([0.0, -ratio, ratio])
    #magmom4 = np.array([0.0, -ratio, -ratio])
    #magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    #new_magmomlist_Ayz = []
   
    #for item in magmomlist:
        #item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 1) * [0.0, 1.0, 0.0])
        #item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        #new_magmomlist_Ayz.append(item)
        
    ############ Ayy-Axx
           
    #magmom1 = np.array([0.0, 1.0, 0.0])
    #magmom3 = np.array([1.0, 0.0, 0.0])
    #magmom4 = np.array([-1.0, 0.0, 0.0])
    #magmom2 = np.array([0.0, -1.0, 0.0])
    #magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    #new_magmomlist_Ayy_xx = []
   
    #for item in magmomlist:
        #item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 1) * [0.0, 1.0, 0.0])
        #item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        #new_magmomlist_Ayy_xx.append(item)
        
    ############ Azz-Axx
           
    #magmom1 = np.array([0.0, 0.0, 1.0])
    #magmom3 = np.array([1.0, 0.0, 0.0])
    #magmom4 = np.array([-1.0, 0.0, 0.0])
    #magmom2 = np.array([0.0, 0.0, -1.0])
    #magmomlist = [magmom1, magmom2, magmom3, magmom4]
   
    #new_magmomlist_Azz_xx = []
   
    #for item in magmomlist:
        #item = np.append(item, (n_magatoms * opt_j_nn1 * opt_k_nn1 - 1) * [0.0, 1.0, 0.0])
        #item = np.append(item, ((n_tot_atoms - n_magatoms) * opt_j_nn1 * opt_k_nn1) * [0.0, 0.0, 0.0])
        #new_magmomlist_Azz_xx.append(item)                           
          
    magmom_dict_all = {
    "J_xx": List([new_magmomlist_xx[0],new_magmomlist_xx[1],new_magmomlist_xx[2],new_magmomlist_xx[3]]),
    "J_yy": List([new_magmomlist_yy[0],new_magmomlist_yy[1],new_magmomlist_yy[2],new_magmomlist_yy[3]]),
    "J_zz": List([new_magmomlist_zz[0],new_magmomlist_zz[1],new_magmomlist_zz[2],new_magmomlist_zz[3]]),
    "J_xy": List([new_magmomlist_xy[0],new_magmomlist_xy[1],new_magmomlist_xy[2],new_magmomlist_xy[3]]),
    "J_yx": List([new_magmomlist_yx[0],new_magmomlist_yx[1],new_magmomlist_yx[2],new_magmomlist_yx[3]]),
    "J_yz": List([new_magmomlist_yz[0],new_magmomlist_yz[1],new_magmomlist_yz[2],new_magmomlist_yz[3]]),
    "J_zy": List([new_magmomlist_zy[0],new_magmomlist_zy[1],new_magmomlist_zy[2],new_magmomlist_zy[3]]),
    "J_xz": List([new_magmomlist_xz[0],new_magmomlist_xz[1],new_magmomlist_xz[2],new_magmomlist_xz[3]]),
    "J_zx": List([new_magmomlist_zx[0],new_magmomlist_zx[1],new_magmomlist_zx[2],new_magmomlist_zx[3]]),
    "J_scalar_2NN": List([new_magmomlist_J2NN[0],new_magmomlist_J2NN[1],new_magmomlist_J2NN[2],new_magmomlist_J2NN[3]]),
    "J_scalar_3NN": List([new_magmomlist_J3NN[0],new_magmomlist_J3NN[1],new_magmomlist_J3NN[2],new_magmomlist_J3NN[3]]),
    "A_scalar": List([new_magmomlist_A_scalar[0],new_magmomlist_A_scalar[1],new_magmomlist_A_scalar[2],new_magmomlist_A_scalar[3]]),
    #"A_xy": List([new_magmomlist_Axy[0],new_magmomlist_Axy[1],new_magmomlist_Axy[2],new_magmomlist_Axy[3]]),
    #"A_xz": List([new_magmomlist_Axz[0],new_magmomlist_Axz[1],new_magmomlist_Axz[2],new_magmomlist_Axz[3]]),
    #"A_yz": List([new_magmomlist_Ayz[0],new_magmomlist_Ayz[1],new_magmomlist_Ayz[2],new_magmomlist_Ayz[3]]),
    #"A_yy_xx": List([new_magmomlist_Ayy_xx[0],new_magmomlist_Ayy_xx[1],new_magmomlist_Ayy_xx[2],new_magmomlist_Ayy_xx[3]]),
    #"A_zz_xx": List([new_magmomlist_Azz_xx[0],new_magmomlist_Azz_xx[1],new_magmomlist_Azz_xx[2],new_magmomlist_Azz_xx[3]])    
    }
   
    magmom_all = Dict(magmom_dict_all)
    
    n_magatoms = Int(n_magatoms)
    magatom = Str(magatom)
    str_out = Str(str_out)
    
    # Return the optimized supercell structures, kpoints.... as a mapping
    return {
        "final_structure_1": final_structure_1,
        "final_structure_2": final_structure_2,
        "final_structure_3": final_structure_3,
        "kmesh_1": kpoints_data_1,
        "kmesh_2": kpoints_data_2,
        "kmesh_3": kpoints_data_3,
        "magmom_uc": magmom_uc,
        "magmom_all": magmom_all,
        "n_magatoms": n_magatoms,
        "magatom": magatom,
        "str_out" : str_out,
        "couples" : couples,
    }










PARAMS_dict = {
'Ac': [1.535, 170.048],
'Ag': [1.503, 249.752],
'Ag_GW': [1.376, 249.752],
'Ag_pv': [1.217, 297.68],
'Ag_sv_GW': [1.111, 354.226],
'Al': [1.402, 240.957],
'Al_GW': [1.058, 240.957],
'Al_sv_GW': [1.312, 411.007],
'Am': [1.588, 255.953],
'Ar': [0.953, 266.101],
'Ar_GW': [1.005, 290.518],
'As': [1.217, 208.87],
'As_d': [1.217, 288.762],
'As_GW': [1.217, 208.87],
'As_sv_GW': [1.111, 415.514],
'At': [1.635, 161.481],
'At_d_GW': [1.323, 266.303],
'At_sv_GW': [1.376, 328.597],
'Au': [1.503, 229.869],
'Au_GW': [1.455, 248.263],
'Au_sv_GW': [1.376, 306.52],
'B': [0.905, 318.762],
'Ba_sv': [1.979, 186.981],
'Ba_sv_GW': [1.482, 237.484],
'Be': [0.995, 247.951],
'Be_GW': [1.005, 247.951],
'Be_sv': [0.995, 308.45],
'Be_sv_GW': [0.847, 536.216],
'B_GW': [0.905, 318.762],
'B_h': [0.582, 700.0],
'Bi': [1.635, 105.071],
'Bi_d': [1.635, 242.856],
'Bi_d_GW': [1.323, 261.89],
'Bi_GW': [1.635, 146.628],
'Bi_sv_GW': [1.376, 323.502],
'Br': [1.164, 216.224],
'Br_GW': [1.164, 216.224],
'Br_sv_GW': [1.111, 475.88],
'B_s': [1.0, 269.251],
'C': [0.863, 400.0],
'Ca_pv': [1.746, 119.552],
'Ca_sv': [1.323, 266.727],
'Ca_sv_GW': [1.482, 281.209],
'Cd': [1.577, 274.265],
'Cd_GW': [1.482, 253.99],
'Cd_sv_GW': [1.111, 361.653],
'Ce': [1.323, 273.088],
'Ce_GW': [1.323, 304.649],
'Ce_h': [1.323, 299.927],
'C_GW': [0.863, 413.992],
'C_GW_new': [0.847, 413.992],
'C_h': [0.863, 700.0],
'C_h_GW': [0.635, 742.464],
'Cl': [1.111, 262.25],
'Cl_GW': [1.005, 262.25],
'Cl_h': [0.953, 409.272],
'Cm': [1.482, 258.027],
'Co': [1.302, 268.056],
'Co_GW': [1.302, 323.447],
'Co_pv': [1.302, 270.871],
'Co_sv': [1.058, 390.343],
'Co_sv_GW': [1.111, 387.407],
'Cr': [1.323, 227.202],
'Cr_pv': [1.323, 265.753],
'Cr_sv': [1.164, 395.443],
'Cr_sv_GW': [1.111, 384.753],
'C_s': [1.0, 273.704],
'Cs_sv': [1.323, 220.727], 
'Cs_sv_GW': [1.482, 198.012],
'Cu': [1.164, 295.521],
'Cu_GW': [1.164, 417.032],
'Cu_pv': [1.058, 368.406],
'Cu_sv_GW': [1.111, 391.702],
'F': [0.794, 400.0],
'Fe': [1.302, 267.969],
'Fe_GW': [1.217, 321.044],
'Fe_pv': [1.164, 293.303],
'Fe_sv': [1.058, 390.513],
'Fe_sv_GW': [1.111, 387.727],
'F_GW': [0.794, 487.335],
'F_GW_new': [0.847, 487.335],
'F_h': [0.635, 772.351],
'F_h_GW': [0.635, 847.822],
'Fr_sv': [1.535, 214.489],
'F_s': [0.9, 289.647],
'Ga': [1.402, 134.8],
'Ga_d': [1.217, 282.829],
'Ga_d_GW': [1.217, 404.723],
'Ga_GW': [1.402, 134.8],
'Ga_h': [1.005, 404.723],
'Ga_sv_GW': [1.111, 404.723],
'Ge': [1.217, 173.969],
'Ge_d': [1.217, 310.448],
'Ge_d_GW': [1.323, 375.614],
'Ge_GW': [1.217, 173.969],
'Ge_h': [1.005, 410.604],
'Ge_sv_GW': [1.111, 410.604],
'H': [0.37, 250.0],
'H1.25': [0.37, 250.0],
'H1.33': [0.37, 250.0],
'H1.5': [0.37, 250.0],
'H1.66': [0.37, 250.0],
'H1.75': [0.37, 250.0],
'H.25': [0.37, 250.0],
'H.33': [0.37, 250.0],
'H.42': [0.582, 250.0],
'H.5': [0.37, 250.0],
'H.58': [0.582, 250.0],
'H.66': [0.37, 250.0],
'H.75': [0.37, 250.0],
'AE': [1000.0],
'He': [0.635, 477.779],
'He_GW': [0.635, 404.806],
'Hf': [1.614, 220.431],
'Hf_pv': [1.614, 220.431],
'GW': [1.614, 237.414],
'Hf_sv_GW': [1.376, 282.716],
'Hg': [1.614, 233.142],
'Hg_sv_GW': [1.376, 311.949],
'H_GW': [0.582, 300.0],
'H_h': [0.37, 700.0],
'H_h_GW': [0.582, 700.0],
'H_s': [0.37, 200.0],
'I': [1.487, 175.712],
'I_GW': [1.487, 175.712],
'In': [1.677, 96.062],
'In_d': [1.677, 239.196],
'In_d_GW': [1.323, 278.582],
'In_sv_GW': [1.164, 366.636],
'Ir': [1.423, 210.837],
'Ir_sv_GW': [1.376, 319.708],
'I_sv_GW': [1.164, 381.757],
'K_pv': [2.275, 116.596],
'Kr': [1.323, 185.392],
'Kr_GW': [1.217, 252.563],
'K_sv': [1.588, 259.279],
'K_sv_GW': [1.482, 248.606],
'La': [1.535, 219.044],
'La_GW': [1.323, 313.728],
'La_s': [1.894, 136.594],
'Li': [1.376, 140.0],
'Li_AE_GW': [0.847, 433.253],
'Li_GW': [1.376, 112.417],
'Li_sv': [1.085, 498.387],
'Li_sv_GW': [0.847, 433.253],
'Mg': [1.524, 200.0],
'Mg_GW': [1.058, 126.671],
'Mg_pv': [1.058, 403.538],
'Mg_pv_GW': [1.058, 403.538],
'Mg_sv': [0.953, 473.54],
'Mg_sv_GW': [1.058, 430.099],
'Mn': [1.323, 269.944],
'Mn_GW': [1.323, 278.537],
'Mn_pv': [1.323, 269.944],
'Mn_sv': [1.164, 387.112],
'Mn_sv_GW': [1.111, 384.488],
'Mo': [1.455, 224.535],
'Mo_pv': [1.455, 224.535],
'Mo_sv': [1.455, 236.514],
'Mo_sv_GW': [1.111, 344.65],
'N': [0.741, 400.0],
'Na': [1.757, 101.956],
'Na_pv': [1.164, 259.494],
'Na_sv': [0.794, 644.874],
'Na_sv_GW': [1.164, 372.86],
'Nb_pv': [1.503, 207.263],
'Nb_sv': [1.27, 293.304],
'Nb_sv_GW': [1.111, 353.857],
'Ne': [0.794, 343.403],
'Ne_GW': [0.847, 431.952],
'Ne_s_GW': [0.9, 317.594],
'N_GW': [0.741, 420.681],
'N_GW_new': [0.847, 420.681],
'N_h': [0.783, 700.0],
'N_h_GW': [0.635, 755.833],
'Ni': [1.286, 269.618],
'Ni_GW': [1.286, 357.352],
'Ni_pv': [1.058, 367.726],
'Ni_sv_GW': [1.111, 389.485],
'Np': [1.588, 254.349],
'Np_s': [1.693, 213.932],
'N_s': [0.979, 279.424],
'N_s_GW': [0.953, 296.222],
'O': [0.82, 400.0],
'O_GW': [0.82, 414.315],
'O_GW_new': [0.847, 433.745],
'O_h': [0.741, 700.0],
'O_h_GW': [0.635, 765.442],
'O_s': [0.9, 282.604],
'Os': [1.413, 228.023],
'O_s_GW': [0.953, 334.366],
'Os_pv': [1.413, 228.023],
'Os_sv_GW': [1.376, 319.628],
'P': [1.233, 255.155],
'Pa': [1.535, 252.31],
'Pa_s': [1.693, 193.642],
'Pb': [1.725, 98.039],
'Pb_d': [1.725, 237.817],
'Pb_d_GW': [1.323, 237.793],
'Pb_sv_GW': [1.376, 317.138],
'Pd': [1.434, 250.832],
'Pd_GW': [1.693, 250.832],
'Pd_pv': [1.217, 250.832],
'Pd_sv_GW': [1.111, 355.88],
'P_GW': [1.005, 255.155],
'P_h': [0.953, 390.903],
'Po': [1.635, 159.801],
'Po_d': [1.635, 264.606],
'Po_d_GW': [1.323, 267.666],
'Po_sv_GW': [1.376, 326.653],
'Pt': [1.455, 230.228],
'Pt_GW': [1.455, 248.657],
'Pt_pv': [1.164, 294.53],
'Pt_sv_GW': [1.376, 323.536],
'Pu': [1.588, 254.436],
'Pu_s': [1.693, 211.493],
'Ra_sv': [1.535, 237.216],
'Rb_pv': [2.418, 122.21],
'Rb_sv': [2.418, 220.215],
'Rb_sv_GW': [1.482, 220.92],
'Re': [1.434, 226.25],
'Re_pv': [1.434, 226.25],
'Re_sv_GW': [1.376, 316.85],
'Rh': [1.402, 228.926],
'Rh_GW': [1.588, 247.321],
'Rh_pv': [1.402, 247.321],
'Rh_sv_GW': [1.111, 350.989],
'Rn': [1.635, 152.086],
'Rn_d_GW': [1.323, 268.546],
'Rn_sv_GW': [1.376, 331.257],
'Ru': [1.402, 213.221],
'Ru_pv': [1.402, 239.907],
'Ru_sv': [1.164, 318.687],
'Ru_sv_GW': [1.111, 347.881],
'S': [1.164, 258.602],
'Sb': [1.577, 172.301],
'Sb_d_GW': [1.429, 263.147],
'Sb_GW': [1.429, 172.301],
'Sb_sv_GW': [1.164, 372.498],
'Sc': [1.588, 155.006],
'Sc_sv': [1.429, 222.7],
'Sc_sv_GW': [1.111, 378.598],
'Se': [1.164, 211.602],
'Se_GW': [1.164, 211.602],
'Se_sv_GW': [1.111, 469.258],
'S_GW': [1.005, 258.602],
'S_h': [0.953, 402.84],
'Si': [1.312, 245.704],
'Si_GW': [1.312, 245.704],
'Si_sv_GW': [1.312, 546.548],
'Sn': [1.566, 103.318],
'Sn_d': [1.566, 241.107],
'Sn_d_GW': [1.323, 260.086],
'Sn_sv_GW': [1.164, 368.704],
'Sr_sv': [2.138, 226.327],
'Sr_sv_GW': [1.482, 224.532],
'Ta': [1.503, 223.759],
'Ta_pv': [1.503, 223.759],
'Ta_sv_GW': [1.376, 285.798],
'Tc': [1.455, 228.636],
'Tc_pv': [1.423, 263.345],
'Tc_sv': [1.423, 318.479],
'Tc_sv_GW': [1.111, 350.798],
'Te': [1.535, 175.144],
'Te_GW': [1.535, 175.144],
'Te_sv_GW': [1.164, 376.686],
'Th': [1.588, 247.389],
'Th_s': [1.693, 169.575],
'Ti': [1.323, 178.52],
'Ti_pv': [1.323, 222.435],
'Ti_sv': [1.217, 274.719],
'Ti_sv_GW': [1.111, 383.48],
'Tl': [1.725, 90.239],
'Tl_d': [1.725, 237.001],
'Tl_d_GW': [1.323, 237.001],
'Tl_sv_GW': [1.376, 316.502],
'U': [1.588, 252.603],
'U_s': [1.693, 209.218],
'V': [1.323, 192.706],
'V_pv': [1.217, 263.722],
'V_sv': [1.217, 263.722],
'V_sv_GW': [1.111, 382.093],
'W': [1.455, 223.126],
'W_sv': [1.455, 223.126],
'W_sv_GW': [1.376, 316.943],
'Xe': [1.323, 153.021],
'Xe_GW': [1.323, 179.528],
'Xe_sv_GW': [1.164, 400.581],
'Y_sv': [1.815, 202.554],
'Y_sv_GW': [1.111, 339.94],
'Zn': [1.27, 276.847],
'Zn_GW': [1.27, 328.312],
'Zn_sv_GW': [1.111, 401.745],
'Zr_sv': [1.625, 230.037],
'Zr_sv_GW': [1.111, 346.437]
}