Skip to content

Commit

Permalink
ENH: jit the 1d quadrature routines (#352)
Browse files Browse the repository at this point in the history 
* ENH: jit the 1d quadrature routines

* ENH: fix some bugs related to jitting 1d quad routines

* DOC: fix docstrings for qnwgamma and _qnwgamma1
  • Loading branch information
@sglyon @mmcky
sglyon authored and mmcky committed Oct 24, 2017
1 parent 6b1f9a8 commit ea900c4
Showing 1 changed file with 66 additions and 52 deletions.
118 changes: 66 additions & 52 deletions quantecon/quad.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
import math
import numpy as np
import scipy.linalg as la
from scipy.special import gammaln
from numba import jit, vectorize
import sympy as sym
from .ce_util import ckron, gridmake
from .util import check_random_state
Expand All @@ -28,6 +28,19 @@
'qnwsimp', 'qnwtrap', 'qnwunif', 'quadrect', 'qnwbeta',
'qnwgamma']

@vectorize(nopython=True)
def gammaln(x):
return math.lgamma(x)


@vectorize(nopython=True)
def fix(x):
if x < 0:
return math.ceil(x)
else:
return math.floor(x)


# ------------------ #
# Exported Functions #
# ------------------ #
Expand Down Expand Up @@ -152,15 +165,15 @@ def qnwequi(n, a, b, kind="N", equidist_pp=None, random_state=None):
if kind.upper() == "N": # Neiderreiter
j = 2.0 ** (np.arange(1, d+1) / (d+1))
nodes = np.outer(i, j)
nodes = (nodes - np.fix(nodes)).squeeze()
nodes = (nodes - fix(nodes)).squeeze()
elif kind.upper() == "W": # Weyl
j = equidist_pp[:d]
nodes = np.outer(i, j)
nodes = (nodes - np.fix(nodes)).squeeze()
nodes = (nodes - fix(nodes)).squeeze()
elif kind.upper() == "H": # Haber
j = equidist_pp[:d]
nodes = np.outer(i * (i+1) / 2, j)
nodes = (nodes - np.fix(nodes)).squeeze()
nodes = (nodes - fix(nodes)).squeeze()
elif kind.upper() == "R": # pseudo-random
nodes = random_state.rand(n, d).squeeze()
else:
Expand Down Expand Up @@ -252,21 +265,21 @@ def qnwnorm(n, mu=None, sig2=None, usesqrtm=False):
Economics and Finance, MIT Press, 2002.
"""
n = np.asarray(n)
n = np.atleast_1d(n)
d = n.size

if mu is None:
mu = np.zeros(d)
else:
mu = np.asarray(mu)
mu = np.atleast_1d(mu)

if sig2 is None:
sig2 = np.eye(d)
else:
sig2 = np.asarray(sig2).reshape(d, d)
sig2 = np.atleast_1d(sig2).reshape(d, d)

if all([x.size == 1 for x in [n, mu, sig2]]):
nodes, weights = _qnwnorm1(n)
nodes, weights = _qnwnorm1(n[0])
else:
nodes = []
weights = []
Expand Down Expand Up @@ -573,7 +586,7 @@ def qnwbeta(n, a=1.0, b=1.0):
return _make_multidim_func(_qnwbeta1, n, a, b)


def qnwgamma(n, a=None):
def qnwgamma(n, a=1.0, b=1.0, tol=3e-14):
"""
Computes nodes and weights for gamma distribution
Expand All @@ -582,14 +595,14 @@ def qnwgamma(n, a=None):
n : int or array_like(float)
A length-d iterable of the number of nodes in each dimension
mu : scalar or array_like(float), optional(default=zeros(d))
The means of each dimension of the random variable. If a scalar
is given, that constant is repeated d times, where d is the
number of dimensions
a : scalar or array_like(float) : optional(default=ones(d))
Shape parameter of the gamma distribution parameter. Must be positive
sig2 : array_like(float), optional(default=eye(d))
A d x d array representing the variance-covariance matrix of the
multivariate normal distribution.
b : scalar or array_like(float) : optional(default=ones(d))
Scale parameter of the gamma distribution parameter. Must be positive
tol : scalar or array_like(float) : optional(default=ones(d) * 3e-14)
Tolerance parameter for newton iterations for each node
Returns
-------
Expand All @@ -610,7 +623,7 @@ def qnwgamma(n, a=None):
Economics and Finance, MIT Press, 2002.
"""
return _make_multidim_func(_qnwgamma1, n, a)
return _make_multidim_func(_qnwgamma1, n, a, b, tol)

# ------------------ #
# Internal Functions #
Expand Down Expand Up @@ -647,12 +660,12 @@ def _make_multidim_func(one_d_func, n, *args):
"""
args = list(args)
n = np.asarray(n)
args = list(map(np.asarray, args))
_args = list(args)
n = np.atleast_1d(n)
args = list(map(np.atleast_1d, _args))

if all([x.size == 1 for x in [n] + args]):
return one_d_func(n, *args)
return one_d_func(n[0], *_args)

d = n.size

Expand All @@ -674,7 +687,7 @@ def _make_multidim_func(one_d_func, n, *args):
nodes = gridmake(*nodes)
return nodes, weights


@jit(nopython=True)
def _qnwcheb1(n, a, b):
"""
Compute univariate Guass-Checbychev quadrature nodes and weights
Expand Down Expand Up @@ -714,15 +727,15 @@ def _qnwcheb1(n, a, b):
# Create temporary arrays to be used in computing weights
t1 = np.arange(1, n+1) - 0.5
t2 = np.arange(0.0, n, 2)
t3 = np.concatenate([np.array([1.0]),
-2.0/(np.arange(1.0, n-1, 2)*np.arange(3.0, n+1, 2))])
t3 = np.concatenate((np.array([1.0]),
-2.0/(np.arange(1.0, n-1, 2)*np.arange(3.0, n+1, 2))))

# compute weights and return
weights = ((b-a)/n)*np.cos(np.pi/n*np.outer(t1, t2)).dot(t3)
weights = ((b-a)/n)*np.cos(np.pi/n*np.outer(t1, t2)) @ t3

return nodes, weights


@jit(nopython=True)
def _qnwlege1(n, a, b):
"""
Compute univariate Guass-Legendre quadrature nodes and weights
Expand Down Expand Up @@ -759,19 +772,19 @@ def _qnwlege1(n, a, b):
"""
# import ipdb; ipdb.set_trace()
maxit = 100
m = np.fix((n + 1) / 2.0).astype(int)
m = int(fix((n + 1) / 2.0))
xm = 0.5 * (b + a)
xl = 0.5 * (b - a)
nodes = np.zeros(n)

weights = nodes.copy()
i = np.arange(m, dtype='int')
i = np.arange(m)

z = np.cos(np.pi * ((i + 1.0) - 0.25) / (n + 0.5))

for its in range(maxit):
p1 = 1.0
p2 = 0.0
p1 = np.ones_like(z)
p2 = np.zeros_like(z)
for j in range(1, n+1):
p3 = p2
p2 = p1
Expand All @@ -780,7 +793,7 @@ def _qnwlege1(n, a, b):
pp = n * (z * p1 - p2)/(z * z - 1.0)
z1 = z.copy()
z = z1 - p1/pp
if all(np.abs(z - z1) < 1e-14):
if np.all(np.abs(z - z1) < 1e-14):
break

if its == maxit - 1:
Expand All @@ -794,7 +807,7 @@ def _qnwlege1(n, a, b):

return nodes, weights


@jit(nopython=True)
def _qnwnorm1(n):
"""
Compute nodes and weights for quadrature of univariate standard
Expand Down Expand Up @@ -826,7 +839,7 @@ def _qnwnorm1(n):
"""
maxit = 100
pim4 = 1 / np.pi**(0.25)
m = np.fix((n + 1) / 2).astype(int)
m = int(fix((n + 1) / 2))
nodes = np.zeros(n)
weights = np.zeros(n)

Expand Down Expand Up @@ -872,7 +885,7 @@ def _qnwnorm1(n):

return nodes, weights


@jit(nopython=True)
def _qnwsimp1(n, a, b):
"""
Compute univariate Simpson quadrature nodes and weights
Expand Down Expand Up @@ -913,14 +926,14 @@ def _qnwsimp1(n, a, b):

nodes = np.linspace(a, b, n)
dx = nodes[1] - nodes[0]
weights = np.tile([2.0, 4.0], (n + 1) // 2)
weights = np.kron(np.ones((n+1) // 2), np.array([2.0, 4.0]))
weights = weights[:n]
weights[0] = weights[-1] = 1
weights = (dx / 3.0) * weights

return nodes, weights


@jit(nopython=True)
def _qnwtrap1(n, a, b):
"""
Compute univariate trapezoid rule quadrature nodes and weights
Expand Down Expand Up @@ -967,7 +980,7 @@ def _qnwtrap1(n, a, b):

return nodes, weights


@jit(nopython=True)
def _qnwbeta1(n, a=1.0, b=1.0):
"""
Computes nodes and weights for quadrature on the beta distribution.
Expand Down Expand Up @@ -1090,21 +1103,24 @@ def _qnwbeta1(n, a=1.0, b=1.0):

return nodes, weights


def _qnwgamma1(n, a=None):
@jit(nopython=True)
def _qnwgamma1(n, a=1.0, b=1.0, tol=3e-14):
"""
Insert docs. Default is a=0
NOTE: For now I am just following compecon; would be much better to
find a different way since I don't know what they are doing.
1d quadrature weights and nodes for Gamma distributed random variable
Parameters
----------
n : scalar : int
The number of quadrature points
a : scalar : float
Gamma distribution parameter
a : scalar : float, optional(default=1.0)
Shape parameter of the gamma distribution parameter. Must be positive
b : scalar : float, optional(default=1.0)
Scale parameter of the gamma distribution parameter. Must be positive
tol : scalar : float, optional(default=3e-14)
Tolerance parameter for newton iterations for each node
Returns
-------
Expand All @@ -1125,12 +1141,9 @@ def _qnwgamma1(n, a=None):
Economics and Finance, MIT Press, 2002.
"""
if a is None:
a = 0
else:
a -= 1
a -= 1

maxit = 10
maxit = 25

factor = -math.exp(gammaln(a+n) - gammaln(n) - gammaln(a+1))
nodes = np.zeros(n)
Expand All @@ -1151,10 +1164,11 @@ def _qnwgamma1(n, a=None):
# root finding iterations
its = 0
z1 = -10000
while abs(z - z1) > 1e-10 and its < maxit:
while abs(z - z1) > tol and its < maxit:
p1 = 1.0
p2 = 0.0
for j in range(1, n+1):
# Recurrance relation for Laguerre polynomials
p3 = p2
p2 = p1
p1 = ((2*j - 1 + a - z)*p2 - (j - 1 + a)*p3) / j
Expand All @@ -1170,4 +1184,4 @@ def _qnwgamma1(n, a=None):
nodes[i] = z
weights[i] = factor / (pp*n*p2)

return nodes, weights
return nodes*b, weights

0 comments on commit ea900c4

Please # to comment.