def preconditioners(shape, ths, pml_params):
'''`(pre, inv_pre)` as 3-tuples of `(1, 1, xx, yy, zz)` arrays.'''
pre = []
for axis in range(3):
p = functools.reduce(onp.multiply, (scpml_coeffs(
shape[i], ths[i], pml_params, axis=i, transpose=(i == axis))
for i in range(3)))
p = onp.reshape(onp.sqrt(p), (1, 1) + shape)
pre.append(p)
pre = vecfield.VecField(*pre)
inv_pre = vecfield.VecField(*(1 / p for p in pre))
return pre.as_array(), inv_pre.as_array()
def test_operator(self):
input = np.eye(375, dtype=np.complex128)
input = np.reshape(input, (375, 3, 5, 5, 5))
input = vecfield.VecField(*np.split(input, 3, axis=1))
z = vecfield.VecField(*(np.ones((1, 1, 5, 5, 5), np.complex128)
for _ in range(3)))
pml_params = ops.PmlParams(w_eff=0.3)
pre, inv_pre = ops.preconditioners((5, 5, 5),
ths=((2, 2), ) * 3,
pml_params=pml_params)
A = ops.operator(input, z, pre, inv_pre, ((2, 2), (2, 2), (2, 2)),
pml_params)
A = np.reshape(np.concatenate(A, axis=1), (375, 375))
onp.testing.assert_array_almost_equal(A, np.transpose(A))
def curl(x, ths, pml_params, transpose=False):
'''Stretched curl of `vecfield.VecField` of `(1, 1, xx, yy, zz)` arrays.'''
diff_fn = functools.partial(stretched_spatial_diff,
pml_params=pml_params,
transpose=transpose)
y = []
for i in range(3):
j, k = (i + 1) % 3, (i + 2) % 3
y.append(
diff_fn(x[k], axis=j, th=ths[j]) - diff_fn(x[j], axis=k, th=ths[k]))
return vecfield.VecField(*y)
def test_loop_fns(self):
b = onp.zeros((1, 1, 10, 10, 10), onp.complex128)
b[0, 0, 5, 5, 5] = 1.
b = vecfield.VecField(0 * b, 0 * b, b)
z = vecfield.zeros((1, 1, 10, 10, 10))
x, errs = jaxwell.solve(z,
b,
ths=((10, 10), ) * 3,
pml_params=operators.PmlParams(w_eff=0.3),
max_iters=1)
self.assertIsInstance(x, vecfield.VecField)
self.assertEqual(x.shape, (1, 1, 10, 10, 10))
self.assertEqual(len(errs), 1)
self.assertAlmostEqual(errs[0], 35.25115523)
def test_loop_fns(self):
loop_init, loop_iter = solver.loop_fns(
shape=(1, 1, 10, 10, 10),
ths=((2, 2), ) * 3,
pml_params=operators.PmlParams(w_eff=0.3),
eps=1e-6)
b = onp.zeros((1, 1, 10, 10, 10), onp.complex128)
b[0, 0, 5, 5, 5] = 1.
b = vecfield.VecField(0 * b, 0 * b, b)
z = vecfield.zeros((1, 1, 10, 10, 10))
p, r, x, term_err = loop_init(z, b)
onp.testing.assert_array_equal(p, r)
onp.testing.assert_array_equal(x, onp.zeros_like(x))
self.assertEqual(term_err, 1e-6)
p, r, x, err = loop_iter(p, r, x, z)
self.assertAlmostEqual(err, 0.8660254)
def test_fns(self):
shape = (1, 1, 10, 10, 10)
ths = ((2, 2), ) * 3
pml_params = operators.PmlParams(w_eff=0.3)
pre, inv_pre = operators.preconditioners(shape[2:], ths, pml_params)
def A(x, z):
return operators.operator(x, z, pre, inv_pre, ths, pml_params)
b = onp.zeros(shape, onp.complex128)
b[0, 0, 5, 5, 5] = 1.
b = vecfield.VecField(0 * b, 0 * b, b)
z = vecfield.zeros(shape)
init, iter = cocg.solver(A, b, eps=1e-6)
p, r, x, term_err = init(z, b)
onp.testing.assert_array_equal(p, r)
onp.testing.assert_array_equal(x, onp.zeros_like(x))
self.assertEqual(term_err, 1e-6)
p, r, x, err = iter(p, r, x, z)
self.assertAlmostEqual(err, 0.8660254)