def _complexcn(z): """Jacobi elliptical function cn with parameter 1/2, valid for complex numbers.""" x = z.real y = z.imag snx, cnx, dnx, phx = ellipj(x, 1 / 2) sny, cny, dny, phy = ellipj(y, 1 / 2) numer = cnx * cny - 1j * snx * dnx * sny * dny denom = (1 - dnx**2 * sny**2) return numer / denom
def complex_jacobi(u, k, tol=10.**(-10.)): if abs(np.imag(u)) < tol: return sp.ellipj(np.real(u), k**2) real_ellip = sp.ellipj(np.real(u), k**2) sn_r, cn_r, dn_r = real_ellip[0], real_ellip[1], real_ellip[2] imag_ellip = sp.ellipj(np.imag(u), 1. - k**2) sn_c, cn_c, dn_c = imag_ellip[0], imag_ellip[1], imag_ellip[2] sn_c, cn_c, dn_c = 1j * sn_c / cn_c, 1. / cn_c, dn_c / cn_c denom = 1. - (k * sn_r * sn_c)**2 sn = (sn_r * cn_c * dn_c + sn_c * cn_r * dn_r) / denom cn = (cn_r * cn_c - sn_r * dn_r * sn_c * dn_c) / denom dn = (dn_r * cn_r - k**2 * sn_r * cn_r * sn_c * dn_c) / denom return sn, cn, dn
def into_jv(self, jgrid, vgrid): sv = np.sqrt(vgrid) jgrid = jgrid / pulsation_canon(jgrid, vgrid) return np.select( [vgrid < 1, vgrid > 1, True], [2 * np.arcsin(sv * scsp.ellipj(jgrid, vgrid)[0]), 2 * np.arcsin(scsp.ellipj(sv * jgrid, 1/vgrid)[0]), 2 * np.arcsin(np.tanh(jgrid)) ] ), np.select( [vgrid < 1, vgrid > 1, True], [2 * sv * scsp.ellipj(jgrid, vgrid)[1], 2 * sv * scsp.ellipj(sv * jgrid, 1/vgrid)[2], 2 / np.cosh(jgrid) ] )
def cplx_cn( z, k): z = numpy.asarray(z) if z.dtype != complex: return ellipj(z,k)[1] # note, 'dn' result of ellipj is buggy, as of Mar 2015 # see https://github.com/scipy/scipy/issues/3904 ss,cc = ellipj( z.real, k )[:2] dd = numpy.sqrt(1-k*ss**2) # make our own dn s1,c1 = ellipj( z.imag, 1-k )[:2] d1 = numpy.sqrt(1-k*s1**2) ds1 = dd*s1 den = (1-ds1**2) rx = cc*c1/den ry = ss*ds1*d1/den return rx - 1j*ry
def test_ellipfun_sn(self): # Oscillating function --- limit range of first argument; the # loss of precision there is an expected numerical feature # rather than an actual bug assert_mpmath_equal(lambda u, m: sc.ellipj(u, m)[0], lambda u, m: mpmath.ellipfun("sn", u=u, m=m), [Arg(-1e6, 1e6), Arg(a=0, b=1)], atol=1e-20)
def performPlot(self): #Plotting waveform plotxL = list(np.arange(-1, 1.001, 0.001)) plottheta = [2*self.plotDict["K"]*\ (i - (self.plotDict["time"]/self.plotDict["T"])) for i in plotxL] pSN, pCN, pDN, pPH = sp.ellipj(plottheta, self.plotDict["m"]) pCSD = pSN * pCN * pDN if self.plotDict["O"] == 1: ploteta = self.plotDict["d"]*\ (self.plotDict["A0"] + self.plotDict["A1"]*pCN**2) plotu = math.sqrt(self.g*self.plotDict["d"])*\ (self.plotDict["B00"] + self.plotDict["B10"]*pCN**2) plotw = math.sqrt(self.g*self.plotDict["d"])*\ (4.0*self.plotDict["K"]*self.plotDict["d"]*pCSD/self.plotDict["L"])*\ ((self.plotDict["z"] + self.plotDict["d"])/self.plotDict["d"])*self.plotDict["B10"] else: ploteta = self.plotDict["d"]*(self.plotDict["A0"] +\ self.plotDict["A1"]*pCN**2 + self.plotDict["A2"]*pCN**4) plotu = math.sqrt(self.g*self.plotDict["d"])*\ ((self.plotDict["B00"] + self.plotDict["B10"]*pCN**2 +\ self.plotDict["B20"]*pCN**4) -\ 0.5*((self.plotDict["z"] + self.plotDict["d"]) / self.plotDict["d"])**2 *\ (self.plotDict["B01"] + self.plotDict["B11"]*pCN**2 + self.plotDict["B21"]*pCN**4)) pw1 = ((self.plotDict["z"] + self.plotDict["d"])/self.plotDict["d"])*\ (self.plotDict["B10"] + 2.0*self.plotDict["B20"]*pCN**2) pw2 = (1.0/6.0)*(((self.plotDict["z"] + self.plotDict["d"])/self.plotDict["d"])**3)*\ (self.plotDict["B11"] + 2.0*self.plotDict["B21"]*pCN**2) plotw = math.sqrt(self.g*self.plotDict["d"])*\ (4.0*self.plotDict["K"]*self.plotDict["d"]*pCSD/self.plotDict["L"])*(pw1 - pw2) # end if plt.figure(1, figsize=(8, 12), dpi=self.plotConfigDict["dpi"]) plt.subplot(3, 1, 1) plt.plot(plotxL, ploteta) plt.axhline(y=0.0, color="r", LineStyle="--") plt.ylabel("Elevation [%s]" % self.labelUnitDist) plt.subplot(3, 1, 2) plt.plot(plotxL, plotu) plt.axhline(y=0.0, color="r", LineStyle="--") plt.ylabel("Velocity, u [%s/s]" % self.labelUnitDist) plt.subplot(3, 1, 3) plt.plot(plotxL, plotw) plt.axhline(y=0.0, color="r", LineStyle="--") plt.ylabel("Velocity, w [%s/s]" % self.labelUnitDist) plt.xlabel("x/L") plt.show() self.plotDict["plotxL"] = plotxL self.plotDict["ploteta"] = ploteta self.plotDict["plotu"] = plotu self.plotDict["plotw"] = plotw self.fileOutputPlotWriteData()
def getnormal(a, b, t): e = sl.ellipj(t, 1 - ((a**2 * (1 - b**2)) / (b**2 * (1 - a**2)))) s = e[0] c = e[1] d = e[2] x = 1 / (math.sqrt(1 - a**2)) * s y = 1 / (math.sqrt(1 - b**2)) * c z = math.sqrt(b**2 / (1 - b**2)) * d return [x, y, z]
def a_v(n, theta_deg): ''' Calculate a[v] and delta values for both odd and even orders of type 'a' cauer lowpass filters. The equations are from equations (37), (38), and (39) of Saal and Ulbrich paper on page 294. n = number of sections or lowpass order p = reflection coefficient or % of reflection theta_deg: theta in degrees such that sin(theta)=wc/ws. See page 292 ''' theta = theta_deg * np.pi / 180 #radians #m calculation, see EQN(37) p294 if n % 2 == 0: #n is even m = n // 2 else: #n is odd m = (n - 1) // 2 # #Complete elliptic integral of the 1st kind K with modulus k=sin(theta) #In scipy library: ellipk(u,m) where m = k^2 K = ss.ellipk(np.sin(theta)**2) a = np.array([1 ]) #a[0] is filled with 1 so indices of a[] complies with SU for v in range(1, n): #1 to n-1 inclusive u = K * v / n #EQN(38) notates sn(u,theta) #In scipy library: sn(u,m) where m = sin^2(theta) jac_ell = ss.ellipj(u, np.sin(theta)**2) sn = jac_ell[0] a = np.append(a, np.sqrt(np.sin(theta)) * sn) a = np.append(a, np.sqrt(np.sin(theta))) #append a[n] wc_pa = a[n] #delta calculation, see EQN(39) delta_a = 1 if n % 2 == 0: m_end = m scale = 1.0 else: m_end = m + 1 scale = 1.0 / a[n] for u in range(1, m_end + 1): delta_a *= a[2 * u - 1]**2 delta_a = scale * delta_a return a, delta_a, wc_pa
def cn_id(gridnum, kappa, k0, L): xspace = np.linspace(-L / 2.0, L / 2.0 - L / gridnum, gridnum) _, cn, _, asd = special.ellipj(kappa * (xspace), k0**2) initialprofile = k0 * kappa / np.sqrt(2.0) * cn k = kvec(gridnum) expconsts = -1j * (2.0 * np.pi / L * k)**2 return (xspace, initialprofile, expconsts)
def __call__(self, t): S = np.sin(0.5 * (self.theta0)) K_S = ellipk(S**2) omega_0 = np.sqrt(9.81) sn, cn, dn, ph = ellipj(K_S - omega_0 * t, S**2) theta = 2.0 * np.arcsin(S * sn) d_sn_du = cn * dn d_sn_dt = -omega_0 * d_sn_du d_theta_dt = 2.0 * S * d_sn_dt / np.sqrt(1.0 - (S * sn)**2) return {'theta': theta, 'v': d_theta_dt}
def sol(t, theta0): S = np.sin(0.5 * (theta0)) K_S = ellipk(S**2) omega_0 = np.sqrt(9.81) sn, cn, dn, ph = ellipj(K_S - omega_0 * t, S**2) theta = 2.0 * np.arcsin(S * sn) d_sn_du = cn * dn d_sn_dt = -omega_0 * d_sn_du d_theta_dt = 2.0 * S * d_sn_dt / np.sqrt(1.0 - (S * sn)**2) return np.stack([theta, d_theta_dt], axis=1)
def compute_rational_approximation(self): # Compute shifts and quadrature weights m, M = 10e-6, 10e6 k2 = m / M kp = special.ellipk(1.0 - k2) t = 1j * np.arange(0.5, self.n_shift) * kp / self.n_shift sn, cn, dn, ph = special.ellipj(t.imag, 1-k2) cn = 1./cn cn *= cn shifts = -m*sn*sn*cn weights = 2.*cn*dn*kp*np.sqrt(m) / (np.pi * self.n_shift) return (shifts, weights)
def find_collision_pts(e, q): """ Finds the collision points for a period q, given an eccentricity e """ collisions_dict = {} a = semi_axes[e][0] b = semi_axes[e][1] l = period_lambda_dict[q][e] k_l_sq = ((a**2) - (b**2)) / ((a**2) - (l**2)) for j in range(int(q)): d_l_q = (4 * (special.ellipk(k_l_sq))) / int(q) t_j = (special.ellipk(k_l_sq)) + j * d_l_q collisions_dict[str(j).zfill(2)] = special.ellipj(t_j, k_l_sq)[3] return (collisions_dict)
def analyticalSolution(x,t,k,a0,a1): """ Returns the cnoidal solution with parameters k,a0,a1 at (x,t) (possibly arrays) """ g = 9.81 kappa = np.sqrt(3*a1)/(2*np.sqrt(a0*(a0+a1)*(a0+(1-k*k)*a1))) h0 = a0+ a1*special.ellipe(k)/special.ellipk(k) c = np.sqrt(g*a0*(a0+a1)*(a0+(1.-k*k)*a1))/h0 sn,cn,dn,ph = special.ellipj(kappa*(x-c*t),k) h = a0+a1*dn**2 u = c*(1-h0/h) return h,u
def sol_an(theta_0, g, l, t, dt): """ Aquí se calcula la solución analítica de theta, donde tambien se calcula la velocidad pero esta utilizando el método de Euler. """ f_sen = np.sin(theta_0 / 2) k_ellip = special.ellipk(f_sen**2) ellip_function = special.ellipj(k_ellip - np.sqrt(g / l) * t, f_sen**2)[0] theta = 2 * np.arcsin(f_sen * ellip_function) #Calcular v F = lambda theta, g, l: -(g / l) * np.sin(theta) v = np.zeros(len(theta)) for i in range(len(theta) - 1): v[i + 1] = v[i] + dt * F(theta[i], g, l) return theta, v
def calc_psi_z(qz, zp, zm, En, aa): """ angle used in polar geodesic calculations Parameters: qz (float) zp (float): polar root zm (float): polar root En (float): energy aa (float): spin Returns: psi_z (float) """ ktheta = (aa**2 * (1 - En**2) * zm**2) / zp**2 u = (2 * (pi / 2.0 + qz) * ellipk(ktheta)) / pi m = ktheta __, __, __, ph = ellipj(u, m) return ph
def calc_psi_r(qr, r1, r2, r3, r4): """ radial geodesic angle Parameters: qr (float) r1 (float): radial root r2 (float): radial root r3 (float): radial root r4 (float): radial root Returns: psi_r (float) """ kr = ((r1 - r2) * (r3 - r4)) / ((r1 - r3) * (r2 - r4)) u = (qr * ellipk(kr)) / pi m = kr __, __, __, ph = ellipj(u, m) return ph
def calc_zq(qz, zp, zm, En, aa): """ function used in computing polar geodesic coordinates Parameters: qz (float) zp (float): polar root zm (float): polar root En (float): energy aa (float): spin Returns: zq (float) """ ktheta = (aa**2 * (1 - En**2) * zm**2) / zp**2 u = (2 * (pi / 2.0 + qz) * ellipk(ktheta)) / pi m = ktheta sn, __, __, __ = ellipj(u, m) return zm * sn
def calc_rq(qr, r1, r2, r3, r4): """ function used in computing radial geodesic coordinates Parameters: qr (float) r1 (float): radial root r2 (float): radial root r3 (float): radial root r4 (float): radial root Returns: rq (float) """ kr = ((r1 - r2) * (r3 - r4)) / ((r1 - r3) * (r2 - r4)) u = (qr * ellipk(kr)) / pi m = kr sn, __, __, __ = ellipj(u, m) return (-(r2 * (r1 - r3)) + (r1 - r2) * r3 * sn**2) / (-r1 + r3 + (r1 - r2) * sn**2)
def Test(): a = np.array([1,2,3,4]) b = np.array((5,6,7,8)) c = np.array([[1,2],[6,7]]) x = np.arange(0,2,0.1) print(a,b,'\n',c) print(a.shape, b.shape, c.shape) print(x) l=[] for i in range(0, 10): l.append(i) y = l y = np.logspace(0, 1, 12, base=2, endpoint=False) print('y= ', y) x = np.linspace(0, 2*np.pi, 10, endpoint=False) y = np.sin(x) x = special.ellipj(1,1) print('x= ', x) print('y= ', y) print(linalg.det(c))
x = 0.1 y = 17.0 * 2.0 * np.pi / 60.0 + 0.1 z = 0.0 E = 0.5 * (Ix * x**2 + Iy * y**2 + Iz * z**2) L = (Ix * x)**2 + (Iy * y)**2 + (Iz * z)**2 l = np.sqrt((Ix - Iy) * (2 * Iz * E - L) / Ix / Iy / Iz) k = np.sqrt((Iz - Ix) / (Iy - Ix) * (2 * Iy * E - L) / (2 * Iz * E - L)) a = np.sqrt((2 * Iy * E - L) / Ix / (Iy - Ix)) b = np.sqrt((2 * Iz * E - L) / Iy / (Iz - Iy)) c = np.sqrt((2 * Iy * E - L) / Iz / (Iy - Iz)) K = ellipk(k) m = k**2 ax131.plot(t, omega_x, color="c", linestyle="-", label="Numerical Solution") ax131.plot(t, a * ellipj(l * t, m)[1], color="m", linestyle="-.", label="Exact Solution") ax131.set_xlim(0, end_time) ax131.set_ylim(-0.1, 0.1) ax131.set_ylabel("$\omega_x\;\;$[rad/s]") ax131.legend(loc="upper right").get_frame().set_alpha(1) ax132.plot(t, omega_y, color="c", linestyle="-", label="Numerical Solution") ax132.plot(t, b * ellipj(l * t - K, m)[2], color="m", linestyle="-.", label="Exact Solution") ax132.set_xlim(0, end_time) ax132.set_ylim(1.8800, 1.8812)
import numpy as np from scipy import io as spio from scipy import special as special print("\n\r") print("====================================") print("特殊函数:scipy.special") print("====================================") print("=== 贝塞尔函数,如scipy.special.jn() (整数n阶贝塞尔函数) ===") data = special.jn(22, 33) print(data) arrA = [11, 22, 33] arrB = [44, 55, 66] data = special.jn(arrA, arrB) print(data) print("=== 椭圆函数: scipy.special.ellipj() (雅可比椭圆函数,……) ===") x1 = [0.4, 0.5, 0.3] x2 = [0.6, 0.7, 0.8] #x1, x2, out1=None, out2=None, out3=None, out4=None data = special.ellipj(x1, x2) print(data) print("=== 伽马函数:scipy.special.gamma() ===") #还要注意scipy.special.gammaln,这个函数给出对数坐标的伽马函数,因此有更高的数值精度 x = [0.4, 0.6, 0.8] data = special.gammaln(x) print(data)
def test_ellipfun_dn(self): # see comment in ellipfun_sn assert_mpmath_equal(lambda u, m: sc.ellipj(u, m)[2], lambda u, m: mpmath.ellipfun("dn", u=u, m=m), [Arg(-1e6, 1e6), Arg(a=0, b=1)], atol=1e-20)
# MODULES import numpy as np import scipy.integrate as integrate import scipy.special as sp import cmath import matplotlib.pyplot as plt # PARAMETERS q = 2 # coefficient of the linear operator N = 64 # number of Fourier modes D = 128 # number of Floquet modes k = 0.9 # only used for elliptic functions! L = 2 * sp.ellipk(k**2) # period f1 = lambda x: -1 f2 = lambda x: 6 * (k**2) * (sp.ellipj(x, k**2)[0] )**2 # sn(x|k); cf. docs for special.ellipj # FUNCTIONS def frange(start, stop, step): i = start while i < stop: yield i i += step def fourier_coeffs(fun, modes, period): cosines = np.zeros(modes + 1, dtype=np.float_) sines = np.zeros(modes + 1, dtype=np.float) output = np.zeros(2 * modes + 1, dtype=np.complex_)
def analytical_dtheta(t): k = sin(theta0/2) T = ellipk(k**2) (sn, cn, dn, ph) = ellipj(T-t, k) return -2*k*cn
def analytical_theta(t): k = sin(theta0/2) T = ellipk(k**2) (sn, cn, dn, ph) = ellipj(T-t, k) return 2 * arcsin(k*sn)
def test_ellipfun_dn(self): assert_mpmath_equal(lambda u, m: sc.ellipj(u, m)[2], lambda u, m: mpmath.ellipfun("dn", u=u, m=m), [Arg(), Arg(a=0, b=1)])
P0 = 2 * np.pi / omega0 # eccentricity parameter for elliptic functions k = np.sin(phi0 / 2) m = k**2 # discrete times t = np.linspace(0, 4 * P0, 4000) # small angle approx phi_approx = phi0 * np.cos(omega0 * t) # exact solution x = omega0 * t P = (4 / omega0) * special.ellipk(m) sn_x, cn_x, dn_x, ph_x = special.ellipj(x, m) phi = 2 * np.arcsin(k * cn_x) # define figure fig = plt.figure() ax = fig.add_subplot(111, aspect='equal', autoscale_on=False, xlim=(-1.25, 1.25), ylim=(-1.25, 1.25)) ax.grid() line1, = ax.plot([], [], 'o--', color='grey', lw=2) line2, = ax.plot([], [], 'o-', color='black', lw=2) ax.text(-1, 0.5, 'grey, dashed: small-angle approx') ax.text(-1, 0.75, 'black, solid: exact solution')
def dnsq(eta,m): sn,cn,dn,ph = ellipj(eta,m) return dn*dn
def ell(u,m): return ssp.ellipj(u,m)
""" Copyright 2009-2012 Karsten Ahnert Copyright 2009-2012 Mario Mulansky Stochastic euler stepper example and Ornstein-Uhlenbeck process Distributed under the Boost Software License, Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt) """ from pylab import * from scipy import special data1 = loadtxt("elliptic1.dat") data2 = loadtxt("elliptic2.dat") data3 = loadtxt("elliptic3.dat") sn1, cn1, dn1, phi1 = special.ellipj(data1[:, 0], 0.51) sn2, cn2, dn2, phi2 = special.ellipj(data2[:, 0], 0.51) sn3, cn3, dn3, phi3 = special.ellipj(data3[:, 0], 0.51) semilogy(data1[:, 0], abs(data1[:, 1] - sn1)) semilogy(data2[:, 0], abs(data2[:, 1] - sn2), 'ro') semilogy(data3[:, 0], abs(data3[:, 1] - sn3), '--') show()
C=I3 - I1 N=np.linspace(math.sqrt(I1*T), math.sqrt(I3*T), 20) for (lnumber, L) in enumerate(N): (ktype, k)=modulus(I, T, L) outfile=open("L{}.dat".format(lnumber), "w") outfile.write("t\t omega1 \t omega2 \t omega3\n") print("#"*32) print("Sequence={}, L={}, ktype={}".format(lnumber, L, ktype)) U=I3*T-L**2 V=L**2 - I2*T if V > 0: for t in np.linspace(0, 0.1, 50): argument=-t*math.sqrt(C*V/I1*I2*I3) if not int(k - 1)==1: (sn, cn, dn, ph) = ellipj(argument, k) else: sn=np.tanh(argument) cs=np.cosh(argument) if int(cs - 1)==0: print("********zero cosh for {}".format(argument)) cn=0 dn=0 else: cn=1/cs dn=1/cs #print(U, I1, C) omega1=math.sqrt(U/I1*C)*sn omega2=math.sqrt(U/I2*B)*cn omega3=math.sqrt(U/I3*B)*dn print("{}\t{}".format(t, omega1))
""" Copyright 2011 Mario Mulansky Copyright 2012 Karsten Ahnert Stochastic euler stepper example and Ornstein-Uhlenbeck process Distributed under the Boost Software License, Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt) """ from pylab import * from scipy import special data1 = loadtxt("elliptic1.dat") data2 = loadtxt("elliptic2.dat") data3 = loadtxt("elliptic3.dat") sn1,cn1,dn1,phi1 = special.ellipj( data1[:,0] , 0.51 ) sn2,cn2,dn2,phi2 = special.ellipj( data2[:,0] , 0.51 ) sn3,cn3,dn3,phi3 = special.ellipj( data3[:,0] , 0.51 ) semilogy( data1[:,0] , abs(data1[:,1]-sn1) ) semilogy( data2[:,0] , abs(data2[:,1]-sn2) , 'ro' ) semilogy( data3[:,0] , abs(data3[:,1]-sn3) , '--' ) show()
def ellipj_(k): return ellipj(k*k)
import numpy as np import scipy.special as sp import matplotlib.pyplot as plt #종속변수 설정하기 k = np.array([0.25, 0.5, 0.75, 0.9, 0.99]) u = np.linspace(-0.5, 0.5, 100) #타원적분 K(k)함수값 구하기 integralk = sp.ellipk(k) print(integralk) #타원함수 모음집에서 sn,cu,dn함수만 꺼내오기, 주기를 1로 normalize 하기 ellipf1 = np.array([sp.ellipj(u * 4 * integralk[0], k[0])]) ellipf2 = np.array([sp.ellipj(u * 4 * integralk[1], k[1])]) ellipf3 = np.array([sp.ellipj(u * 4 * integralk[2], k[2])]) ellipf4 = np.array([sp.ellipj(u * 4 * integralk[3], k[3])]) ellipf5 = np.array([sp.ellipj(u * 4 * integralk[4], k[4])]) y1 = ellipf1[0, 0, :] y2 = ellipf2[0, 0, :] y3 = ellipf3[0, 0, :] y4 = ellipf4[0, 0, :] y5 = ellipf5[0, 0, :] #그래프 그리기 plt.plot(u, y1) plt.plot(u, y2) plt.plot(u, y3) plt.plot(u, y4) plt.plot(u, y5) plt.show()
def JacobiAmplitude(q, k): k2 = np.power(k, 2) JA = sp.ellipj(q, k2) return JA
def cnsoln(xspace,t,kappa,k0): _,cn,_,asd=special.ellipj(kappa*(xspace),k0**2) return (kappa*k0*cn/np.sqrt(2))*np.exp(1j*(kappa**2)*((2*k0**2)-1)*t)
def KdVcnSoln(x, t, kappa, k): _, cn, _, _ = special.ellipj( kappa * (x - 4.0 * kappa**2 * (2.0 * k**2 - 1.0) * t), k**2) return 6.0 * kappa**2 * k**2 * cn**2
def contour_integral_quad( lazy_tensor, rhs, inverse=False, weights=None, shifts=None, max_lanczos_iter=20, num_contour_quadrature=None, shift_offset=0, ): r""" Performs :math:`\mathbf K^{1/2} \mathbf b` or `\mathbf K^{-1/2} \mathbf b` using contour integral quadrature. :param gpytorch.lazy.LazyTensor lazy_tensor: LazyTensor representing :math:`\mathbf K` :param torch.Tensor rhs: Right hand side tensor :math:`\mathbf b` :param bool inverse: (default False) whether to compute :math:`\mathbf K^{1/2} \mathbf b` (if False) or `\mathbf K^{-1/2} \mathbf b` (if True) :param int max_lanczos_iter: (default 10) Number of Lanczos iterations to run (to estimate eigenvalues) :param int num_contour_quadrature: How many quadrature samples to use for approximation. Default is in settings. :rtype: torch.Tensor :return: Approximation to :math:`\mathbf K^{1/2} \mathbf b` or :math:`\mathbf K^{-1/2} \mathbf b`. """ import numpy as np from scipy.special import ellipj, ellipk if num_contour_quadrature is None: num_contour_quadrature = settings.num_contour_quadrature.value() output_batch_shape = _mul_broadcast_shape(lazy_tensor.batch_shape, rhs.shape[:-2]) preconditioner, preconditioner_lt, _ = lazy_tensor._preconditioner() def sqrt_precond_matmul(rhs): if preconditioner_lt is not None: solves, weights, _, _ = contour_integral_quad(preconditioner_lt, rhs, inverse=False) return (solves * weights).sum(0) else: return rhs # if not inverse: rhs = sqrt_precond_matmul(rhs) if shifts is None: # Determine if init_vecs has extra_dimensions num_extra_dims = max(0, rhs.dim() - lazy_tensor.dim()) lanczos_init = rhs.__getitem__( (*([0] * num_extra_dims), Ellipsis, slice(None, None, None), slice(None, 1, None))).expand(*lazy_tensor.shape[:-1], 1) with warnings.catch_warnings(), torch.no_grad(): warnings.simplefilter( "ignore", NumericalWarning) # Supress CG stopping warning _, lanczos_mat = linear_cg( lambda v: lazy_tensor._matmul(v), rhs=lanczos_init, n_tridiag=1, max_iter=max_lanczos_iter, tolerance=1e-5, max_tridiag_iter=max_lanczos_iter, preconditioner=preconditioner, ) lanczos_mat = lanczos_mat.squeeze( 0 ) # We have an extra singleton batch dimension from the Lanczos init """ K^{-1/2} b = 2/pi \int_0^\infty (K - t^2 I)^{-1} dt We'll approximate this integral as a sum using quadrature We'll determine the appropriate values of t, as well as their weights using elliptical integrals """ # Compute an approximate condition number # We'll do this with Lanczos try: if settings.verbose_linalg.on(): settings.verbose_linalg.logger.debug( f"Running symeig on a matrix of size {lanczos_mat.shape}.") approx_eigs = lanczos_mat.symeig()[0] if approx_eigs.min() <= 0: raise RuntimeError except RuntimeError: approx_eigs = lazy_tensor.diag() max_eig = approx_eigs.max(dim=-1)[0] min_eig = approx_eigs.min(dim=-1)[0] k2 = min_eig / max_eig # Compute the shifts needed for the contour flat_shifts = torch.zeros(num_contour_quadrature + 1, k2.numel(), dtype=k2.dtype, device=k2.device) flat_weights = torch.zeros(num_contour_quadrature, k2.numel(), dtype=k2.dtype, device=k2.device) # For loop because numpy for i, (sub_k2, sub_min_eig) in enumerate( zip(k2.flatten().tolist(), min_eig.flatten().tolist())): # Compute shifts Kp = ellipk(1 - sub_k2) # Elliptical integral of the first kind N = num_contour_quadrature t = 1j * (np.arange(1, N + 1) - 0.5) * Kp / N sn, cn, dn, _ = ellipj(np.imag(t), 1 - sub_k2) # Jacobi elliptic functions cn = 1.0 / cn dn = dn * cn sn = 1j * sn * cn w = np.sqrt(sub_min_eig) * sn w_pow2 = np.real(np.power(w, 2)) sub_shifts = torch.tensor(w_pow2, dtype=rhs.dtype, device=rhs.device) # Compute weights constant = -2 * Kp * np.sqrt(sub_min_eig) / (math.pi * N) dzdt = torch.tensor(cn * dn, dtype=rhs.dtype, device=rhs.device) dzdt.mul_(constant) sub_weights = dzdt # Store results flat_shifts[1:, i].copy_(sub_shifts) flat_weights[:, i].copy_(sub_weights) weights = flat_weights.view(num_contour_quadrature, *k2.shape, 1, 1) shifts = flat_shifts.view(num_contour_quadrature + 1, *k2.shape) shifts.sub_(shift_offset) # Make sure we have the right shape if k2.shape != output_batch_shape: weights = torch.stack( [w.expand(*output_batch_shape, 1, 1) for w in weights], 0) shifts = torch.stack( [s.expand(output_batch_shape) for s in shifts], 0) # Compute the solves at the given shifts # Do one more matmul if we don't want to include the inverse with torch.no_grad(): solves = minres(lambda v: lazy_tensor._matmul(v), rhs, value=-1, shifts=shifts, preconditioner=preconditioner) no_shift_solves = solves[0] solves = solves[1:] if not inverse: solves = lazy_tensor._matmul(solves) return solves, weights, no_shift_solves, shifts