gridlock/gridlock/test/test_grid.py

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_get_slice_uses_shifted_grid_bounds() -> None:
    grid = Grid([[0, 1, 2], [0, 1, 2], [0, 1, 2]], shifts=[[0.5, 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(x=2.0), which_shifts=0)

    assert_allclose(grid_slice, cell_data[0, 1, :, :])

    with pytest.raises(GridError):
        grid.get_slice(cell_data, Plane(x=2.1), which_shifts=0)


def test_draw_extrude_rectangle_uses_boundary_slice() -> None:
    grid = Grid([[0, 1, 2], [0, 1, 2], [0, 1, 2]], shifts=[[0, 0, 0]])
    cell_data = grid.allocate(0)
    source = numpy.array([[1, 2],
                          [3, 4]], dtype=float)
    cell_data[0, :, :, 1] = source

    grid.draw_extrude_rectangle(
        cell_data,
        rectangle=[[0, 0, 2], [2, 2, 2]],
        direction=2,
        polarity=-1,
        distance=2,
    )

    assert_allclose(cell_data[0, :, :, 0], source)
    assert_allclose(cell_data[0, :, :, 1], source)


def test_sampled_preview_exyz_tracks_nonuniform_centers() -> None:
    grid = Grid([[0, 1, 3, 6, 10], [0, 1, 2], [0, 1, 2]], shifts=[[0, 0, 0]])

    sampled_exyz = grid._sampled_exyz(0, 2)

    assert_allclose(sampled_exyz[0], [-1.5, 2.5, 6.5])


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)