From 6c54848105b767f6145f06f384468619b4fb163c Mon Sep 17 00:00:00 2001
From: sebastiankloth <kloth@mailbox.org>
Date: Mon, 2 May 2022 17:20:34 +0200
Subject: [PATCH] Reimplemented index reader

---
 mdevaluate/__init__.py |  4 ++--
 mdevaluate/reader.py   | 44 +++++++++++++++++++++++++++++++++---------
 2 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/mdevaluate/__init__.py b/mdevaluate/__init__.py
index 0aa3347..6600f5d 100644
--- a/mdevaluate/__init__.py
+++ b/mdevaluate/__init__.py
@@ -15,7 +15,7 @@ __version__ = '22.5.dev1'
 
 
 def open(directory='', topology='*.tpr', trajectory='*.xtc', cached=False,
-         nojump=False):
+         nojump=False, index_file=None):
     """
     Open a simulation from a directory.
 
@@ -64,7 +64,7 @@ def open(directory='', topology='*.tpr', trajectory='*.xtc', cached=False,
         raise FileNotFoundError('Trajectory file could not be identified.')
         
     atom_set, frames = reader.open_with_mdanalysis(
-        top_file, traj_file, cached=cached
+        top_file, traj_file, cached=cached, index_file=index_file
     )    
     coords = coordinates.Coordinates(frames, atom_subset=atom_set)
     if nojump:
diff --git a/mdevaluate/reader.py b/mdevaluate/reader.py
index d386d9c..692c547 100755
--- a/mdevaluate/reader.py
+++ b/mdevaluate/reader.py
@@ -18,12 +18,14 @@ import builtins
 import warnings
 import subprocess 
 import re
+import itertools
 
 import numpy as np
 import MDAnalysis as mdanalysis
 from scipy import sparse
 from dask import delayed, __version__ as DASK_VERSION
 import pandas as pd
+import re
 
 
 class NojumpError(Exception):
@@ -32,7 +34,7 @@ class NojumpError(Exception):
 class WrongTopologyError(Exception):
     pass
 
-def open_with_mdanalysis(topology, trajectory, cached=False):
+def open_with_mdanalysis(topology, trajectory, cached=False, index_file=None):
     """Open a the topology and trajectory with mdanalysis."""
     uni = mdanalysis.Universe(topology, trajectory, convert_units=False)
     if cached is not False:
@@ -45,19 +47,43 @@ def open_with_mdanalysis(topology, trajectory, cached=False):
         reader = BaseReader(uni.trajectory)
     reader.universe = uni
     if topology.endswith('.tpr'):
-        atms = atoms.Atoms(
-            np.stack((uni.atoms.resids, uni.atoms.resnames, uni.atoms.names), axis=1),
-            charges=uni.atoms.charges, masses=uni.atoms.masses
-        ).subset()
+        charges = uni.atoms.charges
+        masses = uni.atoms.masses
     elif topology.endswith('.gro'):
-        atms = atoms.Atoms(
-            np.stack((uni.atoms.resids, uni.atoms.resnames, uni.atoms.names), axis=1),
-            charges=None, masses=None
-        ).subset()
+        charges = None
+        masses = None
     else:
         raise WrongTopologyError('Topology file should end with ".tpr" or ".gro"')
+    indices = None
+    if index_file:
+        indices = load_indices(indexfile)
+    
+    atms = atoms.Atoms(
+            np.stack((uni.atoms.resids, uni.atoms.resnames, uni.atoms.names), axis=1),
+            charges=charges, masses=masses, indices=indices
+        ).subset()
+    
     return atms, reader
 
+group_re = re.compile('\[ ([-+\w]+) \]')
+
+def load_indices(indexfile):
+    indices = {}
+    index_array = None
+    with open(indexfile) as idx_file:
+        for line in idx_file:
+            m = group_re.search(line)
+            if m is not None:
+                group_name = m.group(1)
+                index_array = indices.get(group_name, [])
+                indices[group_name] = index_array
+            else:
+                elements = line.strip().split('\t')
+                elements = [x.split(' ') for x in elements]
+                elements = itertools.chain(*elements)  # make a flat iterator
+                elements = [x for x in elements if x != '']
+                index_array += [int(x) - 1 for x in elements]
+    return indices
 
 def is_writeable(fname):
     """Test if a directory is actually writeable, by writing a temporary file."""