def test_unit_grid_2d():
"""test 2D grids"""
# test special case
grid = UnitGrid([4, 4], periodic=True)
assert grid.dim == 2
assert grid.numba_type == "f8[:, :]"
assert grid.volume == 16
np.testing.assert_array_equal(grid.discretization, np.ones(2))
assert grid.get_image_data(np.zeros(grid.shape))["extent"] == [0, 4, 0, 4]
for _ in range(10):
p = np.random.randn(2)
assert np.all(grid.polar_coordinates_real(p) < np.sqrt(8))
large_enough = grid.polar_coordinates_real((0, 0)) > np.sqrt(4)
assert np.any(large_enough)
periodic = random.choices([True, False], k=2)
grid = UnitGrid([4, 4], periodic=periodic)
assert grid.dim == 2
assert grid.volume == 16
assert grid.polar_coordinates_real((1, 1)).shape == (4, 4)
grid = UnitGrid([4, 8], periodic=periodic)
assert grid.dim == 2
assert grid.volume == 32
assert grid.polar_coordinates_real((1, 1)).shape == (4, 8)
# test conversion between polar and Cartesian coordinates
c1 = grid.cell_coords
p = np.random.random(2) * grid.shape
d, a = grid.polar_coordinates_real(p, ret_angle=True)
c2 = grid.from_polar_coordinates(d, a, p)
assert np.allclose(grid.distance_real(c1, c2), 0)
# test boundary points
np.testing.assert_equal(
grid._boundary_coordinates(0, False),
np.c_[np.full(8, 0), np.linspace(0.5, 7.5, 8)],
)
np.testing.assert_equal(
grid._boundary_coordinates(0, True),
np.c_[np.full(8, 4), np.linspace(0.5, 7.5, 8)],
)
np.testing.assert_equal(
grid._boundary_coordinates(1, False),
np.c_[np.linspace(0.5, 3.5, 4),
np.full(4, 0)],
)
np.testing.assert_equal(
grid._boundary_coordinates(1, True),
np.c_[np.linspace(0.5, 3.5, 4),
np.full(4, 8)],
)
def test_unit_grid_1d(periodic):
"""test 1D grids"""
grid = UnitGrid(4, periodic=periodic)
assert grid.dim == 1
assert grid.numba_type == "f8[:]"
assert grid.volume == 4
np.testing.assert_array_equal(grid.discretization, np.ones(1))
dist, angle = grid.polar_coordinates_real(0, ret_angle=True)
if periodic:
np.testing.assert_allclose(dist, [0.5, 1.5, 1.5, 0.5])
else:
np.testing.assert_allclose(dist, np.arange(4) + 0.5)
assert angle.shape == (4, )
grid = UnitGrid(8, periodic=periodic)
assert grid.dim == 1
assert grid.volume == 8
norm_numba = grid.make_normalize_point_compiled(reflect=False)
def norm_numba_wrap(x):
y = np.array([x])
norm_numba(y)
return y
for normalize in [
partial(grid.normalize_point, reflect=False), norm_numba_wrap
]:
if periodic:
np.testing.assert_allclose(normalize(-1e-10), 8 - 1e-10)
np.testing.assert_allclose(normalize(1e-10), 1e-10)
np.testing.assert_allclose(normalize(8 - 1e-10), 8 - 1e-10)
np.testing.assert_allclose(normalize(8 + 1e-10), 1e-10)
else:
for x in [-1e-10, 1e-10, 8 - 1e-10, 8 + 1e-10]:
np.testing.assert_allclose(normalize(x), x)
grid = UnitGrid(8, periodic=periodic)
# test conversion between polar and Cartesian coordinates
c1 = grid.cell_coords
p = np.random.random(1) * grid.shape
d, a = grid.polar_coordinates_real(p, ret_angle=True)
c2 = grid.from_polar_coordinates(d, a, p)
assert np.allclose(grid.distance_real(c1, c2), 0)
# test boundary points
np.testing.assert_equal(grid._boundary_coordinates(0, False),
np.array([0]))
np.testing.assert_equal(grid._boundary_coordinates(0, True), np.array([8]))
def test_unit_grid_3d():
"""test 3D grids"""
grid = UnitGrid([4, 4, 4])
assert grid.dim == 3
assert grid.numba_type == "f8[:, :, :]"
assert grid.volume == 64
np.testing.assert_array_equal(grid.discretization, np.ones(3))
assert grid.get_image_data(np.zeros(grid.shape))["extent"] == [0, 4, 0, 4]
assert grid.polar_coordinates_real((1, 1, 3)).shape == (4, 4, 4)
periodic = random.choices([True, False], k=3)
grid = UnitGrid([4, 6, 8], periodic=periodic)
assert grid.dim == 3
assert grid.volume == 192
assert grid.polar_coordinates_real((1, 1, 2)).shape == (4, 6, 8)
grid = UnitGrid([4, 4, 4], periodic=True)
assert grid.dim == 3
assert grid.volume == 64
for _ in range(10):
p = np.random.randn(3)
not_too_large = grid.polar_coordinates_real(p) < np.sqrt(12)
assert np.all(not_too_large)
large_enough = grid.polar_coordinates_real((0, 0, 0)) > np.sqrt(6)
assert np.any(large_enough)
# test boundary points
for bndry in grid._iter_boundaries():
assert grid._boundary_coordinates(*bndry).shape == (4, 4, 3)
