def test_CartesianGrid():
grid = grids.CartesianGrid(np.array([-1, -1, -1]) * u.cm,
np.array([1, 1, 1]) * u.cm,
num=(10, 10, 10))
x_arr, y_arr, z_arr = grid.grids
x_axis, y_axis, z_axis = grid.ax0, grid.ax1, grid.ax2
d_x, d_y, d_z = grid.dax0, grid.dax1, grid.dax2
is_uniform = grid.is_uniform
shape = grid.shape
unit = grid.units
# Grid should be uniform
assert grid.is_uniform
# Test initializing with a provided grid and a quantity
q = np.zeros(grid.shape)
grid2 = grids.CartesianGrid(grid.grids[0],
grid.grids[1],
grid.grids[2],
test_quantity=q)
# Test that input with the wrong units will raise an exception
L0 = [-1 * u.mm, 0 * u.rad, -1 * u.mm]
L1 = [1 * u.mm, 2 * np.pi * u.rad, 1 * u.mm]
with pytest.raises(ValueError):
grid = grids.CartesianGrid(L0, L1, num=10)
def test_CartesianGrid():
grid = grids.CartesianGrid(np.array([-1, -1, -1]) * u.cm,
np.array([1, 1, 1]) * u.cm,
num=(10, 10, 10))
x_arr, y_arr, z_arr = grid.grids
x_axis, y_axis, z_axis = grid.ax0, grid.ax1, grid.ax2
d_x, d_y, d_z = grid.dax0, grid.dax1, grid.dax2
is_uniform = grid.is_uniform
shape = grid.shape
unit = grid.units
# Grid should be uniform
assert grid.is_uniform == True
# Test initializing with a provided grid
grid2 = grids.CartesianGrid(
grid.grids[0],
grid.grids[1],
grid.grids[2],
)
# Units not all consistent
with pytest.raises(ValueError):
grid = grids.CartesianGrid([-1 * u.m, -1 * u.rad, -1 * u.m],
[1 * u.m, 1 * u.rad, 1 * u.m])
def test_CartesianGrid_creation(args, kwargs, shape, error):
# If no exception is expected, create the grid and check its shape
if error is None:
grid = grids.CartesianGrid(*args, **kwargs)
assert grid.shape == shape
# If an exception is expected, verify that it is raised
else:
with pytest.raises(error):
grid = grids.CartesianGrid(*args, **kwargs)
def test_print_summary(abstract_grid_uniform, abstract_grid_nonuniform):
"""
Verify that both __str__ methods can be called without errors
"""
print(abstract_grid_uniform)
print(abstract_grid_nonuniform)
# Test printing a grid with no quantities
grid = grids.CartesianGrid(-1 * u.cm, 1 * u.cm, num=3)
print(grid)
# Test printing a grid with unrecognized quantities
grid = grids.CartesianGrid(-1 * u.cm, 1 * u.cm, num=3)
grid.add_quantities(unrecognized_quantity=np.ones([3, 3, 3]) * u.T)
print(grid)
def test_AbstractGrid_creation(args, kwargs, shape, error):
"""
Test the creation of AbstractGrids
Use CartesianGrid as the test example
"""
# If no exception is expected, create the grid and check its shape
if error is None:
grid = grids.CartesianGrid(*args, **kwargs)
assert grid.shape == shape
# If an exception is expected, verify that it is raised
else:
with pytest.raises(error):
grid = grids.CartesianGrid(*args, **kwargs)
def test_volume_averaged_interpolator():
# Create grid
grid = grids.CartesianGrid(-1 * u.cm, 1 * u.cm, num=25)
# Add some data to the grid
grid.add_quantities(x=grid.grids[0])
grid.add_quantities(y=grid.grids[1])
# One position
pos = np.array([0.1, -0.3, 0]) * u.cm
pout = grid.volume_averaged_interpolator(pos, "x")
assert np.allclose(pos[0], pout, atol=0.1)
# Test quantity key not present in dataset
with pytest.raises(KeyError):
pout = grid.volume_averaged_interpolator(pos, "B_x")
# Two positions, two quantities
pos = np.array([[0.1, -0.3, 0], [0.1, -0.3, 0]]) * u.cm
pout = grid.volume_averaged_interpolator(pos, "x", "y")
# Contains out-of-bounds values (must handle NaNs correctly)
pos = np.array([5, -0.3, 0]) * u.cm
pout = grid.volume_averaged_interpolator(pos, "x")
assert np.allclose(pout, 0 * u.cm, atol=0.1)
# Try running with persistance
pos = np.array([[0.1, -0.3, 0], [0.1, -0.3, 0]]) * u.cm
p1, p2 = grid.volume_averaged_interpolator(pos, "x", "y", persistent=True)
p1, p2 = grid.volume_averaged_interpolator(pos, "x", "y", persistent=True)
# Try changing the arg list, make sure it catchs this and auto-reverts
# to non-persistent interpolation in that case
p1, p2 = grid.volume_averaged_interpolator(pos, "x", persistent=True)
assert p1.size == 1
def test_grid_methods():
# ************ UNIFORM CARTESIAN ****************************
grid = grids.CartesianGrid(np.array([-1, -1, -1]) * u.cm,
np.array([1, 1, 1]) * u.cm,
num=(10, 10, 10))
# Test on-grid
pos = np.array([[0.1, -0.3, 0], [3, -0.3, 0]]) * u.cm
out = grid.on_grid(pos)
assert np.all(out == np.array([True, False]))
# Test vector_intersects
# This vector passes through the grid
p1, p2 = np.array([0, -5, 0]) * u.cm, np.array([0, 5, 0]) * u.cm
assert grid.vector_intersects(p1, p2)
# Test going backwards yields the same result
assert grid.vector_intersects(p2, p1)
# This one doesn't
p1, p2 = np.array([0, -5, 0]) * u.cm, np.array([0, -5, 10]) * u.cm
assert not grid.vector_intersects(p1, p2)
assert not grid.vector_intersects(p2, p1)
# ************ NON-UNIFORM CARTESIAN ****************************
grid = grids.NonUniformCartesianGrid(-1 * u.cm, 1 * u.cm, num=10)
pos = np.array([[0.1, -0.3, 0], [3, -0.3, 0]]) * u.cm
out = grid.on_grid(pos)
assert np.all(out == np.array([True, False]))
def test_interpolate_indices():
# Create grid
grid = grids.CartesianGrid(-1 * u.cm, 1 * u.cm, num=25)
# One position
pos = np.array([0.1, -0.3, 0]) * u.cm
i = grid.interpolate_indices(pos)[0]
# Assert that nearest grid cell was found
pout = grid.grid[int(i[0]), int(i[1]), int(i[2])]
assert np.allclose(pos, pout, atol=0.1)
# Two positions
pos = np.array([[0.1, -0.3, 0], [0.1, -0.3, 0]]) * u.cm
i = grid.interpolate_indices(pos)[0]
# Contains out-of-bounds values (index array should contain NaNs)
pos = np.array([5, -0.3, 0]) * u.cm
i = grid.interpolate_indices(pos)[0]
assert np.sum(np.isnan(i)) > 0
# ***********************************************************************
# Create a non-uniform grid
grid = grids.NonUniformCartesianGrid(-1 * u.cm, 1 * u.cm, num=100)
# One position
pos = np.array([0.1, -0.3, 0]) * u.cm
i = grid.interpolate_indices(pos)[0]
# Assert that nearest grid cell was found
pout = grid.grid[int(i)]
assert np.allclose(pos, pout, atol=0.5)
def test_nearest_neighbor_interpolator():
# Create grid
grid = grids.CartesianGrid(-1 * u.cm, 1 * u.cm, num=25)
# Add some data to the grid
grid.add_quantities(x=grid.grids[0])
grid.add_quantities(y=grid.grids[1])
# One position
pos = np.array([0.1, -0.3, 0]) * u.cm
pout = grid.nearest_neighbor_interpolator(pos, "x")
assert np.allclose(pos[0], pout, atol=0.1)
# Test quantity key not present in dataset
with pytest.raises(KeyError):
pout = grid.nearest_neighbor_interpolator(pos, "B_x")
# Two positions, two quantities
pos = np.array([[0.1, -0.3, 0], [0.1, -0.3, 0]]) * u.cm
pout = grid.nearest_neighbor_interpolator(pos, "x", "y")
# Contains out-of-bounds values (must handle NaNs correctly)
pos = np.array([5, -0.3, 0]) * u.cm
pout = grid.nearest_neighbor_interpolator(pos, "x")
assert np.allclose(pout, 0 * u.cm, atol=0.1)
# ***********************************************************************
# Create a non-uniform grid
grid = grids.NonUniformCartesianGrid(-1 * u.cm, 1 * u.cm, num=100)
grid.add_quantities(x=grid.grids[0] * u.cm)
# One position
pos = np.array([0.1, -0.3, 0]) * u.cm
pout = grid.nearest_neighbor_interpolator(pos, "x")
assert np.allclose(pos[0], pout, atol=0.5)
def debug_volume_avg_interpolator():
"""
Plot the comparison of the nearest neighbor interpolator and volume
averaged interpolator for `~plasmapy.plasma.grids.CartesianGrid`.
"""
import matplotlib.pyplot as plt
grid = grids.CartesianGrid(-1 * u.cm, 1 * u.cm, num=24)
# add x and y positions to grid
grid.add_quantities(x=grid.grids[0])
grid.add_quantities(y=grid.grids[1])
# add a mass density to grid
radius = np.sqrt(grid.pts0 ** 2 + grid.pts1 ** 2 + grid.pts2 ** 2)
rho = radius.to(u.mm).value ** 4 * u.kg * u.m ** -3
grid.add_quantities(rho=rho)
# Create a low resolution test grid
npts = 150
interp_pts = (
np.array([np.linspace(-0.99, 1, num=npts), np.zeros(npts), np.zeros(npts)])
* u.cm
)
interp_pts = np.moveaxis(interp_pts, 0, -1)
interp_hax = interp_pts[:, 0].to(u.mm).value
va_rho = grid.volume_averaged_interpolator(interp_pts, "rho")
nn_rho = grid.nearest_neighbor_interpolator(interp_pts, "rho")
analytic = interp_hax ** 4
raw_hax = grid.ax0.to(u.mm).value
half = 12
raw_rho = grid["rho"][:, half, half]
plt.plot(raw_hax, raw_rho, marker="*", label="Interpolated points")
plt.plot(interp_hax, nn_rho, label="Nearest neighbor")
plt.plot(interp_hax, va_rho, marker="o", label="Volume averaged")
plt.plot(interp_hax, analytic, label="Analytic")
plt.legend()
plt.xlim(-11, 11)
def test_grid_methods():
grid = grids.CartesianGrid(np.array([-1, -1, -1]) * u.cm,
np.array([1, 1, 1]) * u.cm,
num=(10, 10, 10))
# Test on-grid
pos = np.array([[0.1, -0.3, 0], [3, -0.3, 0]]) * u.cm
out = grid.on_grid(pos)
assert np.all(out == np.array([True, False]))
# Test vector_intersects
# This vector passes through the grid
p1, p2 = np.array([0, -5, 0]) * u.cm, np.array([0, 5, 0]) * u.cm
assert grid.vector_intersects(p1, p2) == True
# Test going backwards yields the same result
assert grid.vector_intersects(p2, p1) == True
# This one doesn't
p1, p2 = np.array([0, -5, 0]) * u.cm, np.array([0, -5, 10]) * u.cm
assert grid.vector_intersects(p1, p2) == False
assert grid.vector_intersects(p2, p1) == False
def test_volume_averaged_interpolator_known_solutions():
# Create a 4x4x4 test grid with positions -3, -1, 1, and 3 cm
# Add a quantity that equals 0 when x=-3, 1 when x=-1,
# 2 when x=1, and 3 when x= 3
grid = grids.CartesianGrid(-3 * u.cm, 3 * u.cm, num=4)
Bz = np.zeros([4, 4, 4]) * u.T
Bz[1, :, :] = 1 * u.T
Bz[2, :, :] = 2 * u.T
Bz[3, :, :] = 3 * u.T
grid.add_quantities(B_z=Bz)
# Interpolate the following points:
xpts = np.linspace(-4, 4, num=9) * u.cm
pos = np.zeros([xpts.size, 3])
pos[:, 0] = xpts
pts = grid.volume_averaged_interpolator(pos, "B_z", persistent=True)
assert np.allclose(
pts.value, np.array([np.nan, 0, 0.5, 1, 1.5, 2, 2.5, 3, np.nan]), equal_nan=True
)
def uniform_cartesian_grid():
"""
A `pytest` fixture that generates a CartesianGrid that spans
-1 cm to 1 cm in all dimensions. Three quantities are added to the
grid:
1. "x" which is the x position at each point in the grid [in cm]
2. "y" which is the y position at each point in the grid [in cm]
3. "z" which is the z position at each point in the grid [in cm]
4. "rho" which is a mass density at each point in the grid [kg/m^-3]
"""
# Create grid
grid = grids.CartesianGrid(-1 * u.cm, 1 * u.cm, num=21)
# Add some data to the grid
grid.add_quantities(x=grid.grids[0])
grid.add_quantities(y=grid.grids[1])
grid.add_quantities(z=grid.grids[2])
radius = np.sqrt(grid.pts0 ** 2 + grid.pts1 ** 2 + grid.pts2 ** 2)
rho = radius.to(u.mm).value ** 4 * u.kg * u.m ** -3
grid.add_quantities(rho=rho)
return grid
def test_interpolators():
# Test interpolator
# Test interpolator
# Create grid
grid = grids.CartesianGrid(-1 * u.cm, 1 * u.cm, num=25)
# Add some data to the grid
data = grid.add_quantity("pts0", grid.grids[0])
data = grid.add_quantity("pts1", grid.grids[1])
pos = np.array([0.1, -0.3, 0]) * u.cm
pos2 = np.array([[0.1, -0.3, 0], [0.1, -0.3, 0]]) * u.cm
# Interpolate indices
i = grid.interpolate_indices(pos)[0]
pout = grid.grid[i[0], i[1], i[2], :] * grid.unit
# Assert that nearest grid cell was found
assert np.allclose(pos, pout, atol=0.1)
# Interpolate grid values using nearest-neighbor interpolator
pout = grid.nearest_neighbor_interpolator(pos, "pts0")
assert np.allclose(pos[0], pout, atol=0.1)
# Interpolate grid values using volume-weighted interpolator
pout = grid.volume_averaged_interpolator(pos, "pts0")
assert np.allclose(pos[0], pout, atol=0.1)
# Test using multiple arguments
pout, pout2 = grid.nearest_neighbor_interpolator(pos, "pts0", "pts1")
# Test multiple grid points
pos = np.array([[0, 0, 0], [0.5, 0.2, 0]]) * u.cm
i = grid.interpolate_indices(pos)
pout = grid.nearest_neighbor_interpolator(pos.value, "pts0")
pout = grid.volume_averaged_interpolator(pos, "pts0")
