cython

1
# mode: run
2
# tag: pythran, numpy, cpp
3
# cython: np_pythran=True
4

5
import numpy as np
6
cimport numpy as cnp
7

8
def diffuse():
9
    """
10
    >>> u = diffuse()
11
    >>> count_non_zero = np.sum(u > 0)
12
    >>> bool(850 < count_non_zero < (2**5) * (2**5)) or count_non_zero
13
    True
14
    """
15
    lx, ly = (2**5, 2**5)
16
    u = np.zeros([lx, ly], dtype=np.double)
17
    u[lx // 2, ly // 2] = 1000.0
18
    _diffuse_numpy(u, 50)
19
    return u
20

21

22
def _diffuse_numpy(cnp.ndarray[double, ndim=2] u, int N):
23
    """
24
    Apply Numpy matrix for the Forward-Euler Approximation
25
    """
26
    cdef cnp.ndarray[double, ndim=2] temp = np.zeros_like(u)
27
    mu = 0.1
28

29
    for n in range(N):
30
        temp[1:-1, 1:-1] = u[1:-1, 1:-1] + mu * (
31
            u[2:, 1:-1] - 2 * u[1:-1, 1:-1] + u[0:-2, 1:-1] +
32
            u[1:-1, 2:] - 2 * u[1:-1, 1:-1] + u[1:-1, 0:-2])
33
        u[:, :] = temp[:, :]
34
        temp[:, :] = 0.0
35

36

37
def calculate_tax(cnp.ndarray[double, ndim=1] d):
38
    """
39
    >>> mu, sigma = 10.64, .35
40
    >>> np.random.seed(1234)
41
    >>> d = np.random.lognormal(mu, sigma, 10000)
42
    >>> avg = calculate_tax(d)
43
    >>> bool(0.243 < avg < 0.244) or avg  # 0.24342652180085891
44
    True
45
    """
46
    tax_seg1 = d[(d > 256303)] * 0.45 - 16164.53
47
    tax_seg2 = d[(d > 54057) & (d <= 256303)] * 0.42 - 8475.44
48
    seg3 = d[(d > 13769) & (d <= 54057)] - 13769
49
    seg4 = d[(d > 8820) & (d <= 13769)] - 8820
50
    prog_seg3 = seg3 * 0.0000022376 + 0.2397
51
    prog_seg4 = seg4 * 0.0000100727 + 0.14
52
    return (
53
        np.sum(tax_seg1) +
54
        np.sum(tax_seg2) +
55
        np.sum(seg3 * prog_seg3 + 939.57) +
56
        np.sum(seg4 * prog_seg4)
57
    ) / np.sum(d)
58

59
def access_shape():
60
    """
61
    >>> access_shape()
62
    10
63
    """
64
    cdef cnp.ndarray[double, ndim=2, mode='c'] array_in = \
65
                    1e10 * np.ones((10, 10))
66

67
    return array_in.shape[0]
68