enable get subgroup_comp function

Author: qzhu2017

pyxtal/symmetry.py

@@ -521,6 +521,58 @@ def get_spg_representation(self):
             one_hot: a (15, 26) numpy (0, 1) array
         """
         return self.lattice_id, self.get_spg_symmetry_object().to_one_hot()
+    
+    def get_subgroup_composition(self, ids=None, g_types=['t', 'k'], max_atoms=100, 
+                                 verbose=False):
+        """
+        Get the composition of the subgroup Wyckoff positions.
+
+        Args:
+            ids (list, optional): List of Wyckoff position indices.
+            verbose (bool): Whether to print debug information.
+            g_types (list): List of subgroup types to consider ('t' for translationengleiche, 'k' for klassengleiche).
+            max_atoms (int): Maximum number of atoms to consider for subgroup search.
+
+        Returns:
+            list: List of multiplicities of the Wyckoff positions.
+        """
+        if verbose:
+            strs = f"{self.number} ({self.symbol}): "
+            for id in ids:
+                wp = self[id]
+                strs += f"{wp.multiplicity}{wp.letter}"
+            print(strs)
+        sub_symmetries = []
+        N_atoms = sum([self[id].multiplicity for id in ids])
+        
+        for g_type in g_types:
+            if g_type == 't':
+                sub = self.get_max_t_subgroup()
+            else:
+                sub = self.get_max_k_subgroup()
+
+            for i, sub_g in enumerate(sub['subgroup']):
+                if N_atoms * sub['index'][i] > max_atoms:
+                    continue
+                sub_gg = Group(sub_g)
+                relation = sub['relations'][i]#; print(relation)
+                sub_ids = []
+                for id in ids:
+                    true_id = len(self)-id-1
+                    relation[true_id].sort()
+                    for r in relation[true_id]:
+                        letter = r[-1]#; print("test letter:", relation[true_id])
+                        sub_ids.append(len(sub_gg) - letters.index(letter) - 1)
+                data = (sub_g, sub_ids)
+                if data not in sub_symmetries:
+                    sub_symmetries.append(data)
+                    if verbose:
+                        strs = f"{sub_gg.number} ({sub_gg.symbol}): "
+                        for id in sub_ids:
+                            wp = sub_gg[id]
+                            strs += f"{wp.multiplicity}{wp.letter} "
+                        print(strs, data)
+        return sub_symmetries
 
     def get_lattice_id(self):
         """

tests/test_group.py

@@ -85,6 +85,14 @@ def test_check_compatible(self):
         assert Group(221).check_compatible([15, 16, 17, 21]) == (True, True)
         assert Group(221).check_compatible([15, 16, 2, 2]) == (False, False)
 
+    def test_subgroup_composition(self):
+        g = Group(154)
+        comps = g.get_subgroup_composition([2, 0], g_types=["t", "k"])
+        assert comps[0] == (5, [2, 0, 0, 0, 0])
+        assert comps[1] == (145, [0, 0, 0])
+        assert comps[2] == (152, [0, 0, 0])
+        assert comps[3] == (152, [2, 1, 0, 0])
+        assert len(comps) == 8
 
     def test_search_supergroup_paths(self):
         paths = Group(59, quick=True).search_supergroup_paths(139, 2)