1from numpy import (
2    logspace, linspace, geomspace, dtype, array, sctypes, arange, isnan,
3    ndarray, sqrt, nextafter, stack, errstate
4    )
5from numpy.testing import (
6    assert_, assert_equal, assert_raises, assert_array_equal, assert_allclose,
7    )
8
9
10class PhysicalQuantity(float):
11    def __new__(cls, value):
12        return float.__new__(cls, value)
13
14    def __add__(self, x):
15        assert_(isinstance(x, PhysicalQuantity))
16        return PhysicalQuantity(float(x) + float(self))
17    __radd__ = __add__
18
19    def __sub__(self, x):
20        assert_(isinstance(x, PhysicalQuantity))
21        return PhysicalQuantity(float(self) - float(x))
22
23    def __rsub__(self, x):
24        assert_(isinstance(x, PhysicalQuantity))
25        return PhysicalQuantity(float(x) - float(self))
26
27    def __mul__(self, x):
28        return PhysicalQuantity(float(x) * float(self))
29    __rmul__ = __mul__
30
31    def __div__(self, x):
32        return PhysicalQuantity(float(self) / float(x))
33
34    def __rdiv__(self, x):
35        return PhysicalQuantity(float(x) / float(self))
36
37
38class PhysicalQuantity2(ndarray):
39    __array_priority__ = 10
40
41
42class TestLogspace:
43
44    def test_basic(self):
45        y = logspace(0, 6)
46        assert_(len(y) == 50)
47        y = logspace(0, 6, num=100)
48        assert_(y[-1] == 10 ** 6)
49        y = logspace(0, 6, endpoint=False)
50        assert_(y[-1] < 10 ** 6)
51        y = logspace(0, 6, num=7)
52        assert_array_equal(y, [1, 10, 100, 1e3, 1e4, 1e5, 1e6])
53
54    def test_start_stop_array(self):
55        start = array([0., 1.])
56        stop = array([6., 7.])
57        t1 = logspace(start, stop, 6)
58        t2 = stack([logspace(_start, _stop, 6)
59                    for _start, _stop in zip(start, stop)], axis=1)
60        assert_equal(t1, t2)
61        t3 = logspace(start, stop[0], 6)
62        t4 = stack([logspace(_start, stop[0], 6)
63                    for _start in start], axis=1)
64        assert_equal(t3, t4)
65        t5 = logspace(start, stop, 6, axis=-1)
66        assert_equal(t5, t2.T)
67
68    def test_dtype(self):
69        y = logspace(0, 6, dtype='float32')
70        assert_equal(y.dtype, dtype('float32'))
71        y = logspace(0, 6, dtype='float64')
72        assert_equal(y.dtype, dtype('float64'))
73        y = logspace(0, 6, dtype='int32')
74        assert_equal(y.dtype, dtype('int32'))
75
76    def test_physical_quantities(self):
77        a = PhysicalQuantity(1.0)
78        b = PhysicalQuantity(5.0)
79        assert_equal(logspace(a, b), logspace(1.0, 5.0))
80
81    def test_subclass(self):
82        a = array(1).view(PhysicalQuantity2)
83        b = array(7).view(PhysicalQuantity2)
84        ls = logspace(a, b)
85        assert type(ls) is PhysicalQuantity2
86        assert_equal(ls, logspace(1.0, 7.0))
87        ls = logspace(a, b, 1)
88        assert type(ls) is PhysicalQuantity2
89        assert_equal(ls, logspace(1.0, 7.0, 1))
90
91
92class TestGeomspace:
93
94    def test_basic(self):
95        y = geomspace(1, 1e6)
96        assert_(len(y) == 50)
97        y = geomspace(1, 1e6, num=100)
98        assert_(y[-1] == 10 ** 6)
99        y = geomspace(1, 1e6, endpoint=False)
100        assert_(y[-1] < 10 ** 6)
101        y = geomspace(1, 1e6, num=7)
102        assert_array_equal(y, [1, 10, 100, 1e3, 1e4, 1e5, 1e6])
103
104        y = geomspace(8, 2, num=3)
105        assert_allclose(y, [8, 4, 2])
106        assert_array_equal(y.imag, 0)
107
108        y = geomspace(-1, -100, num=3)
109        assert_array_equal(y, [-1, -10, -100])
110        assert_array_equal(y.imag, 0)
111
112        y = geomspace(-100, -1, num=3)
113        assert_array_equal(y, [-100, -10, -1])
114        assert_array_equal(y.imag, 0)
115
116    def test_boundaries_match_start_and_stop_exactly(self):
117        # make sure that the boundaries of the returned array exactly
118        # equal 'start' and 'stop' - this isn't obvious because
119        # np.exp(np.log(x)) isn't necessarily exactly equal to x
120        start = 0.3
121        stop = 20.3
122
123        y = geomspace(start, stop, num=1)
124        assert_equal(y[0], start)
125
126        y = geomspace(start, stop, num=1, endpoint=False)
127        assert_equal(y[0], start)
128
129        y = geomspace(start, stop, num=3)
130        assert_equal(y[0], start)
131        assert_equal(y[-1], stop)
132
133        y = geomspace(start, stop, num=3, endpoint=False)
134        assert_equal(y[0], start)
135
136    def test_nan_interior(self):
137        with errstate(invalid='ignore'):
138            y = geomspace(-3, 3, num=4)
139
140        assert_equal(y[0], -3.0)
141        assert_(isnan(y[1:-1]).all())
142        assert_equal(y[3], 3.0)
143
144        with errstate(invalid='ignore'):
145            y = geomspace(-3, 3, num=4, endpoint=False)
146
147        assert_equal(y[0], -3.0)
148        assert_(isnan(y[1:]).all())
149
150    def test_complex(self):
151        # Purely imaginary
152        y = geomspace(1j, 16j, num=5)
153        assert_allclose(y, [1j, 2j, 4j, 8j, 16j])
154        assert_array_equal(y.real, 0)
155
156        y = geomspace(-4j, -324j, num=5)
157        assert_allclose(y, [-4j, -12j, -36j, -108j, -324j])
158        assert_array_equal(y.real, 0)
159
160        y = geomspace(1+1j, 1000+1000j, num=4)
161        assert_allclose(y, [1+1j, 10+10j, 100+100j, 1000+1000j])
162
163        y = geomspace(-1+1j, -1000+1000j, num=4)
164        assert_allclose(y, [-1+1j, -10+10j, -100+100j, -1000+1000j])
165
166        # Logarithmic spirals
167        y = geomspace(-1, 1, num=3, dtype=complex)
168        assert_allclose(y, [-1, 1j, +1])
169
170        y = geomspace(0+3j, -3+0j, 3)
171        assert_allclose(y, [0+3j, -3/sqrt(2)+3j/sqrt(2), -3+0j])
172        y = geomspace(0+3j, 3+0j, 3)
173        assert_allclose(y, [0+3j, 3/sqrt(2)+3j/sqrt(2), 3+0j])
174        y = geomspace(-3+0j, 0-3j, 3)
175        assert_allclose(y, [-3+0j, -3/sqrt(2)-3j/sqrt(2), 0-3j])
176        y = geomspace(0+3j, -3+0j, 3)
177        assert_allclose(y, [0+3j, -3/sqrt(2)+3j/sqrt(2), -3+0j])
178        y = geomspace(-2-3j, 5+7j, 7)
179        assert_allclose(y, [-2-3j, -0.29058977-4.15771027j,
180                            2.08885354-4.34146838j, 4.58345529-3.16355218j,
181                            6.41401745-0.55233457j, 6.75707386+3.11795092j,
182                            5+7j])
183
184        # Type promotion should prevent the -5 from becoming a NaN
185        y = geomspace(3j, -5, 2)
186        assert_allclose(y, [3j, -5])
187        y = geomspace(-5, 3j, 2)
188        assert_allclose(y, [-5, 3j])
189
190    def test_dtype(self):
191        y = geomspace(1, 1e6, dtype='float32')
192        assert_equal(y.dtype, dtype('float32'))
193        y = geomspace(1, 1e6, dtype='float64')
194        assert_equal(y.dtype, dtype('float64'))
195        y = geomspace(1, 1e6, dtype='int32')
196        assert_equal(y.dtype, dtype('int32'))
197
198        # Native types
199        y = geomspace(1, 1e6, dtype=float)
200        assert_equal(y.dtype, dtype('float_'))
201        y = geomspace(1, 1e6, dtype=complex)
202        assert_equal(y.dtype, dtype('complex'))
203
204    def test_start_stop_array_scalar(self):
205        lim1 = array([120, 100], dtype="int8")
206        lim2 = array([-120, -100], dtype="int8")
207        lim3 = array([1200, 1000], dtype="uint16")
208        t1 = geomspace(lim1[0], lim1[1], 5)
209        t2 = geomspace(lim2[0], lim2[1], 5)
210        t3 = geomspace(lim3[0], lim3[1], 5)
211        t4 = geomspace(120.0, 100.0, 5)
212        t5 = geomspace(-120.0, -100.0, 5)
213        t6 = geomspace(1200.0, 1000.0, 5)
214
215        # t3 uses float32, t6 uses float64
216        assert_allclose(t1, t4, rtol=1e-2)
217        assert_allclose(t2, t5, rtol=1e-2)
218        assert_allclose(t3, t6, rtol=1e-5)
219
220    def test_start_stop_array(self):
221        # Try to use all special cases.
222        start = array([1.e0, 32., 1j, -4j, 1+1j, -1])
223        stop = array([1.e4, 2., 16j, -324j, 10000+10000j, 1])
224        t1 = geomspace(start, stop, 5)
225        t2 = stack([geomspace(_start, _stop, 5)
226                    for _start, _stop in zip(start, stop)], axis=1)
227        assert_equal(t1, t2)
228        t3 = geomspace(start, stop[0], 5)
229        t4 = stack([geomspace(_start, stop[0], 5)
230                    for _start in start], axis=1)
231        assert_equal(t3, t4)
232        t5 = geomspace(start, stop, 5, axis=-1)
233        assert_equal(t5, t2.T)
234
235    def test_physical_quantities(self):
236        a = PhysicalQuantity(1.0)
237        b = PhysicalQuantity(5.0)
238        assert_equal(geomspace(a, b), geomspace(1.0, 5.0))
239
240    def test_subclass(self):
241        a = array(1).view(PhysicalQuantity2)
242        b = array(7).view(PhysicalQuantity2)
243        gs = geomspace(a, b)
244        assert type(gs) is PhysicalQuantity2
245        assert_equal(gs, geomspace(1.0, 7.0))
246        gs = geomspace(a, b, 1)
247        assert type(gs) is PhysicalQuantity2
248        assert_equal(gs, geomspace(1.0, 7.0, 1))
249
250    def test_bounds(self):
251        assert_raises(ValueError, geomspace, 0, 10)
252        assert_raises(ValueError, geomspace, 10, 0)
253        assert_raises(ValueError, geomspace, 0, 0)
254
255
256class TestLinspace:
257
258    def test_basic(self):
259        y = linspace(0, 10)
260        assert_(len(y) == 50)
261        y = linspace(2, 10, num=100)
262        assert_(y[-1] == 10)
263        y = linspace(2, 10, endpoint=False)
264        assert_(y[-1] < 10)
265        assert_raises(ValueError, linspace, 0, 10, num=-1)
266
267    def test_corner(self):
268        y = list(linspace(0, 1, 1))
269        assert_(y == [0.0], y)
270        assert_raises(TypeError, linspace, 0, 1, num=2.5)
271
272    def test_type(self):
273        t1 = linspace(0, 1, 0).dtype
274        t2 = linspace(0, 1, 1).dtype
275        t3 = linspace(0, 1, 2).dtype
276        assert_equal(t1, t2)
277        assert_equal(t2, t3)
278
279    def test_dtype(self):
280        y = linspace(0, 6, dtype='float32')
281        assert_equal(y.dtype, dtype('float32'))
282        y = linspace(0, 6, dtype='float64')
283        assert_equal(y.dtype, dtype('float64'))
284        y = linspace(0, 6, dtype='int32')
285        assert_equal(y.dtype, dtype('int32'))
286
287    def test_start_stop_array_scalar(self):
288        lim1 = array([-120, 100], dtype="int8")
289        lim2 = array([120, -100], dtype="int8")
290        lim3 = array([1200, 1000], dtype="uint16")
291        t1 = linspace(lim1[0], lim1[1], 5)
292        t2 = linspace(lim2[0], lim2[1], 5)
293        t3 = linspace(lim3[0], lim3[1], 5)
294        t4 = linspace(-120.0, 100.0, 5)
295        t5 = linspace(120.0, -100.0, 5)
296        t6 = linspace(1200.0, 1000.0, 5)
297        assert_equal(t1, t4)
298        assert_equal(t2, t5)
299        assert_equal(t3, t6)
300
301    def test_start_stop_array(self):
302        start = array([-120, 120], dtype="int8")
303        stop = array([100, -100], dtype="int8")
304        t1 = linspace(start, stop, 5)
305        t2 = stack([linspace(_start, _stop, 5)
306                    for _start, _stop in zip(start, stop)], axis=1)
307        assert_equal(t1, t2)
308        t3 = linspace(start, stop[0], 5)
309        t4 = stack([linspace(_start, stop[0], 5)
310                    for _start in start], axis=1)
311        assert_equal(t3, t4)
312        t5 = linspace(start, stop, 5, axis=-1)
313        assert_equal(t5, t2.T)
314
315    def test_complex(self):
316        lim1 = linspace(1 + 2j, 3 + 4j, 5)
317        t1 = array([1.0+2.j, 1.5+2.5j,  2.0+3j, 2.5+3.5j, 3.0+4j])
318        lim2 = linspace(1j, 10, 5)
319        t2 = array([0.0+1.j, 2.5+0.75j, 5.0+0.5j, 7.5+0.25j, 10.0+0j])
320        assert_equal(lim1, t1)
321        assert_equal(lim2, t2)
322
323    def test_physical_quantities(self):
324        a = PhysicalQuantity(0.0)
325        b = PhysicalQuantity(1.0)
326        assert_equal(linspace(a, b), linspace(0.0, 1.0))
327
328    def test_subclass(self):
329        a = array(0).view(PhysicalQuantity2)
330        b = array(1).view(PhysicalQuantity2)
331        ls = linspace(a, b)
332        assert type(ls) is PhysicalQuantity2
333        assert_equal(ls, linspace(0.0, 1.0))
334        ls = linspace(a, b, 1)
335        assert type(ls) is PhysicalQuantity2
336        assert_equal(ls, linspace(0.0, 1.0, 1))
337
338    def test_array_interface(self):
339        # Regression test for https://github.com/numpy/numpy/pull/6659
340        # Ensure that start/stop can be objects that implement
341        # __array_interface__ and are convertible to numeric scalars
342
343        class Arrayish:
344            """
345            A generic object that supports the __array_interface__ and hence
346            can in principle be converted to a numeric scalar, but is not
347            otherwise recognized as numeric, but also happens to support
348            multiplication by floats.
349
350            Data should be an object that implements the buffer interface,
351            and contains at least 4 bytes.
352            """
353
354            def __init__(self, data):
355                self._data = data
356
357            @property
358            def __array_interface__(self):
359                return {'shape': (), 'typestr': '<i4', 'data': self._data,
360                        'version': 3}
361
362            def __mul__(self, other):
363                # For the purposes of this test any multiplication is an
364                # identity operation :)
365                return self
366
367        one = Arrayish(array(1, dtype='<i4'))
368        five = Arrayish(array(5, dtype='<i4'))
369
370        assert_equal(linspace(one, five), linspace(1, 5))
371
372    def test_denormal_numbers(self):
373        # Regression test for gh-5437. Will probably fail when compiled
374        # with ICC, which flushes denormals to zero
375        for ftype in sctypes['float']:
376            stop = nextafter(ftype(0), ftype(1)) * 5  # A denormal number
377            assert_(any(linspace(0, stop, 10, endpoint=False, dtype=ftype)))
378
379    def test_equivalent_to_arange(self):
380        for j in range(1000):
381            assert_equal(linspace(0, j, j+1, dtype=int),
382                         arange(j+1, dtype=int))
383
384    def test_retstep(self):
385        for num in [0, 1, 2]:
386            for ept in [False, True]:
387                y = linspace(0, 1, num, endpoint=ept, retstep=True)
388                assert isinstance(y, tuple) and len(y) == 2
389                if num == 2:
390                    y0_expect = [0.0, 1.0] if ept else [0.0, 0.5]
391                    assert_array_equal(y[0], y0_expect)
392                    assert_equal(y[1], y0_expect[1])
393                elif num == 1 and not ept:
394                    assert_array_equal(y[0], [0.0])
395                    assert_equal(y[1], 1.0)
396                else:
397                    assert_array_equal(y[0], [0.0][:num])
398                    assert isnan(y[1])
399
400    def test_object(self):
401        start = array(1, dtype='O')
402        stop = array(2, dtype='O')
403        y = linspace(start, stop, 3)
404        assert_array_equal(y, array([1., 1.5, 2.]))
405
406    def test_round_negative(self):
407        y = linspace(-1, 3, num=8, dtype=int)
408        t = array([-1, -1, 0, 0, 1, 1, 2, 3], dtype=int)
409        assert_array_equal(y, t)
410