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)