def test_legendre_interpolation(show=False):
"""Test the inperpolation routine of LegendreExtremaGrid"""
from psecas import LegendreExtremaGrid
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 = LegendreExtremaGrid(N, zmin, zmax)
grid_fine = LegendreExtremaGrid(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 Legendre")
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_legendre_differentation(show=False):
"""Test the differentation routine of LegendreExtremaGrid"""
from psecas import LegendreExtremaGrid
import numpy as np
N = 20
zmin = -1
zmax = 1
grid = LegendreExtremaGrid(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 (LegendreExtremaGrid)")
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)
"""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
system = System(grid, variables='phi', eigenvalue='E')
system.hbar = hbar
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):
# Create system
system = System(grid, variables='f', eigenvalue='sigma')
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
# Create figure
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()
dvdz_sym = diff(v_sym, z)
d2vdz_sym = diff(dvdz_sym, z)
self.v = lambdify(z, v_sym)(zg)
self.dvdz = lambdify(z, dvdz_sym)(zg)
self.d2vdz = lambdify(z, d2vdz_sym)(zg)
if __name__ == '__main__':
import numpy as np
from psecas import Solver, LegendreExtremaGrid
import matplotlib.pyplot as plt
from psecas import plot_solution
grid = LegendreExtremaGrid(N=200, zmin=-0.5, zmax=0.5)
system = KelvinHelmholtzMiura(grid, kx=2*np.pi, B=0.2, z1=0, a=1/25.0)
solver = Solver(grid, system)
Ns = np.arange(1, 8) * 64
omega, v, err = solver.iterate_solver(Ns, verbose=True, tol=1e-8)
print('\nB={:1.2f} has gammma*2a = {:1.6f}'.format(system.B,
omega.real*system.a*2))
plot_solution(system, num=1, filename='Ryu_and_Frank1996.pdf')
# Make figure 4 in Miura & Pritchett (1982)
# This calculation takes some time.
if True: