def test_restrict_weights(njit):
if njit:
restrict_weights = njitted.restrict_weights
else:
restrict_weights = njitted.restrict_weights.py_func
# 1. Simple example following equation 9, [Muld06]_.
edges = np.array([0., 500, 1200, 2000, 3000])
width = (edges[1:] - edges[:-1])
centr = edges[:-1] + width / 2
c_edges = edges[::2]
c_width = (c_edges[1:] - c_edges[:-1])
c_centr = c_edges[:-1] + c_width / 2
# Result
wtl = np.array([350 / 250, 250 / 600, 400 / 900])
wt0 = np.array([1., 1., 1.])
wtr = np.array([350 / 600, 500 / 900, 400 / 500])
# Result from implemented function
wl, w0, wr = restrict_weights(edges, centr, width, c_edges, c_centr,
c_width)
assert_allclose(wtl, wl)
assert_allclose(wt0, w0)
assert_allclose(wtr, wr)
# 2. Test with stretched grid and compare with alternative formulation
# Create a highly stretched, non-centered grid
hx = get_h(2, 2, 200, 1.8)
hy = get_h(0, 8, 800, 1.2)
hz = get_h(0, 4, 400, 1.4)
grid = utils.TensorMesh([hx, hy, hz], np.array([-100000, 3000, 100]))
# Create coarse grid thereof
ch = [
np.diff(grid.vectorNx[::2]),
np.diff(grid.vectorNy[::2]),
np.diff(grid.vectorNz[::2])
]
cgrid = utils.TensorMesh(ch, x0=grid.x0)
# Calculate the weights in a numpy-way, instead of numba-way
wl, w0, wr = alternatives.alt_restrict_weights(grid.vectorNx,
grid.vectorCCx, grid.hx,
cgrid.vectorNx,
cgrid.vectorCCx, cgrid.hx)
# Get the implemented numba-result
wxl, wx0, wxr = restrict_weights(grid.vectorNx, grid.vectorCCx, grid.hx,
cgrid.vectorNx, cgrid.vectorCCx, cgrid.hx)
# Compare
assert_allclose(wxl, wl)
assert_allclose(wx0, w0)
assert_allclose(wxr, wr)
def test_solver_homogeneous_laplace():
# Regression test for homogeneous halfspace in Laplace domain.
# Not very sophisticated; replace/extend by more detailed tests.
dat = REGRES['lap'][()]
grid = utils.TensorMesh(**dat['input_grid'])
model = utils.Model(**dat['input_model'])
sfield = utils.get_source_field(**dat['input_source'])
# F-cycle
efield = solver.solve(grid, model, sfield, verb=1)
# Check all fields (ex, ey, and ez)
assert_allclose(dat['Fresult'], efield)
# BiCGSTAB with some print checking.
efield = solver.solve(grid, model, sfield, verb=1, sslsolver=True)
# Check all fields (ex, ey, and ez)
assert_allclose(dat['bicresult'], efield)
# If efield is complex, assert it fails.
efield = utils.Field(grid, dtype=complex)
with pytest.raises(ValueError):
efield = solver.solve(grid, model, sfield, efield=efield, verb=1)
def test_krylov(capsys):
# Everything should be tested just fine in `test_solver`.
# Just check here for bicgstab-error.
# Load any case.
dat = REGRES['res'][()]
grid = utils.TensorMesh(**dat['input_grid'])
model = utils.Model(**dat['input_model'])
sfield = utils.get_source_field(**dat['input_source'])
vmodel = utils.VolumeModel(grid, model, sfield)
efield = utils.Field(grid) # Initiate e-field.
# Get var-instance
var = solver.MGParameters(
cycle=None,
sslsolver=True,
semicoarsening=False,
linerelaxation=False,
vnC=grid.vnC,
verb=3,
maxit=-1, # Set stupid input to make bicgstab fail.
)
var.l2_refe = njitted.l2norm(sfield)
# Call krylov and ensure it fails properly.
solver.krylov(grid, vmodel, sfield, efield, var)
out, _ = capsys.readouterr()
assert '* ERROR :: Error in bicgstab' in out
def test_solver_heterogeneous(capsys):
# Regression test for heterogeneous case.
dat = REGRES['reg_2'][()]
grid = dat['grid']
model = dat['model']
sfield = dat['sfield']
inp = dat['inp']
inp['verb'] = 4
efield = solver.solve(grid, model, sfield, **inp)
assert_allclose(dat['result'], efield.field)
# Check with provided e-field; 2x2 iter should yield the same as 4 iter.
efield2 = solver.solve(grid, model, sfield, maxit=4, verb=1)
efield3 = solver.solve(grid, model, sfield, maxit=2, verb=1)
solver.solve(grid, model, sfield, efield3, maxit=2, verb=1)
assert_allclose(efield2, efield3)
out, _ = capsys.readouterr() # Clean up
# One test without post-smoothing to check if it runs.
efield4 = solver.solve(grid,
model,
sfield,
sslsolver=True,
semicoarsening=True,
linerelaxation=True,
maxit=20,
nu_pre=0,
nu_post=4,
verb=3)
efield5 = solver.solve(grid,
model,
sfield,
sslsolver=True,
semicoarsening=True,
linerelaxation=True,
maxit=20,
nu_pre=4,
nu_post=0,
verb=3)
# They don't converge, and hence don't agree. Just a lazy test.
assert_allclose(efield4, efield5, atol=1e-15, rtol=1e-5)
# Check the QC plot if it is too long.
# Coincidently, this one also diverges if nu_pre=0!
# Mesh: 2-cells in y- and z-direction; 2**9 in x-direction
mesh = utils.TensorMesh(
[np.ones(2**9) / np.ones(2**9).sum(),
np.ones(2),
np.ones(2)],
x0=np.array([-0.5, -1, -1]))
sfield = utils.get_source_field(mesh, [0, 0, 0, 0, 0], 1)
model = utils.Model(mesh)
_ = solver.solve(mesh, model, sfield, verb=3, nu_pre=0)
out, _ = capsys.readouterr()
assert "(Cycle-QC restricted to first 70 steps of 72 steps.)" in out
assert "DIVERGED" in out
def test_solver_backwards(capsys):
grid = utils.TensorMesh(
[np.ones(8), np.ones(8), np.ones(8)], x0=np.array([0, 0, 0]))
model = utils.Model(grid, res_x=1.5, res_y=1.8, res_z=3.3)
sfield = utils.get_source_field(grid, src=[4, 4, 4, 0, 0], freq=10.0)
out, _ = capsys.readouterr()
_ = solver.solver(grid, model, sfield, verb=0)
out, _ = capsys.readouterr()
assert "* WARNING :: ``emg3d.solver.solver()`` is renamed to " in out
def test_volume_avg_weights(njit):
if njit:
volume_avg_weights = njitted._volume_avg_weights
else:
volume_avg_weights = njitted._volume_avg_weights.py_func
grid_in = utils.TensorMesh(
[np.ones(11), np.ones(10) * 2,
np.ones(3) * 10],
x0=np.array([0, 0, 0]))
grid_out = utils.TensorMesh(
[np.arange(4) + 1,
np.arange(5) + 1,
np.arange(6) + 1],
x0=np.array([0.5, 3.33, 5]))
wx, ix_in, ix_out = volume_avg_weights(grid_in.vectorNx, grid_out.vectorNx)
assert_allclose(
wx, [0, 0.5, 0.5, 0.5, 1, 0.5, 0.5, 1, 1, 0.5, 0.5, 1, 1, 1, 0.5, 0])
assert_allclose(ix_in, [0, 0, 1, 1, 2, 3, 3, 4, 5, 6, 6, 7, 8, 9, 10, 10])
assert_allclose(ix_out, [0, 0, 0, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3])
wy, iy_in, iy_out = volume_avg_weights(grid_in.vectorNy, grid_out.vectorNy)
assert_allclose(wy, [
0, 0, 0.67, 0.33, 1.67, 0.33, 1.67, 1.33, 0.67, 2., 1.33, 0.67, 2, 2,
0.33, 0.
])
assert_allclose(iy_in, [0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 6, 6, 7, 8, 9, 9])
assert_allclose(iy_out, [0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 3, 4, 4, 4, 4, 4])
wz, iz_in, iz_out = volume_avg_weights(grid_in.vectorNz, grid_out.vectorNz)
assert_allclose(wz, [0, 1, 2, 2, 1, 4, 5, 6, 0])
assert_allclose(iz_in, [0, 0, 0, 0, 1, 1, 1, 2, 2])
assert_allclose(iz_out, [0, 0, 1, 2, 2, 3, 4, 5, 5])
w, inp, out = volume_avg_weights(np.array([0., 5, 7, 10]),
np.array([-1., 1, 4, 6, 7, 11]))
assert_allclose(w, [1, 1, 3, 1, 1, 1, 3, 1])
assert_allclose(inp, [0, 0, 0, 0, 1, 1, 2, 2])
assert_allclose(out, [0, 0, 1, 2, 2, 3, 4, 4])
def test_restriction():
# Simple test with restriction followed by prolongation.
src = [0, 0, 0, 0, 45]
grid = utils.TensorMesh(
[np.ones(4) * 100,
np.ones(4) * 100,
np.ones(4) * 100], x0=np.zeros(3))
# Create dummy model and fields, parameters don't matter.
model = utils.Model(grid, 1, 1, 1, 1)
sfield = utils.get_source_field(grid, src, 1)
# Get volume-averaged model parameters.
vmodel = utils.VolumeModel(grid, model, sfield)
rx = np.arange(sfield.fx.size, dtype=complex).reshape(sfield.fx.shape)
ry = np.arange(sfield.fy.size, dtype=complex).reshape(sfield.fy.shape)
rz = np.arange(sfield.fz.size, dtype=complex).reshape(sfield.fz.shape)
rr = utils.Field(rx, ry, rz)
# Restrict it
cgrid, cmodel, csfield, cefield = solver.restriction(grid,
vmodel,
sfield,
rr,
sc_dir=0)
assert_allclose(csfield.fx[:, 1:-1, 1],
np.array([[196. + 0.j], [596. + 0.j]]))
assert_allclose(csfield.fy[1:-1, :, 1], np.array([[356. + 0.j,
436. + 0.j]]))
assert_allclose(csfield.fz[1:-1, 1:-1, :],
np.array([[[388. + 0.j, 404. + 0.j]]]))
assert cgrid.nNx == cgrid.nNy == cgrid.nNz == 3
assert cmodel.eta_x[0, 0, 0] / 8. == vmodel.eta_x[0, 0, 0]
assert np.sum(grid.hx) == np.sum(cgrid.hx)
assert np.sum(grid.hy) == np.sum(cgrid.hy)
assert np.sum(grid.hz) == np.sum(cgrid.hz)
# Add pi to the coarse e-field
efield = utils.Field(grid)
cefield += np.pi
# Prolong it
solver.prolongation(grid, efield, cgrid, cefield, sc_dir=0)
assert np.all(efield.fx[:, 1:-1, 1:-1] == np.pi)
assert np.all(efield.fy[1:-1, :, 1:-1] == np.pi)
assert np.all(efield.fz[1:-1, 1:-1, :] == np.pi)
def test_amat_x(njit):
if njit:
amat_x = njitted.amat_x
else:
amat_x = njitted.amat_x.py_func
# 1. Compare to alternative amat_x
# Create a grid
src = [200, 300, -50., 5, 60]
hx = get_h(8, 4, 100, 1.2)
hy = np.ones(8) * 800
hz = np.ones(4) * 500
grid = utils.TensorMesh([hx, hy, hz],
np.array(
[-hx.sum() / 2, -hy.sum() / 2, -hz.sum() / 2]))
# Create some resistivity model
x = np.arange(1, grid.nCx + 1) * 2
y = 1 / np.arange(1, grid.nCy + 1)
z = np.arange(1, grid.nCz + 1)[::-1] / 10
res_x = np.outer(np.outer(x, y), z).ravel()
freq = 0.319
model = utils.Model(grid, res_x, 0.8 * res_x, 2 * res_x)
# Create a source field
sfield = utils.get_source_field(grid=grid, src=src, freq=freq)
# Get volume-averaged model parameters.
vmodel = utils.VolumeModel(grid, model, sfield)
# Run two iterations to get a e-field
efield = solver.solve(grid, model, sfield, maxit=2, verb=1)
# amat_x
rr1 = utils.Field(grid)
amat_x(rr1.fx, rr1.fy, rr1.fz, efield.fx, efield.fy, efield.fz,
vmodel.eta_x, vmodel.eta_y, vmodel.eta_z, vmodel.zeta, grid.hx,
grid.hy, grid.hz)
# amat_x - alternative
rr2 = utils.Field(grid)
alternatives.alt_amat_x(rr2.fx, rr2.fy, rr2.fz, efield.fx, efield.fy,
efield.fz, vmodel.eta_x, vmodel.eta_y,
vmodel.eta_z, vmodel.zeta, grid.hx, grid.hy,
grid.hz)
# Check all fields (ex, ey, and ez)
assert_allclose(-rr1, rr2, atol=1e-23)
def test_volume_average(njit):
if njit:
volume_average = njitted.volume_average
else:
volume_average = njitted.volume_average.py_func
# Comparison to alt_version.
grid_in = utils.TensorMesh(
[np.ones(30), np.ones(20) * 5,
np.ones(10) * 10],
x0=np.array([0, 0, 0]))
grid_out = utils.TensorMesh(
[np.arange(7) + 1,
np.arange(13) + 1,
np.arange(13) + 1],
x0=np.array([0.5, 3.33, 5]))
values = np.arange(grid_in.nC, dtype=float).reshape(grid_in.vnC, order='F')
points = (grid_in.vectorNx, grid_in.vectorNy, grid_in.vectorNz)
new_points = (grid_out.vectorNx, grid_out.vectorNy, grid_out.vectorNz)
# Calculate volume.
vol = np.outer(np.outer(grid_out.hx, grid_out.hy).ravel('F'), grid_out.hz)
vol = vol.ravel('F').reshape(grid_out.vnC, order='F')
# New solution.
new_values = np.zeros(grid_out.vnC, dtype=values.dtype)
volume_average(*points, values, *new_points, new_values, vol)
# Old solution.
new_values_alt = np.zeros(grid_out.vnC, dtype=values.dtype)
alternatives.alt_volume_average(*points, values, *new_points,
new_values_alt)
assert_allclose(new_values, new_values_alt)
def test_one_liner(capsys):
grid = utils.TensorMesh(
[np.ones(8), np.ones(8), np.ones(8)], x0=np.array([0, 0, 0]))
model = utils.Model(grid, res_x=1.5, res_y=1.8, res_z=3.3)
sfield = utils.get_source_field(grid, src=[4, 4, 4, 0, 0], freq=10.0)
out, _ = capsys.readouterr()
_ = solver.solve(grid, model, sfield, verb=-1)
out, _ = capsys.readouterr()
assert '6; 0:00:' in out
assert '; CONVERGED' in out
out, _ = capsys.readouterr()
_ = solver.solve(grid, model, sfield, sslsolver=True, verb=-1)
out, _ = capsys.readouterr()
assert '3(5); 0:00:' in out
assert '; CONVERGED' in out
def test_residual():
# The only thing to test here is that residual returns the same as
# sfield-amat_x. Basically a copy of the function itself.
# Create a grid
src = [90, 1600, 25., 45, 45]
grid = utils.TensorMesh(
[get_h(4, 2, 20, 1.2),
np.ones(16) * 200,
np.ones(2) * 25],
x0=np.zeros(3))
# Create some resistivity model
x = np.arange(1, grid.nCx + 1) * 2
y = 1 / np.arange(1, grid.nCy + 1)
z = np.arange(1, grid.nCz + 1)[::-1] / 10
res_x = np.outer(np.outer(x, y), z).ravel()
freq = 0.319
model = utils.Model(grid, res_x, 0.8 * res_x, 2 * res_x)
# Create a source field
sfield = utils.get_source_field(grid=grid, src=src, freq=freq)
# Get volume-averaged model parameters.
vmodel = utils.VolumeModel(grid, model, sfield)
# Run two iterations to get an e-field
efield = solver.solve(grid, model, sfield, maxit=2, verb=1)
# Use directly amat_x
rfield = sfield.copy()
njitted.amat_x(rfield.fx, rfield.fy, rfield.fz, efield.fx, efield.fy,
efield.fz, vmodel.eta_x, vmodel.eta_y, vmodel.eta_z,
vmodel.zeta, grid.hx, grid.hy, grid.hz)
# Calculate residual
out = solver.residual(grid, vmodel, sfield, efield)
outnorm = solver.residual(grid, vmodel, sfield, efield, True)
# Compare
assert_allclose(out, rfield)
assert_allclose(outnorm, np.linalg.norm(out))
def test_gauss_seidel(njit):
if njit:
gauss_seidel = njitted.gauss_seidel
gauss_seidel_x = njitted.gauss_seidel_x
gauss_seidel_y = njitted.gauss_seidel_y
gauss_seidel_z = njitted.gauss_seidel_z
else:
gauss_seidel = njitted.gauss_seidel.py_func
gauss_seidel_x = njitted.gauss_seidel_x.py_func
gauss_seidel_y = njitted.gauss_seidel_y.py_func
gauss_seidel_z = njitted.gauss_seidel_z.py_func
# At the moment we only compare `gauss_seidel_x/y/z` to `gauss_seidel`.
# Better tests should be implemented.
# Rotate the source, so we have a strong enough signal in all directions
src = [0, 0, 0, 45, 45]
freq = 0.9
nu = 2 # One back-and-forth
for lr_dir in range(1, 4):
# `gauss_seidel`/`_x/y/z` loop over z, then y, then x. Together with
# `lr_dir`, we have to keep the dimension at 2 in order that they
# agree.
nx = [1, 4, 4][lr_dir - 1]
ny = [4, 1, 4][lr_dir - 1]
nz = [4, 4, 1][lr_dir - 1]
# Get this grid.
hx = get_h(0, nx, 80, 1.1)
hy = get_h(0, ny, 100, 1.3)
hz = get_h(0, nz, 200, 1.2)
grid = utils.TensorMesh(
[hx, hy, hz],
np.array([-hx.sum() / 2, -hy.sum() / 2, -hz.sum() / 2]))
# Initialize model with some resistivities.
res_x = np.arange(grid.nC) + 1
res_y = 0.5 * np.arange(grid.nC) + 1
res_z = 2 * np.arange(grid.nC) + 1
model = utils.Model(grid, res_x, res_y, res_z)
# Initialize source field.
sfield = utils.get_source_field(grid, src, freq)
# Get volume-averaged model parameters.
vmodel = utils.VolumeModel(grid, model, sfield)
# Run two iterations to get some e-field.
efield = solver.solve(grid, model, sfield, maxit=2, verb=1)
inp = (sfield.fx, sfield.fy, sfield.fz, vmodel.eta_x, vmodel.eta_y,
vmodel.eta_z, vmodel.zeta, grid.hx, grid.hy, grid.hz, nu)
# Get result from `gauss_seidel`.
cfield = utils.Field(grid, efield.copy())
gauss_seidel(cfield.fx, cfield.fy, cfield.fz, *inp)
# Get result from `gauss_seidel_x/y/z`.
if lr_dir == 1:
gauss_seidel_x(efield.fx, efield.fy, efield.fz, *inp)
elif lr_dir == 2:
gauss_seidel_y(efield.fx, efield.fy, efield.fz, *inp)
elif lr_dir == 3:
gauss_seidel_z(efield.fx, efield.fy, efield.fz, *inp)
# Check the resulting field.
assert_allclose(efield.field, cfield.field)
def test_restrict(njit):
if njit:
restrict = njitted.restrict
restrict_weights = njitted.restrict_weights
else:
restrict = njitted.restrict.py_func
restrict_weights = njitted.restrict_weights.py_func
# Simple comparison using the most basic mesh.
h = np.array([1, 1, 1, 1, 1, 1])
# Fine grid.
fgrid = utils.TensorMesh([h, h, h], x0=np.array([-3, -3, -3]))
# Create fine field.
ffield = utils.Field(fgrid)
# Put in values.
ffield.fx[:, :, :] = 1
ffield.fy[:, :, :] = 2
ffield.fz[:, :, :] = 4
# Ensure PEC.
ffield.ensure_pec
# Get weigths
wlr = np.zeros(fgrid.nNx, dtype=float)
w0 = np.ones(fgrid.nNx, dtype=float)
fw = (wlr, w0, wlr)
# # CASE 0 -- regular # #
# Coarse grid.
cgrid = utils.TensorMesh([
np.diff(fgrid.vectorNx[::2]),
np.diff(fgrid.vectorNy[::2]),
np.diff(fgrid.vectorNz[::2])
], fgrid.x0)
# Regular grid, so all weights (wx, wy, wz) are the same...
w = restrict_weights(fgrid.vectorNx, fgrid.vectorCCx, fgrid.hx,
cgrid.vectorNx, cgrid.vectorCCx, cgrid.hx)
# Create coarse field.
cfield = utils.Field(cgrid)
# Restrict it.
restrict(cfield.fx, cfield.fy, cfield.fz, ffield.fx, ffield.fy, ffield.fz,
w, w, w, 0)
# Check sum of fine and coarse fields.
assert cfield.fx.sum() == ffield.fx.sum()
assert cfield.fy.sum() == ffield.fy.sum()
assert cfield.fz.sum() == ffield.fz.sum()
# Assert fields are multiples from each other.
np.allclose(cfield.fx[0, :, :] * 2, cfield.fy[:, 0, :])
np.allclose(cfield.fy[:, 0, :] * 2, cfield.fz[:, :, 0])
# # CASE 1 -- y & z # #
# Coarse grid.
cgrid = utils.TensorMesh(
[fgrid.hx,
np.diff(fgrid.vectorNy[::2]),
np.diff(fgrid.vectorNz[::2])], fgrid.x0)
# Create coarse field.
cfield = utils.Field(cgrid)
# Restrict field.
restrict(cfield.fx, cfield.fy, cfield.fz, ffield.fx, ffield.fy, ffield.fz,
fw, w, w, 1)
# Check sum of fine and coarse fields.
assert cfield.fx.sum() == ffield.fx.sum()
assert cfield.fy.sum() == ffield.fy.sum()
assert cfield.fz.sum() == ffield.fz.sum()
# Assert fields are multiples from each other.
np.allclose(cfield.fy[:, 0, :] * 2, cfield.fz[:, :, 0])
# # CASE 2 -- x & z # #
# Coarse grid.
cgrid = utils.TensorMesh(
[np.diff(fgrid.vectorNx[::2]), fgrid.hy,
np.diff(fgrid.vectorNz[::2])], fgrid.x0)
# Create coarse field.
cfield = utils.Field(cgrid)
# Restrict field.
restrict(cfield.fx, cfield.fy, cfield.fz, ffield.fx, ffield.fy, ffield.fz,
w, fw, w, 2)
# Check sum of fine and coarse fields.
assert cfield.fx.sum() == ffield.fx.sum()
assert cfield.fy.sum() == ffield.fy.sum()
assert cfield.fz.sum() == ffield.fz.sum()
# Assert fields are multiples from each other.
np.allclose(cfield.fx[0, :, :].T * 2, cfield.fz[:, :, 0])
# # CASE 3 -- x & y # #
# Coarse grid.
cgrid = utils.TensorMesh(
[np.diff(fgrid.vectorNx[::2]),
np.diff(fgrid.vectorNy[::2]), fgrid.hz], fgrid.x0)
# Create coarse field.
cfield = utils.Field(cgrid)
# Restrict field.
restrict(cfield.fx, cfield.fy, cfield.fz, ffield.fx, ffield.fy, ffield.fz,
w, w, fw, 3)
# Check sum of fine and coarse fields.
assert cfield.fx.sum() == ffield.fx.sum()
assert cfield.fy.sum() == ffield.fy.sum()
assert cfield.fz.sum() == ffield.fz.sum()
# Assert fields are multiples from each other.
np.allclose(cfield.fx[0, :, :] * 2, cfield.fy[:, 0, :])
# # CASE 4 -- x # #
# Coarse grid.
cgrid = utils.TensorMesh(
[np.diff(fgrid.vectorNx[::2]), fgrid.hy, fgrid.hz], fgrid.x0)
# Create coarse field.
cfield = utils.Field(cgrid)
# Restrict field.
restrict(cfield.fx, cfield.fy, cfield.fz, ffield.fx, ffield.fy, ffield.fz,
w, fw, fw, 4)
# Check sum of fine and coarse fields.
assert cfield.fx.sum() == ffield.fx.sum()
assert cfield.fy.sum() == ffield.fy.sum()
assert cfield.fz.sum() == ffield.fz.sum()
# Assert fields are multiples from each other.
np.allclose(cfield.fy[:, 0, :] * 2, cfield.fz[:, :, 0])
# # CASE 5 -- y # #
# Coarse grid.
cgrid = utils.TensorMesh(
[fgrid.hx, np.diff(fgrid.vectorNy[::2]), fgrid.hz], fgrid.x0)
# Create coarse field.
cfield = utils.Field(cgrid)
# Restrict field.
restrict(cfield.fx, cfield.fy, cfield.fz, ffield.fx, ffield.fy, ffield.fz,
fw, w, fw, 5)
# Check sum of fine and coarse fields.
assert cfield.fx.sum() == ffield.fx.sum()
assert cfield.fy.sum() == ffield.fy.sum()
assert cfield.fz.sum() == ffield.fz.sum()
# Assert fields are multiples from each other.
np.allclose(cfield.fx[0, :, :].T * 4, cfield.fz[:, :, 0])
# # CASE 6 -- z # #
# Coarse grid.
cgrid = utils.TensorMesh(
[fgrid.hx, fgrid.hy, np.diff(fgrid.vectorNz[::2])], fgrid.x0)
# Create coarse field.
cfield = utils.Field(cgrid)
# Restrict field.
restrict(cfield.fx, cfield.fy, cfield.fz, ffield.fx, ffield.fy, ffield.fz,
fw, fw, w, 6)
# Check sum of fine and coarse fields.
assert cfield.fx.sum() == ffield.fx.sum()
assert cfield.fy.sum() == ffield.fy.sum()
assert cfield.fz.sum() == ffield.fz.sum()
# Assert fields are multiples from each other.
np.allclose(cfield.fx[0, :, :] * 4, cfield.fy[:, 0, :])
from discretize import TensorMesh
from emg3d import utils, solver
# # # # # # # # # # 1. Homogeneous VTI fullspace # # # # # # # # # #
freq = 1.
hx_min, xdomain = utils.get_domain(x0=0, freq=freq)
hy_min, ydomain = utils.get_domain(x0=0, freq=freq)
hz_min, zdomain = utils.get_domain(x0=250, freq=freq)
nx = 2**3
hx = utils.get_stretched_h(hx_min, xdomain, nx, 0)
hy = utils.get_stretched_h(hy_min, ydomain, nx, 0)
hz = utils.get_stretched_h(hz_min, zdomain, nx, 250)
input_grid = {'h': [hx, hy, hz], 'x0': (xdomain[0], ydomain[0], zdomain[0])}
grid = utils.TensorMesh(**input_grid)
input_model = {
'grid': grid,
'res_x': 1.5,
'res_y': 2.0,
'res_z': 3.3,
}
model = utils.Model(**input_model)
input_source = {
'grid': grid,
'src': [0, 0, 250., 30, 10], # A rotated source to include all
'freq': freq
}
def test_RegularGridProlongator():
def prolon_scipy(grid, cgrid, efield, cefield, yz_points):
"""Calculate SciPy alternative."""
for ixc in range(cgrid.nCx):
# Bilinear interpolation in the y-z plane
fn = si.RegularGridInterpolator((cgrid.vectorNy, cgrid.vectorNz),
cefield.fx[ixc, :, :],
bounds_error=False,
fill_value=None)
hh = fn(yz_points).reshape(grid.vnEx[1:], order='F')
# Piecewise constant interpolation in x-direction
efield[2 * ixc, :, :] += hh
efield[2 * ixc + 1, :, :] += hh
return efield
def prolon_emg3d(grid, cgrid, efield, cefield, yz_points):
"""Calculate emg3d alternative."""
fn = solver.RegularGridProlongator(cgrid.vectorNy, cgrid.vectorNz,
yz_points)
for ixc in range(cgrid.nCx):
# Bilinear interpolation in the y-z plane
hh = fn(cefield.fx[ixc, :, :]).reshape(grid.vnEx[1:], order='F')
# Piecewise constant interpolation in x-direction
efield[2 * ixc, :, :] += hh
efield[2 * ixc + 1, :, :] += hh
return efield
# Create fine grid.
nx = 2**7
hx = 50 * np.ones(nx)
hx = np.array([4, 1.1, 2, 3])
hy = np.array([2, 0.1, 20, np.pi])
hz = np.array([1, 2, 5, 1])
grid = utils.TensorMesh([hx, hy, hz], x0=np.array([0, 0, 0]))
# Create coarse grid.
chx = np.diff(grid.vectorNx[::2])
cgrid = utils.TensorMesh([chx, chx, chx], x0=np.array([0, 0, 0]))
# Create empty fine grid fields.
efield1 = utils.Field(grid)
efield2 = utils.Field(grid)
# Create coarse grid field with some values.
cefield = utils.Field(cgrid)
cefield.fx = np.arange(cefield.fx.size)
cefield.fx = 1j * np.arange(cefield.fx.size) / 10
# Required interpolation points.
yz_points = solver._get_prolongation_coordinates(grid, 'y', 'z')
# Compare
out1 = prolon_scipy(grid, cgrid, efield1.fx, cefield, yz_points)
out2 = prolon_emg3d(grid, cgrid, efield2.fx, cefield, yz_points)
assert_allclose(out1, out2)
def test_smoothing():
# 1. The only thing to test here is that smoothing returns the same as
# the corresponding jitted functions. Basically a copy of the function
# itself.
nu = 2
widths = [np.ones(2) * 100, get_h(10, 27, 10, 1.1), get_h(2, 1, 50, 1.2)]
x0 = [-w.sum() / 2 for w in widths]
src = [0, -10, -10, 43, 13]
# Loop and move the 2-cell dimension (100, 2) from x to y to z.
for xyz in range(3):
# Create a grid
grid = utils.TensorMesh(
[widths[xyz % 3], widths[(xyz + 1) % 3], widths[(xyz + 2) % 3]],
x0=np.array([x0[xyz % 3], x0[(xyz + 1) % 3], x0[(xyz + 2) % 3]]))
# Create some resistivity model
x = np.arange(1, grid.nCx + 1) * 2
y = 1 / np.arange(1, grid.nCy + 1)
z = np.arange(1, grid.nCz + 1)[::-1] / 10
res_x = np.outer(np.outer(x, y), z).ravel()
freq = 0.319
model = utils.Model(grid, res_x, 0.8 * res_x, 2 * res_x)
# Create a source field
sfield = utils.get_source_field(grid=grid, src=src, freq=freq)
# Get volume-averaged model parameters.
vmodel = utils.VolumeModel(grid, model, sfield)
# Run two iterations to get an e-field
field = solver.solve(grid, model, sfield, maxit=2, verb=1)
# Collect Gauss-Seidel input (same for all routines)
inp = (sfield.fx, sfield.fy, sfield.fz, vmodel.eta_x, vmodel.eta_y,
vmodel.eta_z, vmodel.zeta, grid.hx, grid.hy, grid.hz, nu)
func = ['', '_x', '_y', '_z']
for lr_dir in range(8):
# Get it directly from njitted
efield = utils.Field(grid, field)
if lr_dir < 4:
getattr(njitted,
'gauss_seidel' + func[lr_dir])(efield.fx, efield.fy,
efield.fz, *inp)
elif lr_dir == 4:
njitted.gauss_seidel_y(efield.fx, efield.fy, efield.fz, *inp)
njitted.gauss_seidel_z(efield.fx, efield.fy, efield.fz, *inp)
elif lr_dir == 5:
njitted.gauss_seidel_x(efield.fx, efield.fy, efield.fz, *inp)
njitted.gauss_seidel_z(efield.fx, efield.fy, efield.fz, *inp)
elif lr_dir == 6:
njitted.gauss_seidel_x(efield.fx, efield.fy, efield.fz, *inp)
njitted.gauss_seidel_y(efield.fx, efield.fy, efield.fz, *inp)
elif lr_dir == 7:
njitted.gauss_seidel_x(efield.fx, efield.fy, efield.fz, *inp)
njitted.gauss_seidel_y(efield.fx, efield.fy, efield.fz, *inp)
njitted.gauss_seidel_z(efield.fx, efield.fy, efield.fz, *inp)
# Use solver.smoothing
ofield = utils.Field(grid, field)
solver.smoothing(grid, vmodel, sfield, ofield, nu, lr_dir)
# Compare
assert_allclose(efield, ofield)
def test_solver_homogeneous(capsys):
# Regression test for homogeneous halfspace.
# Not very sophisticated; replace/extend by more detailed tests.
dat = REGRES['res'][()]
grid = utils.TensorMesh(**dat['input_grid'])
model = utils.Model(**dat['input_model'])
sfield = utils.get_source_field(**dat['input_source'])
# F-cycle
efield = solver.solve(grid, model, sfield, verb=4)
out, _ = capsys.readouterr()
assert ' emg3d START ::' in out
assert ' [hh:mm:ss] ' in out
assert ' MG cycles ' in out
assert ' Final rel. error ' in out
assert ' emg3d END :: ' in out
# Experimental:
# Check if norms are also the same, at least for first two cycles.
assert "1.509e-01 after 1 F-cycles [9.161e-07, 0.151] 0 0" in out
assert "1.002e-01 after 2 F-cycles [6.082e-07, 0.664] 0 0" in out
# Check all fields (ex, ey, and ez)
assert_allclose(dat['Fresult'], efield)
# W-cycle
wfield = solver.solve(grid, model, sfield, cycle='W', verb=1)
# Check all fields (ex, ey, and ez)
assert_allclose(dat['Wresult'], wfield)
# V-cycle
vfield = solver.solve(grid, model, sfield, cycle='V', verb=1)
_, _ = capsys.readouterr() # clear output
# Check all fields (ex, ey, and ez)
assert_allclose(dat['Vresult'], vfield)
# BiCGSTAB with some print checking.
efield = solver.solve(grid, model, sfield, verb=3, sslsolver=True)
out, _ = capsys.readouterr()
assert ' emg3d START ::' in out
assert ' [hh:mm:ss] ' in out
assert ' CONVERGED' in out
assert ' Solver steps ' in out
assert ' MG prec. steps ' in out
assert ' Final rel. error ' in out
assert ' emg3d END :: ' in out
# Check all fields (ex, ey, and ez)
assert_allclose(dat['bicresult'], efield)
# Same as previous, without BiCGSTAB, but some print checking.
efield = solver.solve(grid, model, sfield, verb=3)
out, _ = capsys.readouterr()
assert ' emg3d START ::' in out
assert ' [hh:mm:ss] ' in out
assert ' CONVERGED' in out
assert ' MG cycles ' in out
assert ' Final rel. error ' in out
assert ' emg3d END :: ' in out
# Max it
maxit = 2
_, info = solver.solve(grid,
model,
sfield,
verb=2,
maxit=maxit,
return_info=True)
out, _ = capsys.readouterr()
assert ' MAX. ITERATION REACHED' in out
assert maxit == info['it_mg']
assert info['exit'] == 1
assert 'MAX. ITERATION REACHED' in info['exit_message']
# BiCGSTAB with lower verbosity, print checking.
_ = solver.solve(grid, model, sfield, verb=2, maxit=1, sslsolver=True)
out, _ = capsys.readouterr()
assert ' MAX. ITERATION REACHED' in out
# Just check if it runs without failing for other solvers.
_ = solver.solve(grid, model, sfield, verb=3, maxit=1, sslsolver='gcrotmk')
# Provide initial field.
_, _ = capsys.readouterr() # empty
efield_copy = efield.copy()
outarray = solver.solve(grid, model, sfield, efield_copy)
out, _ = capsys.readouterr()
# Ensure there is no output.
assert outarray is None
assert "NOTHING DONE (provided efield already good enough)" in out
# Ensure the field did not change.
assert_allclose(efield, efield_copy)
# Provide initial field and return info.
info = solver.solve(grid, model, sfield, efield_copy, return_info=True)
assert info['it_mg'] == 0
assert info['it_ssl'] == 0
assert info['exit'] == 0
assert info['exit_message'] == 'CONVERGED'
# Provide initial field, ensure one initial multigrid is carried out
# without linerelaxation nor semicoarsening.
_, _ = capsys.readouterr() # empty
efield = utils.Field(grid)
outarray = solver.solve(grid,
model,
sfield,
efield,
sslsolver=True,
semicoarsening=True,
linerelaxation=True,
maxit=2,
verb=3)
out, _ = capsys.readouterr()
assert "after 1 F-cycles 4 1" in out
assert "after 2 F-cycles 5 2" in out
# Provide an initial source-field without frequency information.
wrong_sfield = utils.Field(grid)
wrong_sfield.field = sfield.field
with pytest.raises(ValueError):
solver.solve(grid, model, wrong_sfield, efield=efield, verb=2)
out, _ = capsys.readouterr()
assert "ERROR :: Source field is missing frequency information" in out
# Check stagnation by providing an almost zero source field.
_ = solver.solve(grid, model, sfield * 0 + 1e-20, maxit=100)
out, _ = capsys.readouterr()
assert "STAGNATED" in out
