1c9cea78b22cde7fbe2fea507aa42a42f91f3b32
[hooke.git] / hooke / util / calculus.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
6 # modify it under the terms of the GNU Lesser General Public
7 # License as published by the Free Software Foundation, either
8 # version 3 of the License, or (at your option) any later version.
9 #
10 # Hooke is distributed in the hope that it will be useful,
11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 # GNU Lesser General 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 `calculus` module provides functions for calculating
20 derivatives and integrals of discrete data.
21 """
22
23 import copy
24
25 import numpy
26
27 from ..curve import Data
28
29
30 def derivative(data, x_col=0, f_col=1, weights={-1:-0.5, 1:0.5}):
31     """Calculate the discrete derivative (finite difference) of
32     data[:,f_col] with respect to data[:,x_col].
33
34     Examples
35     --------
36
37     >>> import pprint
38     >>> d = Data((5,2), dtype=numpy.float,
39     ...          info={'columns':['x', 'x**2']})
40     >>> for i in range(5):
41     ...     d[i,0] = i
42     ...     d[i,1] = i**2
43     >>> d
44     Data([[  0.,   0.],
45            [  1.,   1.],
46            [  2.,   4.],
47            [  3.,   9.],
48            [  4.,  16.]])
49     >>> dd = derivative(d)
50     >>> dd
51     Data([[ 0.,  1.],
52            [ 1.,  2.],
53            [ 2.,  4.],
54            [ 3.,  6.],
55            [ 4.,  7.]])
56     >>> pprint.pprint(dd.info)
57     {'columns': ['x', 'deriv x**2 with respect to x']}
58
59     Notes
60     -----
61
62     *Weights*
63
64     The returned :class:`Data` block shares its x vector with the
65     input data.  The ith df/dx value in the returned data is
66     caclulated with::
67
68         (df/dx)[i] = (SUM_j w[j] f[i+j]) / h
69
70     where ``h = x[i+1]-x[i]`` is the x coordinate spacing (assumed
71     constant) and ``j`` ranges over the keys of `weights`.
72
73     There standard schemes translate as follows:
74
75     ========  ======================  ===================
76     scheme    formula                 weights       
77     ========  ======================  ===================
78     forward   ``(f[i+1]-f[i])/h``     ``{0:-1,1:1}``
79     backward  ``(f[i]-f[i-1])/h``     ``{0:1,-1:-1}``
80     central   ``(f[i+1]-f[i-1])/2h``  ``{-1:-0.5,1:0.5}``
81     ========  ======================  ===================
82
83     The default scheme is central differencing.
84
85     *Boundary conditions*
86
87     The boundary conditions could be configurable in principle.  The
88     current scheme just extrapolates virtual points out to negative
89     `i` following::
90
91         f[i<0] = 2*f[0] - f[-i]
92
93     With analogous treatment for `i > data.shape[0]`.  This ensures that
94     `f[i]-f[0]` is odd about `i=0`, which keeps derivatives smooth.::
95
96         f[i] - f[0] = f[0] - f[-i] == -(f[-i] - f[0])    
97     """
98     output = Data((data.shape[0],2), dtype=data.dtype)
99     output.info = copy.copy(data.info)
100     output.info['columns'] = [
101         data.info['columns'][x_col],
102         'deriv %s with respect to %s' \
103         % (data.info['columns'][f_col], data.info['columns'][x_col]),
104         ]
105     h = data[1,x_col] - data[0,x_col]
106     chunks = []
107     for i,w in weights.items():
108         chunk = numpy.roll(w*data[:,f_col], -i)
109         if i > 0: # chunk shifted down, replace the high `i`s
110             zero = len(chunk) - 1 - i
111             for j in range(1,i+1):
112                 chunk[zero+j] = 2*chunk[zero] - chunk[zero-j]
113         elif i < 0: # chunk shifted up, replace the low `i`s
114             zero = -i
115             for j in range(1,zero+1):
116                 chunk[zero-j] = 2*chunk[zero] - chunk[zero+j]
117         chunks.append(chunk)
118     print chunks
119     output[:,0] = data[:,x_col]
120     output[:,1] = sum(chunks)
121     return output