Added hooke.util.yaml fixing YAML/NumPy type issues (by dropping data).
[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     The data-type is also YAMLable (see :mod:`hooke.util.yaml`).
92
93     >>> import yaml
94     >>> print yaml.dump(d)
95     null
96     ...
97     <BLANKLINE>
98     """
99     def __new__(subtype, shape, dtype=numpy.float, buffer=None, offset=0,
100                 strides=None, order=None, info=None):
101         """Create the ndarray instance of our type, given the usual
102         input arguments.  This will call the standard ndarray
103         constructor, but return an object of our type.
104         """
105         obj = numpy.ndarray.__new__(
106             subtype, shape, dtype, buffer, offset, strides, order)
107         # add the new attribute to the created instance
108         if info == None:
109             info = {}
110         obj.info = info
111         # Finally, we must return the newly created object:
112         return obj
113
114     def __array_finalize__(self, obj):
115         """Set any extra attributes from the original object when
116         creating a new view object."""
117         # reset the attribute from passed original object
118         self.info = getattr(obj, 'info', {})
119         # We do not need to return anything
120
121     def __reduce__(self):
122         """Collapse an instance for pickling.
123
124         Returns
125         -------
126         reconstruct : callable
127             Called to create the initial version of the object.
128         args : tuple
129             A tuple of arguments for `reconstruct`
130         state : (optional)
131             The state to be passed to __setstate__, if present.
132         iter : iterator (optional)
133             Yielded items will be appended to the reconstructed
134             object.
135         dict : iterator (optional)
136             Yielded (key,value) tuples pushed back onto the
137             reconstructed object.
138         """
139         base_reduce = list(numpy.ndarray.__reduce__(self))
140         # tack our stuff onto ndarray's setstate portion.
141         base_reduce[2] = (base_reduce[2], (self.info,))
142         return tuple(base_reduce)
143
144     def __setstate__(self, state):
145         base_class_state,own_state = state
146         numpy.ndarray.__setstate__(self, base_class_state)
147         self.info, = own_state
148
149
150 class Curve (object):
151     """A grouped set of :class:`Data` runs from the same file with metadata.
152
153     For an approach/retract force spectroscopy experiment, the group
154     would consist of the approach data and the retract data.  Metadata
155     would be the temperature, cantilever spring constant, etc.
156
157     Two important :attr:`info` settings are `filetype` and
158     `experiment`.  These are two strings that can be used by Hooke
159     commands/plugins to understand what they are looking at.
160
161     * :attr:`info['filetype']` should contain the name of the exact
162       filetype defined by the driver (so that filetype-speciofic
163       commands can know if they're dealing with the correct filetype).
164     * :attr:`info['experiment']` should contain an instance of a
165       :class:`hooke.experiment.Experiment` subclass to identify the
166       experiment type.  For example, various
167       :class:`hooke.driver.Driver`\s can read in force-clamp data, but
168       Hooke commands could like to know if they're looking at force
169       clamp data, regardless of their origin.
170
171     Another important attribute is :attr:`command_stack`, which holds
172     a :class:`~hooke.command_stack.CommandStack` listing the commands
173     that have been applied to the `Curve` since loading.
174
175     The data-type is pickleable, to ensure we can move it between
176     processes with :class:`multiprocessing.Queue`\s.
177
178     >>> import pickle
179     >>> import yaml
180     >>> from .engine import CommandMessage
181     >>> c = Curve(path='some/path')
182
183     We add a recursive reference to `c` as you would get from
184     :meth:`hooke.plugin.curve.CurveCommand._add_to_command_stack`.
185
186     >>> c.command_stack.append(CommandMessage('curve info', {'curve':c}))
187
188     >>> s = pickle.dumps(c)
189     >>> z = pickle.loads(s)
190     >>> z
191     <Curve path>
192     >>> z.command_stack[-1].arguments['curve'] == z
193     True
194     >>> print yaml.dump(c)  # doctest: +REPORT_UDIFF
195     &id001 !!python/object:hooke.curve.Curve
196     command_stack: !!python/object/new:hooke.command_stack.CommandStack
197       listitems:
198       - !!python/object:hooke.engine.CommandMessage
199         arguments:
200           curve: *id001
201         command: curve info
202     data: null
203     driver: null
204     info: {}
205     name: path
206     path: some/path
207     <BLANKLINE>
208
209     However, if we try and serialize the command stack first, we run
210     into `Python issue 1062277`_.
211
212     .. _Python issue 1062277: http://bugs.python.org/issue1062277
213
214     >>> pickle.dumps(c.command_stack)
215     Traceback (most recent call last):
216       ...
217         assert id(obj) not in self.memo
218     AssertionError
219
220     YAML still works, though.
221
222     >>> print yaml.dump(c.command_stack)  # doctest: +REPORT_UDIFF
223     &id001 !!python/object/new:hooke.command_stack.CommandStack
224     listitems:
225     - !!python/object:hooke.engine.CommandMessage
226       arguments:
227         curve: !!python/object:hooke.curve.Curve
228           command_stack: *id001
229           data: null
230           driver: null
231           info: {}
232           name: path
233           path: some/path
234       command: curve info
235     <BLANKLINE>
236     """
237     def __init__(self, path, info=None):
238         self.__setstate__({'path':path, 'info':info})
239
240     def __str__(self):
241         return str(self.__unicode__())
242
243     def __unicode__(self):
244         return u'<%s %s>' % (self.__class__.__name__, self.name)
245
246     def __repr__(self):
247         return self.__str__()
248
249     def set_path(self, path):
250         self.path = path
251         if self.name == None and path != None:
252             self.name = os.path.basename(path)
253
254     def __getstate__(self):
255         state = dict(self.__dict__)      # make a copy of the attribute dict.
256         del(state['_hooke'])
257         return state
258
259     def __setstate__(self, state):
260         # .data contains: {name of data: list of data sets [{[x], [y]}]
261         # ._hooke contains a Hooke instance for Curve.load()
262         self.name = self.driver = self.data = self._hooke = None
263         self.info = {}
264         self.command_stack = CommandStack()
265         for key,value in state.items():
266             setattr(self, key, value)
267         if self.info == None:
268             self.info = {}
269         self.set_path(state.get('path', None))
270
271     def set_hooke(self, hooke=None):
272         if hooke != None:
273             self._hooke = hooke
274
275     def identify(self, drivers):
276         """Identify the appropriate :class:`hooke.driver.Driver` for
277         the curve file (`.path`).
278         """
279         if 'filetype' in self.info:
280             driver = [d for d in drivers if d.name == self.info['filetype']]
281             if len(driver) == 1:
282                 driver = driver[0]
283                 if driver.is_me(self.path):
284                     self.driver = driver
285                     return
286         for driver in drivers:
287             if driver.is_me(self.path):
288                 self.driver = driver # remember the working driver
289                 return
290         raise NotRecognized(self)
291
292     def load(self, hooke=None):
293         """Use the driver to read the curve into memory.
294
295         Also runs any commands in :attr:`command_stack`.  All
296         arguments are passed through to
297         :meth:`hooke.command_stack.CommandStack.execute`.
298         """
299         self.set_hooke(hooke)
300         log = logging.getLogger('hooke')
301         log.debug('loading curve %s with driver %s' % (self.name, self.driver))
302         data,info = self.driver.read(self.path, self.info)
303         self.data = data
304         for key,value in info.items():
305             self.info[key] = value
306         if self._hooke != None:
307             self.command_stack.execute(self._hooke)
308         elif len(self.command_stack) > 0:
309             log.warn(
310                 'could not execute command stack for %s without Hooke instance'
311                 % self.name)
312
313     def unload(self):
314         """Release memory intensive :attr:`.data`.
315         """
316         log = logging.getLogger('hooke')
317         log.debug('unloading curve %s' % self.name)
318         self.data = None