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()