def test_transform_coeff_with_x_and_r(self):
"""Test coefficient transform between x and r."""
coeff = np.array([2, 3, 4])
ltf = LinearTF(1, 10) # (-1, 1) -> (r0, rmax)
inv_tf = InverseTF(ltf) # (r0, rmax) -> (-1, 1)
x = np.linspace(-1, 1, 20)
r = ltf.transform(x)
assert r[0] == 1
assert r[-1] == 10
coeff_b = ODE._transformed_coeff_ode(coeff, inv_tf, x)
derivs_fun = [inv_tf.deriv, inv_tf.deriv2, inv_tf.deriv3]
coeff_b_ref = ODE._transformed_coeff_ode_with_r(coeff, derivs_fun, r)
assert_allclose(coeff_b, coeff_b_ref)
def test_error_raises(self):
"""Test proper error raises."""
x = np.linspace(-0.999, 0.999, 20)
# r = btf.transform(x)
def fx(x):
return 1 / x ** 2
coeffs = [-1, -2, 1]
bd_cond = [(0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0)]
with self.assertRaises(NotImplementedError):
ODE.solve_ode(x, fx, coeffs, bd_cond[3:])
with self.assertRaises(NotImplementedError):
ODE.solve_ode(x, fx, coeffs, bd_cond)
with self.assertRaises(NotImplementedError):
test_coeff = [1, 2, 3, 4, 5]
ODE.solve_ode(x, fx, test_coeff, bd_cond)
with self.assertRaises(ValueError):
test_coeff = [1, 2, 3, 3]
tf = BeckeTF(0.1, 1)
def fx(x):
return x
ODE.solve_ode(x, fx, test_coeff, bd_cond[:3], tf)
def test_solver_ode_coeff_a_f_x_with_tf(self):
"""Test ode with a(x) and f(x) involved."""
x = np.linspace(-0.999, 0.999, 20)
btf = BeckeTF(0.1, 5)
r = btf.transform(x)
ibtf = InverseTF(btf)
def fx(x):
return 0 * x
coeffs = [lambda x: x ** 2, lambda x: 1 / x ** 2, 0.5]
bd_cond = [(0, 0, 0), (1, 0, 0)]
# calculate diff equation wt/w tf.
res = ODE.solve_ode(x, fx, coeffs, bd_cond, ibtf)
res_ref = ODE.solve_ode(r, fx, coeffs, bd_cond)
assert_allclose(res(x)[0], res_ref(r)[0], atol=1e-4)
def test_construct_coeffs(self):
"""Test construct coefficients."""
# first test
x = np.linspace(-0.9, 0.9, 20)
coeff = [2, 1.5, lambda x: x ** 2]
coeff_a = ODE._construct_coeff_array(x, coeff)
assert_allclose(coeff_a[0], np.ones(20) * 2)
assert_allclose(coeff_a[1], np.ones(20) * 1.5)
assert_allclose(coeff_a[2], x ** 2)
# second test
coeff = [lambda x: 1 / x, 2, lambda x: x ** 3, lambda x: np.exp(x)]
coeff_a = ODE._construct_coeff_array(x, coeff)
assert_allclose(coeff_a[0], 1 / x)
assert_allclose(coeff_a[1], np.ones(20) * 2)
assert_allclose(coeff_a[2], x ** 3)
assert_allclose(coeff_a[3], np.exp(x))
def test_solver_ode_bvp_with_tf(self):
"""Test result for high level api solve_ode with fx term."""
x = np.linspace(-0.999, 0.999, 20)
btf = BeckeTF(0.1, 5)
r = btf.transform(x)
ibtf = InverseTF(btf)
def fx(x):
return 1 / x ** 2
coeffs = [-1, 1, 1]
bd_cond = [(0, 0, 0), (1, 0, 0)]
# calculate diff equation wt/w tf.
res = ODE.solve_ode(x, fx, coeffs, bd_cond, ibtf)
res_ref = ODE.solve_ode(r, fx, coeffs, bd_cond)
assert_allclose(res(x)[0], res_ref(r)[0], atol=1e-4)
def test_transform_coeff(self):
"""Test coefficient transform with r."""
# d^2y / dx^2 = 1
itf = IdentityRTransform()
inv_tf = InverseTF(itf)
derivs_fun = [inv_tf.deriv, inv_tf.deriv2, inv_tf.deriv3]
coeff = np.array([0, 0, 1])
x = np.linspace(0, 1, 10)
coeff_b = ODE._transformed_coeff_ode_with_r(coeff, derivs_fun, x)
# compute transformed coeffs
assert_allclose(coeff_b, np.zeros((3, 10), dtype=float) + coeff[:, None])
def test_linear_transform_coeff(self):
"""Test coefficient with linear transformation."""
x = GaussLaguerre(10).points
ltf = LinearTF(1, 10)
inv_ltf = InverseTF(ltf)
derivs_fun = [inv_ltf.deriv, inv_ltf.deriv2, inv_ltf.deriv3]
coeff = np.array([2, 3, 4])
coeff_b = ODE._transformed_coeff_ode_with_r(coeff, derivs_fun, x)
# assert values
assert_allclose(coeff_b[0], np.ones(len(x)) * coeff[0])
assert_allclose(coeff_b[1], 1 / 4.5 * coeff[1])
assert_allclose(coeff_b[2], (1 / 4.5) ** 2 * coeff[2])
def test_solve_ode_bvp(self):
"""Test result for high level api solve_ode."""
x = np.linspace(0, 2, 10)
def fx(x):
return 1 if isinstance(x, Number) else np.ones(x.size)
coeffs = [0, 0, 1]
bd_cond = [[0, 0, 0], [1, 0, 0]]
res = ODE.solve_ode(x, fx, coeffs, bd_cond)
assert_almost_equal(res(0)[0], 0)
assert_almost_equal(res(1)[0], -0.5)
assert_almost_equal(res(2)[0], 0)
assert_almost_equal(res(0)[1], -1)
assert_almost_equal(res(1)[1], 0)
assert_almost_equal(res(2)[1], 1)
def solve_poisson_bv(fx, x_range, boundary, m_l=(0, 0), tfm=None):
"""Solve poisson equation for given function.
.. math::
Parameters
----------
fx : Callable
Callable function on the right hand side of equation
x_range : np.narray(K,)
An array with points for numerically solve the ODE
The boundary point should be within radial points range
boundary : float
Boundary value when x to go infinite
m_l : tuple, optional
m and l value of given spherical harmonics
tfm : None, optional
Transformation for given x variable
Returns
-------
scipy.PPline
A callable spline for result.
"""
m_value, l_value = m_l
def f_x(r):
return fx(r) * -4 * np.pi * r
def coeff_0(r):
return -l_value * (l_value + 1) / r**2
coeffs = [coeff_0, 0, 1]
if l_value == 0 and m_value == 0:
bd_cond = [(0, 0, 0), (1, 0, boundary)]
else:
bd_cond = [(0, 0, 0), (1, 0, 0)]
return ODE.solve_ode(x_range, f_x, coeffs, bd_cond, transform=tfm)
def func(x, y):
dy_dx = ODE._rearrange_ode(x, y, coeff_b, fx(x))
return np.vstack((*y[1:], dy_dx))
def func(x, y):
dy_dx = ODE._rearrange_trans_ode(x, y, coeff, ibtf, fx)
return np.vstack((*y[1:], dy_dx))