1# mode: run 2# tag: pythran, numpy, cpp 3# cython: np_pythran=True 4 5import numpy as np 6cimport numpy as cnp 7 8def diffuse(): 9 """ 10 >>> u = diffuse() 11 >>> count_non_zero = np.sum(u > 0) 12 >>> 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 22def _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 37def 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 >>> 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