def test_FinDiff_1d(self):
h, f, fxe, fxxe = self._prep_1d_func()
d_dx = fd.FinDiff(h, dims=[0], acc=4)
d2_dx2 = fd.FinDiff(h, dims=[0, 0], acc=4)
fx = d_dx(f)
self._assertAlmostEqual(fx, fxe)
fxx = d2_dx2(f)
self._assertAlmostEqual(fxx, fxxe)
def test_matrix_repr_with_different_accs(self):
# issue 28
shape = (11, )
d1 = findiff.FinDiff(0, 1, 2).matrix(shape)
d2 = findiff.FinDiff(0, 1, 2, acc=4).matrix(shape)
self.assertTrue(np.max(np.abs((d1 - d2).toarray())) > 1)
x = np.linspace(0, 10, 11)
f = x**2
df = d2.dot(f)
np.testing.assert_almost_equal(2 * np.ones_like(f), df)
def test_diff_2d_ord2_acc2(self):
x = np.linspace(-1, 1, 500)
y = np.linspace(-1, 1, 500)
dx = dy = x[1] - x[0]
X, Y = np.meshgrid(x, y, indexing='ij')
f = Y**3 * X**3
fxxe = 6 * X * Y**3
fxx = fd.FinDiff(h=[dx, dy], dims=[0, 0], acc=2)(f)
err = np.max(np.abs(fxxe - fxx))
self.assertAlmostEqual(0, err, 4)
fyye = 6 * X**3 * Y
fyy = fd.FinDiff(acc=2, h=[dx, dy], dims=[1, 1])(f)
err = np.max(np.abs(fyye - fyy))
self.assertAlmostEqual(0, err, 6)
def derivative(self, *, order=1):
r"""Differentiate a FDataGrid object.
It is calculated using central finite differences when possible. In
the extremes, forward and backward finite differences with accuracy
2 are used.
Args:
order (int, optional): Order of the derivative. Defaults to one.
Examples:
First order derivative
>>> fdata = FDataGrid([1,2,4,5,8], range(5))
>>> fdata.derivative()
FDataGrid(
array([[[ 0.5],
[ 1.5],
[ 1.5],
[ 2. ],
[ 4. ]]]),
sample_points=[array([0, 1, 2, 3, 4])],
domain_range=array([[0, 4]]),
...)
Second order derivative
>>> fdata = FDataGrid([1,2,4,5,8], range(5))
>>> fdata.derivative(order=2)
FDataGrid(
array([[[ 3.],
[ 1.],
[-1.],
[ 2.],
[ 5.]]]),
sample_points=[array([0, 1, 2, 3, 4])],
domain_range=array([[0, 4]]),
...)
"""
order_list = np.atleast_1d(order)
if order_list.ndim != 1 or len(order_list) != self.dim_domain:
raise ValueError("The order for each partial should be specified.")
operator = findiff.FinDiff(
*[(1 + i, p, o)
for i, (p, o) in enumerate(zip(self.sample_points, order_list))])
data_matrix = operator(self.data_matrix.astype(float))
if self.dataset_name:
dataset_name = "{} - {} derivative".format(self.dataset_name,
order)
else:
dataset_name = None
fdatagrid = self.copy(data_matrix=data_matrix,
dataset_name=dataset_name)
return fdatagrid
def test_FinDiff_2d(self):
xy0 = [-5, -5]
xy1 = [5, 5]
nxy = [100, 100]
x = np.linspace(xy0[0], xy1[0], nxy[0])
y = np.linspace(xy0[1], xy1[1], nxy[1])
dxy = [x[1] - x[0], y[1] - y[0]]
X, Y = np.meshgrid(x, y, indexing='ij')
f = X**3 * Y**3
fxe = 3 * X**2 * Y**3
fyye = 6 * X**3 * Y
d_dx = fd.FinDiff(dxy, dims=[0], acc=4)
fx = d_dx(f)
d2_dy2 = fd.FinDiff(dxy, dims=[1, 1], acc=4)
fyy = d2_dy2(f)
self._assertAlmostEqual(fxe, fx)
self._assertAlmostEqual(fyye, fyy)
def test_diff_1d_ord1_acc4(self):
x = np.linspace(-3, 3, 50)
h = x[1] - x[0]
y = x**2
y1e = 2 * x
y1 = fd.FinDiff(h=h, dims=[0], acc=4)(y)
err = np.max(np.abs(y1e - y1))
self.assertAlmostEqual(0, err)
def test_diff_1d_ord2_acc4(self):
x = np.linspace(-3, 3, 50)
h = x[1] - x[0]
y = x**3
y2e = 6 * x
y2 = fd.FinDiff(h=h, dims=[0, 0], acc=4)(y)
err = np.max(np.abs(y2e - y2))
self.assertAlmostEqual(0, err)
def system_fit_error(t, U, system, order="inf"):
dt = np.mean(np.diff(t))
ddt = fd.FinDiff(0, dt, acc=6)
df = rc.SYSTEMS[system]['df']
err = ddt(U) - df(t, U)
if order == "inf":
return np.max(np.abs(err))
if order == 2:
return np.mean(np.sum(err**2, axis=1)**0.5)
def get_pder_loop1(pi, dx, filename):
t0 = time.time()
lenpar = len(Grid.valueref)
idx = Grid.center
t1 = time.time()
p0 = pi[idx, idx, idx, idx, idx, idx, :, :, :]
print("Done p0 in %s sec" % str(t1 - t0))
dpdx = np.array([
findiff.FinDiff((i, dx[i], 1), acc=4)(pi)[idx, idx, idx, idx, idx,
idx, :, :, :]
for i in range(lenpar)
])
t0 = time.time()
print("Done dpdx in %s sec" % str(t1 - t0))
# Second derivatives
d2pdx2 = np.array([
findiff.FinDiff((i, dx[i], 2), acc=2)(pi)[idx, idx, idx, idx, idx,
idx, :, :, :]
for i in range(lenpar)
])
t1 = time.time()
print("Done d2pdx2 in %s sec" % str(t1 - t0))
d2pdxdy = np.array([[
i, j,
findiff.FinDiff((i, dx[i], 1), (j, dx[j], 1),
acc=2)(pi)[idx, idx, idx, idx, idx, idx, :, :, :]
] for (i, j) in combinations(range(lenpar), 2)])
t0 = time.time()
print("Done d2pdxdy in %s sec" % str(t1 - t0))
# Third derivatives: we only need it for A_s, so I do this by hand
d3pdx3 = np.array([
findiff.FinDiff((i, dx[i], 3))(pi)[idx, idx, idx, idx, idx,
idx, :, :, :] for i in range(lenpar)
])
t1 = time.time()
print("Done d3pdx3 in %s sec" % str(t1 - t0))
allder = (p0, dpdx, d2pdx2, d2pdxdy, d3pdx3)
np.save(filename, allder)
return allder
def test_diff_3d_ord2_acc2(self):
x = np.linspace(-0.5, 0.5, 100)
y = np.linspace(-0.5, 0.5, 100)
z = np.linspace(-0.5, 0.5, 100)
dx = dy = dz = x[1] - x[0]
X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
f = Y**3 * X**3 * Z**3
fxxe = 6 * X * Y**3 * Z**3
fxx = fd.FinDiff(acc=2, h=[dx, dy, dz], dims=[0, 0])(f)
err = np.max(np.abs(fxxe - fxx))
self.assertAlmostEqual(0, err, 4)
fyye = 6 * X**3 * Y * Z**3
fyy = fd.FinDiff(acc=2, h=[dx, dy, dz], dims=[1, 1])(f)
err = np.max(np.abs(fyye - fyy))
self.assertAlmostEqual(0, err, 6)
fzze = 6 * X**3 * Y**3 * Z
fzz = fd.FinDiff(acc=2, h=[dx, dy, dz], dims=[2, 2])(f)
err = np.max(np.abs(fzze - fzz))
self.assertAlmostEqual(0, err, 6)
def test_diff_3d_df_dxdy_acc6(self):
x = np.linspace(-12.5, 12.5, 25)
y = np.linspace(-12.5, 12.5, 25)
z = np.linspace(-12.5, 12.5, 25)
dx = dy = dz = x[1] - x[0]
X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
f = Y**3 * X**3 * Z**3
fxye = 9 * X**2 * Y**2 * Z**3
fxy = fd.FinDiff(acc=4, h=[dx, dy, dz], dims=[0, 1])(f)
err = np.max(np.abs(fxye - fxy))
self.assertAlmostEqual(0, err, 4)
def block1():
x, y, z = [np.linspace(0, 3, 10)] * 3
dx, dy, dz = x[1] - x[0], y[1] - y[0], z[1] - z[0]
X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
f = X**2 * Y * Z #np.sin(X) * np.cos(Y) * np.sin(Z)
D = fd.FinDiff((0, dx, 1), (1, dy, 1), (2, dz, 1))
res = D(f)
print(res)
print("")
def block2():
x, y = [np.random.normal(size=30)] * 2
dx, dy = 0.01, 0.01 #x[1] - x[0], y[1] - y[0]
X, Y = np.meshgrid(x, y, indexing='ij')
f = X**2 * Y #np.sin(X) * np.cos(Y) * np.sin(Z)
D = fd.FinDiff((0, dx, 1), (1, dy, 1))
res = D(f)
print(res)
print("")
def block3():
x = np.linspace(1, 10, 10000)
dx = 0.001 #x[1] - x[0], y[1] - y[0]
# X, Y = np.meshgrid(x, indexing='ij')
f = x**2 #np.sin(X) * np.cos(Y) * np.sin(Z)
D = fd.FinDiff((0, dx, 1))
res = D(f)
plt.plot(x, res)
print(res)
print("")
def get_pder_loop2b(pi, dx, filename):
# Findiff syntax is Findiff((axis, delta of uniform grid along the axis, order of derivative, accuracy))
lenpar = len(Grid.valueref)
idx = Grid.center
t0 = time.time()
d3pdxdy2 = np.array([[
i, j,
findiff.FinDiff((i, dx[i], 1), (j, dx[j], 2))(pi)[idx, idx, idx, idx,
idx, idx, :, :, :]
] for (i, j) in combinations(range(lenpar), 2)])
t1 = time.time()
print("Done d3pdxdy2 in %s sec" % str(t1 - t0))
allder = (d3pdxdy2, )
np.save(filename, allder)
return allder
def get_pder_loop2a(pi, dx, filename):
""" Calculates the derivative aroud the Grid.valueref points. Do this only once.
gridshape is 2 * order + 1, times the number of free parameters
pi is of shape gridshape, n multipoles, k length, P columns (zeroth being k's)"""
# Findiff syntax is Findiff((axis, delta of uniform grid along the axis, order of derivative, accuracy))
lenpar = len(Grid.valueref)
idx = Grid.center
t1 = time.time()
d3pdx2dy = np.array([[
i, j,
findiff.FinDiff((i, dx[i], 2), (j, dx[j], 1))(pi)[idx, idx, idx, idx,
idx, idx, :, :, :]
] for (i, j) in combinations(range(lenpar), 2)])
t0 = time.time()
print("Done d3pdx2dy in %s sec" % str(t1 - t0))
allder = (d3pdx2dy, )
np.save(filename, allder)
return allder
def ApplyFindiff(u, dx, d, acc):
"""
Takes dth derivative of data using the findiff module.
:param u: 1D-data vector to be differentiated
:param dx: Grid spacing; assumes uniform spacing
:param d: Order
:param acc: Order of accuracy of FD scheme
:return: Vector of the same dimensions as u containing the derivative approximations
"""
n = u.size
ux = np.zeros(n, dtype=np.float64)
dd_dxd = findiff.FinDiff(
0, dx, d,
acc=acc) # initialization FD-Operator, for axis 0, and order acc
ux[:] = dd_dxd(u) # apply operator
return ux
def ddx(groupSameYZ, nBranch, dx):
trans0 = np.array([x[0] for x in groupSameYZ])
trans1 = np.array([x[1] for x in groupSameYZ])
trans2 = np.array([x[2] for x in groupSameYZ])
trans3 = np.array([x[3] for x in groupSameYZ])
# tran = np.vstack((trans0,trans1,trans2,trans3))
l = [[]] * nBranch
for i in range(nBranch):
# 将频率的单位 cm^-1 转化成 THZ
y = np.array([x[4 + i] * 0.02998 for x in groupSameYZ])
# print (y)
d_dx = fi.FinDiff(0, dx, 1)
res1 = np.array(d_dx(y))
l[i] = res1
l = np.array(l)
tran = np.vstack((trans0, trans1, trans2, trans3, l))
tran = tran.transpose(1, 0)
return tran
def sum_rule(params, f_xc, density_gs, vxc_gs, i):
""" Check various exact conditions, see Ullrich TDDFT book """
N = params.Nspace
# One dimensional derivative operator
d_dx = findiff.FinDiff(0, params.dx, acc=10)
# Check Zero-Force theorem on ground-state XC potential
grad_vxc = d_dx(vxc_gs)
error = np.dot(grad_vxc, density_gs) * params.dx
print('Error in the ground-state zero-force theorem: {}'.format(error))
# Check weakest ZFSR: (8.12) in Ullrich TTDFT.
error_weak = 0
grad_den = d_dx(density_gs)
for j in range(N):
for k in range(N):
error_weak += f_xc[j,
k] * grad_den[j] * density_gs[k] * params.dx**2
print('Weak ZFSR error is {}'.format(error))
# Check strong ZFSR
tmp = np.zeros(N, dtype=complex)
for k in range(N):
tmp[k] = np.sum(f_xc[k, :] * grad_den[:]) * params.dx
error_strong = np.sum(abs(tmp - grad_vxc)) * params.dx
print('Strong ZFSR error is {}'.format(error))
# Plot...
plt.clf()
plt.title('Weak ZFSR error: {:0.3e}. Strong ZFSR error: {:1.3e}'.format(
error_weak, error_strong))
plt.plot(tmp.real, label='int grad(n) f_xc')
plt.plot(grad_vxc.real, label='grad(v_xc)')
plt.legend()
plt.savefig(f'ZFSR_{i}.png')
def get_pder_lin(pi, dx, filename):
""" Calculates the derivative aroud the Grid.valueref points. Do this only once.
gridshape is 2 * order + 1, times the number of free parameters
pi is of shape gridshape, n multipoles, k length, P columns (zeroth being k's)"""
# Findiff syntax is Findiff((axis, delta of uniform grid along the axis, order of derivative, accuracy))
t0 = time.time()
lenpar = len(Grid.valueref)
idx = Grid.center
t1 = time.time()
p0 = pi[idx, idx, idx, idx, idx, :, :, :]
print("Done p0 in %s sec" % str(t1 - t0))
dpdx = np.array([
findiff.FinDiff((i, dx[i], 1), acc=4)(pi)[idx, idx, idx, idx,
idx, :, :, :]
for i in range(lenpar)
])
t0 = time.time()
print("Done dpdx in %s sec" % str(t1 - t0))
# Second derivatives
d2pdx2 = np.array([
findiff.FinDiff((i, dx[i], 2), acc=2)(pi)[idx, idx, idx, idx,
idx, :, :, :]
for i in range(lenpar)
])
t1 = time.time()
print("Done d2pdx2 in %s sec" % str(t1 - t0))
d2pdxdy = np.array([[
i, j,
findiff.FinDiff((i, dx[i], 1), (j, dx[j], 1),
acc=2)(pi)[idx, idx, idx, idx, idx, :, :, :]
] for (i, j) in combinations(range(lenpar), 2)])
t0 = time.time()
print("Done d2pdxdy in %s sec" % str(t1 - t0))
# Third derivatives: we only need it for A_s, so I do this by hand
d3pdx3 = np.array([
findiff.FinDiff((i, dx[i], 3))(pi)[idx, idx, idx, idx, idx, :, :, :]
for i in range(lenpar)
])
t1 = time.time()
print("Done d3pdx3 in %s sec" % str(t1 - t0))
t1 = time.time()
d3pdx2dy = np.array([[
i, j,
findiff.FinDiff((i, dx[i], 2), (j, dx[j], 1))(pi)[idx, idx, idx, idx,
idx, :, :, :]
] for (i, j) in combinations(range(lenpar), 2)])
t0 = time.time()
print("Done d3pdx2dy in %s sec" % str(t1 - t0))
d3pdxdy2 = np.array([[
i, j,
findiff.FinDiff((i, dx[i], 1), (j, dx[j], 2))(pi)[idx, idx, idx, idx,
idx, :, :, :]
] for (i, j) in combinations(range(lenpar), 2)])
t1 = time.time()
print("Done d3pdxdy2 in %s sec" % str(t1 - t0))
d3pdxdydz = np.array([[
i, j, k,
findiff.FinDiff((i, dx[i], 1), (j, dx[j], 1),
(k, dx[k], 1))(pi)[idx, idx, idx, idx, idx, :, :, :]
] for (i, j, k) in combinations(range(lenpar), 3)])
t0 = time.time()
print("Done d3pdxdydz in %s sec" % str(t1 - t0))
#allder = (p0, dpdx, d2pdx2, d2pdxdy, d3pdx3)
#allder = (d3pdx2dy, d3pdxdy2)
#allder = (d3pdxdydz, )
allder = (p0, dpdx, d2pdx2, d2pdxdy, d3pdx3, d3pdx2dy, d3pdxdy2, d3pdxdydz)
np.save(filename, allder)
return allder
# =============================================================================
# Read data, convert speed to km/s
# =============================================================================
telemetry = pd.read_csv('Telemetry_Data.csv').drop(columns='Unnamed: 0')
telemetry['Speed'] = telemetry['Speed'] / 1000 # In km/s
# =============================================================================
# Define conversion between frames and seconds, useful for computing
# derivatives as well as for defining the time coordinate
# =============================================================================
k = 2301 / (9 * 60 + 3)
# =============================================================================
# Define operators for first and second derivatives
# =============================================================================
d_df = fd.FinDiff(0, 1, acc=8) # axis zero, first derivative, accuracy = 8
d2_df2 = fd.FinDiff(0, 1, 2, acc=8)
# =============================================================================
# Smooth the scraped data, replace every point with the polynomial fit to all
# data in a window centered at that point
# =============================================================================
window_size = 101
poly_order = 1
Alt = savgol_filter(telemetry['Alt'], window_size, poly_order) # In km
Speed = savgol_filter(telemetry['Speed'], window_size, poly_order) # In km/s
# =============================================================================
# Compute both components of speed vector, then smooth.
# =============================================================================
def vTG(
z,
u,
v,
b,
k,
l,
Kv,
Kb,
BCv_upper="rigid",
BCv_lower="rigid",
BCb_upper="constant",
BCb_lower="constant",
):
"""
Solver for the viscous Taylor Goldstein equation for the case of constant viscous/diffusive coefficients.
Parameters
----------
z : array
Height [m]. Must be equally spaced.
u : array
Zonal velocity [m s-1].
v : array
Meridional velocity [m s-1].
b : array
Buoyancy [m s-2].
k : float
Zonal wavenumber (angular) [rad m-1].
l : float
Meridional wavenumber (angular) [rad m-1].
Kv : float
Momentum diffusivity or viscosity [m2 s-1].
Kv : float
Buoyancy diffusivity [m2 s-1].
BCv_upper : string
Upper boundary condition on velocity, either "rigid" (default) or "frictionless".
BCv_lower : string
Lower boundary condition on velocity, either "rigid" (default) or "frictionless".
BCb_upper : string
Upper boundary condition on buoyancy, either "constant" (default) or "insulating".
BCb_lower : string
Lower boundary condition on buoyancy, either "constant" (default) or "insulating".
Returns
-------
om : array
Complex frequency where the real part is the growth rate and imaginary part is the frequency (angular) [rad s-1].
wvec : 2d array
Eigenvectors of vertical velocity.
bvec : 2d array
Eigenvectors of buoyancy.
"""
z = np.asarray(z)
u = np.asarray(u)
v = np.asarray(v)
b = np.asarray(b)
if not (z.size == u.size == v.size == b.size):
raise ValueError(
"Size of z, u, v and b must be equal, z.size = {}, u.size = {}, v.size = {}, b.size = {}."
.format(z.size, u.size, v.size, b.size))
dz = z[1] - z[0]
# check for equally spaced z
if not np.all(np.diff(z) == dz):
raise ValueError("z values are not equally spaced.")
flip_data = False
if dz < 0:
flip_data = True
dz *= -1
u = np.flipud(u)
v = np.flipud(v)
b = np.flipud(b)
N = u.size
kh = np.sqrt(k**2 + l**2) # Absolute horizontal wavenumber
# Velocity component parallel to the wave vector (k, l)
U = u * k / kh + v * l / kh
# Derivative matrices
# 1st derivative
dz1 = fd.FinDiff(0, dz, 1).matrix(u.shape).toarray()
# 2nd derivative
dz2 = fd.FinDiff(0, dz, 2).matrix(u.shape).toarray()
# 4th derivative
dz4 = fd.FinDiff(0, dz, 4).matrix(u.shape).toarray()
# Shear and buoyancy frequency.
bz = dz1 @ b
Uzz = dz2 @ U
# Add boundary conditions to matrix
# Impermeable boundary
dz2[0, :] = 0
dz2[0, 0] = -2 / dz**2
dz2[0, 1] = 1 / dz**2
dz2[-1, :] = 0
dz2[-1, -1] = -2 / dz**2
dz2[-1, -2] = 1 / dz**2
# % Asymptotic boundary
# % D2(1,:)=0;
# % D2(1,1)=2*(-del*kt-1)/del^2;
# % D2(1,2)=2/del^2;
# % D2(N,:)=0;
# % D2(N,N)=2*(-del*kt-1)/del^2;
# % D2(N,N-1)=2/del^2;
if BCv_upper == "rigid":
BCv1 = 0
elif BCv_upper == "frictionless":
BCv1 = 1
else:
raise ValueError(
"BCv_upper incorrectly specified, it must be either 'rigid' or 'frictionless'."
)
if BCv_lower == "rigid":
BCvN = 0
elif BCv_lower == "frictionless":
BCvN = 1
else:
raise ValueError(
"BCv_lower incorrectly specified, it must be either 'rigid' or 'frictionless'."
)
# % Rigid or frictionless BCs for 4th derivative
dz4[0, :] = 0
dz4[0, 0] = (5 + 2 * BCv1) / dz**4
dz4[0, 1] = -4 / dz**4
dz4[0, 2] = 1 / dz**4
dz4[1, :] = 0
dz4[1, 0] = -4 / dz**4
dz4[1, 1] = 6 / dz**4
dz4[1, 2] = -4 / dz**4
dz4[1, 3] = 1 / dz**4
dz4[-1, :] = 0
dz4[-1, -1] = (5 + 2 * BCvN) / dz**4
dz4[-1, -2] = -4 / dz**4
dz4[-1, -3] = 1 / dz**4
dz4[-2, :] = 0
dz4[-2, -1] = -4 / dz**4
dz4[-2, -2] = 6 / dz**4
dz4[-2, -3] = -4 / dz**4
dz4[-2, -4] = 1 / dz**4
# Boundary conditions for the second derivative of buoyancy.
dz2b = dz2.copy()
if BCb_upper == "constant":
dz2b[0, :] = 0
dz2b[0, 0] = -2 / dz**2
dz2b[0, 1] = 1 / dz**2
elif BCb_upper == "insulating":
dz2b[0, :] = 0
dz2b[0, 0] = -2 / (3 * dz**2)
dz2b[0, 1] = 2 / (3 * dz**2)
else:
raise ValueError(
"BCb_upper incorrectly specified, it must be either 'constant' or 'insulating'."
)
if BCb_lower == "constant":
dz2b[-1, :] = 0
dz2b[-1, -1] = -2 / dz**2
dz2b[-1, -2] = 1 / dz**2
elif BCb_lower == "insulating":
dz2b[-1, :] = 0
dz2b[-1, -1] = -2 / (3 * dz**2)
dz2b[-1, -2] = 2 / (3 * dz**2)
else:
raise ValueError(
"BCb_lower incorrectly specified, it must be either 'constant' or 'insulating'."
)
# Assemble stability matrices for eigenvalue computation
Id = np.eye(N)
L = dz2 - Id * kh**2 # Laplacian
Lb = dz2b - Id * kh**2 # Laplacian for buoyancy
LL = dz4 - 2 * dz2 * kh**2 + Id * kh**4 # Laplacian of laplacian
A = np.block([[L, np.zeros_like(L)], [np.zeros_like(L), Id]])
b11 = -1j * k * np.diag(U) @ L + 1j * k * np.diag(Uzz) + Kv * LL
b21 = -np.diag(bz)
b12 = -Id * kh**2
b22 = -1j * k * np.diag(U) + Kb * Lb
B = np.block([[b11, b12], [b21, b22]])
# Solve system using eig which returns om the imaginary frequency and vec the eigenvectors.
om, vec = eig(B, A)
# Prepare output
# Sort output by phase speed
cp = -om.imag / kh
idxs = np.argsort(cp)
cp = cp[idxs]
om = om[idxs]
vec = vec[:, idxs]
wvec = vec[:N, :]
bvec = vec[N:, :]
d_dz = fd.FinDiff(0, dz, 1, acc=5)
uvec = 1j * d_dz(wvec) / k
d2_dz2 = fd.FinDiff(0, dz, 2, acc=5)
X1 = d_dz(U)[:, np.newaxis] * wvec
X2 = (cp[np.newaxis, :] - U[:, np.newaxis]) * d_dz(wvec)
X3 = 1j * Kv * (d2_dz2(d_dz(wvec)) - k * d_dz(wvec)) / k
pvec = 1j * (X1 + X2 - X3) / k
if flip_data:
wvec = np.flipud(wvec)
bvec = np.flipud(bvec)
uvec = np.flipud(uvec)
pvec = np.flipud(pvec)
return om, wvec, bvec, uvec, pvec
