Add doctests proving recursive curve->command_stack->arguments->curve serializing...
[hooke.git] / hooke / curve.py
1 # Copyright (C) 2010 W. Trevor King <wking@drexel.edu>
2 #
3 # This file is part of Hooke.
4 #
5 # Hooke is free software: you can redistribute it and/or modify it
6 # under the terms of the GNU Lesser General Public License as
7 # published by the Free Software Foundation, either version 3 of the
8 # License, or (at your option) any later version.
9 #
10 # Hooke is distributed in the hope that it will be useful, but WITHOUT
11 # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
12 # or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General
13 # Public License for more details.
14 #
15 # You should have received a copy of the GNU Lesser General Public
16 # License along with Hooke.  If not, see
17 # <http://www.gnu.org/licenses/>.
18
19 """The `curve` module provides :class:`Curve` and :class:`Data` for
20 storing force curves.
21 """
22
23 import logging
24 import os.path
25
26 import numpy
27
28 from .command_stack import CommandStack
29 from . import experiment
30
31
32 class NotRecognized (ValueError):
33     def __init__(self, curve):
34         self.__setstate__(curve)
35
36     def __getstate__(self):
37         return self.curve
38
39     def __setstate__(self, data):
40         if isinstance(data, Curve):
41             msg = 'Not a recognizable curve format: %s' % data.path
42             super(NotRecognized, self).__init__(msg)
43             self.curve = data
44
45
46 class Data (numpy.ndarray):
47     """Stores a single, continuous data set.
48
49     Adds :attr:`info` :class:`dict` to the standard :class:`numpy.ndarray`.
50
51     See :mod:`numpy.doc.subclassing` for the peculiarities of
52     subclassing :class:`numpy.ndarray`.
53
54     Examples
55     --------
56
57     >>> d = Data(shape=(3,2), info={'columns':['distance (m)', 'force (N)']})
58     >>> type(d)
59     <class 'hooke.curve.Data'>
60     >>> for i in range(3): # initialize d
61     ...    for j in range(2):
62     ...        d[i,j] = i*10 + j
63     >>> d
64     Data([[  0.,   1.],
65            [ 10.,  11.],
66            [ 20.,  21.]])
67     >>> d.info
68     {'columns': ['distance (m)', 'force (N)']}
69
70     The information gets passed on to slices.
71
72     >>> row_a = d[:,0]
73     >>> row_a
74     Data([  0.,  10.,  20.])
75     >>> row_a.info
76     {'columns': ['distance (m)', 'force (N)']}
77
78     The data-type is also pickleable, to ensure we can move it between
79     processes with :class:`multiprocessing.Queue`\s.
80
81     >>> import pickle
82     >>> s = pickle.dumps(d)
83     >>> z = pickle.loads(s)
84     >>> z
85     Data([[  0.,   1.],
86            [ 10.,  11.],
87            [ 20.,  21.]])
88     >>> z.info
89     {'columns': ['distance (m)', 'force (N)']}
90     """
91     def __new__(subtype, shape, dtype=numpy.float, buffer=None, offset=0,
92                 strides=None, order=None, info=None):
93         """Create the ndarray instance of our type, given the usual
94         input arguments.  This will call the standard ndarray
95         constructor, but return an object of our type.
96         """
97         obj = numpy.ndarray.__new__(
98             subtype, shape, dtype, buffer, offset, strides, order)
99         # add the new attribute to the created instance
100         if info == None:
101             info = {}
102         obj.info = info
103         # Finally, we must return the newly created object:
104         return obj
105
106     def __array_finalize__(self, obj):
107         """Set any extra attributes from the original object when
108         creating a new view object."""
109         # reset the attribute from passed original object
110         self.info = getattr(obj, 'info', {})
111         # We do not need to return anything
112
113     def __reduce__(self):
114         """Collapse an instance for pickling.
115
116         Returns
117         -------
118         reconstruct : callable
119             Called to create the initial version of the object.
120         args : tuple
121             A tuple of arguments for `reconstruct`
122         state : (optional)
123             The state to be passed to __setstate__, if present.
124         iter : iterator (optional)
125             Yielded items will be appended to the reconstructed
126             object.
127         dict : iterator (optional)
128             Yielded (key,value) tuples pushed back onto the
129             reconstructed object.
130         """
131         base_reduce = list(numpy.ndarray.__reduce__(self))
132         # tack our stuff onto ndarray's setstate portion.
133         base_reduce[2] = (base_reduce[2], (self.info,))
134         return tuple(base_reduce)
135
136     def __setstate__(self, state):
137         base_class_state,own_state = state
138         numpy.ndarray.__setstate__(self, base_class_state)
139         self.info, = own_state
140
141
142 class Curve (object):
143     """A grouped set of :class:`Data` runs from the same file with metadata.
144
145     For an approach/retract force spectroscopy experiment, the group
146     would consist of the approach data and the retract data.  Metadata
147     would be the temperature, cantilever spring constant, etc.
148
149     Two important :attr:`info` settings are `filetype` and
150     `experiment`.  These are two strings that can be used by Hooke
151     commands/plugins to understand what they are looking at.
152
153     * :attr:`info['filetype']` should contain the name of the exact
154       filetype defined by the driver (so that filetype-speciofic
155       commands can know if they're dealing with the correct filetype).
156     * :attr:`info['experiment']` should contain an instance of a
157       :class:`hooke.experiment.Experiment` subclass to identify the
158       experiment type.  For example, various
159       :class:`hooke.driver.Driver`\s can read in force-clamp data, but
160       Hooke commands could like to know if they're looking at force
161       clamp data, regardless of their origin.
162
163     Another important attribute is :attr:`command_stack`, which holds
164     a :class:`~hooke.command_stack.CommandStack` listing the commands
165     that have been applied to the `Curve` since loading.
166
167     The data-type is pickleable, to ensure we can move it between
168     processes with :class:`multiprocessing.Queue`\s.
169
170     >>> import pickle
171     >>> import yaml
172     >>> from .engine import CommandMessage
173     >>> c = Curve(path='some/path')
174
175     We add a recursive reference to `c` as you would get from
176     :meth:`hooke.plugin.curve.CurveCommand._add_to_command_stack`.
177
178     >>> c.command_stack.append(CommandMessage('curve info', {'curve':c}))
179
180     >>> s = pickle.dumps(c)
181     >>> z = pickle.loads(s)
182     >>> z
183     <Curve path>
184     >>> z.command_stack[-1].arguments['curve'] == z
185     True
186     >>> print yaml.dump(c)
187     &id001 !!python/object:hooke.curve.Curve
188     command_stack: !!python/object/new:hooke.command_stack.CommandStack
189       listitems:
190       - !!python/object:hooke.engine.CommandMessage
191         arguments:
192           curve: *id001
193         command: curve info
194     data: null
195     driver: null
196     info: {}
197     name: path
198     path: some/path
199     <BLANKLINE>
200     """
201     def __init__(self, path, info=None):
202         #the data dictionary contains: {name of data: list of data sets [{[x], [y]}]
203         self.name = None
204         self.set_path(path)
205         self.driver = None
206         self.data = None
207         if info == None:
208             info = {}
209         self.info = info
210         self.command_stack = CommandStack()
211         self._hooke = None  # Hooke instance for Curve.load()
212
213     def __str__(self):
214         return str(self.__unicode__())
215
216     def __unicode__(self):
217         return u'<%s %s>' % (self.__class__.__name__, self.name)
218
219     def __repr__(self):
220         return self.__str__()
221
222     def set_path(self, path):
223         self.path = path
224         if self.name == None and path != None:
225             self.name = os.path.basename(path)
226
227     def __getstate__(self):
228         state = dict(self.__dict__)      # make a copy of the attribute dict.
229         del(state['_hooke'])
230         return state
231
232     def __setstate__(self, state):
233         self.name = self._hooke = None
234         for key,value in state.items():
235             setattr(self, key, value)
236         self.set_path(state.get('path', None))
237
238     def set_hooke(self, hooke=None):
239         if hooke != None:
240             self._hooke = hooke
241
242     def identify(self, drivers):
243         """Identify the appropriate :class:`hooke.driver.Driver` for
244         the curve file (`.path`).
245         """
246         if 'filetype' in self.info:
247             driver = [d for d in drivers if d.name == self.info['filetype']]
248             if len(driver) == 1:
249                 driver = driver[0]
250                 if driver.is_me(self.path):
251                     self.driver = driver
252                     return
253         for driver in drivers:
254             if driver.is_me(self.path):
255                 self.driver = driver # remember the working driver
256                 return
257         raise NotRecognized(self)
258
259     def load(self, hooke=None):
260         """Use the driver to read the curve into memory.
261
262         Also runs any commands in :attr:`command_stack`.  All
263         arguments are passed through to
264         :meth:`hooke.command_stack.CommandStack.execute`.
265         """
266         self.set_hooke(hooke)
267         log = logging.getLogger('hooke')
268         log.debug('loading curve %s with driver %s' % (self.name, self.driver))
269         data,info = self.driver.read(self.path, self.info)
270         self.data = data
271         for key,value in info.items():
272             self.info[key] = value
273         if self._hooke != None:
274             self.command_stack.execute(self._hooke)
275         elif len(self.command_stack) > 0:
276             log.warn(
277                 'could not execute command stack for %s without Hooke instance'
278                 % self.name)
279
280     def unload(self):
281         """Release memory intensive :attr:`.data`.
282         """
283         log = logging.getLogger('hooke')
284         log.debug('unloading curve %s' % self.name)
285         self.data = None