Aluode/PerceptionLabPortable

1"""
2Implement the cmath module functions.
3"""
4
5
6import cmath
7import math
8
9from numba.core.imputils import Registry, impl_ret_untracked
10from numba.core import types, cgutils
11from numba.core.typing import signature
12from numba.cpython import builtins, mathimpl
13from numba.core.extending import overload
14
15registry = Registry('cmathimpl')
16lower = registry.lower
17
18
19def is_nan(builder, z):
20    return builder.fcmp_unordered('uno', z.real, z.imag)
21
22def is_inf(builder, z):
23    return builder.or_(mathimpl.is_inf(builder, z.real),
24                       mathimpl.is_inf(builder, z.imag))
25
26def is_finite(builder, z):
27    return builder.and_(mathimpl.is_finite(builder, z.real),
28                        mathimpl.is_finite(builder, z.imag))
29
30
31@lower(cmath.isnan, types.Complex)
32def isnan_float_impl(context, builder, sig, args):
33    [typ] = sig.args
34    [value] = args
35    z = context.make_complex(builder, typ, value=value)
36    res = is_nan(builder, z)
37    return impl_ret_untracked(context, builder, sig.return_type, res)
38
39@lower(cmath.isinf, types.Complex)
40def isinf_float_impl(context, builder, sig, args):
41    [typ] = sig.args
42    [value] = args
43    z = context.make_complex(builder, typ, value=value)
44    res = is_inf(builder, z)
45    return impl_ret_untracked(context, builder, sig.return_type, res)
46
47
48@lower(cmath.isfinite, types.Complex)
49def isfinite_float_impl(context, builder, sig, args):
50    [typ] = sig.args
51    [value] = args
52    z = context.make_complex(builder, typ, value=value)
53    res = is_finite(builder, z)
54    return impl_ret_untracked(context, builder, sig.return_type, res)
55
56
57@overload(cmath.rect)
58def impl_cmath_rect(r, phi):
59    if all([isinstance(typ, types.Float) for typ in [r, phi]]):
60        def impl(r, phi):
61            if not math.isfinite(phi):
62                if not r:
63                    # cmath.rect(0, phi={inf, nan}) = 0
64                    return abs(r)
65                if math.isinf(r):
66                    # cmath.rect(inf, phi={inf, nan}) = inf + j phi
67                    return complex(r, phi)
68            real = math.cos(phi)
69            imag = math.sin(phi)
70            if real == 0. and math.isinf(r):
71                # 0 * inf would return NaN, we want to keep 0 but xor the sign
72                real /= r
73            else:
74                real *= r
75            if imag == 0. and math.isinf(r):
76                # ditto
77                imag /= r
78            else:
79                imag *= r
80            return complex(real, imag)
81        return impl
82
83
84def intrinsic_complex_unary(inner_func):
85    def wrapper(context, builder, sig, args):
86        [typ] = sig.args
87        [value] = args
88        z = context.make_complex(builder, typ, value=value)
89        x = z.real
90        y = z.imag
91        # Same as above: math.isfinite() is unavailable on 2.x so we precompute
92        # its value and pass it to the pure Python implementation.
93        x_is_finite = mathimpl.is_finite(builder, x)
94        y_is_finite = mathimpl.is_finite(builder, y)
95        inner_sig = signature(sig.return_type,
96                              *(typ.underlying_float,) * 2 + (types.boolean,) * 2)
97        res = context.compile_internal(builder, inner_func, inner_sig,
98                                        (x, y, x_is_finite, y_is_finite))
99        return impl_ret_untracked(context, builder, sig, res)
100    return wrapper
101
102
103NAN = float('nan')
104INF = float('inf')
105
106@lower(cmath.exp, types.Complex)
107@intrinsic_complex_unary
108def exp_impl(x, y, x_is_finite, y_is_finite):
109    """cmath.exp(x + y j)"""
110    if x_is_finite:
111        if y_is_finite:
112            c = math.cos(y)
113            s = math.sin(y)
114            r = math.exp(x)
115            return complex(r * c, r * s)
116        else:
117            return complex(NAN, NAN)
118    elif math.isnan(x):
119        if y:
120            return complex(x, x)  # nan + j nan
121        else:
122            return complex(x, y)  # nan + 0j
123    elif x > 0.0:
124        # x == +inf
125        if y_is_finite:
126            real = math.cos(y)
127            imag = math.sin(y)
128            # Avoid NaNs if math.cos(y) or math.sin(y) == 0
129            # (e.g. cmath.exp(inf + 0j) == inf + 0j)
130            if real != 0:
131                real *= x
132            if imag != 0:
133                imag *= x
134            return complex(real, imag)
135        else:
136            return complex(x, NAN)
137    else:
138        # x == -inf
139        if y_is_finite:
140            r = math.exp(x)
141            c = math.cos(y)
142            s = math.sin(y)
143            return complex(r * c, r * s)
144        else:
145            r = 0
146            return complex(r, r)
147
148@lower(cmath.log, types.Complex)
149@intrinsic_complex_unary
150def log_impl(x, y, x_is_finite, y_is_finite):
151    """cmath.log(x + y j)"""
152    a = math.log(math.hypot(x, y))
153    b = math.atan2(y, x)
154    return complex(a, b)
155
156
157@lower(cmath.log, types.Complex, types.Complex)
158def log_base_impl(context, builder, sig, args):
159    """cmath.log(z, base)"""
160    [z, base] = args
161
162    def log_base(z, base):
163        return cmath.log(z) / cmath.log(base)
164
165    res = context.compile_internal(builder, log_base, sig, args)
166    return impl_ret_untracked(context, builder, sig, res)
167
168
169@overload(cmath.log10)
170def impl_cmath_log10(z):
171    if not isinstance(z, types.Complex):
172        return
173
174    LN_10 = 2.302585092994045684
175
176    def log10_impl(z):
177        """cmath.log10(z)"""
178        z = cmath.log(z)
179        # This formula gives better results on +/-inf than cmath.log(z, 10)
180        # See http://bugs.python.org/issue22544
181        return complex(z.real / LN_10, z.imag / LN_10)
182
183    return log10_impl
184
185
186@overload(cmath.phase)
187def phase_impl(x):
188    """cmath.phase(x + y j)"""
189
190    if not isinstance(x, types.Complex):
191        return
192
193    def impl(x):
194        return math.atan2(x.imag, x.real)
195    return impl
196
197
198@overload(cmath.polar)
199def polar_impl(x):
200    if not isinstance(x, types.Complex):
201        return
202
203    def impl(x):
204        r, i = x.real, x.imag
205        return math.hypot(r, i), math.atan2(i, r)
206    return impl
207
208
209@lower(cmath.sqrt, types.Complex)
210def sqrt_impl(context, builder, sig, args):
211    # We risk spurious overflow for components >= FLT_MAX / (1 + sqrt(2)).
212
213    SQRT2 = 1.414213562373095048801688724209698079E0
214    ONE_PLUS_SQRT2 = (1. + SQRT2)
215    theargflt = sig.args[0].underlying_float
216    # Get a type specific maximum value so scaling for overflow is based on that
217    MAX = mathimpl.DBL_MAX if theargflt.bitwidth == 64 else mathimpl.FLT_MAX
218    # THRES will be double precision, should not impact typing as it's just
219    # used for comparison, there *may* be a few values near THRES which
220    # deviate from e.g. NumPy due to rounding that occurs in the computation
221    # of this value in the case of a 32bit argument.
222    THRES = MAX / ONE_PLUS_SQRT2
223
224    def sqrt_impl(z):
225        """cmath.sqrt(z)"""
226        # This is NumPy's algorithm, see npy_csqrt() in npy_math_complex.c.src
227        a = z.real
228        b = z.imag
229        if a == 0.0 and b == 0.0:
230            return complex(abs(b), b)
231        if math.isinf(b):
232            return complex(abs(b), b)
233        if math.isnan(a):
234            return complex(a, a)
235        if math.isinf(a):
236            if a < 0.0:
237                return complex(abs(b - b), math.copysign(a, b))
238            else:
239                return complex(a, math.copysign(b - b, b))
240
241        # The remaining special case (b is NaN) is handled just fine by
242        # the normal code path below.
243
244        # Scale to avoid overflow
245        if abs(a) >= THRES or abs(b) >= THRES:
246            a *= 0.25
247            b *= 0.25
248            scale = True
249        else:
250            scale = False
251        # Algorithm 312, CACM vol 10, Oct 1967
252        if a >= 0:
253            t = math.sqrt((a + math.hypot(a, b)) * 0.5)
254            real = t
255            imag = b / (2 * t)
256        else:
257            t = math.sqrt((-a + math.hypot(a, b)) * 0.5)
258            real = abs(b) / (2 * t)
259            imag = math.copysign(t, b)
260        # Rescale
261        if scale:
262            return complex(real * 2, imag)
263        else:
264            return complex(real, imag)
265
266    res = context.compile_internal(builder, sqrt_impl, sig, args)
267    return impl_ret_untracked(context, builder, sig, res)
268
269
270@lower(cmath.cos, types.Complex)
271def cos_impl(context, builder, sig, args):
272    def cos_impl(z):
273        """cmath.cos(z) = cmath.cosh(z j)"""
274        return cmath.cosh(complex(-z.imag, z.real))
275
276    res = context.compile_internal(builder, cos_impl, sig, args)
277    return impl_ret_untracked(context, builder, sig, res)
278
279@overload(cmath.cosh)
280def impl_cmath_cosh(z):
281    if not isinstance(z, types.Complex):
282        return
283
284    def cosh_impl(z):
285        """cmath.cosh(z)"""
286        x = z.real
287        y = z.imag
288        if math.isinf(x):
289            if math.isnan(y):
290                # x = +inf, y = NaN => cmath.cosh(x + y j) = inf + Nan * j
291                real = abs(x)
292                imag = y
293            elif y == 0.0:
294                # x = +inf, y = 0 => cmath.cosh(x + y j) = inf + 0j
295                real = abs(x)
296                imag = y
297            else:
298                real = math.copysign(x, math.cos(y))
299                imag = math.copysign(x, math.sin(y))
300            if x < 0.0:
301                # x = -inf => negate imaginary part of result
302                imag = -imag
303            return complex(real, imag)
304        return complex(math.cos(y) * math.cosh(x),
305                    math.sin(y) * math.sinh(x))
306    return cosh_impl
307
308
309@lower(cmath.sin, types.Complex)
310def sin_impl(context, builder, sig, args):
311    def sin_impl(z):
312        """cmath.sin(z) = -j * cmath.sinh(z j)"""
313        r = cmath.sinh(complex(-z.imag, z.real))
314        return complex(r.imag, -r.real)
315
316    res = context.compile_internal(builder, sin_impl, sig, args)
317    return impl_ret_untracked(context, builder, sig, res)
318
319@overload(cmath.sinh)
320def impl_cmath_sinh(z):
321    if not isinstance(z, types.Complex):
322        return
323
324    def sinh_impl(z):
325        """cmath.sinh(z)"""
326        x = z.real
327        y = z.imag
328        if math.isinf(x):
329            if math.isnan(y):
330                # x = +/-inf, y = NaN => cmath.sinh(x + y j) = x + NaN * j
331                real = x
332                imag = y
333            else:
334                real = math.cos(y)
335                imag = math.sin(y)
336                if real != 0.:
337                    real *= x
338                if imag != 0.:
339                    imag *= abs(x)
340            return complex(real, imag)
341        return complex(math.cos(y) * math.sinh(x),
342                       math.sin(y) * math.cosh(x))
343    return sinh_impl
344
345
346@lower(cmath.tan, types.Complex)
347def tan_impl(context, builder, sig, args):
348    def tan_impl(z):
349        """cmath.tan(z) = -j * cmath.tanh(z j)"""
350        r = cmath.tanh(complex(-z.imag, z.real))
351        return complex(r.imag, -r.real)
352
353    res = context.compile_internal(builder, tan_impl, sig, args)
354    return impl_ret_untracked(context, builder, sig, res)
355
356
357@overload(cmath.tanh)
358def impl_cmath_tanh(z):
359    if not isinstance(z, types.Complex):
360        return
361
362    def tanh_impl(z):
363        """cmath.tanh(z)"""
364        x = z.real
365        y = z.imag
366        if math.isinf(x):
367            real = math.copysign(1., x)
368            if math.isinf(y):
369                imag = 0.
370            else:
371                imag = math.copysign(0., math.sin(2. * y))
372            return complex(real, imag)
373        # This is CPython's algorithm (see c_tanh() in cmathmodule.c).
374        # XXX how to force float constants into single precision?
375        tx = math.tanh(x)
376        ty = math.tan(y)
377        cx = 1. / math.cosh(x)
378        txty = tx * ty
379        denom = 1. + txty * txty
380        return complex(
381            tx * (1. + ty * ty) / denom,
382            ((ty / denom) * cx) * cx)
383
384    return tanh_impl
385
386
387@lower(cmath.acos, types.Complex)
388def acos_impl(context, builder, sig, args):
389    LN_4 = math.log(4)
390    THRES = mathimpl.FLT_MAX / 4
391
392    def acos_impl(z):
393        """cmath.acos(z)"""
394        # CPython's algorithm (see c_acos() in cmathmodule.c)
395        if abs(z.real) > THRES or abs(z.imag) > THRES:
396            # Avoid unnecessary overflow for large arguments
397            # (also handles infinities gracefully)
398            real = math.atan2(abs(z.imag), z.real)
399            imag = math.copysign(
400                math.log(math.hypot(z.real * 0.5, z.imag * 0.5)) + LN_4,
401                -z.imag)
402            return complex(real, imag)
403        else:
404            s1 = cmath.sqrt(complex(1. - z.real, -z.imag))
405            s2 = cmath.sqrt(complex(1. + z.real, z.imag))
406            real = 2. * math.atan2(s1.real, s2.real)
407            imag = math.asinh(s2.real * s1.imag - s2.imag * s1.real)
408            return complex(real, imag)
409
410    res = context.compile_internal(builder, acos_impl, sig, args)
411    return impl_ret_untracked(context, builder, sig, res)
412
413@overload(cmath.acosh)
414def impl_cmath_acosh(z):
415    if not isinstance(z, types.Complex):
416        return
417
418    LN_4 = math.log(4)
419    THRES = mathimpl.FLT_MAX / 4
420
421    def acosh_impl(z):
422        """cmath.acosh(z)"""
423        # CPython's algorithm (see c_acosh() in cmathmodule.c)
424        if abs(z.real) > THRES or abs(z.imag) > THRES:
425            # Avoid unnecessary overflow for large arguments
426            # (also handles infinities gracefully)
427            real = math.log(math.hypot(z.real * 0.5, z.imag * 0.5)) + LN_4
428            imag = math.atan2(z.imag, z.real)
429            return complex(real, imag)
430        else:
431            s1 = cmath.sqrt(complex(z.real - 1., z.imag))
432            s2 = cmath.sqrt(complex(z.real + 1., z.imag))
433            real = math.asinh(s1.real * s2.real + s1.imag * s2.imag)
434            imag = 2. * math.atan2(s1.imag, s2.real)
435            return complex(real, imag)
436        # Condensed formula (NumPy)
437        #return cmath.log(z + cmath.sqrt(z + 1.) * cmath.sqrt(z - 1.))
438
439    return acosh_impl
440
441
442@lower(cmath.asinh, types.Complex)
443def asinh_impl(context, builder, sig, args):
444    LN_4 = math.log(4)
445    THRES = mathimpl.FLT_MAX / 4
446
447    def asinh_impl(z):
448        """cmath.asinh(z)"""
449        # CPython's algorithm (see c_asinh() in cmathmodule.c)
450        if abs(z.real) > THRES or abs(z.imag) > THRES:
451            real = math.copysign(
452                math.log(math.hypot(z.real * 0.5, z.imag * 0.5)) + LN_4,
453                z.real)
454            imag = math.atan2(z.imag, abs(z.real))
455            return complex(real, imag)
456        else:
457            s1 = cmath.sqrt(complex(1. + z.imag, -z.real))
458            s2 = cmath.sqrt(complex(1. - z.imag, z.real))
459            real = math.asinh(s1.real * s2.imag - s2.real * s1.imag)
460            imag = math.atan2(z.imag, s1.real * s2.real - s1.imag * s2.imag)
461            return complex(real, imag)
462
463    res = context.compile_internal(builder, asinh_impl, sig, args)
464    return impl_ret_untracked(context, builder, sig, res)
465
466@lower(cmath.asin, types.Complex)
467def asin_impl(context, builder, sig, args):
468    def asin_impl(z):
469        """cmath.asin(z) = -j * cmath.asinh(z j)"""
470        r = cmath.asinh(complex(-z.imag, z.real))
471        return complex(r.imag, -r.real)
472
473    res = context.compile_internal(builder, asin_impl, sig, args)
474    return impl_ret_untracked(context, builder, sig, res)
475
476@lower(cmath.atan, types.Complex)
477def atan_impl(context, builder, sig, args):
478    def atan_impl(z):
479        """cmath.atan(z) = -j * cmath.atanh(z j)"""
480        r = cmath.atanh(complex(-z.imag, z.real))
481        if math.isinf(z.real) and math.isnan(z.imag):
482            # XXX this is odd but necessary
483            return complex(r.imag, r.real)
484        else:
485            return complex(r.imag, -r.real)
486
487    res = context.compile_internal(builder, atan_impl, sig, args)
488    return impl_ret_untracked(context, builder, sig, res)
489
490@lower(cmath.atanh, types.Complex)
491def atanh_impl(context, builder, sig, args):
492    LN_4 = math.log(4)
493    THRES_LARGE = math.sqrt(mathimpl.FLT_MAX / 4)
494    THRES_SMALL = math.sqrt(mathimpl.FLT_MIN)
495    PI_12 = math.pi / 2
496
497    def atanh_impl(z):
498        """cmath.atanh(z)"""
499        # CPython's algorithm (see c_atanh() in cmathmodule.c)
500        if z.real < 0.:
501            # Reduce to case where z.real >= 0., using atanh(z) = -atanh(-z).
502            negate = True
503            z = -z
504        else:
505            negate = False
506
507        ay = abs(z.imag)
508        if math.isnan(z.real) or z.real > THRES_LARGE or ay > THRES_LARGE:
509            if math.isinf(z.imag):
510                real = math.copysign(0., z.real)
511            elif math.isinf(z.real):
512                real = 0.
513            else:
514                # may be safe from overflow, depending on hypot's implementation...
515                h = math.hypot(z.real * 0.5, z.imag * 0.5)
516                real = z.real/4./h/h
517            imag = -math.copysign(PI_12, -z.imag)
518        elif z.real == 1. and ay < THRES_SMALL:
519            # C99 standard says:  atanh(1+/-0.) should be inf +/- 0j
520            if ay == 0.:
521                real = INF
522                imag = z.imag
523            else:
524                real = -math.log(math.sqrt(ay) /
525                                 math.sqrt(math.hypot(ay, 2.)))
526                imag = math.copysign(math.atan2(2., -ay) / 2, z.imag)
527        else:
528            sqay = ay * ay
529            zr1 = 1 - z.real
530            real = math.log1p(4. * z.real / (zr1 * zr1 + sqay)) * 0.25
531            imag = -math.atan2(-2. * z.imag,
532                               zr1 * (1 + z.real) - sqay) * 0.5
533
534        if math.isnan(z.imag):
535            imag = NAN
536        if negate:
537            return complex(-real, -imag)
538        else:
539            return complex(real, imag)
540
541    res = context.compile_internal(builder, atanh_impl, sig, args)
542    return impl_ret_untracked(context, builder, sig, res)
543