Refactoring the way data is returned in panedr

panedr/panedr.py

@@ -45,7 +45,7 @@
 import sys
 import itertools
 import time
-import pandas
+
 
 #Index for the IDs of additional blocks in the energy file.
 #Blocks can be added without sacrificing backward and forward
@@ -75,7 +75,7 @@
 Enxnm = collections.namedtuple('Enxnm', 'name unit')
 ENX_VERSION = 5
 
-__all__ = ['edr_to_df']
+__all__ = ['edr_to_df', 'edr_to_dict']
 
 
 class EDRFile(object):
@@ -395,21 +395,21 @@ def edr_strings(data, file_version, n):
 
 def is_frame_magic(data):
     """Unpacks an int and checks whether it matches the EDR frame magic number
-    
+
     Does not roll the reading position back.
     """
     magic = data.unpack_int()
     return magic == -7777777
 
 
-def edr_to_df(path, verbose=False):
+def read_edr(path, verbose_set=False):
     begin = time.time()
     edr_file = EDRFile(str(path))
     all_energies = []
     all_names = [u'Time'] + [nm.name for nm in edr_file.nms]
     times = []
     for ifr, frame in enumerate(edr_file):
-        if verbose:
+        if verbose_set:
             if ((ifr < 20 or ifr % 10 == 0) and
                     (ifr < 200 or ifr % 100 == 0) and
                     (ifr < 2000 or ifr % 1000 == 0)):
@@ -421,11 +421,28 @@ def edr_to_df(path, verbose=False):
             all_energies.append([frame.t] + [ener.e for ener in frame.ener])
 
     end = time.time()
-    if verbose:
+    if verbose_set:
         print('\rLast Frame read : {}, time : {} ps'
               .format(ifr, frame.t),
               end='', file=sys.stderr)
         print('\n{} frame read in {:.2f} seconds'.format(ifr, end - begin),
               file=sys.stderr)
+
+    return all_energies, all_names, times
+
+
+def edr_to_df(path: str, verbose: bool = False):
+    import pandas
+    all_energies, all_names, times = read_edr(path, verbose_set=verbose)
     df = pandas.DataFrame(all_energies, columns=all_names, index=times)
     return df
+
+
+def edr_to_dict(path: str, verbose: bool = False):
+    import numpy as np
+    all_energies, all_names, times = read_edr(path, verbose_set=verbose)
+    energy_dict = {}
+    for idx, name in enumerate(all_names):
+        energy_dict[name] = np.array(
+            [all_energies[frame][idx] for frame in range(len(times))])
+    return energy_dict

tests/test_edr.py

@@ -163,6 +163,14 @@ def _assert_progress_range(self, progress, dt, start, stop, step):
             assert ref_line == progress_line
 
 
+def test_edr_to_dict():
+    array_dict = panedr.edr_to_dict(EDR)
+    ref_df = panedr.edr_to_df(EDR)
+    array_df = pandas.DataFrame.from_dict(array_dict).set_index(
+        "Time", drop=False)
+    assert array_df.equals(ref_df)
+
+
 def read_xvg(path):
     """
     Reads XVG file, returning the data, names, and precision.