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