add an util method to expand symmetry

SFC_Torch/Fmodel.py

@@ -288,6 +288,23 @@ def orth2frac(self, orth_pos: torch.Tensor) -> torch.Tensor:
         frac_pos = torch.einsum("n...x,yx->n...y", orth_pos, self.orth2frac_tensor)
         return frac_pos
 
+    def exp_sym(self, frac_pos: torch.Tensor | None = None) -> torch.Tensor:
+        """
+        Apply all symmetry operations to the fractional coordinates
+
+        Args:
+            frac_pos, torch.Tensor, [n_points, 3]
+                fractional coordinates of model in single ASU. 
+                If not given, will use self.atom_pos_frac
+        Returns:
+            torch.Tensor, [n_points, n_ops, 3]
+                fractional coordinates of symmetry operated models        
+        """
+        if frac_pos is None:
+            frac_pos = self.atom_pos_frac
+        sym_oped_frac_pos = torch.einsum("oxy,ay->aox", self.R_G_tensor_stack, frac_pos) + self.T_G_tensor_stack
+        return sym_oped_frac_pos
+
     def init_mtz(self, mtzdata, N_bins, expcolumns, set_experiment, freeflag, testset_value):
         """
         set mtz file for HKL list, resolution and experimental related properties

SFC_Torch/io.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import gemmi
 import urllib.request, os
 
@@ -180,6 +182,41 @@ def to_gemmi(self, include_header=True):
     @property
     def atom_pos_frac(self):
         return self.orth2frac(self.atom_pos)
+
+    @property
+    def operations(self):
+        return self.spacegroup.operations() 
+
+    @property
+    def R_G_stack(self):
+        """
+        [n_ops, 3, 3], np.ndarray
+        """
+        return np.array([np.array(sym_op.rot) / sym_op.DEN for sym_op in self.operations])
+
+    @property
+    def T_G_stack(self):
+        """
+        [n_ops, 3], np.ndarray
+        """
+        return np.array([np.array(sym_op.tran) / sym_op.DEN for sym_op in self.operations])
+
+    def exp_sym(self, frac_pos: np.ndarray | None = None) -> np.ndarray:
+        """
+        Apply all symmetry operations to the fractional coordinates
+
+        Args:
+            frac_pos, np.ndarray, [n_points, 3]
+                fractional coordinates of model in single ASU. 
+                If not given, will use self.atom_pos_frac
+        Returns:
+            np.ndarray, [n_points, n_ops, 3]
+                fractional coordinates of symmetry operated models
+        """
+        if frac_pos is None:
+            frac_pos = self.atom_pos_frac
+        sym_oped_frac_pos = np.einsum("oxy,ay->aox", self.R_G_stack, frac_pos) + self.T_G_stack
+        return sym_oped_frac_pos
 
     def set_spacegroup(self, spacegroup):
         """