Update app.py

app.py

@@ -7,7 +7,7 @@ import xlsxwriter
 from io import BytesIO
 from collections import defaultdict
 
-# Detect homo repeats like AA, AAA, etc.
+# Utility to check homo repeat
 def is_homo_repeat(s):
     return all(c == s[0] for c in s)
 
@@ -63,9 +63,8 @@ def find_new_boundary_repeats(fragments, final_repeats, overlap=50):
 def process_protein_sequence(sequence, analysis_type, overlap=50):
     fragments = fragment_protein_sequence(sequence)
     final_repeats = defaultdict(int)
-    homo_repeats = {}
 
-    if analysis_type in ["Hetero", "Both"]:
+    if analysis_type == "Hetero":
         for fragment in fragments:
             fragment_repeats = find_hetero_amino_acid_repeats(fragment)
             for k, v in fragment_repeats.items():
@@ -74,21 +73,28 @@ def process_protein_sequence(sequence, analysis_type, overlap=50):
         new_repeats = find_new_boundary_repeats(fragments, final_repeats, overlap)
         for k, v in new_repeats.items():
             final_repeats[k] += v
+        final_repeats = {k: v for k, v in final_repeats.items() if not is_homo_repeat(k)]
 
-    if analysis_type in ["Homo", "Both"]:
-        homo_repeats = find_homorepeats(sequence)
-        for k, v in homo_repeats.items():
-            final_repeats[k] += v
+    elif analysis_type == "Homo":
+        final_repeats = find_homorepeats(sequence)
 
-    # Remove homo from hetero-only results
-    if analysis_type == "Hetero":
-        for k in list(final_repeats.keys()):
-            if is_homo_repeat(k):
-                del final_repeats[k]
+    elif analysis_type == "Both":
+        hetero_repeats = defaultdict(int)
+        for fragment in fragments:
+            fragment_repeats = find_hetero_amino_acid_repeats(fragment)
+            for k, v in fragment_repeats.items():
+                hetero_repeats[k] += v
+        hetero_repeats = check_boundary_repeats(fragments, hetero_repeats, overlap)
+        new_repeats = find_new_boundary_repeats(fragments, hetero_repeats, overlap)
+        for k, v in new_repeats.items():
+            hetero_repeats[k] += v
+        hetero_repeats = {k: v for k, v in hetero_repeats.items() if not is_homo_repeat(k)]
 
-    # Keep only homo repeats for homo-only results
-    if analysis_type == "Homo":
-        final_repeats = {k: v for k, v in final_repeats.items() if is_homo_repeat(k)}
+        homo_repeats = find_homorepeats(sequence)
+
+        final_repeats = homo_repeats.copy()
+        for k, v in hetero_repeats.items():
+            final_repeats[k] += v
 
     return final_repeats
 
@@ -134,7 +140,6 @@ def create_excel(sequences_data, repeats, filenames):
     output.seek(0)
     return output
 
-# Streamlit UI
 st.title("Protein Repeat Analysis")
 analysis_type = st.radio("Select analysis type:", ["Homo", "Hetero", "Both"], index=2)
 uploaded_files = st.file_uploader("Upload Excel files", accept_multiple_files=True, type=["xlsx"])