Merge pull request #2245 from prody/copilot/fix-system-size-limitation

Fix writeNMD residue number overflow for large systems (>9999 residues)

Author: jamesmkrieger

prody/dynamics/nmdfile.py

@@ -422,10 +422,40 @@ def writeNMD(filename, modes, atoms):
     try:
         data = atoms.getResnums()
         if data is not None:
+            if data.max() > 9999:
+                LOGGER.warn('Residue numbers exceed 9999. Renumbering residues '
+                            'in NMD output to be compatible with NMWiz/VMD PDB '
+                            'format (max 9999 per chain).')
+                try:
+                    chids = atoms.getChids()
+                except Exception:
+                    chids = None
+                new_data = np.empty_like(data)
+                if chids is not None:
+                    for chid in np.unique(chids):
+                        mask = (chids == chid)
+                        chain_resnums = data[mask]
+                        unique_res = {}
+                        counter = 0
+                        for r in chain_resnums:
+                            if r not in unique_res:
+                                counter += 1
+                                unique_res[r] = (counter - 1) % 9999 + 1
+                        new_data[mask] = np.array([unique_res[r]
+                                                   for r in chain_resnums])
+                else:
+                    unique_res = {}
+                    counter = 0
+                    for i, r in enumerate(data):
+                        if r not in unique_res:
+                            counter += 1
+                            unique_res[r] = (counter - 1) % 9999 + 1
+                        new_data[i] = unique_res[r]
+                data = new_data
             out.write('resids ')
             data.tofile(out, ' ')
             out.write('\n')
-    except:
+    except Exception:
         pass
     try:
         data = atoms.getChids()