|
Packit |
562c7a |
# mode: run
|
|
Packit |
562c7a |
# tag: pythran, numpy, cpp
|
|
Packit |
562c7a |
# cython: np_pythran=True
|
|
Packit |
562c7a |
|
|
Packit |
562c7a |
import numpy as np
|
|
Packit |
562c7a |
cimport numpy as cnp
|
|
Packit |
562c7a |
|
|
Packit |
562c7a |
def diffuse():
|
|
Packit |
562c7a |
"""
|
|
Packit |
562c7a |
>>> u = diffuse()
|
|
Packit |
562c7a |
>>> count_non_zero = np.sum(u > 0)
|
|
Packit |
562c7a |
>>> 15000 < count_non_zero < (2**7) * (2**7) or count_non_zero
|
|
Packit |
562c7a |
True
|
|
Packit |
562c7a |
"""
|
|
Packit |
562c7a |
lx, ly = (2**7, 2**7)
|
|
Packit |
562c7a |
u = np.zeros([lx, ly], dtype=np.double)
|
|
Packit |
562c7a |
u[lx // 2, ly // 2] = 1000.0
|
|
Packit |
562c7a |
_diffuse_numpy(u, 500)
|
|
Packit |
562c7a |
return u
|
|
Packit |
562c7a |
|
|
Packit |
562c7a |
|
|
Packit |
562c7a |
def _diffuse_numpy(cnp.ndarray[double, ndim=2] u, int N):
|
|
Packit |
562c7a |
"""
|
|
Packit |
562c7a |
Apply Numpy matrix for the Forward-Euler Approximation
|
|
Packit |
562c7a |
"""
|
|
Packit |
562c7a |
cdef cnp.ndarray[double, ndim=2] temp = np.zeros_like(u)
|
|
Packit |
562c7a |
mu = 0.1
|
|
Packit |
562c7a |
|
|
Packit |
562c7a |
for n in range(N):
|
|
Packit |
562c7a |
temp[1:-1, 1:-1] = u[1:-1, 1:-1] + mu * (
|
|
Packit |
562c7a |
u[2:, 1:-1] - 2 * u[1:-1, 1:-1] + u[0:-2, 1:-1] +
|
|
Packit |
562c7a |
u[1:-1, 2:] - 2 * u[1:-1, 1:-1] + u[1:-1, 0:-2])
|
|
Packit |
562c7a |
u[:, :] = temp[:, :]
|
|
Packit |
562c7a |
temp[:, :] = 0.0
|
|
Packit |
562c7a |
|
|
Packit |
562c7a |
|
|
Packit |
562c7a |
def calculate_tax(cnp.ndarray[double, ndim=1] d):
|
|
Packit |
562c7a |
"""
|
|
Packit |
562c7a |
>>> mu, sigma = 10.64, .35
|
|
Packit |
562c7a |
>>> np.random.seed(1234)
|
|
Packit |
562c7a |
>>> d = np.random.lognormal(mu, sigma, 10000)
|
|
Packit |
562c7a |
>>> avg = calculate_tax(d)
|
|
Packit |
562c7a |
>>> 0.243 < avg < 0.244 or avg # 0.24342652180085891
|
|
Packit |
562c7a |
True
|
|
Packit |
562c7a |
"""
|
|
Packit |
562c7a |
tax_seg1 = d[(d > 256303)] * 0.45 - 16164.53
|
|
Packit |
562c7a |
tax_seg2 = d[(d > 54057) & (d <= 256303)] * 0.42 - 8475.44
|
|
Packit |
562c7a |
seg3 = d[(d > 13769) & (d <= 54057)] - 13769
|
|
Packit |
562c7a |
seg4 = d[(d > 8820) & (d <= 13769)] - 8820
|
|
Packit |
562c7a |
prog_seg3 = seg3 * 0.0000022376 + 0.2397
|
|
Packit |
562c7a |
prog_seg4 = seg4 * 0.0000100727 + 0.14
|
|
Packit |
562c7a |
return (
|
|
Packit |
562c7a |
np.sum(tax_seg1) +
|
|
Packit |
562c7a |
np.sum(tax_seg2) +
|
|
Packit |
562c7a |
np.sum(seg3 * prog_seg3 + 939.57) +
|
|
Packit |
562c7a |
np.sum(seg4 * prog_seg4)
|
|
Packit |
562c7a |
) / np.sum(d)
|