From 15e8156f1ccd74e0e824b553cb011205b2d17de2 Mon Sep 17 00:00:00 2001
From: Oskar Weser <oskar.weser@gmail.com>
Date: Tue, 21 Jan 2025 15:27:36 -0500
Subject: [PATCH] added get_fs function

---
 src/quemb/molbe/chemfrag.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/src/quemb/molbe/chemfrag.py b/src/quemb/molbe/chemfrag.py
index d118bb71..4fadac96 100644
--- a/src/quemb/molbe/chemfrag.py
+++ b/src/quemb/molbe/chemfrag.py
@@ -153,6 +153,25 @@ def AO_indices_of_fragment(fragment: set[AtomIdx]) -> set[AOIdx]:
     )
 
 
+def get_fs(
+    mol: Mole, fragments: AtomPerFrag
+) -> dict[CenterIdx, dict[AtomIdx, set[AOIdx]]]:
+    atom_to_AO = [
+        range(AO_offsets[2], AO_offsets[3]) for AO_offsets in mol.aoslice_by_atom()
+    ]
+
+    def AO_indices_of_fragment(fragment: set[AtomIdx]) -> dict[AtomIdx, set[AOIdx]]:
+        return {
+            i_atom: {AOIdx(AO_idx) for AO_idx in atom_to_AO[i_atom]}
+            for i_atom in fragment
+        }
+
+    return {
+        i_center: AO_indices_of_fragment(fragment)
+        for i_center, fragment in fragments.items()
+    }
+
+
 def get_edge_idx(fragments: AtomPerFrag) -> dict[CenterIdx, set[CenterIdx]]:
     return {
         i_center: {