def interpolation(interP,knotP=None):
"""
Interpolates the given points and returns an object of the Spline class
Arguments:
interP: interpolation points, (L x 2) matrix
knotP: knot points, (L+4 x 1) matrix
default: equidistant on [0,1]
"""
nip=len(interP)
ctrlP=zeros((nip,2))
knotP=fixKnotPoints(nip, knotP)
xi=(knotP[:-2]+knotP[1:-1]+knotP[2:])/3
nMatrix=zeros((nip,nip))
for i in xrange(nip):
fun=basisFunction(i,knotP)
for k in xrange(nip):
nMatrix[k,i]=fun(xi[k],3)
print nMatrix
print knotP
ctrlP[:,0]=sl.solve_banded((nip-4,nip-4),nMatrix,interP[:,0])
ctrlP[:,1]=sl.solve_banded((nip-4,nip-4),nMatrix,interP[:,1])
return Spline(ctrlP,knotP)
def calculate_vd_vL_vR(self): for index in xrange(self.neuron.BPcount) : if self.test_list[index] : i = self.i_list[index] j = self.j_list[index] self.vL[i:j+1] = solve_banded((1,1),self.ab[:,i:j+1],self.bL[i:j+1],overwrite_ab=False,overwrite_b=False) self.vR[i:j+1] = solve_banded((1,1),self.ab[:,i:j+1],self.bR[i:j+1],overwrite_ab=False,overwrite_b=False) self.d[i:j+1] = solve_banded((1,1),self.ab[:,i:j+1],self.bd[i:j+1],overwrite_ab=False,overwrite_b=False)
def solve(self,f): # Solves u'=f for u(-1)=0
fhat = ifcgt(f*(1-self.x)) # Compute (psi_j,f)
fhat[0] *= 2
y = solve_banded((1,0),self.Lbands,fhat,\
overwrite_ab=False,overwrite_b=True)
uhat = solve_banded((0,1),self.Ubands,y,\
overwrite_ab=False,overwrite_b=True)
u = fcgt(uhat)*(1+self.x) # u=phi*uhat
return u
def fit_grad_descent(self,r_init=0.001,learning_rate=0.1):
''' Calculate fit using gradient decent, as opposed to iteratively
computing exact solutions
'''
delta_r = 1
r = r_init
r_list = []
error_list = []
F_C = solve_banded((1,1),self.ab,self.F_M - r*self.F_N)
itmax = 10000
it = 0
exceed_bounds = False
while (delta_r > 0.0001 and it < itmax):
F_C = solve_banded((1,1),self.ab,self.F_M - r*self.F_N)
r_new = r + learning_rate*np.mean((self.F_M - F_C - r*self.F_N)*self.F_N)
'''compute error on cross-validation set'''
F_C_crossval = solve_banded((1,1),self.ab,self.F_M_crossval - r*self.F_N_crossval)
# error_it = error_calc_outlier(self.F_M_crossval, self.F_N_crossval, F_C_crossval, r)
error_it = abs(error_calc(self.F_M_crossval, self.F_N_crossval, F_C_crossval, r))
delta_r = np.abs(r_new - r)/r
r = r_new
r_list.append(r)
error_list.append(error_it)
it+=1
if error_it > error_list[it-1]: # early stopping
break
# if r or error_it go out of acceptable bounds, break
if r < 0.0 or r > 1.0:
exceed_bounds = True
if error_it > 0.2:
exceed_bounds = True
F_C = solve_banded((1,1),self.ab,self.F_M - r*self.F_N)
r_list = np.array(r_list)
error_list = np.array(error_list)
self.r_vals = r_list
self.error_vals = error_list
self.r = self.r_vals[-1]
self.F_C = F_C
self.F_C_crossval = F_C_crossval
self.it = it
return exceed_bounds
def impl(V, L1, R1x, L2, R2x, dt, n, initial=None, callback=None):
V = V.copy()
# L1i = flatten_tensor(L1)
L1i = L1.copy()
R1 = np.array(R1x).T
# L2i = flatten_tensor(L2)
L2i = L2.copy()
R2 = np.array(R2x)
m = 2
# L = (As + Ass - H.interest_rate*np.eye(nspots))*-dt + np.eye(nspots)
L1i.data *= -dt
L1i.data[m, :] += 1
R1 *= -dt
L2i.data *= -dt
L2i.data[m, :] += 1
R2 *= -dt
offsets1 = (abs(min(L1i.offsets)), abs(max(L1i.offsets)))
offsets2 = (abs(min(L2i.offsets)), abs(max(L2i.offsets)))
dx = np.gradient(spots)[:, np.newaxis]
dy = np.gradient(vars)
Y, X = np.meshgrid(vars, spots)
gradgrid = dt * coeffs[(0, 1)](0, X, Y) / (dx * dy)
gradgrid[:, 0] = 0
gradgrid[:, -1] = 0
gradgrid[0, :] = 0
gradgrid[-1, :] = 0
print_step = max(1, int(n / 10))
to_percent = 100.0 / n
utils.tic("Impl:\t")
for k in xrange(n):
if not k % print_step:
if np.isnan(V).any():
print "Impl fail @ t = %f (%i steps)" % (dt * k, k)
return crumbs
print int(k * to_percent),
if callback is not None:
callback(V, ((n - k) * dt))
Vsv = np.gradient(np.gradient(V)[0])[1] * gradgrid
# Vsv = 0.0
V = spl.solve_banded(offsets2, L2i.data, (V + Vsv - R2).flat, overwrite_b=True).reshape(V.shape)
V = spl.solve_banded(offsets1, L1i.data, (V - R1).T.flat, overwrite_b=True).reshape(V.shape[::-1]).T
crumbs.append(V.copy())
utils.toc(": \t")
return crumbs
def crank(V, L1, R1x, L2, R2x, dt, n, crumbs=[], callback=None):
V = V.copy()
dt *= 0.5
L1e = flatten_tensor(L1)
L1i = L1e.copy()
R1 = np.array(R1x).T
L2e = flatten_tensor(L2)
L2i = L2e.copy()
R2 = np.array(R2x)
m = 2
# L = (As + Ass - r*np.eye(nspots))*-dt + np.eye(nspots)
L1e.data *= dt
L1e.data[m, :] += 1
L1i.data *= -dt
L1i.data[m, :] += 1
R1 *= dt
L2e.data *= dt
L2e.data[m, :] += 1
L2i.data *= -dt
L2i.data[m, :] += 1
R2 *= dt
offsets1 = (abs(min(L1i.offsets)), abs(max(L1i.offsets)))
offsets2 = (abs(min(L2i.offsets)), abs(max(L2i.offsets)))
print_step = max(1, int(n / 10))
to_percent = 100.0 / n
utils.tic("Crank:")
R = R1 + R2
normal_shape = V.shape
transposed_shape = normal_shape[::-1]
for k in xrange(n):
if not k % print_step:
if isnan(V).any():
print "Crank fail @ t = %f (%i steps)" % (dt * k, k)
return crumbs
print int(k * to_percent),
if callback is not None:
callback(V, ((n - k) * dt))
V = (L2e.dot(V.flat).reshape(normal_shape) + R).T
V = spl.solve_banded(offsets1, L1i.data, V.flat, overwrite_b=True)
V = (L1e.dot(V).reshape(transposed_shape).T) + R
V = spl.solve_banded(offsets2, L2i.data, V.flat, overwrite_b=True).reshape(normal_shape)
crumbs.append(V.copy())
utils.toc()
return crumbs
def __init__(self, a, b, y, alpha=0, beta=0):
y = np.asarray(y)
n = y.shape[0] - 1
h = (b - a)/n
coeff = np.zeros(n + 3, dtype=y.dtype)
# Solutions to boundary coeffcients of spline
coeff[1] = 1/6. * (y[0] - (alpha * h**2)/6) #C2 in paper
coeff[n + 1] = 1/6. * (y[n] - (beta * h**2)/6) #cn+2 in paper
# Compressed tridiagonal matrix
ab = np.ones((3, n - 1), dtype=float)
ab[0,0] = 0 # Because top row is upper diag with one less elem
ab[1, :] = 4
ab[-1,-1] = 0 # Because bottom row is lower diag with one less elem
B = y[1:-1].copy() #grabs elements y[1] - > y[n-2] for reduced array
B[0] -= coeff[1]
B[-1] -= coeff[n + 1]
coeff[2:-2] = la.solve_banded((1, 1), ab, B, overwrite_ab=True,
overwrite_b=True, check_finite=False)
coeff[0] = alpha * h**2/6. + 2 * coeff[1] - coeff[2]
coeff[-1] = beta * h**2/6. + 2 * coeff[-2] - coeff[-3]
self.a = a # Lower-bound of domain
self.b = b # Uppser-bound of domain
self.coeffs = coeff # Spline coefficients
self.is_complex = (y.dtype == complex) #Tells which dtype solver to use
def beamwarming(u, nt, dt, dx):
##Tridiagonal setup##
a = numpy.zeros_like(u)
b = numpy.ones_like(u)
c = numpy.zeros_like(u)
d = numpy.zeros_like(u)
un = numpy.zeros((nt,len(u)))
un[:]=u.copy()
for n in range(1,nt):
u[0] = 1
E = utoE(u)
au = utoA(u)
a[0] = -dt/(4*dx)*u[0]
a[1:] = -dt/(4*dx)*u[0:-1]
a[-1] = -dt/(4*dx)*u[-1]
#b is all ones
c[:-1] = dt/(4*dx)*u[1:]
print str(-.5*dt/dx*(E[2:]-E[0:-2])+dt/(4*dx)*(au[2:]-au[:-2]))
#print str(dt/(4*dx)*(au[2:]-au[:-2]))
d[1:-1] = u[1:-1]-.5*dt/dx*(E[2:]-E[0:-2])+dt/(4*dx)*(au[2:]-au[:-2])
###subtract a[0]*LHS B.C to 'fix' thomas algorithm
d[0] = u[0] - .5*dt/dx*(E[1]-E[0])+dt/(4*dx)*(au[1]-au[0]) - a[0]
ab = numpy.matrix([c,b,a])
u = linalg.solve_banded((1,1), ab, d)
u[0]=1
un[n] = u.copy()
return un
def __init__(self, xp, yp):
if np.any(np.diff(xp) < 0):
raise TypeError("xp must be in ascending order")
# n = len(x)
self.m = len(xp)
xk = xp[1:-1]
yk = yp[1:-1]
xkp = xp[2:]
ykp = yp[2:]
xkm = xp[:-2]
ykm = yp[:-2]
b = (ykp - yk) / (xkp - xk) - (yk - ykm) / (xk - xkm)
l = (xk - xkm) / 6.0
d = (xkp - xkm) / 3.0
u = (xkp - xk) / 6.0
# u[0] = 0.0 # non-existent entries
# l[-1] = 0.0
# solve for second derivatives
fpp = solve_banded((1, 1), np.matrix([u, d, l]), b)
self.fpp = np.concatenate([[0.0], fpp, [0.0]]) # natural spline
self.xp = xp
self.yp = yp
def __init__(self, pixbound, flux):
#- If pixbound and flux have same number of elements, assume they
#- are centers rather than edges
if len(pixbound) == len(flux):
pixbound = cen2bound(pixbound)
npix = len(flux)
# Test for correct argument dimensions:
if (len(pixbound) - npix) != 1:
raise PixSplineError('Need one more element in pixbound than in flux!')
# The array of "delta-x" values:
dxpix = pixbound[1:] - pixbound[:-1]
# Test for monotonic increase:
if dxpix.min() <= 0.:
raise PixSplineError('Pixel boundaries not monotonically increasing!')
self.npix = npix
self.pixbound = pixbound.copy()
self.dxpix = dxpix.copy()
self.xcen = 0.5 * (pixbound[1:] + pixbound[:-1]).copy()
self.flux = flux.copy()
maindiag = (dxpix[:-1] + dxpix[1:]) / 3.
offdiag = dxpix[1:-1] / 6.
upperdiag = np.append(0., offdiag)
lowerdiag = np.append(offdiag, 0.)
band_matrix = np.vstack((upperdiag, maindiag, lowerdiag))
# The right-hand side:
rhs = flux[1:] - flux[:-1]
# Solve the banded matrix for the slopes at the ducks:
acoeff = la.solve_banded((1,1), band_matrix, rhs)
self.duckslopes = np.append(np.append(0., acoeff), 0.)
def phi_np1(phi_n, x, dx, t, dt, M, eps, a, Df, verbose=False):
""" returns phi(x, t=n+1) given phi(x, t=n) using finite difference method.
Uses *backward* finite difference in time (implicit), central finite
difference in x
"""
mu = M * dt / (dx * dx)
eps2 = eps * eps
# Set up matrix from 1 to n-2, and handle x = 0 and x = n-1 via boundary
# conditions
alpha, beta = boundary(t)
boundary_vec = zeros_like(phi_n[1:-1])
boundary_vec[0] = -mu * eps2 * alpha
boundary_vec[-1] = -mu * eps2 * beta
fp = fprime(phi_n[1:-1], a, Df)
b = phi_n[1:-1] - boundary_vec + M * dt * fp
d = ones_like(phi_n[1:-1]) * (1 + 2 * mu * eps2)
u = ones_like(phi_n[1:-1]) * (-mu * eps2)
u[0] = 0
l = ones_like(phi_n[1:-1]) * (-mu * eps2)
l[-1] = 0
ab = np.matrix([u, d, l])
res = zeros_like(phi_n)
res[1:-1] = linalg.solve_banded((1, 1), ab, b)
res[0] = alpha
res[-1] = beta
return res
def _solve_for_angular_rates(self, dt, angular_rates, rotvecs):
angular_rate_first = angular_rates[0].copy()
A = _angular_rate_to_rotvec_dot_matrix(rotvecs)
A_inv = _rotvec_dot_to_angular_rate_matrix(rotvecs)
M = _create_block_3_diagonal_matrix(
2 * A_inv[1:-1] / dt[1:-1, None, None],
2 * A[1:-1] / dt[1:-1, None, None],
4 * (1 / dt[:-1] + 1 / dt[1:]))
b0 = 6 * (rotvecs[:-1] * dt[:-1, None] ** -2 +
rotvecs[1:] * dt[1:, None] ** -2)
b0[0] -= 2 / dt[0] * A_inv[0].dot(angular_rate_first)
b0[-1] -= 2 / dt[-1] * A[-1].dot(angular_rates[-1])
for iteration in range(self.MAX_ITER):
rotvecs_dot = _matrix_vector_product_of_stacks(A, angular_rates)
delta_beta = _angular_acceleration_nonlinear_term(
rotvecs[:-1], rotvecs_dot[:-1])
b = b0 - delta_beta
angular_rates_new = solve_banded((5, 5), M, b.ravel())
angular_rates_new = angular_rates_new.reshape((-1, 3))
delta = np.abs(angular_rates_new - angular_rates[:-1])
angular_rates[:-1] = angular_rates_new
if np.all(delta < self.TOL * (1 + np.abs(angular_rates_new))):
break
rotvecs_dot = _matrix_vector_product_of_stacks(A, angular_rates)
angular_rates = np.vstack((angular_rate_first, angular_rates[:-1]))
return angular_rates, rotvecs_dot
def solve_banded(self, rhs): rhsex = np.zeros(self.M) rhsex[::self.nBlockSize] = rhs if not hasattr(self, 'Aex_diag'): Aex_diag = sp.dia_matrix(self.Aex) solex = solve_banded((self.m + 1,self.m + 1), Aex_diag.data[::-1,:], rhsex) return solex[::self.nBlockSize]
def fit_block_coordinate_desc(self, r_init=5.0, min_delta_r=0.00000001):
F_M = np.concatenate(self.F_M)
F_N = np.concatenate(self.F_N)
r_vals = []
error_vals = []
r = r_init
delta_r = None
it = 0
ab = ab_from_T(self.T, self.lam, self.dt)
while delta_r is None or delta_r > min_delta_r:
F_C = solve_banded((1, 1), ab, F_M - r * F_N)
new_r = - np.sum((F_C - F_M) * F_N) / np.sum(np.square(F_N))
error = self.estimate_error(new_r)
error_vals.append(error)
r_vals.append(new_r)
if r is not None:
delta_r = np.abs(r - new_r) / r
r = new_r
it += 1
self.r_vals = r_vals
self.error_vals = error_vals
self.r = r_vals[-1]
self.error = error_vals.min()
def test_check_finite(self):
a = array([[1.0, 20, 0, 0], [-30, 4, 6, 0], [2, 1, 20, 2], [0, -1, 7, 14]])
ab = array([[0.0, 20, 6, 2], [1, 4, 20, 14], [-30, 1, 7, 0], [2, -1, 0, 0]])
l, u = 2, 1
b4 = array([10.0, 0.0, 2.0, 14.0])
x = solve_banded((l, u), ab, b4, check_finite=False)
assert_array_almost_equal(dot(a, x), b4)
def integrate_outwards(self, u_g, rgd, vr_g, g0, e,
scalar_relativistic=False, pt_g=None):
l = self.l
r_g = rgd.r_g
cm1_g, c0_g, cp1_g = coefs(rgd, l, vr_g, e, scalar_relativistic)
c_xg = np.zeros((3, g0 + 2))
c_xg[0, :2] = 1.0
c_xg[0, 2:] = cp1_g[1:g0 + 1]
c_xg[1, 1:-1] = c0_g[1:g0 + 1]
c_xg[2, :-2] = cm1_g[1:g0 + 1]
b_g = np.zeros(g0 + 2)
if pt_g is not None:
b_g[2:] = -2 * pt_g[1:g0 + 1] * r_g[1:g0 + 1]**(1 - l)
a0 = pt_g[1] / r_g[1]**l / (vr_g[1] / r_g[1] - e)
else:
a0 = 1
a1 = a0 + vr_g[0] * rgd.dr_g[0]
b_g[:2] = [a0, a1]
a_g = solve_banded((2, 0), c_xg, b_g,
overwrite_ab=True, overwrite_b=True)
r = r_g[g0]
dr = rgd.dr_g[g0]
da = 0.5 * (a_g[g0 + 1] - a_g[g0 - 1])
dudr = (l + 1) * r**l * a_g[g0] + r**(l + 1) * da / dr
u_g[:g0 + 2] = a_g * r_g[:g0 + 2]**(l + 1)
return dudr
def diffuseImplicit(gr, phi, k, dt):
""" diffuse phi implicitly through timestep dt """
phinew = gr.scratchArray()
alpha = k*dt/gr.dx**2
# create the RHS of the matrix
R = phi[gr.ilo:gr.ihi+1]
# create the diagonal, d+1 and d-1 parts of the matrix
d = (1.0 + 2.0*alpha)*numpy.ones(gr.nx)
u = -alpha*numpy.ones(gr.nx)
u[0] = 0.0
l = -alpha*numpy.ones(gr.nx)
l[gr.nx-1] = 0.0
# set the boundary conditions by changing the matrix elements
# homogeneous neumann
d[0] = 1.0 + alpha
d[gr.nx-1] = 1.0 + alpha
# solve
A = numpy.matrix([u,d,l])
phinew[gr.ilo:gr.ihi+1] = linalg.solve_banded((1,1), A, R)
return phinew
def radau(alpha,beta,xr):
"""
Compute the Radau nodes and weights with the preassigned node xr
Inputs:
alpha - recursion coefficients
beta - recursion coefficients
xr - assigned node location
Outputs:
x - quadrature nodes
w - quadrature weights
Based on the section 7 of the paper "Some modified matrix eigenvalue
problems" by Gene Golub, SIAM Review Vol 15, No. 2, April 1973, pp.318--334
"""
from scipy.linalg import solve_banded
n = len(alpha)-1
f = np.zeros(n)
f[-1] = beta[-1]
A = np.vstack((np.sqrt(beta),alpha-xr))
J = np.vstack((A[:,0:-1],A[0,1:]))
delta = solve_banded((1,1),J,f)
alphar = alpha
alphar[-1] = xr+delta[-1]
x,w = gauss(alphar,beta)
return x,w
def _forward_control_points(self, dps):
num_cps = dps.shape[0] - 1;
ab = np.matrix([
np.tile(1, num_cps),
np.tile(4, num_cps),
np.tile(1, num_cps)
])
# Fixup first and last row
ab[0,0] = 0
ab[1,0] = 2
ab[1, num_cps - 1] = 7
ab[2, num_cps - 1] = 0
b = np.empty([num_cps, dps.shape[1]])
for i in range(0, num_cps - 1):
b[i] = 4 * dps[i] + 2 * dps[i + 1]
# Fixup first and last element
b[0] = dps[0] + 2 * dps[1]
b[num_cps - 1] = 8 * dps[num_cps - 1] + dps[num_cps]
return solve_banded((1, 1), ab, b)
def __call__(self, neuron):
'''
Updates the state variables.
'''
if not self._isprepared:
self.prepare()
self._isprepared=True
# Update the membrane currents (should it be after (no real difference though))
self._state_updater(neuron)
'''
x=solve_banded((lower,upper),ab,b)
lower = number of lower diagonals = 1
upper = number of upper diagonals = 1
ab = array(l+u+1,M)
each row is one diagonal
a[i,j]=ab[u+i-j,j]
'''
b=-neuron.Cm/neuron.clock.dt*neuron.v-neuron._I0
ab = zeros((3,len(neuron))) # part of it could be precomputed
# a[i,j]=ab[1+i-j,j]
# ab[1,:]=main diagonal
# ab[0,1:]=upper diagonal
# ab[2,:-1]=lower diagonal
ab[0,1:]=self.Aplus
ab[2,:-1]=self.Aminus
ab[1,:]=-neuron.Cm/neuron.clock.dt-neuron._gtot
ab[1,1:]-=self.Aminus
ab[1,:-1]-=self.Aplus
neuron.v=solve_banded((1,1),ab,b,overwrite_ab=True,overwrite_b=True)
def solve(fk):
N = len(fk)+2
k = ST.wavenumbers(N)
if solver == "banded":
A = np.zeros((N-2, N-2))
A[-1, :] = -2*np.pi*(k+1)*(k+2)
for i in range(2, N-2, 2):
A[-i-1, i:] = -4*np.pi*(k[:-i]+1)
uk_hat = solve_banded((0, N-3), A, fk)
elif solver == "sparse":
aij = [-2*np.pi*(k+1)*(k+2)]
for i in range(2, N-2, 2):
aij.append(np.array(-4*np.pi*(k[:-i]+1)))
A = diags(aij, range(0, N-2, 2))
uk_hat = la.spsolve(A, fk)
elif solver == "bs":
fc = fk.copy()
uk_hat = np.zeros(N-2)
uk_hat = SFTc.BackSubstitution_1D(uk_hat, fc)
#for i in range(N-3, -1, -1):
#for l in range(i+2, N-2, 2):
#fc[i] += (4*np.pi*(i+1)uk_hat[l])
#uk_hat[i] = -fc[i] / (2*np.pi*(i+1)*(i+2))
return uk_hat
def _get_natural_f(knots):
"""Returns mapping of natural cubic spline values to 2nd derivatives.
.. note:: See 'Generalized Additive Models', Simon N. Wood, 2006, pp 145-146
:param knots: The 1-d array knots used for cubic spline parametrization,
must be sorted in ascending order.
:return: A 2-d array mapping natural cubic spline values at
knots to second derivatives.
:raise ImportError: if scipy is not found, required for
``linalg.solve_banded()``
"""
try:
from scipy import linalg
except ImportError: # pragma: no cover
raise ImportError("Cubic spline functionality requires scipy.")
h = knots[1:] - knots[:-1]
diag = (h[:-1] + h[1:]) / 3.
ul_diag = h[1:-1] / 6.
banded_b = np.array([np.r_[0., ul_diag], diag, np.r_[ul_diag, 0.]])
d = np.zeros((knots.size - 2, knots.size))
for i in range(knots.size - 2):
d[i, i] = 1. / h[i]
d[i, i + 2] = 1. / h[i + 1]
d[i, i + 1] = - d[i, i] - d[i, i + 2]
fm = linalg.solve_banded((1, 1), banded_b, d)
return np.vstack([np.zeros(knots.size), fm, np.zeros(knots.size)])
def solve_tridag(A, D):
# Find the diagonals
ud = insert(diag(A, 1), 0, 0) # upper diagonal
d = diag(A) # main diagonal
ld = insert(diag(A, -1), len(d) - 1, 0) # lower diagonal
# simplified matrix
ab = matrix([ud, d, ld])
return solve_banded((1, 1), ab, D, overwrite_ab=True, overwrite_b=True)
def deriv(d,f,df):
bcast = np.ones(len(np.shape(f)),int); bcast[0]=5
for iy in range(0,ny+3):
iyc = iy*(iy>1 and iy < ny+1) + ny*(iy>=ny+1) + 2*(iy<=1)
df[iy,...] += np.sum(der[iy,:][d].reshape(bcast)*f[iyc-2:iyc+3,...],axis=0)
df[2]-=der[2,1]['d0']*df[1]+der[2,0]['d0']*df[0]; df[3]-=der[3,0]['d0']*df[1]
df[ny]-=der[ny,3]['d0']*df[ny+1]+der[ny,4]['d0']*df[ny+2]; df[ny-1]-=der[ny-1,4]['d0']*df[ny+1]
df[2:ny+1]=solve_banded((2,2),d0mat,df[2:ny+1])
def symmetric_solve_simple(P, L, tridiagonal, cell_sizes, free_values, alpha=(1. + sqrt(17)) / 8):
dtype=tridiagonal.dtype
z = linalg.solve_triangular(L, P.T.dot(free_values), lower=True, unit_diagonal=True)
#print tridiagonal
w = linalg.solve_banded((1, 1), tridiagonal, z)
tri_inv = tridiagonal_inversion(tridiagonal, cell_sizes, dtype=dtype)
w1 = tridiagonal_dot(tri_inv, z, dtype=dtype)
y = linalg.solve_triangular(np.matrix(L, dtype=dtype).getH(), w, lower=False, unit_diagonal=True)
return P.dot(y)
def __call__(self, neuron):
'''
Updates the state variables.
'''
if not self._isprepared:
self.prepare()
self._isprepared=True
# Update the membrane currents (should it be after (no real difference though))
self._state_updater(neuron)
'''
x=solve_banded((lower,upper),ab,b)
lower = number of lower diagonals = 1
upper = number of upper diagonals = 1
ab = array(l+u+1,M)
each row is one diagonal
a[i,j]=ab[u+i-j,j]
That is:
a[i,j]=ab[1+i-j,j]
ab[1,:]=main diagonal
ab[0,1:]=upper diagonal
ab[2,:-1]=lower diagonal '''
# Particular solution
b=-neuron.Cm/neuron.clock.dt*neuron.v-neuron._I0
ab = zeros((3,len(neuron)))
ab[:]=self.ab_star
ab[1,:]-=neuron._gtot
self.v_star[:]=solve_banded((1,1),ab,b,overwrite_ab=True,overwrite_b=True)
# Homogeneous solutions
b[:]=self.b_plus
ab[:]=self.ab_plus
ab[1,:]-=neuron._gtot
self.u_plus[:]=solve_banded((1,1),ab,b,overwrite_ab=True,overwrite_b=True)
b[:]=self.b_minus
ab[:]=self.ab_minus
ab[1,:]-=neuron._gtot
self.u_minus[:]=solve_banded((1,1),ab,b,overwrite_ab=True,overwrite_b=True)
# Solve the linear system connecting branches
self.P[:]=0
self.B[:]=0
self.fill_matrix(self.neuron.morphology)
self.V = solve(self.P,self.B)
# Calculate solutions by linear combination
self.linear_combination(self.neuron.morphology)
def solve_neumann(f):
A = np.zeros((N-2, N-2))
A[-1, :] = -2*np.pi*(k+1)*k**2/(k+2)
for i in range(2, N-3, 2):
A[-i-1, i:] = -4*np.pi*(k[:-i]+i)**2*(k[:-i]+1)/(k[:-i]+2)**2
uk_hat = solve_banded((0, N-4), A[1:, 1:], f)
return uk_hat
def test_tridiag_complex(self):
ab = array([[0.0, 20, 6, 2j], [1, 4, 20, 14], [-30, 1, 7, 0]])
a = np.diag(ab[0, 1:], 1) + np.diag(ab[1, :], 0) + np.diag(ab[2, :-1], -1)
b4 = array([10.0, 0.0, 2.0, 14.0j])
b4by1 = b4.reshape(-1, 1)
b4by2 = array([[2, 1], [-30, 4], [2, 3], [1, 3]])
b4by4 = array([[1, 0, 0, 0], [0, 0, 0, 1], [0, 1, 0, 0], [0, 1, 0, 0]])
for b in [b4, b4by1, b4by2, b4by4]:
x = solve_banded((1, 1), ab, b)
assert_array_almost_equal(dot(a, x), b)
def _solve_implicit_banded(current, banded_matrix):
# can improve performance by storing the banded form once and not
# recalculating it...
# but whatever
J = banded_matrix.shape[0]
diag = np.zeros((3, J))
diag[1, :] = np.diag(banded_matrix, k=0)
diag[0, 1:] = np.diag(banded_matrix, k=1)
diag[2, :-1] = np.diag(banded_matrix, k=-1)
return solve_banded((1, 1), diag, current)
def test_complex(self):
a = array([[1.0, 20, 0, 0], [-30, 4, 6, 0], [2j, 1, 20, 2j], [0, -1, 7, 14]])
ab = array([[0.0, 20, 6, 2j], [1, 4, 20, 14], [-30, 1, 7, 0], [2j, -1, 0, 0]])
l, u = 2, 1
b4 = array([10.0, 0.0, 2.0, 14.0j])
b4by1 = b4.reshape(-1, 1)
b4by2 = array([[2, 1], [-30, 4], [2, 3], [1, 3]])
b4by4 = array([[1, 0, 0, 0], [0, 0, 0, 1j], [0, 1, 0, 0], [0, 1, 0, 0]])
for b in [b4, b4by1, b4by2, b4by4]:
x = solve_banded((l, u), ab, b)
assert_array_almost_equal(dot(a, x), b)
def milcshake_a(dt, bond, r_old, r, v):
"""First part of velocity Verlet algorithm with constraints."""
# This subroutine iteratively adjusts the positions stored in the array r
# and the velocities stored in the array v, to satisfy the bond constraints
# using a tri-diagonal solver
# See AG Bailey, CP Lowe, and AP Sutton, J Comput Phys, 227, 8949 (2008)
# and AG Bailey, CP Lowe, and AP Sutton, Comput Phys Commun, 180, 594 (2009)
# On entry to this routine we assume:
# r_old stores the positions at the start of the step
# r stores the positions following the unconstrained drift and
# v stores the velocities following the first unconstrained half-kick
# The returned arrays r and v will hold the constrained values
import numpy as np
import scipy.linalg as la
n, d = r.shape
assert d == 3, 'r dimension error in milcshake_a'
k = n - 1 # Number of constraints
tol = 1.0e-9
iter_max = 500
r_new = r.copy() # Saves unconstrained positions
# Old and new (non-constrained) bond vectors
rij_old = r_old[:-1, :] - r_old[1:, :]
rij_new = r_new[:-1, :] - r_new[1:, :]
# Elements of tridiagonal matrix (dot products of old and new bond vectors)
# In this example, all masses are equal to unity. Let k=n-1 be number of constraints
tridiag = np.zeros((3, k), dtype=np.float_)
tridiag[0, 1:] = -2.0 * np.sum(
rij_old[1:, :] * rij_new[:-1, :],
axis=1) # leading zero to pad, then k-1 elements of upper-diagonal
tridiag[1, :] = 2.0 * np.sum(rij_old[:, :] * rij_new[:, :],
axis=1) / 0.5 # k elements of diagonal
tridiag[2, :-1] = -2.0 * np.sum(
rij_old[:-1, :] * rij_new[1:, :],
axis=1) # k-1 elements of lower-diagonal, then trailing zero to pad
# Set up rhs of constraint equation
rijsq = np.sum(rij_new**2, axis=1)
rhs = bond**2 - rijsq
rhsold = rhs.copy()
iter = 0
while True: # Iterative loop until done
# Test for done
max_error = np.max(np.fabs(rijsq - bond**2)) / (2.0 * bond**2)
if max_error <= tol:
break
# Reset tridiagonal elements (may have been over-written by solver)
tridiag_tmp = tridiag.copy()
lam = la.solve_banded((1, 1), tridiag_tmp, rhs)
# Constraint effects on position from lambda multipliers
r = r_new.copy()
r[:-1, :] = r[:-1, :] + lam[:, np.newaxis] * rij_old
r[1:, :] = r[1:, :] - lam[:, np.newaxis] * rij_old
# New bond vectors
rij = r[:-1, :] - r[1:, :]
# Prepare for next iteration
rijsq = np.sum(rij**2, axis=1)
rhs = bond**2 - rijsq + rhsold
rhsold = rhs.copy()
iter = iter + 1
assert iter <= iter_max, "{}{:15d}{:15d}".format(
'Too many iterations', iter, iter_max)
# Effect of constraints on velocities
v[:-1, :] = v[:-1, :] + lam[:, np.newaxis] * rij_old / dt
v[1:, :] = v[1:, :] - lam[:, np.newaxis] * rij_old / dt
return r, v
def FDM_Vanilla_Implicit(self, Ns, Nt, m):
'''
Abstract:
--------
Finite difference method for vanilla option.
Trivial implicit method.
Parameters:
----------
Ns: Number of points in price axis
Nt: Number of points in time axis
m : monitoring times
'''
# discretize Nt-1 points between every two monitoring time, total Nt*m + 1 gird in time axis
step = Nt
Nt = Nt * m
# set up parameters
mu = self.r - self.q
_range = 3 * self.sig * np.sqrt(self.maturity)
Smax = self.spot_price * np.exp(
(mu - self.sig**2 / 2.0) * self.maturity + _range)
Smin = 0
# totally Nt + 1 in row grid
dt = self.maturity / float(Nt)
ds = (Smax - Smin) / float(Ns) # totally Ns + 1 in column grid
# initialize the payoff
sGrid = np.linspace(Smin, Smax, Ns + 1)
if self.position.upper() == "CALL":
V_Nt = np.maximum(sGrid - self.strike, 0)
elif self.position.upper() == "PUT":
V_Nt = np.maximum(self.strike - sGrid, 0)
s = np.linspace(Smin, Smax, Ns + 1)
payoff = np.maximum(s - self.strike, 0)
# initialize the Dirichlet boundary condition
if self.position.upper() == "CALL":
f_0 = np.linspace(0, 0, Nt + 1)
f_Ns = Smax * np.exp(-self.q * np.linspace(
0, self.maturity, Nt + 1)) - self.strike * np.exp(
-self.r * (np.linspace(0, self.maturity, Nt + 1)))
elif self.position.upper() == "PUT":
f_0 = self.strike * np.exp(
-self.r * np.linspace(0, self.maturity, Nt + 1))
f_Ns = np.linspace(0, 0, Nt + 1)
else:
raise ValueError('Invalid option_type!!')
# initialize the tridiagonal matrix by scalar-form
i = np.linspace(1, Ns - 1, Ns - 1)
# from a_2 to a_I-1 are in the calculation matrix
a = -(self.sig**2 * i**2 - (self.r - self.q) * i) * dt / 2.0
# from b_1 to b_I-1 are in the calculation matrix
b = 1 + self.sig**2 * i**2 * dt + self.r * dt
# from c_1 to c_I-2 are in the calculation matrix
c = -(self.sig**2 * i**2 + (self.r - self.q) * i) * dt / 2.0
# From Nt to 1, calculate V_Nt-1, V_Nt-2, ..., V_0 (vectors)
V_Nplus = V_Nt[1:Ns]
for k in range(Nt, 0, -1):
V_Nplus[0] = V_Nplus[0] - a[0] * f_0[k]
V_Nplus[Ns - 2] = V_Nplus[Ns - 2] - c[Ns - 2] * f_Ns[k]
ab = self.tri_bound_to_ab(a, b, c)
V_N = linalg.solve_banded((1, 1), ab, V_Nplus)
#American process
if self.exercise_type == 'AMERICAN':
V_N = self.Projected_SOR(a[1:], b, c[:-1], V_Nplus, V_N,
payoff[1:-1], k, step, s[1:Ns])
V_Nplus = V_N
# linear interpolation
w = (self.spot_price - sGrid[int(self.spot_price / ds)]) / (
sGrid[int(self.spot_price / ds) + 1] -
sGrid[int(self.spot_price / ds)])
v_0 = V_N[int(self.spot_price /
ds)] * (1 - w) + w * V_N[int(self.spot_price / ds) + 1]
return v_0
def __init__(self, x, y, axis=0, bc_type='not-a-knot', extrapolate=None):
x, dx, y, axis, _ = prepare_input(x, y, axis)
n = len(x)
bc, y = self._validate_bc(bc_type, y, y.shape[1:], axis)
if extrapolate is None:
if bc[0] == 'periodic':
extrapolate = 'periodic'
else:
extrapolate = True
dxr = dx.reshape([dx.shape[0]] + [1] * (y.ndim - 1))
slope = np.diff(y, axis=0) / dxr
# If bc is 'not-a-knot' this change is just a convention.
# If bc is 'periodic' then we already checked that y[0] == y[-1],
# and the spline is just a constant, we handle this case in the same
# way by setting the first derivatives to slope, which is 0.
if n == 2:
if bc[0] in ['not-a-knot', 'periodic']:
bc[0] = (1, slope[0])
if bc[1] in ['not-a-knot', 'periodic']:
bc[1] = (1, slope[0])
# This is a very special case, when both conditions are 'not-a-knot'
# and n == 3. In this case 'not-a-knot' can't be handled regularly
# as the both conditions are identical. We handle this case by
# constructing a parabola passing through given points.
if n == 3 and bc[0] == 'not-a-knot' and bc[1] == 'not-a-knot':
A = np.zeros((3, 3)) # This is a standard matrix.
b = np.empty((3, ) + y.shape[1:], dtype=y.dtype)
A[0, 0] = 1
A[0, 1] = 1
A[1, 0] = dx[1]
A[1, 1] = 2 * (dx[0] + dx[1])
A[1, 2] = dx[0]
A[2, 1] = 1
A[2, 2] = 1
b[0] = 2 * slope[0]
b[1] = 3 * (dxr[0] * slope[1] + dxr[1] * slope[0])
b[2] = 2 * slope[1]
s = solve(A,
b,
overwrite_a=True,
overwrite_b=True,
check_finite=False)
else:
# Find derivative values at each x[i] by solving a tridiagonal
# system.
A = np.zeros((3, n)) # This is a banded matrix representation.
b = np.empty((n, ) + y.shape[1:], dtype=y.dtype)
# Filling the system for i=1..n-2
# (x[i-1] - x[i]) * s[i-1] +\
# 2 * ((x[i] - x[i-1]) + (x[i+1] - x[i])) * s[i] +\
# (x[i] - x[i-1]) * s[i+1] =\
# 3 * ((x[i+1] - x[i])*(y[i] - y[i-1])/(x[i] - x[i-1]) +\
# (x[i] - x[i-1])*(y[i+1] - y[i])/(x[i+1] - x[i]))
A[1, 1:-1] = 2 * (dx[:-1] + dx[1:]) # The diagonal
A[0, 2:] = dx[:-1] # The upper diagonal
A[-1, :-2] = dx[1:] # The lower diagonal
b[1:-1] = 3 * (dxr[1:] * slope[:-1] + dxr[:-1] * slope[1:])
bc_start, bc_end = bc
if bc_start == 'periodic':
# In case when number of points is 3 we should count derivatives
# manually
if n == 3:
t = (slope / dx).sum() / (1. / dx).sum()
s = [t, t, t]
else:
# Due to the periodicity, and because y[-1] = y[0], the linear
# system has (n-1) unknowns/equations instead of n:
A = A[:, 0:-1]
A[1, 0] = 2 * (dx[-1] + dx[0])
A[0, 1] = dx[-1]
b = b[:-1]
# Also, due to the periodicity, the system is not tri-diagonal.
# We need to compute a "condensed" matrix of shape (n-2, n-2).
# See https://web.archive.org/web/20151220180652/http://www.cfm.brown.edu/people/gk/chap6/node14.html
# for more explanations.
# The condensed matrix is obtained by removing the last column
# and last row of the (n-1, n-1) system matrix. The removed
# values are saved in scalar variables with the (n-1, n-1)
# system matrix indices forming their names:
a_m1_0 = dx[-2] # lower left corner value: A[-1, 0]
a_m1_m2 = dx[-1]
a_m1_m1 = 2 * (dx[-1] + dx[-2])
a_m2_m1 = dx[-3]
a_0_m1 = dx[0]
b[0] = 3 * (dxr[0] * slope[-1] + dxr[-1] * slope[0])
b[-1] = 3 * (dxr[-1] * slope[-2] + dxr[-2] * slope[-1])
Ac = A[:, :-1]
b1 = b[:-1]
b2 = np.zeros_like(b1)
b2[0] = -a_0_m1
b2[-1] = -a_m2_m1
# s1 and s2 are the solutions of (n-2, n-2) system
s1 = solve_banded((1, 1),
Ac,
b1,
overwrite_ab=False,
overwrite_b=False,
check_finite=False)
s2 = solve_banded((1, 1),
Ac,
b2,
overwrite_ab=False,
overwrite_b=False,
check_finite=False)
# computing the s[n-2] solution:
s_m1 = ((b[-1] - a_m1_0 * s1[0] - a_m1_m2 * s1[-1]) /
(a_m1_m1 + a_m1_0 * s2[0] + a_m1_m2 * s2[-1]))
# s is the solution of the (n, n) system:
s = np.empty((n, ) + y.shape[1:], dtype=y.dtype)
s[:-2] = s1 + s_m1 * s2
s[-2] = s_m1
s[-1] = s[0]
else:
if bc_start == 'not-a-knot':
A[1, 0] = dx[1]
A[0, 1] = x[2] - x[0]
d = x[2] - x[0]
b[0] = ((dxr[0] + 2 * d) * dxr[1] * slope[0] +
dxr[0]**2 * slope[1]) / d
elif bc_start[0] == 1:
A[1, 0] = 1
A[0, 1] = 0
b[0] = bc_start[1]
elif bc_start[0] == 2:
A[1, 0] = 2 * dx[0]
A[0, 1] = dx[0]
b[0] = -0.5 * bc_start[1] * dx[0]**2 + 3 * (y[1] - y[0])
if bc_end == 'not-a-knot':
A[1, -1] = dx[-2]
A[-1, -2] = x[-1] - x[-3]
d = x[-1] - x[-3]
b[-1] = ((dxr[-1]**2 * slope[-2] +
(2 * d + dxr[-1]) * dxr[-2] * slope[-1]) / d)
elif bc_end[0] == 1:
A[1, -1] = 1
A[-1, -2] = 0
b[-1] = bc_end[1]
elif bc_end[0] == 2:
A[1, -1] = 2 * dx[-1]
A[-1, -2] = dx[-1]
b[-1] = 0.5 * bc_end[1] * dx[-1]**2 + 3 * (y[-1] - y[-2])
s = solve_banded((1, 1),
A,
b,
overwrite_ab=True,
overwrite_b=True,
check_finite=False)
super(CubicSpline, self).__init__(x,
y,
s,
axis=0,
extrapolate=extrapolate)
self.axis = axis
import numpy as np
import time
from scipy.sparse.linalg import bicg
from scipy.linalg import solve_banded
from scipy.sparse import dia_matrix
A = dia_matrix((np.array([[13, 14, 15, 28, 33, 44], [15, 16, 18, 22, 32, 42],
[17, 18, 19, 44, 71, 81]]), [0, -1, 1]),
shape=(6, 6)).toarray()
b = [1, 2, 3, 4, 5, 6]
start_time_1 = time.time()
print(bicg(A, b))
end_time_1 = (time.time() - start_time_1) * 1000
ud = np.insert(np.diag(A, 1), 0, 0)
d = np.diag(A)
ld = np.insert(np.diag(A, -1), len(d) - 1, 0)
ab = [ud, d, ld]
start_time_2 = time.time()
print(solve_banded((1, 1), ab, b))
end_time_2 = (time.time() - start_time_2) * 1000
print('bicg:', end_time_1, 'solve_banded:', end_time_2)
def __init__(self, x, y, axis=0, bc_type='not-a-knot', extrapolate=None):
x, y = map(np.asarray, (x, y))
if np.issubdtype(x.dtype, np.complexfloating):
raise ValueError("`x` must contain real values.")
if np.issubdtype(y.dtype, np.complexfloating):
dtype = complex
else:
dtype = float
y = y.astype(dtype, copy=False)
axis = axis % y.ndim
if x.ndim != 1:
raise ValueError("`x` must be 1-dimensional.")
if x.shape[0] < 2:
raise ValueError("`x` must contain at least 2 elements.")
if x.shape[0] != y.shape[axis]:
raise ValueError("The length of `y` along `axis`={0} doesn't "
"match the length of `x`".format(axis))
if not np.all(np.isfinite(x)):
raise ValueError("`x` must contain only finite values.")
if not np.all(np.isfinite(y)):
raise ValueError("`y` must contain only finite values.")
dx = np.diff(x)
if np.any(dx <= 0):
raise ValueError("`x` must be strictly increasing sequence.")
n = x.shape[0]
y = np.rollaxis(y, axis)
bc, y = self._validate_bc(bc_type, y, y.shape[1:], axis)
if extrapolate is None:
if bc[0] == 'periodic':
extrapolate = 'periodic'
else:
extrapolate = True
dxr = dx.reshape([dx.shape[0]] + [1] * (y.ndim - 1))
slope = np.diff(y, axis=0) / dxr
# If bc is 'not-a-knot' this change is just a convention.
# If bc is 'periodic' then we already checked that y[0] == y[-1],
# and the spline is just a constant, we handle this case in the same
# way by setting the first derivatives to slope, which is 0.
if n == 2:
if bc[0] in ['not-a-knot', 'periodic']:
bc[0] = (1, slope[0])
if bc[1] in ['not-a-knot', 'periodic']:
bc[1] = (1, slope[0])
# This is a very special case, when both conditions are 'not-a-knot'
# and n == 3. In this case 'not-a-knot' can't be handled regularly
# as the both conditions are identical. We handle this case by
# constructing a parabola passing through given points.
if n == 3 and bc[0] == 'not-a-knot' and bc[1] == 'not-a-knot':
A = np.zeros((3, 3)) # This is a standard matrix.
b = np.empty((3, ) + y.shape[1:], dtype=y.dtype)
A[0, 0] = 1
A[0, 1] = 1
A[1, 0] = dx[1]
A[1, 1] = 2 * (dx[0] + dx[1])
A[1, 2] = dx[0]
A[2, 1] = 1
A[2, 2] = 1
b[0] = 2 * slope[0]
b[1] = 3 * (dxr[0] * slope[1] + dxr[1] * slope[0])
b[2] = 2 * slope[1]
s = solve(A,
b,
overwrite_a=True,
overwrite_b=True,
check_finite=False)
else:
# Find derivative values at each x[i] by solving a tridiagonal
# system.
A = np.zeros((3, n)) # This is a banded matrix representation.
b = np.empty((n, ) + y.shape[1:], dtype=y.dtype)
# Filling the system for i=1..n-2
# (x[i-1] - x[i]) * s[i-1] +\
# 2 * ((x[i] - x[i-1]) + (x[i+1] - x[i])) * s[i] +\
# (x[i] - x[i-1]) * s[i+1] =\
# 3 * ((x[i+1] - x[i])*(y[i] - y[i-1])/(x[i] - x[i-1]) +\
# (x[i] - x[i-1])*(y[i+1] - y[i])/(x[i+1] - x[i]))
A[1, 1:-1] = 2 * (dx[:-1] + dx[1:]) # The diagonal
A[0, 2:] = dx[:-1] # The upper diagonal
A[-1, :-2] = dx[1:] # The lower diagonal
b[1:-1] = 3 * (dxr[1:] * slope[:-1] + dxr[:-1] * slope[1:])
bc_start, bc_end = bc
if bc_start == 'periodic':
# Due to the periodicity, and because y[-1] = y[0], the linear
# system has (n-1) unknowns/equations instead of n:
A = A[:, 0:-1]
A[1, 0] = 2 * (dx[-1] + dx[0])
A[0, 1] = dx[-1]
b = b[:-1]
# Also, due to the periodicity, the system is not tri-diagonal.
# We need to compute a "condensed" matrix of shape (n-2, n-2).
# See https://web.archive.org/web/20151220180652/http://www.cfm.brown.edu/people/gk/chap6/node14.html
# for more explanations.
# The condensed matrix is obtained by removing the last column
# and last row of the (n-1, n-1) system matrix. The removed
# values are saved in scalar variables with the (n-1, n-1)
# system matrix indices forming their names:
a_m1_0 = dx[-2] # lower left corner value: A[-1, 0]
a_m1_m2 = dx[-1]
a_m1_m1 = 2 * (dx[-1] + dx[-2])
a_m2_m1 = dx[-2]
a_0_m1 = dx[0]
b[0] = 3 * (dxr[0] * slope[-1] + dxr[-1] * slope[0])
b[-1] = 3 * (dxr[-1] * slope[-2] + dxr[-2] * slope[-1])
Ac = A[:, :-1]
b1 = b[:-1]
b2 = np.zeros_like(b1)
b2[0] = -a_0_m1
b2[-1] = -a_m2_m1
# s1 and s2 are the solutions of (n-2, n-2) system
s1 = solve_banded((1, 1),
Ac,
b1,
overwrite_ab=False,
overwrite_b=False,
check_finite=False)
s2 = solve_banded((1, 1),
Ac,
b2,
overwrite_ab=False,
overwrite_b=False,
check_finite=False)
# computing the s[n-2] solution:
s_m1 = ((b[-1] - a_m1_0 * s1[0] - a_m1_m2 * s1[-1]) /
(a_m1_m1 + a_m1_0 * s2[0] + a_m1_m2 * s2[-1]))
# s is the solution of the (n, n) system:
s = np.empty((n, ) + y.shape[1:], dtype=y.dtype)
s[:-2] = s1 + s_m1 * s2
s[-2] = s_m1
s[-1] = s[0]
else:
if bc_start == 'not-a-knot':
A[1, 0] = dx[1]
A[0, 1] = x[2] - x[0]
d = x[2] - x[0]
b[0] = ((dxr[0] + 2 * d) * dxr[1] * slope[0] +
dxr[0]**2 * slope[1]) / d
elif bc_start[0] == 1:
A[1, 0] = 1
A[0, 1] = 0
b[0] = bc_start[1]
elif bc_start[0] == 2:
A[1, 0] = 2 * dx[0]
A[0, 1] = dx[0]
b[0] = -0.5 * bc_start[1] * dx[0]**2 + 3 * (y[1] - y[0])
if bc_end == 'not-a-knot':
A[1, -1] = dx[-2]
A[-1, -2] = x[-1] - x[-3]
d = x[-1] - x[-3]
b[-1] = ((dxr[-1]**2 * slope[-2] +
(2 * d + dxr[-1]) * dxr[-2] * slope[-1]) / d)
elif bc_end[0] == 1:
A[1, -1] = 1
A[-1, -2] = 0
b[-1] = bc_end[1]
elif bc_end[0] == 2:
A[1, -1] = 2 * dx[-1]
A[-1, -2] = dx[-1]
b[-1] = 0.5 * bc_end[1] * dx[-1]**2 + 3 * (y[-1] - y[-2])
s = solve_banded((1, 1),
A,
b,
overwrite_ab=True,
overwrite_b=True,
check_finite=False)
# Compute coefficients in PPoly form.
t = (s[:-1] + s[1:] - 2 * slope) / dxr
c = np.empty((4, n - 1) + y.shape[1:], dtype=t.dtype)
c[0] = t / dxr
c[1] = (slope - s[:-1]) / dxr - t
c[2] = s[:-1]
c[3] = y[:-1]
super(CubicSpline, self).__init__(c, x, extrapolate=extrapolate)
self.axis = axis
def test_1x1(self):
x = solve_banded((1, 1), [[0], [1], [0]], [[1, 2, 3]])
assert_array_equal(x, [[1.0, 2.0, 3.0]])
assert_equal(x.dtype, np.dtype('f8'))
def cubic_with_deriv(x, xp, yp):
"""deprecated"""
x, n = _checkIfFloat(x)
if np.any(np.diff(xp) < 0):
raise TypeError("xp must be in ascending order")
# n = len(x)
m = len(xp)
y = np.zeros(n)
dydx = np.zeros(n)
dydxp = np.zeros((n, m))
dydyp = np.zeros((n, m))
xk = xp[1:-1]
yk = yp[1:-1]
xkp = xp[2:]
ykp = yp[2:]
xkm = xp[:-2]
ykm = yp[:-2]
b = (ykp - yk) / (xkp - xk) - (yk - ykm) / (xk - xkm)
l = (xk - xkm) / 6.0
d = (xkp - xkm) / 3.0
u = (xkp - xk) / 6.0
# u[0] = 0.0 # non-existent entries
# l[-1] = 0.0
# solve for second derivatives
fpp = solve_banded((1, 1), np.array([u, d, l]), b)
fpp = np.concatenate([[0.0], fpp, [0.0]]) # natural spline
# find location in vector
for i in range(n):
if x[i] < xp[0]:
j = 0
elif x[i] > xp[-1]:
j = m - 2
else:
for j in range(m - 1):
if xp[j + 1] > x[i]:
break
x1 = xp[j]
y1 = yp[j]
x2 = xp[j + 1]
y2 = yp[j + 1]
A = (x2 - x[i]) / (x2 - x1)
B = 1 - A
C = 1.0 / 6 * (A ** 3 - A) * (x2 - x1) ** 2
D = 1.0 / 6 * (B ** 3 - B) * (x2 - x1) ** 2
y[i] = A * y1 + B * y2 + C * fpp[j] + D * fpp[j + 1]
dAdx = -1.0 / (x2 - x1)
dBdx = -dAdx
dCdx = 1.0 / 6 * (3 * A ** 2 - 1) * dAdx * (x2 - x1) ** 2
dDdx = 1.0 / 6 * (3 * B ** 2 - 1) * dBdx * (x2 - x1) ** 2
dydx[i] = dAdx * y1 + dBdx * y2 + dCdx * fpp[j] + dDdx * fpp[j + 1]
if n == 1:
y = y[0]
dydx = dydx[0]
return y
def mono(tau, ssa, hg, rsurf, MU0, F0PI=1.0, BTOP=0.0):
"""Monochromatic two-stream radiative transfer solver with:
tau - optical depths per layer
ssa - single scattering albedo
hg - scattering asymetry parameter
rsurf - surface reflectivity
MU0 - cos(theta) where theta is both the
incidence an emission angle.
F0PI - top of atmosphere incident flux * pi
BTOP - top of the atmosphere diffuse flux.
"""
import numpy as np
nlay = len(tau)
# Cumulative optical depth
taub = np.cumsum(tau)
taut = np.append(0, taub[:-1])
# Surface reflectance and lower boundary condition
bsurf = rsurf * MU0[-1] * F0PI * np.exp(-np.sum(tau) / MU0[-1])
twopi = np.pi + np.pi
g1 = 0.86602540378 * (2. - ssa * (1 + hg))
g2 = (1.7320508075688772 * ssa / 2.) * (1 - hg)
g3 = (1. - 1.7320508075688772 * hg * MU0) / 2
g4 = 1. - g3
lam = np.sqrt(g1 * g1 - g2 * g2)
gamma = (g1 - lam) / g2
alpha = np.sqrt((1. - ssa) / (1. - ssa * hg))
Am = F0PI * ssa * (g4 * (g1 + 1. / MU0) + g2 * g3) / (lam * lam - 1. /
(MU0 * MU0))
Ap = F0PI * ssa * (g3 * (g1 - 1. / MU0) + g2 * g4) / (lam * lam - 1. /
(MU0 * MU0))
# Cpm1 and Cmm1 are the C+ and C- terms evaluated at the top of the layer.
Cpm1 = Ap * np.exp(-taut / MU0)
Cmm1 = Am * np.exp(-taut / MU0)
# Cp and Cm are the C+ and C- terms evaluated at the bottom of the layer.
Cp = Ap * np.exp(-taub / MU0)
Cm = Am * np.exp(-taub / MU0)
# Solve for the coefficients of system of equations using boundary conditions
# Exponential terms:
exptrm = lam * tau
exptrm[exptrm > 35] = 35 # clipped so that exponential doesn't explode
Ep = np.exp(exptrm)
Em = 1. / Ep
E1 = Ep + gamma * Em
E2 = Ep - gamma * Em
E3 = gamma * Ep + Em
E4 = gamma * Ep - Em
L = nlay + nlay
Af = np.empty(L)
Bf = np.empty(L)
Cf = np.empty(L)
Df = np.empty(L)
# First Term
Af[0] = 0.0
Bf[0] = gamma[0] + 1.
Cf[0] = gamma[0] - 1.
Df[0] = BTOP - Cmm1[0]
AA = (E1[:-1] + E3[:-1]) * (gamma[1:] - 1)
BB = (E2[:-1] + E4[:-1]) * (gamma[1:] - 1)
CC = 2. * (1. - gamma[1:] * gamma[1:])
DD = (gamma[1:] - 1) * (Cpm1[1:] - Cp[:-1]) + (1 - gamma[1:]) * (Cm[:-1] -
Cmm1[1:])
Af[1:-1:2] = AA
Bf[1:-1:2] = BB
Cf[1:-1:2] = CC
Df[1:-1:2] = DD
AA = 2. * (1. - gamma[:-1] * gamma[:-1])
BB = (E1[:-1] - E3[:-1]) * (gamma[1:] + 1.)
CC = (E1[:-1] + E3[:-1]) * (gamma[1:] - 1.)
DD = E3[:-1] * (Cpm1[1:] - Cp[:-1]) + E1[:-1] * (Cm[:-1] - Cmm1[1:])
Af[2::2] = AA
Bf[2::2] = BB
Cf[2::2] = CC
Df[2::2] = DD
# Last term:
Af[-1] = E1[-1] - rsurf * E3[-1]
Bf[-1] = E2[-1] - rsurf * E4[-1]
Cf[-1] = 0.0
Df[-1] = bsurf - Cp[-1] + rsurf * Cm[-1]
# Match to IDL
Af2 = np.append(Af[1:], 0)
Cf2 = np.append(0, Cf[0:-1])
# Solve tridiagonal matrix
ab = np.matrix([
Cf2,
Bf,
Af2,
])
k = solve_banded((1, 1), ab, Df)
# Unmix coefficients
even = np.arange(0, 2 * nlay, 2)
odd = even + 1
k1 = k[even] + k[odd]
k2 = k[even] - k[odd]
# Iteratively determine upward intensity (along mu) in each layer
Fsurf = k1[-1] * Ep[-1] + gamma[-1] * k2[-1] * Em[-1] + Cp[
-1] # flux at surface
Isurf = Fsurf / np.pi # Lambert intensity distribution
Am = Am * np.exp(-taut / MU0)
Ap = Ap * np.exp(-taut / MU0)
F1 = 1 + 1.5 * hg * MU0
F2 = 1 - 1.5 * hg * MU0
G = (ssa * k1 * (F1 + gamma * F2)) / (twopi)
H = (ssa * k2 * (gamma * F1 + F2)) / (twopi)
A = (ssa * (F1 * Ap + F2 * Am)) / (twopi)
# Integrate upward intensity starting from surface
Imu = np.empty(nlay + 1) # upward intensity at each level (not layer)
Imu[-1] = Isurf # upward intensity from surface
addterms = ((ssa * F0PI / (8 * np.pi)) * (1 - hg) * np.exp(-taut / MU0) *
(1 - np.exp(-2. * tau / MU0)) + A *
(1 - np.exp(-2 * tau / MU0)) / 2 + G *
(np.exp(exptrm - tau / MU0) - 1) / (lam * MU0 - 1) + H *
(1 - np.exp(-exptrm - tau / MU0)) / (lam * MU0 + 1))
multterms = np.exp(-tau / MU0)
for j in np.arange(nlay - 1, -1, -1):
Imu[j] = Imu[j + 1] * multterms[j] + addterms[j]
return Imu
error_counter = 0
plot_counter = 0
t = tmax
while t > dt:
for i in range(1, N[k] - 1):
rhs[i - 1] = alpha * u[i + 1] + (
1 - 2 * alpha) * u[i] + alpha * u[i - 1] + Q[i] * dt
Ab = np.zeros((3, N[k] - 2))
Ab[0, 1:] = upper_diag
Ab[1, 0:] = central_diag
Ab[2, :-1] = lower_diag
sol = solve_banded((1, 1), Ab, rhs)
for i in range(1, N[k] - 1):
u[i] = sol[i - 1]
t -= dt
if error_counter == 1000 - 1:
for i in range(N[k]):
error += pow(u[i] - S[i], 2)
difference.append(math.sqrt(error / (N[k])))
error = 0
error_counter = 0
error_counter += 1
def main():
# For timing purposes
start = time.time()
pr = cProfile.Profile()
pr.enable()
# Initial parameters
a1 = 1 + TIME_STEP * ((1j * H_BAR) / (2 * ELECTRON_MASS * (A_DIST**2)))
a2 = -TIME_STEP * H_BAR * 1j / (4 * ELECTRON_MASS * A_DIST**2)
b1 = 1 - 1j * TIME_STEP * H_BAR / (2 * ELECTRON_MASS * A_DIST**2)
b2 = -a2
# initial conditions
x_points = linspace(0, L, N_SLICES + 1, endpoint=False)
psi = array(list(map(psi_function, x_points)), complex)
psi[0] = psi[N_SLICES - 1] = 0
# Create the matrix A
A2 = empty([3, N_SLICES + 1], complex)
A2[0, 0] = 0
A2[0, 1:] = a2
A2[1, :] = a1
A2[2, 0:N_SLICES] = a2
A2[2, N_SLICES] = 0
# store the wave function at each time step in a list
solution = [psi]
for i in range(500):
psi[1:N_SLICES] = b1 * psi[1:N_SLICES] + b2 * (psi[2:] +
psi[0:N_SLICES - 1])
psi = solve_banded((1, 1), A2, psi)
solution.append(psi)
plt.plot(x_points, real(solution[25]))
plt.plot(x_points, real(solution[49])**2)
plt.plot(x_points, real(solution[250])**2)
plt.xlabel("x (m)")
plt.ylabel("$\psi(x)$")
plt.show()
# Animation
def init():
line.set_data(x_points, psi_function(xpoints, L, SIGMA, KAPPA))
return line,
A_2 = zeros([N_SLICES, N_SLICES], dtype=complex)
A_2 = TriDiag_Matrix(A_2, a1, a2)
B = zeros([N_SLICES, N_SLICES], dtype=complex)
B = TriDiag_Matrix(B, b1, b2, N_SLICES)
xpoints = (L / N_SLICES) * linspace(
0, N_SLICES, num=N_SLICES + 1, endpoint=False)
xpoints = delete(xpoints, 0)
def animate(i):
v = dot(B, psi[i])
psi_2.append(solve(A_2, v))
x = xpoints
y = psi[i]
i += 1
line.set_data(x, y)
return line,
# calling the animation function
# creating a blank window
# for the animation
psi_2 = [psi_function(xpoints, L, SIGMA, KAPPA)]
fig = plt.figure()
axis = plt.axes(xlim=(-1e-9, 11e-9), ylim=(-1, 1))
axis.set_xlabel('x')
axis.set_ylabel('$\psi(x)$')
axis.set_title('Crank-Nicolson Wave')
line, = axis.plot([], [], lw=.5)
animate_wave = animation.FuncAnimation(fig,
animate,
init_func=init,
frames=100000,
interval=500,
blit=True,
save_count=50)
my_writer = animation.PillowWriter(fps=30, codec='libx264', bitrate=2)
plt.show()
pr.disable()
# pr.print_stats(sort='time')
end = time.time()
print(f"Program took :{end - start} seconds to run")
# - Block matrices diagonal, triangular, periodic,...
# %% {"slideshow": {"slide_type": "fragment"}}
import scipy.linalg as spl
b=sp.ones(5)
A=sp.array([[1.,3.,0., 0.,0.],
[ 2.,1.,-4, 0.,0.],
[ 6.,1., 2,-3.,0.],
[ 0.,1., 4.,-2.,-3.],
[ 0.,0., 6.,-3., 2.]])
print("x=",spl.solve(A,b,sym_pos=False)) # LAPACK ( gesv ou posv )
AB=sp.array([[0.,3.,-4.,-3.,-3.],
[1.,1., 2.,-2., 2.],
[2.,1., 4.,-3., 0.],
[6.,1., 6., 0., 0.]])
print("x=",spl.solve_banded((2,1),AB,b)) # LAPACK ( gbsv )
# %% {"slideshow": {"slide_type": "slide"}}
P,L,U = spl.lu(A) # P A = L U
np.set_printoptions(precision=3)
for M in (P,L,U):
print(M, end="\n"+20*"-"+"\n")
# %% [markdown] {"slideshow": {"slide_type": "slide"}}
# ## CSC (Compressed Sparse Column)
#
# - All operations are optimized
# - Efficient "slicing" along axis=1.
# - Fast Matrix-vector product.
# - Conversion to other format could be costly.
#print ht
scheme = ImplicitEulerScheme(M,N, ht)
scheme.roll_back()
#print scheme
#scheme.U
import scipy as sp
from scipy.linalg import solve_banded
A = np.zeros((3, 5))
A[0, :] = np.ones(5) * 1. # 上对角线
A[1, :] = np.ones(5) * 3. # 对角线
A[2, :] = np.ones(5) * (-1.) # 下对角线
b = [1.,2.,3.,4.,5.]
x = solve_banded((1,1), A,b)
#print 'x = A^-1b = ',x
import scipy as sp
from scipy.linalg import solve_banded
for k in range(0,M):
udiag = - np.ones(N-1) * rho
ldiag = - np.ones(N-1) * rho
cdiag = np.ones(N-1) * (1.0 + 2. * rho)
mat = np.zeros((3,N-1))
mat[0,:] = udiag
#print mat[0,:]
mat[1,:] = cdiag
#print mat[1,:]
# create column vector {bb}
bb = np.zeros(m)
bb[0] = Ti
bb[1:m-1] = Ti
bb[m-1] = Ti + 2*Fo*Bi*(1 + b/(2*m))*Tinf
# check: ab and b
#print('ab \n', ab) # banded matrix of tridiagonal coefficient matrix [A]
#print('bb \n', bb) # vector {b}
# Solve using scipy.sparse.linalg.lsqr
# -------------------------------------------------------------------------
for i in range(1, nt+1):
T = sp.solve_banded((1, 1), ab, bb)
bb = T.copy()
bb[m-1] = T[m-1] + 2*Fo*Bi*(1 + b/(2*m))*Tinf
TT[i, :] = T
# check: display final T, best if nr is small number like nr=5
#print('T \n', T)
# plot results
py.figure(1)
py.plot(t, TT[:, m-1], '-k', label='surface')
py.plot(t, TT[:, 0], '--k', label='center')
py.axhline(Tinf, color='r', linestyle='-.')
py.ylabel('Temperature (K)')
py.xlabel('Time (s)')
py.legend(loc='best', numpoints=1)
def solve_func(b):
return linalg.solve_banded((l, u), diag, b)
def test_solve_banded(self):
assert_no_overwrite(lambda ab, b: solve_banded((2, 1), ab, b), [(4, 6),
(6, )])
def solve_func(b: np.ndarray) -> np.ndarray:
return linalg.solve_banded((l, u), diag, b)
ab = trid(Nx, ksi) # with NEUMANN CONDITIONS
side = np.zeros((Nx + 1, 2)) # right-hand vector for matrix eq
A = 0.1 * np.ones(Nx + 1) # FHN main parameter
A[:5] = -0.2 # Reverb
for timeStep in range(Nt):
runge = RK4(q)
# Neumann conditions: derivatives at the edges are null.
side[0] = runge[:, 0] + ksi * 2 * (q[:, 1] - q[:, 0])
for i in range(1, Nx):
side[i] = runge[:, i] + ksi * (q[:, i - 1] - 2 * q[:, i] + q[:, i + 1])
# side = np.array([RK4(q[:, i]) + ksi*(q[:, i-1] - 2*q[:, i] + q[:, i+1]) for i in range(Nx)])
# side = np.append(side, [RK4(q[:, -1]) + ksi*2*(q[:, -2] - q[:, -1])], axis=0)
side[-1] = runge[:, -1] + ksi * 2 * (q[:, -2] - q[:, -1])
# Solve the equation a x = b for x, assuming a is banded matrix using the matrix diagonal ordered form.
# q = np.array(list(map(lambda x: slin.solve_banded((1, 1), ab[x], side.T[x]), [0, 1])))
q[0] = slin.solve_banded((1, 1), ab[0], side.T[0])
q[1] = slin.solve_banded((1, 1), ab[1], side.T[1])
R.append(q.copy())
if timeStep % 100 == 0:
print(timeStep, "/", Nt, time.strftime("%H:%M:%S", time.localtime()))
R = np.array(R)
#R = 0.5 - 0.5*R
# =============================================================================
# Plot block
# plt.ioff()
# plt.ion()
plt.close('all')
#plt.style.use('fivethirtyeight')
fig = plt.figure()
ax1 = fig.add_subplot(1, 1, 1)
anim = animation.FuncAnimation(fig, animate, range(0, Nt + 1, 5), interval=10)
def calculate(self):
try:
Therm.albedo = float(self.entryAlbedo.get())
Therm.Q_prime = float(self.entrySW.get())
Therm.Q_prime_var = float(self.entrySW_var.get())
Therm.l_star = float(self.entryLW.get())
Therm.l_star_var = float(self.entryLW_var.get())
Therm.T_air = float(self.entryAir.get())
Therm.T_air_var = float(self.entryAir_var.get())
Therm.u_wind = float(self.entryWind.get())
Therm.RH = float(self.entryRH.get())
Therm.RH_var = float(self.entryRH_var.get())
Therm.depth = float(self.entryDebris.get())
Therm.k_deb = float(self.entryKdeb.get())
Therm.max_time = 86400 * float(self.entryTime.get())
temp = float(self.entryFreq.get())
Therm.plot_frequency = (Therm.delta_t / 600.) * np.round(
float(self.entryFreq.get()) / (Therm.delta_t / 600.))
if (temp != Therm.plot_frequency):
print('Plot frequency rounded to ' +
str(Therm.plot_frequency) + ' hours')
Therm.plot_frequency = Therm.plot_frequency * 3600.
Therm.meas_begin = 86400 * float(self.entryPlotStart.get())
T = np.zeros(Therm.n_levels + 1)
time = np.arange(0, Therm.max_time + Therm.delta_t, Therm.delta_t)
melt = np.zeros(np.shape(time))
QS = np.zeros(np.shape(time))
QLA = np.zeros(np.shape(time))
QLS = np.zeros(np.shape(time))
QH = np.zeros(np.shape(time))
QC = np.zeros(np.shape(time))
QE = np.zeros(np.shape(time))
n_plot = int(np.round(self.max_time / self.plot_frequency))
T_plot = np.zeros([self.n_levels + 2, n_plot + 1])
Therm.A_bl = (0.4 / np.log(Therm.z_eddy / Therm.z0))**2
Therm.delta_z = Therm.depth / Therm.n_levels
Therm.mu = Therm.kappa_deb * Therm.delta_t / Therm.delta_z**2
Z = np.linspace(0, Therm.depth, Therm.n_levels + 1)
Therm.fig.clf()
aa = Therm.fig.add_subplot(121)
bb = Therm.fig.add_subplot(122)
self.str.set("RUNNING")
Therm.canvas.draw()
for n in range(1, len(time)):
Therm.Q_prime_t = Therm.Q_prime - np.cos(
2 * np.pi * time[n] / 86400) * Therm.Q_prime_var
Therm.l_star_t = Therm.l_star - np.cos(
2 * np.pi * time[n] / 86400) * Therm.l_star_var
Therm.T_air_t = Therm.T_air - np.cos(
2 * np.pi * time[n] / 86400) * Therm.T_air_var
Therm.RH_t = Therm.RH - np.cos(
2 * np.pi * time[n] / 86400) * Therm.RH_var
if (time[n] % 86400 == 0):
print(str(time[n]) + ' seconds')
T_current = T[:-1]
T_new = T[:-1]
for k in range(Therm.max_newton_iter):
F, Qs, Ql_atm, Ql_snow, Qh, Qc, Qe = self.residual(
T_new, T_current)
if ((LA.norm(F) / np.sqrt(Therm.n_levels)) <
Therm.newton_tol):
break
if (k == (Therm.max_newton_iter - 1)):
print('did not converge: ' + str(LA.norm(F)))
Jab = self.residual_gradient(T_new)
T_new = T_new - LA2.solve_banded((1, 1), Jab, F)
T[:-1] = T_new
mr = self.k_deb * (
T_new[-2] -
T_new[-1]) / self.delta_z / Therm.rho_i / Therm.Lf * 86400.
melt[n] = max(mr, 0)
QS[n] = Qs
QLA[n] = Ql_atm
QLS[n] = Ql_snow
QH[n] = Qh
QC[n] = Qc
QE[n] = Qe
if ((time[n] % self.plot_frequency) == 0):
T_plot[1:,
int(np.round(time[n] / self.plot_frequency))] = T
T_plot[0, int(np.round(time[n] /
self.plot_frequency))] = time[n]
bb.clear()
bb.plot(time[0:n], melt[0:n])
if (time[n] > Therm.meas_begin):
aa.plot(T, np.flipud(Z))
Therm.canvas.draw()
# tk.Tk.update(self)
avg_melt = np.mean(melt[time > Therm.meas_begin])
# for k in range(0,n_plot+1):
# if (k*Therm.plot_frequency > Therm.meas_begin):
# aa.plot(T_plot[1:,k],np.flipud(Z))
aa.set_xlabel('Temperature (C)')
aa.set_ylabel('distance from ice (m)')
bb.clear()
bb.plot(time / 86400, melt)
bb.plot([Therm.meas_begin / 86400, Therm.meas_begin / 86400],
[0, avg_melt * 1.2], 'r')
self.str.set(str(avg_melt) + ' m/day')
bb.set_xlabel('time(days)')
bb.set_ylabel('melt rate (m/day)')
Therm.canvas.draw()
savevar = np.concatenate(
(time[1:, None], melt[1:, None], QS[1:, None], QLA[1:, None],
QLS[1:, None], QH[1:, None], QE[1:, None], QC[1:, None]), 1)
str1 = 'Time (s)'
for n in range(1, self.n_levels + 1):
str1 = str1 + ',' + str((n - 1) * Therm.delta_z) + 'm depth'
str1 = str1 + ',base\n'
str2 = 'Time (s),Melt,Qs,QL_down,QL_up,Qh,Qe,Qc'
str2 = str2 + '\n'
overw = self.chekvar.get()
files = list(
filter(os.path.isfile,
glob.glob("TemperatureProfiles[0-9]*csv")))
if ((len(files) > 0) & (overw == 0)):
files.sort(key=lambda x: os.path.getmtime(x))
filenum = int(files[-1].split('.')[0][19:])
profiles_filename = 'TemperatureProfiles' + str(filenum +
1) + '.csv'
else:
profiles_filename = 'TemperatureProfiles' + str(0) + '.csv'
files = list(
filter(os.path.isfile, glob.glob("MeltHistory[0-9]*csv")))
if ((len(files) > 0) & (overw == 0)):
files.sort(key=lambda x: os.path.getmtime(x))
filenum = int(files[-1].split('.')[0][11:])
hist_filename = 'MeltHistory' + str(filenum + 1) + '.csv'
else:
hist_filename = 'MeltHistory' + str(0) + '.csv'
with open(profiles_filename, 'w') as f:
try:
f.write('albedo,' + str(Therm.albedo) + '\n')
f.write('Shortwave Average,' + str(self.Q_prime) + '\n')
f.write('Shortwave Variation,' + str(self.Q_prime_var) +
'\n')
f.write('Atm. Longwave Average,' + str(self.l_star) + '\n')
f.write('Atm. Longwave Variation,' + str(self.l_star_var) +
'\n')
f.write('Air temp average,' + str(self.T_air) + '\n')
f.write('Air temp variaion,' + str(self.T_air_var) + '\n')
f.write('Rel Humidity average,' + str(self.RH) + '\n')
f.write('Rel Humidity variaion,' + str(self.RH_var) + '\n')
f.write('Debris depth,' + str(self.depth) + '\n')
f.write('Wind speed,' + str(self.u_wind) + '\n')
f.write('Thermal Conductivity,' + str(self.k_deb) + '\n')
f.write('Plot Frequency,' + str(self.plot_frequency) +
'\n')
f.write(str1)
T_plot = np.transpose(T_plot)
xxx = np.shape(T_plot)
for i in range(0, xxx[0]):
str3 = str(T_plot[i, 0])
for n in range(1, xxx[1]):
str3 = str3 + ',' + str(T_plot[i, n])
str3 = str3 + '\n'
f.write(str3)
f.close()
except:
self.str.set('ERROR: CLOSE EXCEL FILES')
print('GOT HERE')
Therm.canvas.draw()
with open(hist_filename, 'w') as f:
try:
f.write('albedo,' + str(Therm.albedo) + '\n')
f.write('Shortwave Average,' + str(self.Q_prime) + '\n')
f.write('Shortwave Variation,' + str(self.Q_prime_var) +
'\n')
f.write('Atm. Longwave Average,' + str(self.l_star) + '\n')
f.write('Atm. Longwave Variation,' + str(self.l_star_var) +
'\n')
f.write('Air temp average,' + str(self.T_air) + '\n')
f.write('Air temp variaion,' + str(self.T_air_var) + '\n')
f.write('Rel Humidity average,' + str(self.RH) + '\n')
f.write('Rel Humidity variaion,' + str(self.RH_var) + '\n')
f.write('Debris depth,' + str(self.depth) + '\n')
f.write('Wind speed,' + str(self.u_wind) + '\n')
f.write('Thermal Conductivity,' + str(self.k_deb) + '\n')
f.write('Plot Frequency,' + str(self.plot_frequency) +
'\n')
f.write(str2)
xxx = np.shape(savevar)
for i in range(0, xxx[0]):
str3 = str(savevar[i, 0])
for n in range(1, xxx[1]):
str3 = str3 + ',' + str(savevar[i, n])
str3 = str3 + '\n'
f.write(str3)
f.close()
except:
self.str.set('ERROR: CLOSE EXCEL FILES')
print('GOT HERE')
Therm.canvas.draw()
Therm.fig.savefig('LatestPlot.png')
print('DONE')
except ValueError:
result = "Please enter numbers only, using . for decimal (e.g. 0.1 instead of 0,1)"
self.str.set(result)
print(result)
# for debugging, use the built-in banded solver from numpy
u = np.zeros(n - 1)
d = np.zeros(n - 1)
l = np.zeros(n - 1)
d[:] = 4.0 * dx
u[:] = dx
u[0] = 0.0
l[:] = dx
l[n - 2] = 0.0
# create a banded matrix -- this doesn't store every element -- just
# the diagonal and one above and below and solve
A = np.matrix([u, d, l])
xsol_np = linalg.solve_banded((1, 1), A, b[1:n])
# xsol_np now hold all the second derivatives for points 1 to n-1.
# Natural boundary conditions are imposed here
xsol_np = np.insert(xsol_np, 0, 0)
xsol_np = np.insert(xsol_np, n, 0) # insert at the end
# report error from our iterative method vs. direct solve
err = np.max(np.abs((xsol[1:n] - xsol_np[1:n]) / xsol_np[1:n]))
print("relative error of iterative solution compared to direct: ", err)
# go ahead with our iterative solution -- natural boundary conditions are
# already in place for this
ppp = xsol
def solve_hemholtz(f, boundary_func=lambda x,y: x*y, lap_f=lambda x: None, n=20,xint=(0,1), solver='spdiags', scheme='5-point', k=20):
h = (xint[1]-xint[0])/float(n+1)
if scheme == 'deferred':
#Get the second-order approximation
u_bar = np.zeros((n+2, n+2))
u_bar[1:-1,1:-1] = solve_hemholtz(f, boundary_func, lap_f, n, xint,k=k).reshape((n,n))
#Fill in the boundary data
grid_1d = np.linspace(xint[0], xint[1], n+2)
u_bar[0] = boundary_func(xint[0], grid_1d)
u_bar[-1] = boundary_func(xint[1], grid_1d)
u_bar[:,0] = boundary_func(grid_1d, xint[0])
u_bar[:,-1] = boundary_func(grid_1d, xint[1])
#Compute the xxyy derivative
u_xx = (u_bar[:-2]+u_bar[2:]-2*u_bar[1:-1])/h**2
u_xxyy = (u_xx[:,2:]+u_xx[:,:-2]-2*u_xx[:,1:-1])/h**2
#First, construct the diagonals and the RHS
#Here we assume that the domain is square
b = np.zeros((n,n))
yvals = np.linspace(xint[0], xint[1], n+2)[1:-1]
# print h, yvals[1]-yvals[0]
for i in xrange(n):
xval = yvals[i]
b[i] = f(xval, yvals)
if scheme == '9-point':
b[i] += h**2/12. * (lap_f(xval, yvals) - k**2*f(xval, yvals))
elif scheme == 'deferred':
b[i] += h**2/12. * (lap_f(xval, yvals) - k**2*f(xval, yvals) + k**4 * u_bar[i+1,1:-1] - 2*u_xxyy[i] )
if scheme == '5-point' or scheme == 'deferred':
b *= h**2
else: b *= 6*h**2
# print b
if scheme=='5-point' or scheme == 'deferred':
diags = make_5_point_diags(n, h)
diags[0] += h**2 * k**2
#Now we need to make use of the boundary conditions
b[0] -= boundary_func(xint[0], yvals)
b[-1] -= boundary_func(xint[1], yvals)
b[:,0] -= boundary_func(yvals, xint[0])
b[:,-1] -= boundary_func(yvals, xint[1])
elif scheme=='9-point':
diags = make_9_point_diags(n, h)
diags[0] += 6 *((h*k)**2 - 1./12* (h*k)**4)
#This is like the 5-point stencil
b[0] -= 4*boundary_func(xint[0], yvals)
b[-1] -= 4*boundary_func(xint[1], yvals)
b[:,0] -= 4*boundary_func(yvals, xint[0])
b[:,-1] -= 4*boundary_func(yvals, xint[1])
#Here we handle the corners of the stencil
vals = np.linspace(xint[0], xint[1], n+2)
#Take care of effects to the top and bottom boundaries
b[0] -= boundary_func(xint[0], vals[:-2]) + boundary_func(xint[0], vals[2:])
b[-1] -= boundary_func(xint[1], vals[:-2]) + boundary_func(xint[1], vals[2:])
#Handle the sides - be careful not to double count the corners
b[1:,0] -= boundary_func(vals[1:-2], xint[0])
b[:-1,0] -= boundary_func(vals[2:-1], xint[0])
b[1:,-1] -= boundary_func(vals[1:-2], xint[1])
b[:-1,-1] -= boundary_func(vals[2:-1], xint[1])
else:
raise ValueError("Unrecognized scheme {}".format(scheme))
#Next, create the sparse matrix and solve the system
# print([(d, sum(diags[d]!=0)) for d in diags if isinstance(diags[d], np.ndarray)])
offsets = sorted(diags.keys())
mat = sparse_diags([diags[d] for d in offsets], offsets, format='csc')
# print mat.todense()
if solver == 'spdiags':
# print b
# print mat.shape
return spsolve(mat, b.flatten())
elif solver == 'solve_banded':
ab, l_and_u = diags_to_banded(diags)
# print b.flatten().shape, ab.shape
return solve_banded(l_and_u,ab, b.flatten())
else:
return solver(mat, b.flatten())
ab[1, 1:m-1] = 8
ab[1, m-1] = 4
# lower diagonal
ab[2, 0:m-2] = 7
ab[2, m-2] = 3
# create row vector b
b = np.zeros(m)
b[0] = 1
b[1:] = 2
b[-1] = 3
# solve using scipy.sparse.linalg.lsqr
ti_sp = timeit.default_timer() # start timer
x_sp = sp.solve_banded((1, 1), ab, b) # solve A*x = B for x
tf_sp = timeit.default_timer() # stop timer
# Print Results
# ------------------------------------------------------------------------------
print('m is', m)
# numpy approach
if m < 7:
print('A is\n', A)
print('B is\n', B)
print('x is\n', x_np)
print('NumPy time is', tf_np - ti_np, 'seconds')
def FDM_SingleBarrier_NonUnifromGrid(self, Ns, Nt, theta, ratio, m):
'''
Abstract:
--------
Finite difference method for barrier option.
Using a non-uniform grid, with shape controlled by input ratio.
Discrete monioring.
Iteration process is only suitable for knock out option.
For knock in, this funcition uses a vanilla option to minus the identical knock out.
Parameters:
----------
Ns: Number of points in price axis
Nt: Number of points in time axis
theta:
0 : Fully implicit method
0.5 : Crank-nicolson method
According to reference, fully implicit method is better than crank-nicolson method
Reference :
Zvan R, Vetzal K R, Forsyth P A. PDE methods for pricing barrier options ☆[J].
Journal of Economic Dynamics & Control, 1997, 24(11-12):1563-1590.
ratio: The parameter use to controll the shape of the grid
m : monitoring times
'''
# discretize Nt-1 points between every two monitoring time, total Nt*m + 1 gird in time axis
step = Nt
Nt = Nt * m
# set up parameters
mu = self.r - self.q
_range = 5 * self.sig * np.sqrt(self.maturity)
if self.option_type == 'DOWN-AND-OUT-BARRIER' or self.option_type == 'DOWN-AND-IN-BARRIER':
Smax = self.spot_price * np.exp(
(mu - self.sig**2 / 2.0) * self.maturity + _range)
Smin = self.barrier * 0.99999999
elif self.option_type == 'UP-AND-OUT-BARRIER' or self.option_type == 'UP-AND-IN-BARRIER':
Smax = max(self.barrier, self.strike) * 1.0000001
Smin = 0
# totally Nt + 1 in row grid
dt = self.maturity / float(Nt)
# generate non-uniform grid
s = np.linspace(Smin, Smax, Ns * (1 - ratio) + 1)
if self.option_type == 'DOWN-AND-OUT-BARRIER':
temp = [self.barrier, self.spot_price]
elif self.option_type == 'UP-AND-OUT-BARRIER':
temp = [self.spot_price, self.barrier]
elif self.option_type == 'DOWN-AND-IN-BARRIER':
temp = [self.barrier, self.spot_price]
else:
temp = [self.spot_price, self.barrier]
lower_index = np.array([sum(s < i) for i in temp]) - 1
upper_index = lower_index + 1
delta_s = -s[lower_index] + s[upper_index]
delta_s = delta_s[0]
if lower_index[1] > lower_index[0]:
count = int(Ns * ratio / 2.0)
else:
count = int(Ns * ratio)
ds = delta_s / (count - 1)
#enforce the grid density around key value
insert_vector = [
np.linspace(s[lower_index[j]] + ds, s[upper_index[j]] - ds, count)
for j in [0, 1]
]
s_temp = np.append(s[:lower_index[0] + 1], insert_vector[0])
s_temp = np.append(s_temp, s[upper_index[0]:lower_index[1] + 1])
if lower_index[1] > lower_index[0]:
s_temp = np.append(s_temp, insert_vector[1])
s_temp = np.append(s_temp, s[upper_index[1]:])
s = s_temp
Ns = len(s) - 1
# initialize the payoff
if self.position == 'CALL':
if self.option_type == 'DOWN-AND-OUT-BARRIER' or self.option_type == 'DOWN-AND-IN-BARRIER':
V_Nt = np.maximum(s - self.strike, 0) * np.where(
s >= self.barrier, 1, 0)
else:
V_Nt = np.maximum(s - self.strike, 0) * np.where(
s <= self.barrier, 1, 0)
payoff = np.maximum(s - self.strike, 0)
elif self.position == 'PUT':
if self.option_type == 'DOWN-AND-OUT-BARRIER' or self.option_type == 'DOWN-AND-IN-BARRIER':
V_Nt = np.maximum(-s + self.strike, 0) * np.where(
s >= self.barrier, 1, 0)
else:
V_Nt = np.maximum(-s + self.strike, 0) * np.where(
s <= self.barrier, 1, 0)
payoff = np.maximum(-s + self.strike, 0)
# initialize the Dirichlet boundary condition
if self.position == "CALL":
f_0 = np.linspace(0, 0, Nt + 1)
if self.option_type == 'DOWN-AND-OUT-BARRIER' or self.option_type == 'DOWN-AND-IN-BARRIER':
f_Ns = Smax * np.exp(-self.r * np.linspace(
0, self.maturity, Nt + 1)) - self.strike * np.exp(
-self.r * (np.linspace(0, self.maturity, Nt + 1)))
else:
f_Ns = np.linspace(0, 0, Nt + 1)
elif self.position == "PUT":
if self.option_type == 'DOWN-AND-OUT-BARRIER' or self.option_type == 'DOWN-AND-IN-BARRIER':
f_0 = np.linspace(0, 0, Nt + 1)
else:
f_0 = self.strike * np.exp(
-self.r * np.linspace(0, self.maturity, Nt + 1))
f_Ns = np.linspace(0, 0, Nt + 1)
# initialize the tridiagonal matrix by scalar-form
delta_s_i = 0.5 * (s[2:] - s[0:Ns - 1])
delta_s_plus = s[2:] - s[1:Ns]
delta_s_minus = s[1:Ns] - s[0:Ns - 1]
# from a_2 to a_I-1 are in the calculation matrix
a = -(1.0 - theta) * self.sig**2 * s[1:Ns]**2 / (
2.0 * delta_s_i *
delta_s_minus) + (1 - theta) * mu * s[1:Ns] / (2 * delta_s_i)
# from b_1 to b_I-1 are in the calculation matrix
b = 1.0 / dt + (1 - theta) * self.r + (
1.0 - theta) * self.sig**2 * s[1:Ns]**2 / (2.0 * delta_s_i *
delta_s_minus)
b = b + (1.0 - theta) * self.sig**2 * s[1:Ns]**2 / (2.0 * delta_s_i *
delta_s_plus)
# from c_1 to c_I-2 are in the calculation matrix
c = -(1.0 - theta) * self.sig**2 * s[1:Ns]**2 / (
2.0 * delta_s_i *
delta_s_plus) - (1 - theta) * mu * s[1:Ns] / (2 * delta_s_i)
# from alpha_2 to alpha_I-1 are in the calculation matrix
alpha = theta * self.sig**2 * s[1:Ns]**2 / (
2.0 * delta_s_i *
delta_s_minus) - theta * mu * s[1:Ns] / (2 * delta_s_i)
# from beta_1 to beta_I-1 are in the calculation matrix
beta = 1.0 / dt - theta * self.sig**2 * s[1:Ns]**2 / (
2.0 * delta_s_i * delta_s_minus) - self.r * theta
beta = beta - theta * self.sig**2 * s[1:Ns]**2 / (2.0 * delta_s_i *
delta_s_plus)
# from gamma_1 to gamma_I-2 are in the calculation matrix
gamma = theta * self.sig**2 * s[1:Ns]**2 / (
2.0 * delta_s_i *
delta_s_plus) + theta * mu * s[1:Ns] / (2 * delta_s_i)
# From Nt to 1, calculate V_Nt-1, V_Nt-2, ..., V_0 (vectors)
V_Nplus = V_Nt[1:Ns]
for k in range(Nt, 0, -1):
#for k in range(1,0,-1):
#V_Nplus : b of Ax=b
V_Nplus = self.my_dot_product(alpha, beta, gamma, V_Nplus)
V_Nplus[0] = V_Nplus[0] - a[0] * f_0[k - 1] + alpha[0] * f_0[k]
V_Nplus[Ns - 2] = V_Nplus[
Ns - 2] - c[Ns - 2] * f_Ns[k - 1] + gamma[Ns - 2] * f_Ns[k]
#V_N : Intial Guess for american case / x of Ax=b for european case
ab = self.tri_bound_to_ab(a, b, c)
V_N = linalg.solve_banded((1, 1), ab, V_Nplus)
#American process
if self.exercise_type == 'AMERICAN':
V_N = self.Projected_SOR(a[1:], b, c[:-1], V_Nplus, V_N,
payoff[1:-1], k, step, s[1:Ns])
V_Nplus = V_N
# linear interpolation
index = sum(s < self.spot_price)
w = (self.spot_price - s[index - 1]) / (s[index] - s[index - 1])
v_0 = V_Nplus[index - 1] * (1 - w) + w * V_Nplus[index]
'''
Above process is only for knock out option
'''
if self.option_type == 'UP-AND-OUT-BARRIER' or self.option_type == 'DOWN-AND-OUT-BARRIER':
return v_0
else:
if self.position == 'CALL':
v_0 = self.Black_Scholes_Call() - v_0
return v_0
else:
if self.exercise_type == 'EUROPEAN':
v_0 = self.Black_Scholes_Put() - v_0
return v_0
else:
v_0 = self.BTM_Vanilla(1200) - v_0
return v_0
def LibRun(mx, f): ud = np.insert(np.diag(mx, 1), 0, 0) d = np.diag(mx) ld = np.insert(np.diag(mx, -1), len(d)-1, 0) return lg.solve_banded((1, 1), np.matrix([ud, d, ld]), f)
def _woodbury_algorithm(A, ur, ll, b, k):
'''
Solve a cyclic banded linear system with upper right
and lower blocks of size ``(k-1) / 2`` using
the Woodbury formula
Parameters
----------
A : 2-D array, shape(k, n)
Matrix of diagonals of original matrix(see
``solve_banded`` documentation).
ur : 2-D array, shape(bs, bs)
Upper right block matrix.
ll : 2-D array, shape(bs, bs)
Lower left block matrix.
b : 1-D array, shape(n,)
Vector of constant terms of the system of linear equations.
k : int
B-spline degree.
Returns
-------
c : 1-D array, shape(n,)
Solution of the original system of linear equations.
Notes
-----
This algorithm works only for systems with banded matrix A plus
a correction term U @ V.T, where the matrix U @ V.T gives upper right
and lower left block of A
The system is solved with the following steps:
1. New systems of linear equations are constructed:
A @ z_i = u_i,
u_i - columnn vector of U,
i = 1, ..., k - 1
2. Matrix Z is formed from vectors z_i:
Z = [ z_1 | z_2 | ... | z_{k - 1} ]
3. Matrix H = (1 + V.T @ Z)^{-1}
4. The system A' @ y = b is solved
5. x = y - Z @ (H @ V.T @ y)
Also, ``n`` should be greater than ``k``, otherwise corner block
elements will intersect with diagonals.
Examples
--------
Consider the case of n = 8, k = 5 (size of blocks - 2 x 2).
The matrix of a system: U: V:
x x x * * a b a b 0 0 0 0 1 0
x x x x * * c 0 c 0 0 0 0 0 1
x x x x x * * 0 0 0 0 0 0 0 0
* x x x x x * 0 0 0 0 0 0 0 0
* * x x x x x 0 0 0 0 0 0 0 0
d * * x x x x 0 0 d 0 1 0 0 0
e f * * x x x 0 0 e f 0 1 0 0
References
----------
.. [1] William H. Press, Saul A. Teukolsky, William T. Vetterling
and Brian P. Flannery, Numerical Recipes, 2007, Section 2.7.3
'''
k_mod = k - k % 2
bs = int((k - 1) / 2) + (k + 1) % 2
n = A.shape[1] + 1
U = np.zeros((n - 1, k_mod))
VT = np.zeros((k_mod, n - 1)) # V transpose
# upper right block
U[:bs, :bs] = ur
VT[np.arange(bs), np.arange(bs) - bs] = 1
# lower left block
U[-bs:, -bs:] = ll
VT[np.arange(bs) - bs, np.arange(bs)] = 1
Z = solve_banded((bs, bs), A, U)
H = solve(np.identity(k_mod) + VT @ Z, np.identity(k_mod))
y = solve_banded((bs, bs), A, b)
c = y - Z @ (H @ (VT @ y))
return c
import matplotlib.pyplot as plt
import numpy as np
from export_img import export_img
from time import perf_counter
nb_resolutions = 10000
with open('matrices.data', 'r+b') as outfile:
dict_save_arrays = np.load(outfile)['arr_0'].ravel((-1, ))[0]
tridiag_for_scipy = dict_save_arrays['tridiag_for_scipy']
seconde_partie_sec_membre = dict_save_arrays['seconde_partie_sec_membre']
offset = perf_counter()
for _ in range(nb_resolutions):
la.solve_banded((1, 1), tridiag_for_scipy, seconde_partie_sec_membre)
times_calcul_u = (perf_counter() - offset) / (nb_resolutions)
print("temps utilisé : ", perf_counter() - offset)
times_calcul_alpha = []
imax = 0
# ---------- CALCUL PERFORMANCE : MATRICE PLEINE DE TAILLE i
try:
for i in range(1, 30):
imax = i + 1
print("i =", i)
Phi_T_M_Phi = dict_save_arrays['Phi_T_M_Phi' + str(i)]
sec_membre = dict_save_arrays['sec_membre' + str(i)]
offset = perf_counter()
for _ in range(nb_resolutions):
np.linalg.solve(Phi_T_M_Phi, sec_membre)
H = D + Sub + Sup
# define the identity matrix
I = np.identity(P - 1)
# define L and R from eq 17, lab manual
L_vec = I + 1j * (h / 2 * h_bar) * H
R = I - 1j * (h / 2 * h_bar) * H
# reformatted L_vec to be solved with solve_banded
L_vec_banded = np.zeros((3, L_vec.shape[0]))
L_vec_banded[0, 1:] = L_vec.diagonal(1)
L_vec_banded[1, :] = L_vec.diagonal()
L_vec_banded[2, :-1] = L_vec.diagonal(-1)
# Main loop
# store the wavefunction at each time step in a list
solution = [psi]
for i in range(N):
x = np.matmul(R, psi)
psi = solve_banded((1, 1), L_vec_banded, x)
solution.append(psi)
# plotting wave function at different posititons depending on time
plt.plot(x_points, (solution[0]))
plt.plot(x_points, (solution[150]))
plt.plot(x_points, (solution[300]))
plt.xlabel("x (m)")
plt.ylabel("$\psi(x)$")
plt.show()
# B[0, 1] = b2
# B[N, N - 1] = b2
# B[N, N] = b1
# for i in range(N):
# B[i, i - 1] = b2
# B[i, i] = b1
# B[i, i + 1] = b2
# Create the matrix A in the form appropriate for the function solve_banded
A2 = empty([3, N + 1], complex)
A2[0, 0] = 0
A2[0, 1:] = a2
A2[1, :] = a1
A2[2, 0:N] = a2
A2[2, N] = 0
# Main loop
# store the wavefunction at each time step in a list
solution = [psi]
for i in range(300):
psi[1:N] = b1 * psi[1:N] + b2 * (psi[2:] + psi[0:N - 1])
psi = solve_banded((1, 1), A2, psi)
solution.append(psi)
plot(x_points, abs(solution[0])**2)
plot(x_points, abs(solution[49])**2)
plot(x_points, abs(solution[250])**2)
xlabel("x (m)")
ylabel("$\psi(x)$")
show()
cost, penalty, lagrange_mult, xres = calc_loss(x, l, rho)
#print(total_cost)
print("hey now")
#print(cost)
total_cost = cost + penalty + lagrange_mult
#total_cost = cost
gradL, hess = getGradHessBand(total_cost, (NVars + NControls) * 3, x)
print(hess)
#print(hess.shape)
gradL = gradL.reshape(-1)
#print(gradL.shape)
#easiest thing might be to put lagrange mutlipleirs into x.
#Alternatively, use second order step in penalty method.
bandn = (NVars + NControls) * 3 // 2
dx = linalg.solve_banded((bandn, bandn), hess, gradL) #
newton_dec = np.dot(dx, gradL)
df0 = dx[:NControls].reshape(-1, NControls)
dx = dx[NControls:].reshape(-1, N - 1, NVars + NControls)
with torch.no_grad():
x[:, 1:, :] -= torch.tensor(dx)
print(x[:, 0, NVars:].shape)
print(df0.shape)
x[:, 0, NVars:] -= torch.tensor(df0)
costval = cost.detach().numpy()
if newton_dec / costval < 1e-10:
break
#print(x)
with torch.no_grad():
l += 2 * rho * xres
def compute_reference_solution_x_1(self, delta_t=0.005, xb=2.5, nx=1000):
self.xb = xb # range of x, [-xb, xb]
self.nx = nx # number of discrete interval
self.dx = 2.0 * self.xb / self.nx
self.delta_t = delta_t
beta = 2
self.xvec = np.linspace(-self.xb, self.xb, self.nx, endpoint=True)
# A = D^{-1} L D
# assumes Neumann boundary conditions
A = np.zeros([self.nx, self.nx])
for i in range(0, self.nx):
x = -self.xb + (i + 0.5) * self.dx
if i > 0:
x0 = -self.xb + (i - 0.5) * self.dx
x1 = -self.xb + i * self.dx
A[i, i - 1] = -exp(
beta * 0.5 *
(self.V(x0) + self.V(x) - 2 * self.V(x1))) / self.dx**2
A[i, i] = exp(beta * (self.V(x) - self.V(x1))) / self.dx**2
if i < self.nx - 1:
x0 = -self.xb + (i + 1.5) * self.dx
x1 = -self.xb + (i + 1) * self.dx
A[i, i + 1] = -exp(
beta * 0.5 *
(self.V(x0) + self.V(x) - 2 * self.V(x1))) / self.dx**2
A[i, i] = A[i, i] + exp(beta *
(self.V(x) - self.V(x1))) / self.dx**2
A = -A / beta
N = int(self.T / self.delta_t)
D = np.diag(exp(beta * self.V(self.xvec) / 2))
D_inv = np.diag(exp(-beta * self.V(self.xvec) / 2))
np.linalg.cond(np.eye(self.nx) - self.delta_t * A)
#w, vv = np.linalg.eigh(np.eye(self.nx) - self.delta_t * A)
self.psi = np.zeros([N + 1, self.nx])
self.psi[N, :] = exp(-self.g_1(self.xvec))
for n in range(N - 1, -1, -1):
band = -self.delta_t * np.vstack([
np.append([0], np.diagonal(A, offset=1)),
np.diagonal(A, offset=0) - N / self.T,
np.append(np.diagonal(A, offset=1), [0])
])
self.psi[n, :] = D.dot(
solve_banded([1, 1], band, D_inv.dot(self.psi[n + 1, :])))
#psi[n, :] = np.dot(D, np.linalg.solve(np.eye(self.nx) - delta_t * A, D_inv.dot(psi[n + 1, :])));
self.u = np.zeros([N + 1, self.nx - 1])
for n in range(N + 1):
for i in range(self.nx - 1):
self.u[n, i] = -2 / beta * self.B[0, 0] * (
-log(self.psi[n, i + 1]) + log(self.psi[n, i])) / self.dx
def cFunctionUncoupled(ws, Kv, F, Fsurf, Fbed, data, n, hasMatrix=False):
####################################################################################################################
# Init
####################################################################################################################
jmax = data.v('grid', 'maxIndex',
'x') # maximum index of x grid (jmax+1 grid points incl. 0)
kmax = data.v('grid', 'maxIndex',
'z') # maximum index of z grid (kmax+1 grid points incl. 0)
OMEGA = data.v('OMEGA')
# Init Ctd
nRHS = F.shape[-1]
cCoef = np.zeros((jmax + 1, kmax + 1, 1, nRHS), dtype=complex)
cMatrix = np.zeros((jmax + 1, 3, kmax + 1), dtype=complex)
try:
z = ny.dimensionalAxis(data.slice('grid'), 'z')[:, :, 0]
except:
z = -np.linspace(0, 1, kmax + 1).reshape((1, kmax + 1)) * np.ones(
(jmax + 1, 1))
####################################################################################################################
# build, save and solve the matrices in every water column
####################################################################################################################
for j in range(0, jmax + 1):
dz = z[j, 1:] - z[j, :-1]
dz = dz
dz_down = dz[:-1]
dz_up = dz[1:]
dz_av = 0.5 * (dz_up + dz_down)
if not hasMatrix:
A = np.zeros((3, kmax + 1), dtype=complex)
##### LEFT HAND SIDE #####
# Build matrix.
# NB. can use general numerical schemes as dz < 0
a = -0.5 * (Kv[j, 0:-2, 0] + Kv[j, 1:-1, 0]) / dz_down + ws[
j, 0:-2, 0] # for k-1: from 1..kmax-1
b = 0.5 * (Kv[j, 0:-2, 0] + Kv[j, 1:-1, 0]) / dz_down + 0.5 * (
Kv[j, 2:, 0] + Kv[j, 1:-1, 0]) / dz_up - ws[
j, 1:-1, 0] # for k: from 1..kmax-1
c = -0.5 * (Kv[j, 2:, 0] +
Kv[j, 1:-1, 0]) / dz_up # for k+1: from 1..kmax
# add inertia later
# Build matrix k=1..kmax-1
A[2, :-2] = a
A[1, 1:-1] = b
A[0, 2:] = c
## surface
b = ws[j, 0, 0] - Kv[j, 0, 0] / dz[0]
c = +Kv[j, 0, 0] / dz[0]
A[1, 0] = b
A[0, 1] = c
## bed
a = -Kv[j, 0, 0] / dz[-1]
b = Kv[j, 0, 0] / dz[-1]
A[2, -2] = a
A[1, -1] = b
# save matrix
cMatrix[j, Ellipsis] = A[Ellipsis]
else:
A = Kv[j, Ellipsis]
A[1, 1:-1] += n * 1j * OMEGA * dz_av
################################################################################################################
# Right hand side
################################################################################################################
RHS = np.zeros([kmax + 1, nRHS], dtype=complex)
RHS[1:-1, :] = F[j, 1:-1, :] * dz_up.reshape((kmax - 1, 1))
RHS[0, :] = Fsurf[j, 0, :]
RHS[-1, :] += Fbed[j, 0, :]
################################################################################################################
# Solve
################################################################################################################
cstag = solve_banded((1, 1),
A,
RHS,
overwrite_ab=True,
overwrite_b=True)
cCoef[j, :, 0, :] = cstag.reshape(kmax + 1, nRHS)
return cCoef, cMatrix
