From a47631367c9f34e226e05e3af9d27c3a876831ad Mon Sep 17 00:00:00 2001
From: David Koes <dkoes@pitt.edu>
Date: Tue, 19 Sep 2023 21:11:00 -0400
Subject: [PATCH] Support multi-model files.

At least as long is they are all the same molecule.
---
 prody/proteins/mmtffile.py | 25 ++++++++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/prody/proteins/mmtffile.py b/prody/proteins/mmtffile.py
index f981b36d3..d465aaabe 100644
--- a/prody/proteins/mmtffile.py
+++ b/prody/proteins/mmtffile.py
@@ -179,12 +179,30 @@ def set_info(atomgroup, mmtf_data,get_bonds=False,altloc_sel='A'):
         "SACCHARIDE"])
 
     asize = mmtf_data.num_atoms
+
+    if len(mmtf_data.chains_per_model) > 1:
+        # get number of atoms in first model
+        asize = 0
+        groupIndex = 0
+        modelChainCount = mmtf_data.chains_per_model[0]
+        for chainGroupCount in mmtf_data.groups_per_chain[:modelChainCount]:
+            #traverse groups
+            for _ in range(chainGroupCount):
+                group = mmtf_data.group_list[mmtf_data.group_type_list[groupIndex]]        
+                asize += len(group['atomNameList'])
+                groupIndex += 1
+
     fields = OrderedDict()
 
     x = mmtf_data.x_coord_list
     y = mmtf_data.y_coord_list
     z = mmtf_data.z_coord_list
-    coords = np.array([x, y, z]).T
+
+    if len(x) != mmtf_data.num_models*asize:
+        LOGGER.warn('Multi-model MMTF files with different molecules not supported.  Keeping only first model')
+        coords = np.array([x, y, z]).T.reshape(mmtf_data.num_models, -1, 3)[0]
+    else:     
+        coords = np.array([x, y, z]).T.reshape(mmtf_data.num_models, asize, 3)
 
     # Initialize arrays for atom properties
     atom_names = np.empty(asize, dtype=ATOMIC_FIELDS['name'].dtype)
@@ -240,6 +258,7 @@ def set_info(atomgroup, mmtf_data,get_bonds=False,altloc_sel='A'):
                 groupIndex += 1
             chainIndex += 1
         modelIndex += 1
+        break
 
     #detect termini based on chain changes
     termini[:-1] = chainids[1:] != chainids[:-1]
@@ -250,7 +269,7 @@ def set_info(atomgroup, mmtf_data,get_bonds=False,altloc_sel='A'):
         #mask out any unwanted alternative locations
         mask = (altlocs == '') | (altlocs == altloc_sel)
         
-    atomgroup.setCoords(coords[mask])
+    atomgroup.setCoords(coords[:,mask])
     atomgroup.setNames(atom_names[mask])
     atomgroup.setResnums(resnums[mask])
     atomgroup.setResnames(resnames[mask])
@@ -277,7 +296,7 @@ def set_info(atomgroup, mmtf_data,get_bonds=False,altloc_sel='A'):
         nonpeptide = []
         #irgnore bonds between C and N in adjacent residues
         for a,b in allbonds:
-            if mask[a] and mask[b]:
+            if mask[a] and mask[b] and a < asize and b < asize:
                 if atom_names[a] != 'N' or atom_names[b] != 'C' or resnums[a]-resnums[b] != 1:
                     nonpeptide.append((remap[a],remap[b]))
         atomgroup.setBonds(nonpeptide)