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


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)