def test_init(self):
# Test invalid input
with self.assertRaises(TypeError):
gd.UniformGridData(1, 0)
# Test invalid input
with self.assertRaises(ValueError):
gd.UniformGridData(self.geom, np.array([2]))
data = np.array([i * np.linspace(1, 5, 51) for i in range(101)])
ug_data = gd.UniformGridData(self.geom, data)
self.assertEqual(ug_data.grid, self.geom)
self.assertIsNot(ug_data.grid, self.geom)
self.assertTrue(np.array_equal(ug_data.data, data))
self.assertIsNot(ug_data.data, data)
# Test from_grid_structure
ug_data_from_grid_structure = gd.UniformGridData.from_grid_structure(
data, x0=[0, 0], x1=[1, 0.5]
)
self.assertEqual(ug_data, ug_data_from_grid_structure)
# Test not equal of UniformGridData
self.assertNotEqual(ug_data, 2)
# Test num_dimensions
self.assertEqual(ug_data.num_dimensions, 2)
self.assertEqual(ug_data.num_extended_dimensions, 2)
self.assertCountEqual(ug_data.extended_dimensions, [True, True])
def test_ghost_zones_remove(self):
geom = gd.UniformGrid(
[101, 201], x0=[0, 0], x1=[100, 200], num_ghost=[1, 3]
)
data = np.array([i * np.linspace(0, 200, 201) for i in range(101)])
ug_data = gd.UniformGridData(geom, data)
ug_data.ghost_zones_remove()
expected_data = np.array(
[i * np.linspace(3, 197, 195) for i in range(1, 100)]
)
expected_grid = gd.UniformGrid(
[99, 195], x0=[1, 3], x1=[99, 197], num_ghost=[0, 0]
)
expected_ug_data = gd.UniformGridData(expected_grid, expected_data)
self.assertEqual(ug_data, expected_ug_data)
# Check invalidation of spline
self.assertTrue(ug_data.invalid_spline)
self.assertCountEqual(ug_data.num_ghost, [0, 0])
# Check with num_ghost = 0
ug_data.ghost_zones_remove()
self.assertEqual(ug_data, ug_data.copy())
def test__apply_binary(self):
data1 = np.array([i * np.linspace(1, 5, 51) for i in range(101)])
data2 = np.array([i ** 2 * np.linspace(1, 5, 51) for i in range(101)])
ug_data1 = gd.UniformGridData(self.geom, data1)
ug_data2 = gd.UniformGridData(self.geom, data2)
expected_ug_data = gd.UniformGridData(self.geom, data1 + data2)
self.assertEqual(ug_data1 + ug_data2, expected_ug_data)
# Test incompatible grids
geom = gd.UniformGrid([101, 1], x0=[0, 0], dx=[1, 1])
data3 = np.array([i * np.linspace(1, 5, 1) for i in range(101)])
ug_data3 = gd.UniformGridData(geom, data3)
with self.assertRaises(ValueError):
ug_data1 + ug_data3
# Add number
self.assertEqual(
ug_data1 + 1, gd.UniformGridData(self.geom, data1 + 1)
)
# Incompatible objects
with self.assertRaises(TypeError):
ug_data1 + geom
def test_flat_dimensions_remove(self):
geom = gd.UniformGrid([101, 1], x0=[0, 0], dx=[0.01, 1])
data = np.array([i * np.linspace(1, 5, 1) for i in range(101)])
ug_data = gd.UniformGridData(geom, data)
ug_data.flat_dimensions_remove()
flat_geom = gd.UniformGrid([101], x0=[0], x1=[1])
self.assertEqual(
ug_data, gd.UniformGridData(flat_geom, np.linspace(0, 100, 101))
)
# Check invalidation of spline
self.assertTrue(ug_data.invalid_spline)
# Test from 3D to 2D
grid_data3d = gdu.sample_function_from_uniformgrid(
lambda x, y, z: x * (y + 2) * (z + 5),
gd.UniformGrid([10, 20, 1], x0=[0, 1, 0], dx=[1, 1, 1]),
)
grid_2d = gd.UniformGrid([10, 20], x0=[0, 1], dx=[1, 1])
expected_data2d = gdu.sample_function_from_uniformgrid(
lambda x, y: x * (y + 2) * (0 + 5), grid_2d
)
self.assertEqual(
grid_data3d.flat_dimensions_removed(), expected_data2d
)
def test__apply_unary(self):
data1 = np.array([i * np.linspace(1, 5, 51) for i in range(101)])
ug_data1 = gd.UniformGridData(self.geom, data1)
self.assertEqual(
np.sin(ug_data1), gd.UniformGridData(self.geom, np.sin(data1))
)
def test_hg(name):
ma_func = getattr(ma, name)
np_func = getattr(np.ma, name)
expected = [
gd.UniformGridData(self.ugd1.grid, np_func(self.ugd1.data)),
gd.UniformGridData(self.ugd2.grid, np_func(self.ugd2.data)),
]
self.assertTrue(ma_func(self.hg), expected)
def test_is_complex(self):
data = np.array([i * np.linspace(1, 5, 51) for i in range(101)])
ug_data = gd.UniformGridData(self.geom, data)
self.assertFalse(ug_data.is_complex())
ug_data_c = gd.UniformGridData(self.geom, 1j * data)
self.assertTrue(ug_data_c.is_complex())
def test_hg(name, *args):
np_func = getattr(np.ma, f"masked_{name}")
# We will edit in place t
t = self.hg.copy()
getattr(t, f"mask_{name}")(*args)
expected = [
gd.UniformGridData(self.ugd1.grid,
np_func(self.ugd1.data, *args)),
gd.UniformGridData(self.ugd2.grid,
np_func(self.ugd2.data, *args)),
]
self.assertTrue(t, expected)
def test_sample_function(self):
# Test not grid as input
with self.assertRaises(TypeError):
gdu.sample_function_from_uniformgrid(np.sin, 0)
# Test 1d
geom = gd.UniformGrid(100, x0=0, x1=2 * np.pi)
data = np.sin(np.linspace(0, 2 * np.pi, 100))
self.assertEqual(
gdu.sample_function(np.sin, 100, 0, 2 * np.pi),
gd.UniformGridData(geom, data),
)
# Test with additional arguments
geom_ref_level = gd.UniformGrid(100, x0=0, x1=2 * np.pi, ref_level=0)
self.assertEqual(
gdu.sample_function(np.sin, 100, 0, 2 * np.pi, ref_level=0),
gd.UniformGridData(geom_ref_level, data),
)
# Test 2d
geom2d = gd.UniformGrid([100, 200], x0=[0, 1], x1=[1, 2])
def square(x, y):
return x * y
# Test function takes too few arguments
with self.assertRaises(TypeError):
gdu.sample_function_from_uniformgrid(lambda x: x, geom2d)
# Test function takes too many arguments
with self.assertRaises(TypeError):
gdu.sample_function_from_uniformgrid(square, geom)
# Test other TypeError
with self.assertRaises(TypeError):
gdu.sample_function_from_uniformgrid(np.sin, geom2d)
data2d = np.vectorize(square)(*geom2d.coordinates(as_same_shape=True))
self.assertEqual(
gdu.sample_function(square, [100, 200], [0, 1], [1, 2]),
gd.UniformGridData(geom2d, data2d),
)
self.assertEqual(
gdu.sample_function_from_uniformgrid(square, geom2d),
gd.UniformGridData(geom2d, data2d),
)
def test_read_hdf5(self):
expected_grid = grid_data.UniformGrid(
[29, 29],
x0=[-14, -14],
x1=[14, 14],
ref_level=0,
num_ghost=[3, 3],
time=0,
iteration=0,
component=0,
)
with h5py.File(self.P_file, "r") as fil:
# Notice the .T
expected_data = fil["ILLINOISGRMHD::P it=0 tl=0 rl=0 c=0"][()].T
expected_grid_data = grid_data.UniformGridData(expected_grid,
expected_data)
self.assertEqual(
self.P._read_component_as_uniform_grid_data(self.P_file, 0, 0, 0),
expected_grid_data,
)
# Test that we can read a grid array
self.assertTrue(
isinstance(self.reader["int_em_T_rph"][0],
grid_data.HierarchicalGridData))
def test_read_hdf5(self):
expected_grid = grid_data.UniformGrid(
[29, 29],
x0=[-14, -14],
x1=[14, 14],
ref_level=0,
num_ghost=[3, 3],
time=0,
iteration=0,
component=0,
)
with h5py.File(self.P_file, "r") as fil:
# Notice the .T
expected_data = fil["ILLINOISGRMHD::P it=0 tl=0 rl=0 c=0"][()].T
expected_grid_data = grid_data.UniformGridData(
expected_grid, expected_data
)
self.assertEqual(
self.P._read_component_as_uniform_grid_data(self.P_file, 0, 0, 0),
expected_grid_data,
)
def test_ugd(name):
ma_func = getattr(ma, name)
np_func = getattr(np.ma, name)
self.assertTrue(
ma_func(self.ugd),
gd.UniformGridData(self.grid_2d, np_func(self.ugd.data)),
)
def test__apply_reduction(self):
data = np.array([i * np.linspace(1, 5, 51) for i in range(101)])
ug_data = gd.UniformGridData(self.geom, data)
self.assertAlmostEqual(ug_data.min(), 0)
self.assertAlmostEqual(ug_data.max(), 500)
def test_ugd(name, *args):
np_func = getattr(np.ma, f"masked_{name}")
# We will edit in place t
t = self.ugd.copy()
getattr(t, f"mask_{name}")(*args)
self.assertTrue(
t,
gd.UniformGridData(self.grid_2d, np_func(self.ugd.data,
*args)),
)
def test_init(self):
self.assertCountEqual(self.rho_star._iterations_to_times, {
0: 0,
1: 0.25,
2: 0.5
})
# TODO: This test doesn't test compressed files nor 3D files,
# and it tests only one component/refinement level/iteration
# The first component ends at line 880 of rho_star.xy.asc
# Let's check this first one
#
# We have to remove the blank lines from the count
expected_data = np.loadtxt(self.rho_star_file,
usecols=(12, ),
max_rows=875)
expected_grid = grid_data.UniformGrid(
[29, 29],
x0=[-14, -14],
x1=[14, 14],
ref_level=0,
num_ghost=[3, 3],
time=0,
iteration=0,
component=0,
)
expected_grid_data = grid_data.UniformGridData(
expected_grid,
expected_data.reshape((29, 29)).T)
self.assertEqual(
self.rho_star._read_component_as_uniform_grid_data(
self.rho_star_file, 0, 0, 0),
expected_grid_data,
)
# Test time_at_iteration
self.assertEqual(self.rho_star.time_at_iteration(2), 0.5)
# Iteration not available
with self.assertRaises(ValueError):
self.rho_star.time_at_iteration(20)
# Test file with wrong name
with self.assertRaises(RuntimeError):
self.rho_star._parse_file("/tmp/wrongname")
def sample_function_from_uniformgrid(function, grid):
"""Create a regular dataset by sampling a scalar function of the form
``f(x, y, z, ...)`` on a grid.
:param function: The function to sample.
:type function: A callable that takes as many arguments as the number
of dimensions (in shape).
:param grid: Grid over which to sample the function.
:type grid: :py:class:`~.UniformGrid`
:returns: Sampled data.
:rtype: :py:class:`~.UniformGridData`
"""
if not isinstance(grid, gd.UniformGrid):
raise TypeError("grid has to be a UniformGrid")
# The try except block checks that the function supplied has the correct
# signature for the grid provided. If you try to pass a function that takes
# too many of too few arguments, you will get a TypeError
try:
ret = gd.UniformGridData(
grid, np.vectorize(function)(*grid.coordinates(as_same_shape=True))
)
except TypeError as type_err:
# Too few arguments, type_err = missing N required positional arguments: ....
# Too many arguments, type_err = takes N positional arguments but M were given
ret = str(type_err)
# TODO (REFACTORING): This is fragile way to do error parsing
#
# We are reading the error message to understand what is going on. This
# can is not the best way to do this.
if isinstance(ret, str):
if "missing" in ret:
raise TypeError(
"Provided function takes too few arguments for requested grid"
)
if "takes" in ret:
raise TypeError(
"Provided function takes too many arguments for requested grid"
)
raise TypeError(ret)
return ret
def test_mean_integral_norm1_norm2(self):
data = np.array([i ** 2 * np.linspace(1, 5, 51) for i in range(101)])
ug_data = gd.UniformGridData(self.geom, data)
self.assertAlmostEqual(ug_data.integral(), np.sum(data) * self.geom.dv)
self.assertAlmostEqual(
ug_data.norm1(), np.sum(np.abs(data)) * self.geom.dv
)
self.assertAlmostEqual(
ug_data.norm2(),
np.sum(np.abs(data) ** 2 * self.geom.dv) ** 0.5,
)
self.assertAlmostEqual(
ug_data.norm_p(3),
np.sum(np.abs(data) ** 3 * self.geom.dv) ** (1 / 3),
)
self.assertAlmostEqual(ug_data.average(), np.mean(data))
def test_fourier_transform(self):
prod_data_complex = gdu.sample_function(
lambda x, y: (1 + 1j) * x * (y + 2), [11, 21], [0, 10], [10, 30]
)
dx = prod_data_complex.dx
fft_c = np.fft.fftshift(np.fft.fftn(prod_data_complex.data))
freqs_c = [
np.fft.fftshift(np.fft.fftfreq(11, d=dx[0])),
np.fft.fftshift(np.fft.fftfreq(21, d=dx[1])),
]
f_min_c = [freqs_c[0][0], freqs_c[1][0]]
delta_f_c = [
freqs_c[0][1] - freqs_c[0][0],
freqs_c[1][1] - freqs_c[1][0],
]
freq_grid_c = gd.UniformGrid(fft_c.shape, x0=f_min_c, dx=delta_f_c)
expected_c = gd.UniformGridData(freq_grid_c, fft_c)
self.assertEqual(expected_c, prod_data_complex.fourier_transform())
