Improve the performance of ParmEd converter. (Fix MDAnalysis#3028) (M…

…DAnalysis#3029) Fixes MDAnalysis#3028 * Improves the performance of the ParmEd converter by using a dictionary lookup for the atomgroup to universe index mapping.
lilyminium · Dec 14, 2020 · a64eed9 · a64eed9
1 parent 7f83f4a
commit a64eed9
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/package/AUTHORS b/package/AUTHORS
@@ -149,6 +149,7 @@ Chronological list of authors
   - Marcello Sega
   - Nicholas Craven
   - Ramon Crehuet
+  - Haochuan Chen
 
 External code
 -------------

diff --git a/package/CHANGELOG b/package/CHANGELOG
@@ -15,7 +15,7 @@ The rules for this file:
 ------------------------------------------------------------------------------
 ??/??/20 richardjgowers, IAlibay, orbeckst, tylerjereddy, jbarnoud,
          yuxuanzhuang, lilyminium, VOD555, p-j-smith, bieniekmateusz,
-         calcraven, ianmkenney, rcrehuet, manuel.nuno.melo
+         calcraven, ianmkenney, rcrehuet, manuel.nuno.melo, hanatok
 
   * 1.0.1
 
@@ -50,6 +50,7 @@ Fixes
     (Issue #2934)
 
 Enhancements
+  * Improved performance of the ParmEd converter (Issue #3028, PR #3029)
   * Improved performances when parsing TPR files (PR #2804)
 
 Changes (not affecting users)

diff --git a/package/MDAnalysis/coordinates/ParmEd.py b/package/MDAnalysis/coordinates/ParmEd.py
@@ -133,7 +133,7 @@ def _read_first_frame(self):
 }
 
 def get_indices_from_subset(i, atomgroup=None, universe=None):
-    return atomgroup.index(universe.atoms[i])
+    return atomgroup[universe.atoms[i]]
 
 class ParmEdConverter(base.ConverterBase):
     """Convert MDAnalysis AtomGroup or Universe to ParmEd :class:`~parmed.structure.Structure`.
@@ -263,8 +263,10 @@ def convert(self, obj):
             struct.box = None
 
         if hasattr(ag_or_ts, 'universe'):
+            atomgroup = {atom: index for index,
+                         atom in enumerate(list(ag_or_ts))}
             get_atom_indices = functools.partial(get_indices_from_subset,
-                                                 atomgroup=list(ag_or_ts),
+                                                 atomgroup=atomgroup,
                                                  universe=ag_or_ts.universe)
         else:
             get_atom_indices = lambda x: x