New trajectory object

chemlab/core/__init__.py

@@ -2,5 +2,6 @@
 from .system import (subsystem_from_molecules,
                      subsystem_from_atoms,
                      merge_systems)
+from .trajectory import Trajectory
 from .spacegroup.crystal import crystal
 from .random import random_lattice_box, random_box

chemlab/core/base.py

@@ -135,6 +135,15 @@ def copy_from(self, other):
         self.maps = {k: m.copy() for k, m in other.maps.items()}
         self.dimensions = other.dimensions.copy()
 
+    def update(self, dictionary):
+        """Update the current chemical entity from a dictionary of attributes"""
+        allowed_attrs = self.__attributes__.keys()
+        allowed_attrs += [a.alias for a in self.__attributes__.values()]
+        for k in dictionary:
+            # We only update existing attributes
+            if k in allowed_attrs:
+                setattr(self, k, dictionary[k])
+
     def initialize_empty(self, **kwargs):
         self.dimensions.update(kwargs)
 

chemlab/core/system.py

@@ -48,7 +48,7 @@ class Molecule(ChemicalEntity):
     __dimension__ = 'molecule'
 
     __attributes__ = {
-        'r_array' : Attribute(shape=(3,), dtype='float', dim='atom'),
+        'r_array' : Attribute(shape=(3,), dtype='float', dim='atom', alias="coords"),
         'type_array' : Attribute(dtype='str', dim='atom'),
         'charge_array' : Attribute(dim='atom'),
         'bond_orders' : Attribute(dtype='int', dim='bond'),
@@ -121,7 +121,7 @@ def make_formula(elements):
 class System(ChemicalEntity):
     __dimension__ = 'system'
     __attributes__ = {
-        'r_array' : Attribute(shape=(3,), dtype='float', dim='atom'),
+        'r_array' : Attribute(shape=(3,), dtype='float', dim='atom', alias="coords"),
         'type_array' : Attribute(dtype='str', dim='atom'),
         'charge_array' : Attribute(dim='atom'),
         'molecule_name' : Attribute(dtype='str', dim='molecule'),

chemlab/io/handlers/xtctrr.py

@@ -2,38 +2,15 @@
 
 """
 from .base import IOHandler
+from ...core import Trajectory
 import numpy as np
 from ...libs.pyxdr import XTCReader
-
-
-class BlockedArray(np.ndarray):
-    def __init__(self, *args, **kwargs):
-        super(BlockedArray, self).__init__(*args, **kwargs)
-        self.block_size = 1000 # It preallocates 1000
-        self._last_i = self.shape[0] - 1
-
-
-    def append(self, item):
-        assert item.shape == self.shape[1:]
-        self._last_i += 1
-        if self._last_i == self.shape[0]:
-            self._enlarge()
-        else:
-            nslices = len(self.shape) - 1
-            slices = (slice(None, None, None),)*nslices
-
-            self[(self._last_i,)+slices] = item
-
-
-    def _enlarge(self):
-        newshape = list(self.shape)
-        newshape[0] +=  self.block_size
-        self.resize(newshape, refcheck=False)
-
-    def trim(self):
-        self.resize((self._last_i,) + self.shape[1:] , refcheck=False)
-
+import os
+import shutil
 import itertools
+import time
+import bcolz
+
 class XtcIO(IOHandler):
     '''Reader for GROMACS XTC trajectories.
     
@@ -56,41 +33,142 @@ class XtcIO(IOHandler):
        correspoiding to each frame.
 
     '''
-    can_read = ['trajectory', 'boxes']
+    can_read = ['trajectory']
     can_write = []
 
     def read(self, feature, **kwargs):
-        import time
         t0 = time.time()
 
         if feature == 'trajectory':
             skipframes = kwargs.get("skip", None)
+            rootdir = kwargs.get("rootdir", None)
 
+            if rootdir is not None:
+                p = os.path
+                rootdir = p.join(p.dirname(p.abspath(self.fd.name)), rootdir)
 
-            times = []
             xtcreader = XTCReader(self.fd.name)
-
-            cursize = 0
-
-            # 1 frame, 10 atoms, 3 coordinates each
-            frames = np.empty((0, 10, 3), dtype='float32')
-            self._boxes = []
-
+
+            if rootdir is not None and os.path.exists(rootdir):
+                f = h5py.File(rootdir, mode='r')
+                if f.attrs["timestamp"] >= os.path.getmtime(self.fd.name):
+                    return Trajectory(f['coords'], f['times'], f['boxes'])
+                else:
+                    f.close()
+
+
+            handler = _NumpyFrameHandler() if rootdir is None \
+                       else _HDF5FrameHandler(rootdir, os.path.getmtime(self.fd.name))
+
             for i, frame in enumerate(xtcreader):
                 if skipframes is None or i%skipframes == 0:
-                    cursize += 1
-                    # Enarge if necessary
-                    if cursize > frames.shape[0]:
-                        frames.resize((cursize * 2, ) + frame.coords.shape)
-
-                    frames[cursize - 1, :] = frame.coords
-                    times.append(frame.time)
-                    self._boxes.append(frame.box)
-            # Shrink if necessary
-            if frames.shape[0] != cursize:
-                frames.resize((cursize, frames.shape[1], frames.shape[2]))
+                    handler.handle_frame(i, frame)
 
-            return np.array(times), frames
+            handler.handle_done(i)
 
-        if feature == 'boxes':
-            return self._boxes
+            return Trajectory(handler.frames, handler.times, handler.boxes)
+
+class _NumpyFrameHandler:
+
+    def __init__(self):
+        self.cursize = 0
+
+    def handle_frame(self, i, frame):
+        self.cursize += 1
+
+        if i == 0:
+            self.frames = np.empty((1,) + frame.coords.shape, dtype='float32')
+            self.boxes = []
+            self.times = []
+
+        # Enlarge if necessary
+        if self.cursize > self.frames.shape[0]:
+            self.frames.resize((self.cursize * 2, ) + frame.coords.shape)
+
+        self.frames[self.cursize - 1, :] = frame.coords
+        self.boxes.append(frame.box)
+        self.times.append(frame.time)
+
+    def handle_done(self, i):
+
+        # Trim the edge
+        if self.frames.shape[0] != self.cursize:
+            self.frames.resize((i + 1, self.frames.shape[1], self.frames.shape[2]))
+
+        self.boxes = np.array(self.boxes)
+        self.times = np.array(self.times)
+
+import h5py
+
+
+class _HDF5FrameHandler:
+
+    def __init__(self, rootdir, timestamp):
+        self.f = h5py.File(rootdir, 'w')
+        self.f.attrs['timestamp'] = timestamp
+        self.rootdir = rootdir
+        self.cursize = 0
+        self.boxes = []
+        self.times = []
+
+    def handle_frame(self, i, frame):
+        self.cursize += 1
+
+        if i == 0:
+            self.frames = self.f.create_dataset("coords", 
+                                                shape=(1,) + frame.coords.shape,
+                                                maxshape=(None, frame.coords.shape[0], 3),
+                                                compression='lzf',
+                                                shuffle=True,
+                                                dtype="float32")
+        # Enlarge if necessary
+        if self.cursize > self.frames.shape[0]:
+            self.frames.resize((self.cursize * 2,) + frame.coords.shape)
+
+        self.frames[self.cursize - 1, :] = frame.coords
+        self.boxes.append(frame.box)
+        self.times.append(frame.time)
+
+    def handle_done(self, i):
+        if self.frames.shape[0] != self.cursize:
+            self.frames.resize((i + 1, self.frames.shape[1], self.frames.shape[2]))
+
+        self.boxes = self.f.create_dataset("boxes", data=self.boxes)
+        self.times = self.f.create_dataset("times", data=self.times)
+        self.f.close()
+        self.f = h5py.File(self.rootdir, 'r')
+
+        self.frames = self.f['coords']
+        self.times = self.f['times']
+        self.boxes = self.f['boxes']
+
+
+class _BcolzFrameHandler:
+
+    def __init__(self, rootdir, timestamp):
+        # Clean rootdir
+        if os.path.exists(rootdir):
+            shutil.rmtree(rootdir)
+
+        os.mkdir(rootdir)
+
+        self.rootdir = rootdir
+        self.timestamp = timestamp
+
+    def handle_frame(self, i, frame):
+        if i == 0:
+            self.frames = bcolz.carray(np.zeros((0,) + frame.coords.shape, dtype="float32"),
+                                 rootdir=os.path.join(self.rootdir, "coords"),
+                                 mode='w')
+            self.frames.attrs['timestamp'] = self.timestamp
+            self.boxes = bcolz.carray(np.zeros((0,) + frame.box.shape, dtype="float32"),
+                                 rootdir=os.path.join(self.rootdir, "boxes"),
+                                 mode='w')
+            self.times = []
+
+        self.frames.append(frame.coords)
+        self.boxes.append(frame.box)
+        self.times.append(frame.time)
+
+    def handle_done(self, i):
+        self.times = bcolz.carray(self.times, dtype='float32', rootdir=os.path.join(self.rootdir, "times"))

chemlab/md/simulation.py

@@ -3,7 +3,8 @@ class Simulation(object):
 
     def __init__(self, system, potential, length=1.0, integrator='md', dt=2e-6, dt_io=2e-4, 
                  cutoff=0.9, pme=True, temperature=300, thermostat='v-rescale',
-                 pressure=1.0, barostat='berendsen', constraints='none'):
+                 pressure=1.0, barostat='berendsen', constraints='none',
+                 annealing=[]):
 
         self.system = system
         self.potential = potential
@@ -22,12 +23,15 @@ def __init__(self, system, potential, length=1.0, integrator='md', dt=2e-6, dt_i
         self.pressure = pressure
         self.barostat = barostat
         self.constraints = constraints
+
+        self.annealing = annealing
 
 def to_mdp(simulation):
 
     length = simulation.length * 1000
     dt = simulation.dt * 1000
     dt_io = simulation.dt_io * 1000
+    annealing = simulation.annealing
 
     r = ''
     r += 'integrator = {}\n'.format(simulation.integrator)
@@ -48,6 +52,14 @@ def to_mdp(simulation):
     r += 'ref_t = {:f}\n'.format(simulation.temperature)
     r += 'tau_t = {:f}\n'.format(0.1)
 
+    if annealing:
+
+        r += "annealing = single\n"
+        r += "annealing-npoints = {}\n".format(len(annealing))
+        # Time from ns to ps
+        r += "annealing-time = {}\n".format(" ".join(str(a[0] * 1000) for a in annealing))
+        r += "annealing-temp = {}\n".format(" ".join(str(a[1]) for a in annealing)) 
+
     r += 'pcoupl = {}\n'.format(simulation.barostat)
     r += 'compressibility = 4.5e-5\n'
     r += 'ref_p = {}\n'.format(simulation.pressure)

tests/test_core.py

@@ -9,7 +9,7 @@
 from chemlab.table import vdw_radius
 
 #from chemlab.graphics import display_system
-from .testtools import assert_allclose, assert_eqbonds, assert_npequal
+from .testtools import assert_allclose, assert_eqbonds, assert_npequal, npeq_
 
 
 def _make_water():
@@ -27,7 +27,7 @@ def test_init(self):
         eq_(a.type_array, 'O')
         eq_(a.get_attribute('type_array').value, 'O')
         assert_npequal(a.r_array, [-4.99, 2.49, 0.0])
-        eq_(a.get_attribute('r_array').value, [-4.99, 2.49, 0.0])
+        npeq_(a.get_attribute('r_array').value, [-4.99, 2.49, 0.0])
 
 
 class TestMolecule(object):