1from __future__ import absolute_import
2from __future__ import print_function
3import numpy
4
5from . import _interpolate
6from six.moves import range
7
8
9def make_array_safe(ary, typecode):
10    ary = numpy.atleast_1d(numpy.asarray(ary, typecode))
11    if not ary.flags['CONTIGUOUS']:
12        ary = ary.copy()
13    return ary
14
15
16def linear(x, y, new_x):
17    """ Linearly interpolates values in new_x based on the values in x and y
18
19        Parameters
20        ----------
21        x
22            1-D array
23        y
24            1-D or 2-D array
25        new_x
26            1-D array
27    """
28    x = make_array_safe(x, numpy.float64)
29    y = make_array_safe(y, numpy.float64)
30    new_x = make_array_safe(new_x, numpy.float64)
31
32    assert len(y.shape) < 3, "function only works with 1D or 2D arrays"
33    if len(y.shape) == 2:
34        new_y = numpy.zeros((y.shape[0], len(new_x)), numpy.float64)
35        for i in range(len(new_y)):
36            _interpolate.linear_dddd(x, y[i], new_x, new_y[i])
37    else:
38        new_y = numpy.zeros(len(new_x), numpy.float64)
39        _interpolate.linear_dddd(x, y, new_x, new_y)
40
41    return new_y
42
43
44def logarithmic(x, y, new_x):
45    """ Linearly interpolates values in new_x based in the log space of y.
46
47        Parameters
48        ----------
49        x
50            1-D array
51        y
52            1-D or 2-D array
53        new_x
54            1-D array
55    """
56    x = make_array_safe(x, numpy.float64)
57    y = make_array_safe(y, numpy.float64)
58    new_x = make_array_safe(new_x, numpy.float64)
59
60    assert len(y.shape) < 3, "function only works with 1D or 2D arrays"
61    if len(y.shape) == 2:
62        new_y = numpy.zeros((y.shape[0], len(new_x)), numpy.float64)
63        for i in range(len(new_y)):
64            _interpolate.loginterp_dddd(x, y[i], new_x, new_y[i])
65    else:
66        new_y = numpy.zeros(len(new_x), numpy.float64)
67        _interpolate.loginterp_dddd(x, y, new_x, new_y)
68
69    return new_y
70
71
72def block_average_above(x, y, new_x):
73    """ Linearly interpolates values in new_x based on the values in x and y
74
75        Parameters
76        ----------
77        x
78            1-D array
79        y
80            1-D or 2-D array
81        new_x
82            1-D array
83    """
84    bad_index = None
85    x = make_array_safe(x, numpy.float64)
86    y = make_array_safe(y, numpy.float64)
87    new_x = make_array_safe(new_x, numpy.float64)
88
89    assert len(y.shape) < 3, "function only works with 1D or 2D arrays"
90    if len(y.shape) == 2:
91        new_y = numpy.zeros((y.shape[0], len(new_x)), numpy.float64)
92        for i in range(len(new_y)):
93            bad_index = _interpolate.block_averave_above_dddd(x, y[i],
94                                                              new_x, new_y[i])
95            if bad_index is not None:
96                break
97    else:
98        new_y = numpy.zeros(len(new_x), numpy.float64)
99        bad_index = _interpolate.block_average_above_dddd(x, y, new_x, new_y)
100
101    if bad_index is not None:
102        msg = "block_average_above cannot extrapolate and new_x[%d]=%f "\
103              "is out of the x range (%f, %f)" % \
104              (bad_index, new_x[bad_index], x[0], x[-1])
105        raise ValueError(msg)
106
107    return new_y
108
109
110def window_average(x, y, new_x, width=10.0):
111    bad_index = None
112    x = make_array_safe(x, numpy.float64)
113    y = make_array_safe(y, numpy.float64)
114    new_x = make_array_safe(new_x, numpy.float64)
115    width = float(width)
116    assert len(y.shape) < 3, "function only works with 1D or 2D arrays"
117    if len(y.shape) == 2:
118        new_y = numpy.zeros((y.shape[0], len(new_x)), numpy.float64)
119        for i in range(len(new_y)):
120            _interpolate.window_average_ddddd(x, y[i], new_x, new_y[i],
121                                              width)
122    else:
123        new_y = numpy.zeros(len(new_x), numpy.float64)
124        _interpolate.window_average_ddddd(x, y, new_x, new_y, width)
125
126    return new_y
127
128
129def main():
130    from scipy import arange, ones
131    import time
132    N = 3000.
133    x = arange(N)
134    y = arange(N)
135    new_x = arange(N) + 0.5
136    t1 = time.clock()
137    new_y = linear(x, y, new_x)
138    t2 = time.clock()
139    print('1d interp (sec):', t2 - t1)
140    print(new_y[:5])
141
142    N = 3000.
143    x = arange(N)
144    y = arange(N)
145
146    new_x = arange(N / 2) * 2
147    t1 = time.clock()
148    new_y = block_average_above(x, y, new_x)
149    t2 = time.clock()
150    print('1d block_average_above (sec):', t2 - t1)
151    print(new_y[:5])
152
153    N = 3000.
154    x = arange(N)
155    y = ones((100, N)) * arange(N)
156    new_x = arange(N) + 0.5
157    t1 = time.clock()
158    new_y = linear(x, y, new_x)
159    t2 = time.clock()
160    print('fast interpolate (sec):', t2 - t1)
161    print(new_y[:5, :5])
162
163    import scipy
164    N = 3000.
165    x = arange(N)
166    y = ones((100, N)) * arange(N)
167    new_x = arange(N)
168    t1 = time.clock()
169    interp = scipy.interpolate.interp1d(x, y)
170    new_y = interp(new_x)
171    t2 = time.clock()
172    print('scipy interp1d (sec):', t2 - t1)
173    print(new_y[:5, :5])
174
175if __name__ == '__main__':
176    main()
177