274 lines
8.6 KiB
Python
274 lines
8.6 KiB
Python
import pytest
|
|
import numpy
|
|
from numpy.testing import assert_allclose #, assert_array_equal
|
|
|
|
from .. import Grid, Extent, GridError, Plane, Slab
|
|
|
|
|
|
def test_draw_oncenter_2x2() -> None:
|
|
xs = [-1, 0, 1]
|
|
ys = [-1, 0, 1]
|
|
zs = [-1, 1]
|
|
grid = Grid([xs, ys, zs], shifts=[[0, 0, 0]])
|
|
arr = grid.allocate(0)
|
|
|
|
grid.draw_cuboid(
|
|
arr,
|
|
x=dict(center=0, span=1),
|
|
y=Extent(center=0, span=1),
|
|
z=dict(center=0, span=10),
|
|
foreground=1,
|
|
)
|
|
|
|
correct = numpy.array([[0.25, 0.25],
|
|
[0.25, 0.25]])[None, :, :, None]
|
|
|
|
assert_allclose(arr, correct)
|
|
|
|
|
|
def test_draw_ongrid_4x4() -> None:
|
|
xs = [-2, -1, 0, 1, 2]
|
|
ys = [-2, -1, 0, 1, 2]
|
|
zs = [-1, 1]
|
|
grid = Grid([xs, ys, zs], shifts=[[0, 0, 0]])
|
|
arr = grid.allocate(0)
|
|
|
|
grid.draw_cuboid(
|
|
arr,
|
|
x=dict(center=0, span=2),
|
|
y=dict(min=-1, max=1),
|
|
z=dict(center=0, min=-5),
|
|
foreground=1,
|
|
)
|
|
|
|
correct = numpy.array([[0, 0, 0, 0],
|
|
[0, 1, 1, 0],
|
|
[0, 1, 1, 0],
|
|
[0, 0, 0, 0]])[None, :, :, None]
|
|
|
|
assert_allclose(arr, correct)
|
|
|
|
|
|
def test_draw_xshift_4x4() -> None:
|
|
xs = [-2, -1, 0, 1, 2]
|
|
ys = [-2, -1, 0, 1, 2]
|
|
zs = [-1, 1]
|
|
grid = Grid([xs, ys, zs], shifts=[[0, 0, 0]])
|
|
arr = grid.allocate(0)
|
|
|
|
grid.draw_cuboid(
|
|
arr,
|
|
x=dict(center=0.5, span=1.5),
|
|
y=dict(min=-1, max=1),
|
|
z=dict(center=0, span=10),
|
|
foreground=1,
|
|
)
|
|
|
|
correct = numpy.array([[0, 0, 0, 0],
|
|
[0, 0.25, 0.25, 0],
|
|
[0, 1, 1, 0],
|
|
[0, 0.25, 0.25, 0]])[None, :, :, None]
|
|
|
|
assert_allclose(arr, correct)
|
|
|
|
|
|
def test_draw_yshift_4x4() -> None:
|
|
xs = [-2, -1, 0, 1, 2]
|
|
ys = [-2, -1, 0, 1, 2]
|
|
zs = [-1, 1]
|
|
grid = Grid([xs, ys, zs], shifts=[[0, 0, 0]])
|
|
arr = grid.allocate(0)
|
|
|
|
grid.draw_cuboid(
|
|
arr,
|
|
x=dict(min=-1, max=1),
|
|
y=dict(center=0.5, span=1.5),
|
|
z=dict(center=0, span=10),
|
|
foreground=1,
|
|
)
|
|
|
|
correct = numpy.array([[0, 0, 0, 0],
|
|
[0, 0.25, 1, 0.25],
|
|
[0, 0.25, 1, 0.25],
|
|
[0, 0, 0, 0]])[None, :, :, None]
|
|
|
|
assert_allclose(arr, correct)
|
|
|
|
|
|
def test_draw_2shift_4x4() -> None:
|
|
xs = [-2, -1, 0, 1, 2]
|
|
ys = [-2, -1, 0, 1, 2]
|
|
zs = [-1, 1]
|
|
grid = Grid([xs, ys, zs], shifts=[[0, 0, 0]])
|
|
arr = grid.allocate(0)
|
|
|
|
grid.draw_cuboid(
|
|
arr,
|
|
x=dict(center=0.5, span=1.5),
|
|
y=dict(min=-0.5, max=0.5),
|
|
z=dict(center=0, span=10),
|
|
foreground=1,
|
|
)
|
|
|
|
correct = numpy.array([[0, 0, 0, 0],
|
|
[0, 0.125, 0.125, 0],
|
|
[0, 0.5, 0.5, 0],
|
|
[0, 0.125, 0.125, 0]])[None, :, :, None]
|
|
|
|
assert_allclose(arr, correct)
|
|
|
|
|
|
def test_ind2pos_round_preserves_float_centers() -> None:
|
|
grid = Grid([[0, 1, 3], [0, 2], [0, 1]], shifts=[[0, 0, 0]])
|
|
|
|
pos = grid.ind2pos(numpy.array([1, 0, 0]), which_shifts=0)
|
|
|
|
assert_allclose(pos, [2.0, 1.0, 0.5])
|
|
|
|
|
|
def test_ind2pos_enforces_bounds_for_rounded_and_fractional_indices() -> None:
|
|
grid = Grid([[0, 1, 3], [0, 2], [0, 1]], shifts=[[0, 0, 0]])
|
|
|
|
with pytest.raises(GridError):
|
|
grid.ind2pos(numpy.array([2, 0, 0]), which_shifts=0, check_bounds=True)
|
|
|
|
edge_pos = grid.ind2pos(numpy.array([1.5, 0.5, 0.5]), which_shifts=0, round_ind=False, check_bounds=True)
|
|
assert_allclose(edge_pos, [3.0, 2.0, 1.0])
|
|
|
|
with pytest.raises(GridError):
|
|
grid.ind2pos(numpy.array([1.6, 0.5, 0.5]), which_shifts=0, round_ind=False, check_bounds=True)
|
|
|
|
|
|
def test_draw_polygon_accepts_coplanar_nx3_vertices() -> None:
|
|
grid = Grid([[0, 1, 2], [0, 1, 2], [0, 1]], shifts=[[0, 0, 0]])
|
|
arr_2d = grid.allocate(0)
|
|
arr_3d = grid.allocate(0)
|
|
slab = dict(axis='z', center=0.5, span=1.0)
|
|
|
|
polygon_2d = numpy.array([[0, 0], [1, 0], [1, 1], [0, 1]], dtype=float)
|
|
polygon_3d = numpy.array([[0, 0, 0.5],
|
|
[1, 0, 0.5],
|
|
[1, 1, 0.5],
|
|
[0, 1, 0.5]], dtype=float)
|
|
|
|
grid.draw_polygon(arr_2d, slab=slab, polygon=polygon_2d, foreground=1)
|
|
grid.draw_polygon(arr_3d, slab=slab, polygon=polygon_3d, foreground=1)
|
|
|
|
assert_allclose(arr_3d, arr_2d)
|
|
|
|
|
|
def test_draw_polygon_rejects_noncoplanar_nx3_vertices() -> None:
|
|
grid = Grid([[0, 1, 2], [0, 1, 2], [0, 1]], shifts=[[0, 0, 0]])
|
|
arr = grid.allocate(0)
|
|
polygon = numpy.array([[0, 0, 0.5],
|
|
[1, 0, 0.5],
|
|
[1, 1, 0.75],
|
|
[0, 1, 0.5]], dtype=float)
|
|
|
|
with pytest.raises(GridError):
|
|
grid.draw_polygon(arr, slab=dict(axis='z', center=0.5, span=1.0), polygon=polygon, foreground=1)
|
|
|
|
|
|
def test_get_slice_supports_sampling() -> None:
|
|
grid = Grid([[0, 1, 2, 3], [0, 1, 2, 3], [0, 1]], shifts=[[0, 0, 0]])
|
|
cell_data = numpy.arange(numpy.prod(grid.cell_data_shape), dtype=float).reshape(grid.cell_data_shape)
|
|
|
|
grid_slice = grid.get_slice(cell_data, Plane(z=0.5), sample_period=2)
|
|
|
|
assert_allclose(grid_slice, cell_data[0, ::2, ::2, 0])
|
|
|
|
|
|
def test_sampled_visualization_helpers_do_not_error() -> None:
|
|
matplotlib = pytest.importorskip('matplotlib')
|
|
matplotlib.use('Agg')
|
|
from matplotlib import pyplot
|
|
|
|
grid = Grid([[0, 1, 2, 3], [0, 1, 2, 3], [0, 1]], shifts=[[0, 0, 0]])
|
|
cell_data = numpy.arange(numpy.prod(grid.cell_data_shape), dtype=float).reshape(grid.cell_data_shape)
|
|
|
|
fig_slice, ax_slice = grid.visualize_slice(cell_data, Plane(z=0.5), sample_period=2, finalize=False)
|
|
fig_edges, ax_edges = grid.visualize_edges(cell_data, Plane(z=0.5), sample_period=2, finalize=False)
|
|
|
|
assert fig_slice is ax_slice.figure
|
|
assert fig_edges is ax_edges.figure
|
|
|
|
pyplot.close(fig_slice)
|
|
pyplot.close(fig_edges)
|
|
|
|
|
|
def test_grid_constructor_rejects_invalid_coordinate_count() -> None:
|
|
with pytest.raises(GridError):
|
|
Grid([[0, 1], [0, 1]], shifts=[[0, 0, 0]])
|
|
|
|
with pytest.raises(GridError):
|
|
Grid([[0, 1], [0, 1], [0, 1], [0, 1]], shifts=[[0, 0, 0]])
|
|
|
|
|
|
def test_grid_constructor_rejects_invalid_periodic_length() -> None:
|
|
with pytest.raises(GridError):
|
|
Grid([[0, 1], [0, 1], [0, 1]], shifts=[[0, 0, 0]], periodic=[True, False])
|
|
|
|
|
|
def test_extent_and_slab_reject_inverted_geometry() -> None:
|
|
with pytest.raises(GridError):
|
|
Extent(center=0, min=1)
|
|
|
|
with pytest.raises(GridError):
|
|
Extent(min=2, max=1)
|
|
|
|
with pytest.raises(GridError):
|
|
Slab(axis='z', center=1, max=0)
|
|
|
|
|
|
def test_extent_accepts_scalar_like_inputs() -> None:
|
|
extent = Extent(min=numpy.array([1.0]), span=numpy.array([4.0]))
|
|
|
|
assert_allclose([extent.center, extent.span, extent.min, extent.max], [3.0, 4.0, 1.0, 5.0])
|
|
|
|
|
|
|
|
|
|
def test_visualize_isosurface_sampling_uses_preview_lattice(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
matplotlib = pytest.importorskip('matplotlib')
|
|
matplotlib.use('Agg')
|
|
skimage_measure = pytest.importorskip('skimage.measure')
|
|
from matplotlib import pyplot
|
|
from mpl_toolkits.mplot3d.axes3d import Axes3D
|
|
|
|
captured: dict[str, numpy.ndarray] = {}
|
|
|
|
def fake_marching_cubes(_grid: numpy.ndarray, _level: float) -> tuple[numpy.ndarray, numpy.ndarray, None, None]:
|
|
verts = numpy.array([[0.5, 0.5, 0.5],
|
|
[0.5, 1.5, 0.5],
|
|
[1.5, 0.5, 0.5]], dtype=float)
|
|
faces = numpy.array([[0, 1, 2]], dtype=int)
|
|
return verts, faces, None, None
|
|
|
|
def fake_plot_trisurf( # noqa: ANN202
|
|
_self: object,
|
|
xs: numpy.ndarray,
|
|
ys: numpy.ndarray,
|
|
faces: numpy.ndarray,
|
|
zs: numpy.ndarray,
|
|
*_args: object,
|
|
**_kwargs: object,
|
|
) -> object:
|
|
captured['xs'] = numpy.asarray(xs)
|
|
captured['ys'] = numpy.asarray(ys)
|
|
captured['faces'] = numpy.asarray(faces)
|
|
captured['zs'] = numpy.asarray(zs)
|
|
return object()
|
|
|
|
monkeypatch.setattr(skimage_measure, 'marching_cubes', fake_marching_cubes)
|
|
monkeypatch.setattr(Axes3D, 'plot_trisurf', fake_plot_trisurf)
|
|
|
|
grid = Grid([numpy.arange(7, dtype=float), numpy.arange(7, dtype=float), numpy.arange(7, dtype=float)], shifts=[[0, 0, 0]])
|
|
cell_data = numpy.zeros(grid.cell_data_shape)
|
|
|
|
fig, _ax = grid.visualize_isosurface(cell_data, level=0.5, sample_period=2, finalize=False)
|
|
|
|
assert_allclose(captured['xs'], [1.5, 1.5, 3.5])
|
|
assert_allclose(captured['ys'], [1.5, 3.5, 1.5])
|
|
assert_allclose(captured['zs'], [1.5, 1.5, 1.5])
|
|
|
|
pyplot.close(fig)
|