Refactored to implement phono3py

aiida_phonopy/common/utils.py

@@ -1,68 +1,38 @@
 import numpy as np
 from aiida.engine import calcfunction
 from aiida.plugins import DataFactory
+from aiida.orm import Float, Bool, Str, Int, load_node
 from phonopy.structure.atoms import PhonopyAtoms
-from aiida.orm import Bool, Str, Int, load_node
+from phonopy.interface.calculator import get_default_displacement_distance
+from phonopy.structure.symmetry import symmetrize_borns_and_epsilon
 
 
-@calcfunction
-def get_phonon_setting_info(phonon_settings,
-                            structure,
-                            symmetry_tolerance,
-                            displacement_dataset=None):
-    return_vals = {}
+Dict = DataFactory('dict')
+ArrayData = DataFactory('array')
+XyData = DataFactory('array.xy')
+StructureData = DataFactory('structure')
+BandsData = DataFactory('array.bands')
 
-    ph_settings = {}
-    ph_settings.update(phonon_settings.get_dict())
-    dim = ph_settings['supercell_matrix']
-    if len(np.ravel(dim)) == 3:
-        smat = np.diag(dim)
-    else:
-        smat = np.array(dim)
-    if not np.issubdtype(smat.dtype, np.integer):
-        raise TypeError("supercell_matrix is not integer matrix.")
-    else:
-        ph_settings['supercell_matrix'] = smat.tolist()
 
-    if 'mesh' not in ph_settings:
-        ph_settings['mesh'] = 100.0
-    if 'distance' not in ph_settings:
-        ph_settings['distance'] = 0.01
-    tolerance = symmetry_tolerance.value
-    ph_settings['symmetry_tolerance'] = tolerance
+@calcfunction
+def generate_phonopy_cells(phonon_settings,
+                           structure,
+                           symmetry_tolerance,
+                           dataset=None):
+    ph_settings = _get_setting_info(phonon_settings,
+                                    structure,
+                                    symmetry_tolerance)
 
     ph = get_phonopy_instance(structure, ph_settings, {})
-    ph_settings['primitive_matrix'] = ph.primitive_matrix
-    ph_settings['symmetry'] = {
-        'number': ph.symmetry.dataset['number'],
-        'international': ph.symmetry.dataset['international']}
-
-    if displacement_dataset is None:
+    if dataset is None:
         ph.generate_displacements(distance=ph_settings['distance'])
     else:
-        ph.dataset = displacement_dataset.get_dict()
-    ph_settings['displacement_dataset'] = ph.dataset
-    settings = DataFactory('dict')(dict=ph_settings)
-    settings.label = 'phonon_setting_info'
-    return_vals['phonon_setting_info'] = settings
-
-    # Supercells and primitive cell
-    for i, scell in enumerate(ph.supercells_with_displacements):
-        structure = phonopy_atoms_to_structure(scell)
-        label = "supercell_%03d" % (i + 1)
-        structure.label = "%s %s" % (
-            structure.get_formula(mode='hill_compact'), label)
-        return_vals[label] = structure
+        ph.dataset = dataset.get_dict()
 
-    supercell_structure = phonopy_atoms_to_structure(ph.supercell)
-    supercell_structure.label = "%s %s" % (
-        supercell_structure.get_formula(mode='hill_compact'), 'supercell')
-    return_vals['supercell'] = supercell_structure
-
-    primitive_structure = phonopy_atoms_to_structure(ph.primitive)
-    primitive_structure.label = "%s %s" % (
-        primitive_structure.get_formula(mode='hill_compact'), 'primitive cell')
-    return_vals['primitive'] = primitive_structure
+    _update_structure_info(ph_settings, ph)
+    structures_dict = _generate_aiida_structures(ph)
+    return_vals = {'ph_settings': Dict(dict=ph_settings)}
+    return_vals.update(structures_dict)
 
     return return_vals
 
@@ -94,30 +64,59 @@ def check_imported_supercell_structure(supercell_ref,
 
 
 @calcfunction
-def get_force_sets(**forces_dict):
+def get_vasp_force_sets_dict(**forces_dict):
     forces = []
     energies = []
-    for i in range(len(forces_dict)):
-        label = "forces_%03d" % (i + 1)
-        if label in forces_dict:
-            forces.append(forces_dict[label].get_array('final'))
-        label = "misc_%03d" % (i + 1)
-        if label in forces_dict:
-            energies.append(
-                forces_dict[label]['total_energies']['energy_no_entropy'])
-
-    assert len(forces) == sum(['forces' in k for k in forces_dict])
-
-    force_sets = DataFactory('array')()
+    forces_0 = None
+    energy_0 = None
+
+    for key in forces_dict:
+        num = int(key.split('_')[-1])
+        if num == 0:
+            continue
+        if 'forces' in key:
+            forces.append(None)
+        elif 'misc' in key:
+            energies.append(None)
+
+    for key in forces_dict:
+        num = int(key.split('_')[-1])  # e.g. "001" --> 1
+        if 'forces' in key:
+            forces_ndarray = forces_dict[key].get_array('final')
+            if num == 0:
+                forces_0 = forces_ndarray
+            else:
+                forces[num - 1] = forces_ndarray
+        elif 'misc' in key:
+            energy = forces_dict[key]['total_energies']['energy_no_entropy']
+            if num == 0:
+                energy_0 = energy
+            else:
+                energies[num - 1] = energy
+
+    if forces_0 is not None:
+        for forces_ndarray in forces:
+            forces_ndarray -= forces_0
+
+    force_sets = ArrayData()
     force_sets.set_array('force_sets', np.array(forces))
     if energies:
         force_sets.set_array('energies', np.array(energies))
     force_sets.label = 'force_sets'
-    return force_sets
+    ret_dict = {'force_sets': force_sets}
+    if forces_0 is not None:
+        forces_0_array = ArrayData()
+        forces_0_array.set_array('forces', forces_0)
+        ret_dict['supercell_forces'] = forces_0_array
+    if energy_0 is not None:
+        ret_dict['supercell_energy'] = Float(energy_0)
+
+    return ret_dict
 
 
 @calcfunction
-def get_nac_params(born_charges, epsilon, nac_structure, **params):
+def get_nac_params(born_charges, epsilon, nac_structure, symmetry_tolerance,
+                   primitive=None):
     """Obtain Born effective charges and dielectric constants in primitive cell
 
     When Born effective charges and dielectric constants are calculated within
@@ -132,22 +131,18 @@ def get_nac_params(born_charges, epsilon, nac_structure, **params):
     symmetry_tolerance : Float
 
     """
-    from phonopy.structure.symmetry import symmetrize_borns_and_epsilon
-
     borns = born_charges.get_array('born_charges')
     eps = epsilon.get_array('epsilon')
 
     nac_cell = phonopy_atoms_from_structure(nac_structure)
-    kargs = {}
-    if 'symmetry_tolerance' in params:
-        kargs['symprec'] = params['symmetry_tolerance'].value
-    if 'primitive' in params:
-        pcell = phonopy_atoms_from_structure(params['primitive'])
-        kargs['primitive'] = pcell
+    kwargs = {}
+    kwargs['symprec'] = symmetry_tolerance.value
+    if primitive is not None:
+        kwargs['primitive'] = phonopy_atoms_from_structure(primitive)
     borns_, epsilon_ = symmetrize_borns_and_epsilon(
-        borns, eps, nac_cell, **kargs)
+        borns, eps, nac_cell, **kwargs)
 
-    nac_params = DataFactory('array')()
+    nac_params = ArrayData()
     nac_params.set_array('born_charges', borns_)
     nac_params.set_array('epsilon', epsilon_)
     nac_params.label = 'born_charges & epsilon'
@@ -162,7 +157,7 @@ def get_force_constants(structure, phonon_settings, force_sets):
     phonon.dataset = phonon_settings['displacement_dataset']
     phonon.forces = force_sets.get_array('force_sets')
     phonon.produce_force_constants()
-    force_constants = DataFactory('array')()
+    force_constants = ArrayData()
     force_constants.set_array('force_constants', phonon.force_constants)
     force_constants.set_array('p2s_map', phonon.primitive.p2s_map)
     force_constants.label = 'force_constants'
@@ -199,18 +194,18 @@ def get_data_from_node_id(node_id):
         raise RuntimeError("Crystal structure could not be found.")
 
     if 'born_charges' in n.outputs and 'dielectrics' in n.outputs:
-        born = DataFactory('array')()
+        born = ArrayData()
         born.set_array(
             'born_charges', n.outputs.born_charges.get_array('born_charges'))
         born.label = 'born_charges'
-        epsilon = DataFactory('array')()
+        epsilon = ArrayData()
         epsilon.set_array(
             'epsilon', n.outputs.dielectrics.get_array('epsilon'))
         epsilon.label = 'epsilon'
         return {'born_charges': born, 'dielectrics': epsilon,
                 'structure': structure}
     elif 'forces' in n.outputs:
-        forces = DataFactory('array')()
+        forces = ArrayData()
         forces.set_array('final', n.outputs.forces.get_array('final'))
         forces.label = 'forces'
         return {'forces': forces, 'structure': structure}
@@ -235,15 +230,15 @@ def get_mesh_property_data(ph, mesh):
 
 
 def get_total_dos(total_dos):
-    dos = DataFactory('array.xy')()
+    dos = XyData()
     dos.set_x(total_dos['frequency_points'], 'Frequency', 'THz')
     dos.set_y(total_dos['total_dos'], 'Total DOS', '1/THz')
     dos.label = 'Total DOS'
     return dos
 
 
 def get_projected_dos(projected_dos):
-    pdos = DataFactory('array.xy')()
+    pdos = XyData()
     pdos_list = [pd for pd in projected_dos['projected_dos']]
     pdos.set_x(projected_dos['frequency_points'], 'Frequency', 'THz')
     pdos.set_y(pdos_list,
@@ -254,7 +249,7 @@ def get_projected_dos(projected_dos):
 
 
 def get_thermal_properties(thermal_properties):
-    tprops = DataFactory('array.xy')()
+    tprops = XyData()
     tprops.set_x(thermal_properties['temperatures'], 'Temperature', 'K')
     tprops.set_y([thermal_properties['free_energy'],
                   thermal_properties['entropy'],
@@ -308,7 +303,7 @@ def get_bands(qpoints, frequencies, labels, path_connections, label=None):
             frequencies_list += list(fs)
     labels_list.append((len(qpoints_list) - 1, labels[-1]))
 
-    bs = DataFactory('array.bands')()
+    bs = BandsData()
     bs.set_kpoints(np.array(qpoints_list))
     bs.set_bands(np.array(frequencies_list), units='THz')
     bs.labels = labels_list
@@ -351,7 +346,7 @@ def get_primitive(structure, ph_settings):
     symbols = primitive_phonopy.get_chemical_symbols()
     positions = primitive_phonopy.get_positions()
 
-    primitive_structure = DataFactory('structure')(cell=primitive_cell)
+    primitive_structure = StructureData(cell=primitive_cell)
     for symbol, position in zip(symbols, positions):
         primitive_structure.append_atom(position=position, symbols=symbol)
 
@@ -361,7 +356,7 @@ def get_primitive(structure, ph_settings):
 def phonopy_atoms_to_structure(cell):
     symbols = cell.get_chemical_symbols()
     positions = cell.get_positions()
-    structure = DataFactory('structure')(cell=cell.get_cell())
+    structure = StructureData(cell=cell.get_cell())
     for symbol, position in zip(symbols, positions):
         structure.append_atom(position=position, symbols=symbol)
     return structure
@@ -382,3 +377,118 @@ def from_node_id_to_aiida_node_id(node_id):
     else:
         raise RuntimeError("%s is not supported in load_node."
                            % type(node_id))
+
+
+def collect_vasp_forces_and_energies(ctx, ctx_supercells):
+    forces_dict = {}
+    for key in ctx_supercells:
+        # key: e.g. "supercell_001", "phonon_supercell_001"
+        num = key.split('_')[-1]  # e.g. "001"
+
+        calc = ctx[key]
+        if type(calc) is dict:
+            calc_dict = calc
+        else:
+            calc_dict = calc.outputs
+        if ('forces' in calc_dict and
+            'final' in calc_dict['forces'].get_arraynames()):
+            forces_dict["forces_%s" % num] = calc_dict['forces']
+        else:
+            raise RuntimeError(
+                "Forces could not be found in calculation %s." % num)
+
+        if ('misc' in calc_dict and
+            'total_energies' in calc_dict['misc'].keys()):  # needs .keys()
+            forces_dict["misc_%s" % num] = calc_dict['misc']
+
+    return forces_dict
+
+
+def _get_setting_info(phonon_settings,
+                      structure,
+                      symmetry_tolerance,
+                      code_name='phonopy'):
+    """Convert AiiDA inputs to a dict
+
+    code_name : 'phonopy' or 'phono3py'
+
+    Note
+    ----
+    Designed to be shared by phonopy and phono3py.
+
+    """
+
+    ph_settings = {}
+    ph_settings.update(phonon_settings.get_dict())
+    dim = ph_settings['supercell_matrix']
+    if len(np.ravel(dim)) == 3:
+        smat = np.diag(dim)
+    else:
+        smat = np.array(dim)
+    if not np.issubdtype(smat.dtype, np.integer):
+        raise TypeError("supercell_matrix is not integer matrix.")
+    else:
+        ph_settings['supercell_matrix'] = smat.tolist()
+
+    if 'mesh' not in ph_settings:
+        ph_settings['mesh'] = 100.0
+    if 'distance' not in ph_settings:
+        if code_name == 'phono3py':
+            from phono3py.interface.calculator import (
+                get_default_displacement_distance as get_phono3py_ddd)
+            distance = get_phono3py_ddd('vasp')
+        else:
+            distance = get_default_displacement_distance('vasp')
+        ph_settings['distance'] = distance
+    ph_settings['symmetry_tolerance'] = symmetry_tolerance.value
+
+    return ph_settings
+
+
+def _update_structure_info(ph_settings, ph):
+    """Update a dict
+
+    Note
+    ----
+    Designed to be shared by phonopy and phono3py.
+    ph is either an instance of Phonopy or Phono3py.
+
+    """
+    ph_settings['displacement_dataset'] = ph.dataset
+    ph_settings['primitive_matrix'] = ph.primitive_matrix
+    ph_settings['symmetry'] = {
+        'number': ph.symmetry.dataset['number'],
+        'international': ph.symmetry.dataset['international']}
+    return ph_settings
+
+
+def _generate_aiida_structures(ph):
+    """Generate AiiDA structures of phonon related cells
+
+    Note
+    ----
+    Designed to be shared by phonopy and phono3py.
+    ph is either an instance of Phonopy or Phono3py.
+
+    """
+
+    structures_dict = {}
+
+    for i, scell in enumerate(ph.supercells_with_displacements):
+        structure = phonopy_atoms_to_structure(scell)
+        label = "supercell_%03d" % (i + 1)
+        structure.label = "%s %s" % (
+            structure.get_formula(mode='hill_compact'), label)
+        structures_dict[label] = structure
+
+    supercell_structure = phonopy_atoms_to_structure(ph.supercell)
+    supercell_structure.label = "%s %s" % (
+        supercell_structure.get_formula(mode='hill_compact'), 'supercell')
+    structures_dict['supercell'] = supercell_structure
+
+    primitive_structure = phonopy_atoms_to_structure(ph.primitive)
+    primitive_structure.label = "%s %s" % (
+        primitive_structure.get_formula(mode='hill_compact'), 'primitive cell')
+    structures_dict['primitive'] = primitive_structure
+
+    return structures_dict