Change generate_deformations to return strain_inds and space group nu…

…mber

abipy/dfpt/deformation_utils.py

@@ -30,15 +30,28 @@ def generate_deformations_volumic(structure, eps_V=0.02, scales=None):
     return structures_new
 
 
-def generate_deformations(structure, eps=0.005):
-    spgrp = AbinitSpaceGroup.from_structure(structure )
-    print (spgrp)
-    spgrp_number=spgrp.spgid
-    rprim= structure.lattice.matrix
+#def generate_deformations(structure, eps=0.005) -> tuple:
+def generate_deformations(structure, eps: float) -> tuple:
+    """
+    """
+    spgrp = AbinitSpaceGroup.from_structure(structure)
+    #print(spgrp)
+
+    spgrp_number = spgrp.spgid
+    rprim = structure.lattice.matrix
 
     rprim2 = np.copy(rprim)
-    rprim_new = {}
     structures_new = {}
+    strain_inds = []
+
+    def _add(name, new_rprim, i, j, k, l, m, n) -> None:
+        """Helper function to register a new structure in internal dict."""
+        new_structure = structure.copy()
+        new_structure.lattice = Lattice(new_rprim)
+        structures_new[name] = new_structure
+        strain_inds.append([i, j, k, l, m, n])
+
+    i, j, k, l, m, n = 6 * [None]
 
     if 1 <= spgrp_number <= 2:
         disp=[[1,1,1,1,1,1],  [0,1,1,1,1,1],  [2,1,1,1,1,1],  [1,0,1,1,1,1],  [1,2,1,1,1,1],  [1,1,0,1,1,1],
@@ -79,7 +92,7 @@ def generate_deformations(structure, eps=0.005):
         print(rprim0)
 
         for pair in disp:
-            i,j,k,l,m,n = pair
+            i, j, k, l, m, n = pair
             rprim2[ :,0] = rprim0[ :,0] * (1.00 + eps * i) + rprim0[ :,1] * (eps * l) +rprim0[ :,2] * (eps * m)
             rprim2[ :,1] = rprim0[ :,1] * (1.00 + eps * j) + rprim0[ :,2] * (eps * n)
             rprim2[ :,2] = rprim0[ :,2] * (1.00 + eps * k)
@@ -92,11 +105,8 @@ def generate_deformations(structure, eps=0.005):
             namen = int(round(1000 * (1.00 + eps * n)))
             formatted_namei = f"{namei:04d}_{namej:04d}_{namek:04d}_{namel:04d}_{namem:04d}_{namen:04d}"
 
-            structure2=structure.copy()
-            structure2.lattice=Lattice(rprim2)
-            structures_new[formatted_namei] = structure2
+            _add(formatted_namei, rprim2, i, j, k, l, m, n)
 
-        return structures_new
     elif 3 <= spgrp_number <= 15:
         disp=[[1,1,1,1], [0,1,1,1], [2,1,1,1], [1,0,1,1], [1,2,1,1], [1,1,0,1], [1,1,2,1], [1,1,1,0],
               [1,1,1,2], [0,0,1,1], [1,0,0,1], [1,1,0,0], [0,1,0,1], [1,0,1,0], [0,1,1,0]]
@@ -128,7 +138,7 @@ def generate_deformations(structure, eps=0.005):
         print(rprim0)
 
         for pair in disp:
-            i,j,k,l = pair
+            i, j, k, l = pair
             rprim2[ :,0] = rprim0[ :,0] * (1.00 + eps * i) +rprim0[ :,2] * (eps * l)
             rprim2[ :,1] = rprim0[ :,1] * (1.00 + eps * j)
             rprim2[ :,2] = rprim0[ :,2] * (1.00 + eps * k)
@@ -139,15 +149,12 @@ def generate_deformations(structure, eps=0.005):
             namel = int(round(1000 * (1.00 + eps * l)))
             formatted_namei = f"{namei:04d}_{namej:04d}_{namek:04d}_{namel:04d}"
 
-            structure2=structure.copy()
-            structure2.lattice=Lattice(rprim2)
-            structures_new[formatted_namei] = structure2
+            _add(formatted_namei, rprim2, i, j, k, l, m, n)
 
-        return structures_new
     elif 16 <= spgrp_number <= 74:
         disp=[[0,0,1],[0,1,0],[1,0,0],[1,1,1],[0,1,1],[2,1,1],[1,0,1],[1,2,1],[1,1,0],[1,1,2]]
         for pair in disp:
-            i,j,k = pair
+            i, j, k = pair
             rprim2[ :,0] = rprim[ :,0] * (1.00 + eps * i)
             rprim2[ :,1] = rprim[ :,1] * (1.00 + eps * j)
             rprim2[ :,2] = rprim[ :,2] * (1.00 + eps * k)
@@ -157,11 +164,8 @@ def generate_deformations(structure, eps=0.005):
             namek = int(round(1000 * (1.00 + eps * k)))
             formatted_namei = f"{namei:04d}_{namej:04d}_{namek:04d}"
 
-            structure2=structure.copy()
-            structure2.lattice=Lattice(rprim2)
-            structures_new[formatted_namei] = structure2
+            _add(formatted_namei, rprim2, i, j, k, l, m, n)
 
-        return structures_new
     elif 75 <= spgrp_number <= 194:
         disp=[[0,0],[1,1],[0,1],[2,1],[1,0],[1,2]]
         for pair in disp:
@@ -173,13 +177,9 @@ def generate_deformations(structure, eps=0.005):
             namei = int(round(1000 * (1.00 + eps * i)))
             namek = int(round(1000 * (1.00 + eps * k)))
             formatted_namei = f"{namei:04d}_{namek:04d}"
-            rprim_new[formatted_namei] = rprim2
 
-            structure2=structure.copy()
-            structure2.lattice=Lattice(rprim2)
-            structures_new[formatted_namei] = structure2
+            _add(formatted_namei, rprim2, i, j, k, l, m, n)
 
-        return structures_new
     elif 195 <= spgrp_number <= 230:
         for i in range(3):
             rprim2[ :,0] = rprim[ :,0] * (1.00 + eps * i)
@@ -188,8 +188,10 @@ def generate_deformations(structure, eps=0.005):
             namei = int(round(1000 * (1.00 + eps * i)))
             formatted_namei = f"{namei:04d}"
 
-            structure2=structure.copy()
-            structure2.lattice=Lattice(rprim2)
-            structures_new[formatted_namei] = structure2
-        return structures_new
+            _add(formatted_namei, rprim2, i, j, k, l, m, n)
+
+    else:
+        raise ValueError(f"Invalid {spgrp_number=}")
+
+    return structures_new, np.array(strain_inds, dtype=object), spgrp_number
 

abipy/dfpt/qha_general_stress.py

@@ -1217,7 +1217,7 @@ def cal_stress(self, temp, pressure=0, case=None):
         elif self.case == "ZSISA_slab_3DOF":
             dtol, stress = self.stress_ZSISA_slab_3DOF(temp, pressure)
         else:
-            raise ValueError(f"Unknown case: {case}")
+            raise ValueError(f"Unknown {case=}")
 
         if all(dtol[i] < 1e-8 for i in range(6)):
             with open("cell.txt", "a") as f:

abipy/flowtk/zsisa.py

@@ -61,11 +61,11 @@ def finalize(self):
         It performs some basic post-processing of the results to facilitate further analysis.
         """
         work = self[0]
-        data = {"eps": work.eps}
+        data = {"eps": work.eps, "spgrp_number": work.spgrp_number}
 
         # Build list of strings with path to the relevant output files ordered by V.
         data["gsr_relax_paths"] = [task.gsr_path for task in work.relax_tasks_deformed]
-        data["cell6_inds"] = work.cell6_inds
+        data["strain_inds"] = work.strain_inds
 
         gsr_relax_entries, gsr_relax_volumes = [], []
         for task in work.relax_tasks_deformed:
@@ -143,8 +143,7 @@ def on_ok(self, sender):
             # Get relaxed structure and build new task for structural relaxation at fixed volume.
             relaxed_structure = sender.get_final_structure()
 
-            self.deformed_structures_dict = generate_deformations(relaxed_structure, eps=self.eps)
-            self.cell6_inds
+            self.deformed_structures_dict, self.strain_inds, self.spgrp_number = generate_deformations(relaxed_structure, self.eps)
 
             relax_template = self.relax_template
             self.relax_tasks_deformed = []
@@ -206,7 +205,7 @@ def from_relax_input(cls,
                          relax_input: AbinitInput,
                          zsisa,
                          temperatures,
-                         pressures) -> ThermalRelaxWork
+                         pressures) -> ThermalRelaxWork:
         """
         Args:
             relax_input: