From 139b242f644daf0e172956491d873c77ceaa9d7d Mon Sep 17 00:00:00 2001
From: minhuanli <minhuanli@g.harvard.edu>
Date: Tue, 30 Apr 2024 16:07:01 -0400
Subject: [PATCH] add an util method to expand symmetry

---
 SFC_Torch/Fmodel.py | 17 +++++++++++++++++
 SFC_Torch/io.py     | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/SFC_Torch/Fmodel.py b/SFC_Torch/Fmodel.py
index 52ff042..07c6e2f 100644
--- a/SFC_Torch/Fmodel.py
+++ b/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
diff --git a/SFC_Torch/io.py b/SFC_Torch/io.py
index 6a88c45..db8a5a6 100644
--- a/SFC_Torch/io.py
+++ b/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):
         """