Update code/add_3Dalignment.py

code/add_3Dalignment.py

@@ -233,8 +233,23 @@ def get_alignments_3D(identifier, model_num, pdb_path, pdbSequence, source, chai
             atomSequence = ''.join([three_to_one(i[3]) for i in atoms])
             coords = [[i[6] ,i[7] ,i[8]] for i in atoms]
             resnums_for_sasa = [i[5] for i in atoms]
-        """
-        else:
+        elif source == 'SWISSMODEL':
+            atomSequence = ''
+            coords = []
+            resnums_for_sasa = []
+            with open(pdb_path, encoding="utf8") as f:
+                for line in f.readlines():
+                    if line[0:4].strip() == 'ATOM' and line[13:15].strip() == 'CA' and line[21].upper() == chain.upper():
+                        atomSequence += threeToOne(line[17:20].strip())
+                        coords.append([line[31:38].strip(), line[39:46].strip(), line[47:54].strip()])
+                        resnums_for_sasa.append(line[22:26].strip())
+                    elif line[0:4].strip() == 'ATOM' and line[13:15].strip() == 'CA' and line[21] == ' ':
+                        atomSequence += threeToOne(line[17:20].strip())
+                        coords.append([line[31:38].strip(), line[39:46].strip(), line[47:54].strip()])
+                        resnums_for_sasa.append(line[22:26].strip())
+     
+            
+        elif source == 'MODBASE':
             atomSequence = ''
             coords = []
             resnums_for_sasa = []
@@ -244,7 +259,7 @@ def get_alignments_3D(identifier, model_num, pdb_path, pdbSequence, source, chai
                         atomSequence += threeToOne(line[17:20].strip())
                         coords.append([line[31:38].strip(), line[39:46].strip(), line[47:54].strip()])
                         resnums_for_sasa.append(line[22:26].strip())
-        """        
+             
         aligner.mode = 'local'
         aligner.substitution_matrix = substitution_matrices.load("BLOSUM62")
         aligner.open_gap_score = -11