Minor fix of IterHa

aiida-phonopy · Jun 19, 2020 · a95d0b2 · a95d0b2
1 parent dd580d2
commit a95d0b2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/aiida_phonopy/workflows/iter_ha.py b/aiida_phonopy/workflows/iter_ha.py
@@ -70,7 +70,7 @@ def _remove_high_energy_snapshots(d, f, e, ratio):
     idx = np.argsort(e)[:num_include]
     d = d[idx]
     f = f[idx]
-    return d, f, e
+    return d, f, e, idx
 
 
 def _modify_force_constants(ph):
@@ -121,7 +121,7 @@ def get_random_displacements(structure,
     if e is not None and len(e) == len(f) and 'include_ratio' in data:
         ratio = data['include_ratio'].value
         if 0 < ratio and ratio < 1:
-            d, f, e = _remove_high_energy_snapshots(d, f, e, ratio)
+            d, f, e, idx = _remove_high_energy_snapshots(d, f, e, ratio)
 
     # Calculate force constants by fitting using ALM
     phonon_setting_info = data['ph_info_1']