def test_chebyshev_interpolation(show=False): """Test the inperpolation routine of ChebyshevExtremaGrid""" from psecas import ChebyshevExtremaGrid import numpy as np def psi(x, c): from numpy.polynomial.hermite import hermval return hermval(x, c) * np.exp(-x**2 / 2) N = 40 zmin = -1.1 zmax = 1.5 grid = ChebyshevExtremaGrid(N, zmin, zmax) grid_fine = ChebyshevExtremaGrid(N * 4, zmin, zmax) z = grid_fine.zg y = np.exp(grid.zg) * np.sin(5 * grid.zg) y_fine = np.exp(z) * np.sin(5 * z) y_interpolated = grid.interpolate(z, y) if show: import matplotlib.pyplot as plt plt.figure(2) plt.clf() plt.title("Interpolation with Chebyshev") plt.plot(z, y_fine, "-") plt.plot(z, y_interpolated, "--") plt.plot(grid.zg, y, "+") plt.show() np.testing.assert_allclose(y_fine, y_interpolated, atol=1e-12) return (y_fine, y_interpolated)
def test_chebyshev_differentation(show=False): """Test the differentation routine of FourierGrid""" from psecas import ChebyshevExtremaGrid import numpy as np N = 20 zmin = -1 zmax = 1 grid = ChebyshevExtremaGrid(N, zmin, zmax) assert grid.zg[0] < grid.zg[-1] z = grid.zg y = np.exp(z) * np.sin(5 * z) yp_exac = np.exp(z) * (np.sin(5 * z) + 5 * np.cos(5 * z)) yp_num = np.matmul(grid.d1, y) if show: import matplotlib.pyplot as plt plt.figure(1) plt.clf() plt.title("Differentation with matrix (ChebyshevExtremaGrid)") plt.plot(z, yp_exac, "-") plt.plot(z, yp_num, "--") plt.show() np.testing.assert_allclose(yp_num, yp_exac, atol=1e-16) return (yp_num, yp_exac)
def test_mti_solution(show=False, verbose=False): """Test eigenvalue solver using ChebyshevExtremaGrid""" import numpy as np from psecas import Solver, ChebyshevExtremaGrid from psecas.systems.mti import MagnetoThermalInstability grid = ChebyshevExtremaGrid(N=64, zmin=0, zmax=1) system = MagnetoThermalInstability(grid, beta=1e5, Kn0=200, kx=4 * np.pi) system.beta = 1e5 system.Kn0 = 200 assert system.Kn0 == 200 assert system.beta == 1e5 solver = Solver(grid, system) Ns = np.hstack(np.arange(1, 10) * 16) mode = 0 omega, vec, err = solver.iterate_solver(Ns, mode=mode, tol=1e-8, verbose=verbose) system.get_bx_and_by() if show: from psecas import plot_solution phi = np.arctan(vec[2].imag / vec[2].real) solver.keep_result(omega, vec * np.exp(-1j * phi), mode=mode) plot_solution(system, smooth=True, num=1) np.testing.assert_allclose(1.7814514515967603, omega, atol=1e-8) return err
def test_solver_methods(verbose=False): """Show how the solver class methods can be called directly, using the MTI as an example. This can be useful when setting up a new problem. """ import numpy as np from psecas import Solver, ChebyshevExtremaGrid from psecas.systems.mti import MagnetoThermalInstability from scipy.linalg import eig grid = ChebyshevExtremaGrid(N=64, zmin=0, zmax=1) system = MagnetoThermalInstability(grid, beta=1e5, Kn0=200, kx=4 * np.pi) solver = Solver(grid, system) solver.get_matrix1(verbose=verbose) solver.get_matrix2(verbose=verbose) E, V = eig(solver.mat1.toarray(), solver.mat2.toarray()) # Sort the eigenvalues E, index = solver.sorting_strategy(E) mode = 0 # Choose the eigenvalue mode value only sigma = E[index[mode]] v = V[:, index[mode]] np.testing.assert_allclose(1.7814514515967603, sigma, atol=1e-8) # Compare with solve without using OPinv solver.solve(useOPinv=False, saveall=True) np.testing.assert_allclose(solver.E[mode], sigma, atol=1e-8)
class Example(Solver): def sorting_strategy(self, E): """Sorting strategy. E is a list of eigenvalues""" # Sort from smallest to largest eigenvalue index = np.argsort(np.abs(E)) return (E, index) L = 1 hbar = 1 m = 1 # Create grids N = 32 zmin = 0 grid1 = ChebyshevExtremaGrid(N, zmin, zmax=L, z='x') grid2 = ChebyshevRootsGrid(N, zmin, zmax=L, z='x') grid3 = LegendreExtremaGrid(N, zmin, zmax=L, z='x') grids = list([grid1, grid2, grid3]) # Number of solutions to plot for each grid modes = 10 # Create figure plt.figure(1) plt.clf() fig, axes = plt.subplots(num=1, ncols=modes, nrows=3, sharey=True, sharex=True) for j, grid in enumerate(grids): # Create system
import numpy as np from psecas import Solver, System from psecas import ChebyshevExtremaGrid import pytest """ We set up systems with errors, and see if Psecas gives a reasonable error, i.e., a NameError. """ # Create grid N = 32 grid = ChebyshevExtremaGrid(N, -1, 1) # Create a solver object def test_parser_findmatrices(verbose=False): """ Here we add the value A to the equation without setting it in the system. This should return a NameError """ # Create system system = System(grid, variables='f', eigenvalue='sigma') # Add the first (and only) equation system.add_equation("sigma*z*f = dz(dz(f)) + 2*A*f") with pytest.raises(NameError) as e_info:
def test_bessel_solutions(show=False): import numpy as np from psecas import Solver, System from psecas import ( ChebyshevExtremaGrid, ChebyshevRootsGrid, LegendreExtremaGrid, ) from scipy.special import spherical_jn """ We solve the spherical Bessel equation r² d²f/dr² + 2 r df/dr +(κ²r² - l(l+1)) f = 0 on the domain r=[0, 1] with the boundary condition f(1) = 0. The exact solutions are the spherical bessel functions of the first kind, j_l(κ r). The problem is taken from the paper Tensor calculus in spherical coordinates using Jacobi polynomials Part-II: Implementation and Examples, https://arxiv.org/pdf/1804.09283.pdf by Daniel Lecoanet, Geoff Vasil, Keaton Burns, Ben Brown and Jeff Oishi, and the results of our script can be compared with their Figure 2. A key difference is that we use shifted Chebyshev and Legendre grids, and that we enforce the boundary f(0)=0 explicitly. Since we do not treat the coordinate singularity at r=0 in using Jacobi polynomials in the clever way that Lecoanet+ do, the r^50 scaling shown in their inset is not recovered in our calculation. Nevertheless, the solutions for the eigenvalues and eigenfunctions are in very good agreement with the exact solution, in particular for the extrema grids. """ # do_gen_evp = True # equation = "sigma*r**2*f = -r**2*dr(dr(f)) -2*r*dr(f) +l*(l+1)*f" do_gen_evp = False equation = "sigma*f = -dr(dr(f)) -2/r*dr(f) +l*(l+1)/r**2*f" # Overwrite the default sorting method in the Solver class class Example(Solver): def sorting_strategy(self, E): """Sorting strategy. E is a list of eigenvalues""" # Sort from smallest to largest eigenvalue index = np.argsort(np.abs(E)) return (E, index) # Create grids N = 500 zmin = 0 # grid = ChebyshevRootsGrid(N, zmin, zmax=1, z='r') grid1 = ChebyshevExtremaGrid(N, zmin, zmax=1.0, z='r') grid2 = ChebyshevRootsGrid(N, zmin, zmax=1.0, z='r') grid3 = LegendreExtremaGrid(N, zmin, zmax=1.0, z='r') grids = list([grid1, grid2, grid3]) if show: # Create figure import matplotlib.pyplot as plt plt.figure(1) plt.clf() fig, axes = plt.subplots(num=1, nrows=3, ncols=2, sharex=True) # https://keisan.casio.com/exec/system/1180573465 kappa_ref = 389.4203848348835221199 tols = [1e-13, 1e-4, 1e-13] for jj, grid in enumerate(grids): # Create system system = System(grid, variables='f', eigenvalue='sigma') system.l = 50 # Add the first (and only) equation system.add_equation(equation, boundary=True) # Create a solver object solver = Example(grid, system, do_gen_evp=do_gen_evp) r = np.linspace(grid.zmin, grid.zmax, 4000) E, vec = solver.solve(mode=99) kappa = np.sqrt(E) # Normalize solution to bessel function exact solution at r = 0.5 amp = grid.interpolate(0.5, system.result['f'].real) / spherical_jn( 50, kappa_ref * 0.5) # Test statements print(grid) np.testing.assert_allclose(grid.interpolate( r, system.result['f'].real / amp), spherical_jn(50, kappa_ref * r), atol=tols[jj]) np.testing.assert_allclose(kappa, kappa_ref, atol=tols[jj], rtol=tols[jj]) if show: axes[jj, 0].plot(r, grid.interpolate(r, system.result['f'].real)) axes[jj, 0].plot(r, grid.interpolate(r, system.result['f'].imag)) axes[jj, 0].set_title( type(grid).__name__ + r", $\kappa$-error={:1.3e}".format(kappa.real - kappa_ref)) axes[jj, 1].plot( r, grid.interpolate(r, system.result['f'].real / amp) - spherical_jn(50, kappa_ref * r), ) # axes[jj, 1].plot(z) axes[jj, 1].set_title(r'$f(r) - j_\mathcal{l}(\kappa r)$, ' + type(grid).__name__) if show: plt.show()
# Overwrite the default sorting method in the Solver class class Example(Solver): def sorting_strategy(self, E): """Sorting strategy. E is a list of eigenvalues""" # Sort from smallest to largest eigenvalue index = np.argsort(np.abs(E)) return (E, index) # Create grids N = 500 zmin = 0 # grid = ChebyshevRootsGrid(N, zmin, zmax=1, z='r') grid1 = ChebyshevExtremaGrid(N, zmin, zmax=1.0, z='r') grid2 = ChebyshevRootsGrid(N, zmin, zmax=1.0, z='r') grid3 = LegendreExtremaGrid(N, zmin, zmax=1.0, z='r') grids = list([grid1, grid2, grid3]) # Create figure plt.figure(1) plt.clf() fig, axes = plt.subplots(num=1, nrows=3, ncols=2, sharex=True) # https://keisan.casio.com/exec/system/1180573465 kappa_ref = 389.4203848348835221199 for jj, grid in enumerate(grids):
# Make a Child of the System class class HallMRI(System): def __init__(self, grid, kz, variables, eigenvalue): # Set parameters self.q = 1.5 self.eta = 0.003 self.lh = 1 self.h = 0.25 self.va = 0.002 self.kz = kz super().__init__(grid, variables, eigenvalue) # Create a grid grid = ChebyshevExtremaGrid(N=128, zmin=1, zmax=2, z='r') # grid = ChebyshevTLnGrid(N=128, z='r') variables = ['rho', 'vr', 'vphi', 'vz', 'Aphi', 'bphi'] kz = 2 * np.pi # Create the system system = HallMRI(grid, kz, variables=variables, eigenvalue='sigma') # The linearized equations system.add_equation("-r*sigma*rho = r*dr(vr) + vr + 1j*kz*r*vz") system.add_equation( "-r*r*sigma*vr = - 2*r**(2-q)*vphi + h**2*r*r*dr(rho) + va**2*(DrAphi)") system.add_equation("-sigma*vphi = + (2-q)*r**(-q)*vr - va**2*1j*kz*bphi") system.add_equation("-sigma*vz = h**2*1j*kz*rho")
def test_mri_solution(show=False, verbose=False): """Test eigenvalue solver using ChebyshevExtremaGrid""" import numpy as np from psecas import Solver, ChebyshevExtremaGrid, System class HallMRI(System): def __init__(self, grid, kz, variables, eigenvalue): # Set parameters self.q = 1.5 self.eta = 0.003 self.lh = 1 self.h = 0.25 self.va = 0.002 self.kz = kz super().__init__(grid, variables, eigenvalue) # Create a grid grid = ChebyshevExtremaGrid(N=128, zmin=1, zmax=2, z='r') variables = ['rho', 'vr', 'vphi', 'vz', 'Aphi', 'bphi'] kz = 2 * np.pi # Create the system system = HallMRI(grid, kz, variables=variables, eigenvalue='sigma') # The linearized equations system.add_equation("-r*sigma*rho = r*dr(vr) + vr + 1j*kz*r*vz") system.add_equation( "-r*r*sigma*vr = - 2*r**(2-q)*vphi + h**2*r*r*dr(rho) + va**2*(DrAphi)" ) system.add_equation("-sigma*vphi = + (2-q)*r**(-q)*vr - va**2*1j*kz*bphi") system.add_equation("-sigma*vz = h**2*1j*kz*rho") system.add_equation( "-r*r*sigma*Aphi = + r*r*vr - eta*(DrAphi) + lh*va*1j*kz*r*r*bphi") system.add_equation( "-r*r*sigma*bphi = - 1j*kz*r*r*vphi - 1j*kz*q*r**(2-q)*Aphi - eta*(Drbphi) - lh*va*1j*kz*(DrAphi)" ) # The boundary conditions Aphi_bound = 'r**2*dr(dr(Aphi)) + r*dr(Aphi) - Aphi = 0' system.add_boundary('vr', 'Dirichlet', 'Dirichlet') system.add_boundary('vphi', 'Dirichlet', 'Dirichlet') system.add_boundary('vz', 'Neumann', 'Neumann') system.add_boundary('Aphi', Aphi_bound, Aphi_bound) system.add_boundary('bphi', 'Dirichlet', 'Dirichlet') # Short hands for long expressions for derivatives system.add_substitution( 'DrAphi = r*r*dr(dr(Aphi)) + r*dr(Aphi) - Aphi - kz**2*r*r*Aphi') system.add_substitution( 'Drbphi = r*r*dr(dr(bphi)) + r*dr(bphi) - bphi - kz**2*r*r*bphi') solver = Solver(grid, system) mode = 0 Ns = np.hstack(np.arange(1, 10) * 32) omega, vec, err = solver.iterate_solver(Ns, mode=mode, tol=1e-8, verbose=verbose) if show: from psecas import plot_solution phi = np.arctan(vec[2].imag / vec[2].real) solver.keep_result(omega, vec * np.exp(-1j * phi), mode=mode) plot_solution(system, smooth=True, num=1) np.testing.assert_allclose(0.09892641, omega, atol=1e-8) return err
# Overwrite the default sorting method in the Solver class class Example(Solver): def sorting_strategy(self, E): """Sorting strategy. E is a list of eigenvalues""" # Sort from smallest to largest eigenvalue # We ignore the solutions with negative σ. E[E.real < 0.] = np.inf index = np.argsort(E) return (E, index) # Create grids N = 100 zmin = -1 zmax = 1 grid1 = ChebyshevExtremaGrid(N, zmin, zmax) grid2 = ChebyshevRootsGrid(N, zmin, zmax) grid3 = LegendreExtremaGrid(N, zmin, zmax) grids = list([grid1, grid2, grid3]) def get_normalizing_constant(f1, f2): """EVP solver return eigenmodes with arbitrary sign""" abs_max1_ind = np.argmax(np.abs(f1)) abs_max2_ind = np.argmax(np.abs(f2)) val1 = f1[abs_max1_ind] val2 = f2[abs_max2_ind] A = val1/val2 return A
def test_infinite_well(show=False): import numpy as np from psecas import Solver, System from psecas import ( ChebyshevExtremaGrid, ChebyshevRootsGrid, LegendreExtremaGrid, ) """ This test solves the Schrödinger equation with three different grids. The problem is outlined as follows: Solve the Schrödinger equation -ħ²/2m ∂²/∂x² Φ + V(x) Φ = E Φ for the inifite well potential given by V(x) = 0 for 0 < x < L V(x) = ∞ otherwise For this problem the eigenmodes are sinuisodal and the energies are given by E = n²ħ²π²/2mL² This problem illustrates that the Gauss-Lobatto grids are better at handling problems with a boundary conditions since they have grid points at z=zmin and z=zmax. The test is performed twice, the second time the equation is multiplied by minus 1 on both sides. In the first case, Psecas evaluates mat2 to be the identity matrix and solves a standard evp. The second time, mat2 is not the identity matrix and Psecas therefore solves a generalized evp. This is more time consuming, and equation 1 is therefore the preferred form. Psecas does not automatically convert equation 2 to equation 1, but simply warns the user that a rewrite of the equations could lead to a speed-up. """ equation1 = "E*phi = hbar/(2*m)*dx(dx(phi))" equation2 = "-E*phi = -hbar/(2*m)*dx(dx(phi))" for equation in [equation1, equation2]: # Overwrite the default sorting method in the Solver class class Example(Solver): def sorting_strategy(self, E): """ Sorting strategy. E is a list of eigenvalues """ # Sort from smallest to largest eigenvalue index = np.argsort(np.abs(E)) return (E, index) L = 1 hbar = 1 m = 1 # Create grids N = 32 zmin = 0 grid1 = ChebyshevExtremaGrid(N, zmin, zmax=L, z='x') grid2 = ChebyshevRootsGrid(2 * N, zmin, zmax=L, z='x') grid3 = LegendreExtremaGrid(N, zmin, zmax=L, z='x') grids = list([grid1, grid2, grid3]) tols = [1e-8, 1e-3, 1e-8] # Number of solutions to plot for each grid modes = 10 if show: import matplotlib.pyplot as plt # Create figure plt.figure(1) plt.clf() fig, axes = plt.subplots(num=1, ncols=modes, nrows=3, sharey=True, sharex=True) for j, grid in enumerate(grids): # Create system system = System(grid, variables='phi', eigenvalue='E') system.hbar = hbar system.m = m # Add the first (and only) equation system.add_equation(equation, boundary=True) # Create a solver object solver = Example(grid, system) z = np.linspace(grid.zmin, grid.zmax, 1000) for mode in range(modes): E, vec = solver.solve(mode=mode) np.testing.assert_allclose( E.real / np.pi**2 * 2, -(mode + 1)**2, atol=tols[j], rtol=tols[j], ) if show: # Plottting axes[j, mode].set_title(r"$E/E_0 = ${:1.5f}".format( E.real / np.pi**2 * 2)) axes[j, mode].plot( z, grid.interpolate(z, system.result['phi'].real)) axes[j, mode].plot( z, grid.interpolate(z, system.result['phi'].imag)) axes[j, 0].set_ylabel(type(grid).__name__) if show: plt.show()