def test_1d_linear_combination(self):
x = np.sqrt(np.linspace(0, 1, 10))
d1 = Coef(x) * FinDiff(0, x, acc=4)
d2 = Coef(x**2) * FinDiff(0, x, 2, acc=4)
assert_array_almost_equal(3 * x**3, d1(x**3))
assert_array_almost_equal(6 * x**3, d2(x**3))
assert_array_almost_equal(9 * x**3, (d1 + d2)(x**3))
def test_3d_linear_combination(self):
x, y, z = [np.sqrt(np.linspace(0, 1, 10))] * 3
X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
d1 = Coef(X) * FinDiff(0, x, acc=4)
d2 = Coef(Y) * FinDiff(1, y, acc=4)
assert_array_almost_equal(3 * X ** 3, d1(X**3 + Y**3 + Z**3))
assert_array_almost_equal(3 * Y ** 3, d2(X**3 + Y**3 + Z**3))
assert_array_almost_equal(3 * (X**3 + Y**3), (d1 + d2)(X**3 + Y**3 + Z**3))
def test_matrix_1d_coeffs_nonuni(self):
shape = 11,
x = np.linspace(0, 10, 11)
L = Coef(x) * FinDiff(0, x, 2)
u = np.random.rand(*shape).reshape(-1)
actual = L.matrix(shape).dot(u)
expected = L(u).reshape(-1)
np.testing.assert_array_almost_equal(expected, actual)
def test_local_stencil_operator_with_coef(self):
x = np.linspace(0, 10, 101)
y = np.linspace(0, 10, 101)
X, Y = np.meshgrid(x, y, indexing='ij')
u = X * Y
dx = x[1] - x[0]
dy = y[1] - y[0]
d1x = Coef(2) * FinDiff(0, dx) * FinDiff(1, dy)
stencil1 = d1x.stencil(u.shape)
du_dx = stencil1.apply_all(u)
np.testing.assert_array_almost_equal(2 * np.ones_like(X), du_dx)
def test_identity(self):
x = np.linspace(-1, 1, 100)
u = x**2
identity = Identity()
assert_array_equal(u, identity(u))
twice_id = Coef(2) * Identity()
assert_array_equal(2 * u, twice_id(u))
x_id = Coef(x) * Identity()
assert_array_equal(x * u, x_id(u))
def test_identity_2d(self):
(X, Y), (x, y), _ = grid(2, 100, -1, 1)
u = X**2 + Y**2
identity = Identity()
assert_array_equal(u, identity(u))
twice_id = Coef(2) * Identity()
assert_array_equal(2 * u, twice_id(u))
x_id = Coef(X) * Identity()
assert_array_equal(X * u, x_id(u))
dx = x[1] - x[0]
d_dx = FinDiff(0, dx)
linop = d_dx + 2 * Identity()
assert_array_almost_equal(2 * X + 2 * u, linop(u))
def test_multiply(self):
(X, Y), _, h = grid(2, 5, 0, 1)
u = X**2 + Y**2
d2_dx2 = FinDiff(0, h[0], 2)
d = Coef(X) * d2_dx2
u1 = d(u)
assert_array_almost_equal(u1, 2 * X)
def test_2d_inhom_var_coefs_dirichlet_all(self):
shape = (41, 50)
(x, y), (dx, dy), (X, Y) = make_grid(shape, edges=[(-1, 1), (-1, 1)])
expected = X**3 + Y**3 + X * Y + 1
L = Coef(3 * X) * FinDiff(0, dx, 2) + Coef(2 * Y) * FinDiff(
(0, dx, 1), (1, dy, 1)) + FinDiff(1, dy, 2)
f = 18 * X**2 + 8 * Y
bc = BoundaryConditions(shape)
bc[0, :] = expected
bc[-1, :] = expected
bc[:, 0] = expected
bc[:, -1] = expected
pde = PDE(L, f, bc)
actual = pde.solve()
np.testing.assert_array_almost_equal(expected, actual, decimal=4)
def test_2d_inhom_var_coefs_with_identity_all_dirichlet(self):
shape = (5, 5)
(x, y), (dx, dy), (X, Y) = make_grid(shape, edges=[(-1, 1), (-1, 1)])
expected = X**3 + Y**3 + X * Y + 1
#L = Coef(3*X) * FinDiff(0, dx, 2) + Coef(2*Y) * FinDiff((0, dx, 1), (1, dy, 1)) + FinDiff(1, dy, 2) + Coef(5*X*Y) * Identity()
L = Coef(5 * X * Y) * FinDiff(0, dx, 2) #Identity()
#f = 18 * X**2 + 8*Y + 5*X*Y*expected
mat = L.matrix(shape)
print(mat)
bc = BoundaryConditions(shape)
bc[0, :] = expected
bc[-1, :] = expected
bc[:, 0] = expected
bc[:, -1] = expected
pde = PDE(L, f, bc)
actual = pde.solve()
np.testing.assert_array_almost_equal(expected, actual, decimal=4)
def test_laplace(self):
(X, Y, Z), _, h = grid(3, 50, 0, 1)
u = X**3 + Y**3 + Z**3
d2_dx2, d2_dy2, d2_dz2 = [FinDiff(i, h[i], 2) for i in range(3)]
laplace = d2_dx2 + d2_dy2 + d2_dz2
lap_u = laplace(u)
assert_array_almost_equal(lap_u, 6 * X + 6 * Y + 6 * Z)
d_dx, d_dy, d_dz = [FinDiff(i, h[i]) for i in range(3)]
d = Coef(X) * d_dx + Coef(Y) * d_dy + Coef(Z) * d_dz
f = d(lap_u)
d2 = d * laplace
f2 = d2(u)
assert_array_almost_equal(f2, f)
assert_array_almost_equal(f2, 6 * (X + Y + Z))
def test_1d_with_coeffs(self):
n = 200
shape = n,
t = np.linspace(0, 1, n)
dt = t[1] - t[0]
L = Coef(t) * FinDiff(0, dt, 2)
f = 6 * t**2
bc = BoundaryConditions(shape)
bc[0] = 0
bc[-1] = 1
eq = PDE(L, f, bc)
u = eq.solve()
expected = t**3
np.testing.assert_array_almost_equal(expected, u, decimal=4)
def test_mixed_equation__with_coeffs_2d(self):
shape = (41, 51)
x, y = np.linspace(0, 1, shape[0]), np.linspace(0, 1, shape[1])
dx, dy = x[1] - x[0], y[1] - y[0]
X, Y = np.meshgrid(x, y, indexing='ij')
L = FinDiff(0, dx, 2) + Coef(X * Y) * FinDiff(
(0, dx, 1), (1, dy, 1)) + FinDiff(1, dy, 2)
expected = X**3 + Y**3 + 1
f = 6 * (X + Y)
bc = BoundaryConditions(shape)
bc[0, :] = expected
bc[-1, :] = expected
bc[:, 0] = expected
bc[:, -1] = expected
pde = PDE(L, f, bc)
u = pde.solve()
np.testing.assert_array_almost_equal(expected, u, decimal=4)
def test_eigs(self):
shape = 8,
x = np.linspace(-10, 10, shape[0])
dx = x[1] - x[0]
hbar = 1
m = 1
omega = 1
T = Coef(-hbar**2 / (2 * m)) * FinDiff(0, dx, 2)
V = Coef(0.5 * omega * x**2) * Identity()
H = T + V
print("\n", T.matrix(shape).toarray())
print("\n", V.matrix(shape).toarray())
print("\n", H.matrix(shape).toarray())
print(H.matrix(shape).toarray())
vals, vecs = H.eigs(shape)
print(vals)
print(vecs)
def test_linear_combinations(self):
(X, Y, Z), _, (dx, dy, dz) = grid(3, 30, 0, 1)
u = X**2 + Y**2 + Z**2
d = Coef(X) * FinDiff(0, dx) + Coef(Y**2) * FinDiff(1, dy, 2)
assert_array_almost_equal(d(u), 2 * X**2 + 2 * Y**2)