def test_init(self):
# Test incorrect arguments
# Not a list
with self.assertRaises(TypeError):
gd.HierarchicalGridData(0)
# Empty list
with self.assertRaises(ValueError):
gd.HierarchicalGridData([])
# Not a list of UniformGridData
with self.assertRaises(TypeError):
gd.HierarchicalGridData([0])
# Inconsistent number of dimensions
def product1(x):
return x
def product2(x, y):
return x * y
prod_data1 = gdu.sample_function(product1, [101], [0], [3])
prod_data2 = gdu.sample_function(product2, [101, 101], [0, 0], [3, 3])
with self.assertRaises(ValueError):
gd.HierarchicalGridData([prod_data1, prod_data2])
# Only one component
one = gd.HierarchicalGridData([prod_data1])
# Test content
self.assertDictEqual(
one.grid_data_dict, {-1: [prod_data1.ghost_zones_removed()]}
)
grid = gd.UniformGrid([101], x0=[0], x1=[3], ref_level=2)
# Two components at two different levels
prod_data1_level2 = gdu.sample_function_from_uniformgrid(
product1, grid
)
two = gd.HierarchicalGridData([prod_data1, prod_data1_level2])
self.assertDictEqual(
two.grid_data_dict,
{
-1: [prod_data1.ghost_zones_removed()],
2: [prod_data1_level2.ghost_zones_removed()],
},
)
# Test a good grid
hg_many_components = gd.HierarchicalGridData(self.grid_data)
self.assertEqual(
hg_many_components.grid_data_dict[0], [self.expected_data]
)
# Test a grid with two separate components
hg3 = gd.HierarchicalGridData(self.grid_data_two_comp)
self.assertEqual(hg3.grid_data_dict[0], self.grid_data_two_comp)
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_histogram(self):
# There should be no reason why the histogram behaves differently for
# different dimensions, so let's test it with 1d
sin_data = gdu.sample_function(np.sin, 1000, 0, 2 * np.pi)
sin_data_complex = sin_data + 1j * sin_data
# Test error weights
with self.assertRaises(TypeError):
sin_data.histogram(weights=1)
# Test error complex
with self.assertRaises(ValueError):
sin_data_complex.histogram()
hist = sin_data.histogram()
expected_hist = np.histogram(sin_data.data, range=(-1, 1), bins=400)
self.assertTrue(np.allclose(expected_hist[0], hist[0]))
self.assertTrue(np.allclose(expected_hist[1], hist[1]))
# Test with weights
weights = sin_data.copy()
weights **= 2
hist = sin_data.histogram(weights)
expected_hist = np.histogram(
sin_data.data, range=(-1, 1), bins=400, weights=weights.data
)
self.assertTrue(np.allclose(expected_hist[0], hist[0]))
self.assertTrue(np.allclose(expected_hist[1], hist[1]))
def test__getitem__(self):
def square(x, y):
return x * (y + 2)
# These are just integers
prod_data = gdu.sample_function(square, [11, 21], [0, 10], [10, 30])
self.assertAlmostEqual(prod_data[2, 2], 2 * 14)
def test_copy(self):
sin_data = gdu.sample_function(np.sin, 1000, 0, 2 * np.pi)
sin_data2 = sin_data.copy()
self.assertEqual(sin_data, sin_data2)
self.assertIsNot(sin_data.data, sin_data2.data)
self.assertIsNot(sin_data.grid, sin_data2.grid)
def test_plot_components_boundaries(self):
# Test invalid data
with self.assertRaises(TypeError):
viz.plot_components_boundaries("bob")
# Not 2D data
ugd1D = gdu.sample_function(lambda x: x, [10], [0], [2])
hg1D = gd.HierarchicalGridData([ugd1D])
# Test invalid data
with self.assertRaises(ValueError):
viz.plot_components_boundaries(hg1D)
# Now with valid data
ugd = gdu.sample_function(lambda x, y: x + y, [10, 20], [0, 2], [2, 5])
hg = gd.HierarchicalGridData([ugd])
viz.plot_components_boundaries(hg)
def test_plot_grid(self):
ugd = gdu.sample_function(
lambda x, y: x + y, [100, 20], [0, 1], [2, 5]
)
# Unknown plot type
with self.assertRaises(ValueError):
viz._plot_grid(ugd, plot_type="bubu")
self.assertTrue(
isinstance(
viz.plot_contourf(
ugd, xlabel="x", ylabel="y", colorbar=True, label="test"
),
matplotlib.contour.QuadContourSet,
)
)
# Here we are not providing the coordinates
with self.assertRaises(ValueError):
viz.plot_contourf(ugd.data_xyz, logscale=True)
# Plot_color
self.assertTrue(
isinstance(
viz.plot_color(
ugd, xlabel="x", ylabel="y", colorbar=True, label="test"
),
matplotlib.image.AxesImage,
)
)
# Plot color, grid = None
self.assertTrue(
isinstance(
viz.plot_color(
ugd.data_xyz,
),
matplotlib.image.AxesImage,
)
)
# Plot contour, no levels
with self.assertRaises(ValueError):
viz._plot_grid(ugd, plot_type="contour")
self.assertTrue(
isinstance(
viz.plot_contour(
ugd, xlabel="x", ylabel="y", colorbar=True, label="test"
),
matplotlib.contour.QuadContourSet,
)
)
def test_plot_colorbar(self):
ugd = gdu.sample_function(lambda x, y: x + y, [100, 20], [0, 1],
[2, 5])
cf = viz.plot_contourf(ugd, xlabel="x", ylabel="y", colorbar=False)
self.assertTrue(
isinstance(
viz.plot_colorbar(cf, label="test"),
matplotlib.colorbar.Colorbar,
))
def test_dx_change(self):
def product_complex(x, y):
return (1 + 1j) * x * (y + 2)
prod_data_complex = gdu.sample_function(
product_complex, [301, 401], [0, 1], [3, 4]
)
prod_data_complex_copy = prod_data_complex.copy()
# Test invalid new dx
# Not a list
with self.assertRaises(TypeError):
prod_data_complex.dx_change(0)
# Not a with the correct dimensions
with self.assertRaises(ValueError):
prod_data_complex.dx_change([0])
# Not a integer multiple/factor
with self.assertRaises(ValueError):
prod_data_complex.dx_change(prod_data_complex.dx * np.pi)
# Same dx
self.assertEqual(
prod_data_complex.dx_changed(prod_data_complex.dx),
prod_data_complex,
)
# Half dx
prod_data_complex.dx_change(prod_data_complex.dx / 2)
self.assertCountEqual(
prod_data_complex.dx, prod_data_complex_copy.dx / 2
)
self.assertCountEqual(
prod_data_complex.shape, prod_data_complex_copy.shape * 2 - 1
)
# The data part should be tested with testing resample
# Twice of the dx, which will bring us back to same dx,
# so, same object we started with
prod_data_complex.dx_change(prod_data_complex.dx * 2)
self.assertEqual(prod_data_complex, prod_data_complex_copy)
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())
def test_save_load(self):
grid_file = "test_save_grid.dat"
grid_file_bz = "test_save_grid.dat.bz2"
grid_file_gz = "test_save_grid.dat.gz"
def square(x, y):
return x * y
grid_data = gdu.sample_function(square, [100, 200], [0, 1], [1, 2])
# Test save uncompressed
grid_data.save(grid_file)
# Load it
loaded = gdu.load_UniformGridData(grid_file)
self.assertEqual(loaded, grid_data)
# Clean up file
os.remove(grid_file)
# Test compressed
grid_data.save(grid_file_bz)
loaded_bz = gdu.load_UniformGridData(grid_file_bz)
self.assertEqual(loaded_bz, grid_data)
# Clean up file
os.remove(grid_file_bz)
grid_data.save(grid_file_gz)
loaded_gz = gdu.load_UniformGridData(grid_file_gz)
self.assertEqual(loaded_gz, grid_data)
# Clean up file
os.remove(grid_file_gz)
def test_percentiles(self):
# There should be no reason why the histogram behaves differently for
# different dimensions, so let's test it with 1d
lin_data = gdu.sample_function(lambda x: 1.0 * x, 1000, 0, 2 * np.pi)
# Scalar input
self.assertAlmostEqual(lin_data.percentiles(0.5), np.pi)
# Vector input
self.assertTrue(
np.allclose(
lin_data.percentiles([0.25, 0.5]), np.array([np.pi / 2, np.pi])
)
)
# Not normalized
self.assertTrue(
np.allclose(
lin_data.percentiles([250, 500], relative=False),
np.array([np.pi / 2, np.pi]),
)
)
def test_preprocess_plot_functions(self):
# We test preprocess_plot with a function that returns the argument, so
# we can check that they are what we expect
def func(data, **kwargs):
return data, kwargs
dec_func = viz.preprocess_plot(func)
# Default
self.assertIs(dec_func("")[1]["axis"], plt.gca())
# Passing axis
self.assertIs(dec_func("", axis=self.ax2)[1]["axis"], self.ax2)
# Default
self.assertIs(dec_func("")[1]["figure"], plt.gcf())
# Passing figure
self.assertIs(dec_func("", figure=self.fig)[1]["figure"], self.fig)
dec_func_grid = viz.preprocess_plot_grid(func)
# Check data not provided
with self.assertRaises(TypeError):
dec_func_grid()
# Data numpy array, but not of dimension 2
with self.assertRaises(ValueError):
dec_func_grid(data=np.linspace(0, 1, 2))
# Data numpy array of dimension 2
self.assertTrue(
np.array_equal(
dec_func_grid(data=np.array([[1, 2], [3, 4]]))[0],
np.array([[1, 2], [3, 4]]),
)
)
# Check with UniformGridData
# 2D
ugd = gdu.sample_function(lambda x, y: x + y, [10, 20], [0, 1], [2, 5])
# Passing coordinates
with self.assertWarns(Warning):
ret = dec_func_grid(
data=ugd, coordinates=ugd.coordinates_from_grid()
)
self.assertTrue(np.array_equal(ret[0], ugd.data_xyz))
self.assertIs(ret[1]["coordinates"], ugd.coordinates_from_grid())
# Passing x0 and x1 but not shape
with self.assertRaises(TypeError):
dec_func_grid(data=ugd, x0=ugd.x0, x1=ugd.x1)
# Now HierachicalGridData or resampling UniformGridData
hg = gd.HierarchicalGridData([ugd])
# Shape not provided
with self.assertRaises(TypeError):
dec_func_grid(data=hg)
# Valid (passing x0 and x1)
ret2 = dec_func_grid(data=hg, shape=ugd.shape, x0=ugd.x0, x1=ugd.x1)
# Check coordinates (which checks the resampling)
self.assertTrue(
np.allclose(
ret2[1]["coordinates"][0], ugd.coordinates_from_grid()[0]
)
)
# Not passing x0 and x1
ret2b = dec_func_grid(data=hg, shape=ugd.shape)
self.assertTrue(
np.allclose(
ret2b[1]["coordinates"][0], ugd.coordinates_from_grid()[0]
)
)
# We create an empty OneGridFunctionASCII and populate it with ugd
cactus_ascii = cgf.OneGridFunctionASCII([], "var_name", (0, 0))
cactus_ascii.allfiles = ["file"]
cactus_ascii.alldata = {"file": {0: {0: {0: ugd}}}}
# Iteration not provided
with self.assertRaises(TypeError):
dec_func_grid(data=cactus_ascii)
# Check coordinates (which checks the resampling)
ret3 = dec_func_grid(
data=cactus_ascii, iteration=0, shape=ugd.shape, xlabel="x"
)
self.assertTrue(
np.allclose(
ret3[1]["coordinates"][0], ugd.coordinates_from_grid()[0]
)
)
# Test xlabel and ylabel
self.assertEqual(ret3[1]["xlabel"], "x")
# Test with resample=True
_ = dec_func_grid(
data=cactus_ascii,
iteration=0,
shape=ugd.shape,
xlabel="x",
ylabel="y",
resample=True,
)
# Test with masked data
# HierarchicalGrid
hg_m = km.arcsin(gd.HierarchicalGridData([ugd]))
with self.assertWarns(Warning):
dec_func_grid(data=hg_m, shape=[10, 10])
ugd_m = km.arcsin(ugd)
with self.assertWarns(Warning):
dec_func_grid(data=ugd_m, shape=[10, 10], x0=[1, 3])
def test_resampled(self):
def product(x, y):
return x * (y + 2)
def product_complex(x, y):
return (1 + 1j) * x * (y + 2)
prod_data = gdu.sample_function(product, [101, 201], [0, 1], [3, 4])
prod_data_complex = gdu.sample_function(
product_complex, [3001, 2801], [0, 1], [3, 4]
)
# Check error
with self.assertRaises(TypeError):
prod_data.resampled(2)
# Check same grid
self.assertEqual(prod_data.resampled(prod_data.grid), prod_data)
new_grid = gd.UniformGrid([51, 101], x0=[1, 2], x1=[2, 3])
resampled = prod_data_complex.resampled(new_grid)
exp_resampled = gdu.sample_function_from_uniformgrid(
product_complex, new_grid
)
self.assertEqual(resampled.grid, new_grid)
self.assertTrue(np.allclose(resampled.data, exp_resampled.data))
# Check that the method of the spline is linear
self.assertEqual(prod_data_complex.spline_imag.method, "linear")
# Test using nearest interpolation
resampled_nearest = prod_data_complex.resampled(
new_grid, piecewise_constant=True
)
self.assertTrue(
np.allclose(resampled_nearest.data, exp_resampled.data, atol=1e-3)
)
# Check that the method of the spline hasn't linear
self.assertEqual(prod_data_complex.spline_imag.method, "linear")
# Check single number
self.assertAlmostEqual(resampled_nearest((2, 2.5)), 9 * (1 + 1j))
# Check with one point
new_grid2 = gd.UniformGrid([11, 1], x0=[1, 2], dx=[0.1, 1])
resampled2 = prod_data_complex.resampled(new_grid2)
prod_data_one_point = gdu.sample_function_from_uniformgrid(
product_complex, new_grid2
)
self.assertEqual(resampled2, prod_data_one_point)
# Resample from 3d to 2d
grid_data3d = gdu.sample_function_from_uniformgrid(
lambda x, y, z: x * (y + 2) * (z + 5),
gd.UniformGrid([10, 20, 11], x0=[0, 1, 0], dx=[1, 2, 0.1]),
)
grid_2d = gd.UniformGrid([10, 20, 1], [0, 1, 0], dx=[1, 2, 0.1])
expected_data2d = gdu.sample_function_from_uniformgrid(
lambda x, y, z: x * (y + 2) * (z + 5), grid_2d
)
self.assertEqual(grid_data3d.resampled(grid_2d), expected_data2d)
def test_splines(self):
# Let's start with 1d.
sin_data = gdu.sample_function(np.sin, 12000, 0, 2 * np.pi)
sin_data_complex = sin_data + 1j * sin_data
# Test unknown ext
with self.assertRaises(ValueError):
sin_data.evaluate_with_spline(1, ext=3)
# Test k!=0!=1
with self.assertRaises(ValueError):
sin_data._make_spline(k=3)
self.assertAlmostEqual(
sin_data_complex.evaluate_with_spline([np.pi / 3]),
(1 + 1j) * np.sin(np.pi / 3),
)
# Test with point in cell but outside boundary, in
# We change the boundary values to be different from 0
sin_data_complex_plus_one = sin_data_complex + 1 + 1j
dx = sin_data_complex.dx[0]
# At the boundary, we do a constant extrapolation, so the value should
# be the boundary value
self.assertAlmostEqual(
sin_data_complex_plus_one.evaluate_with_spline([0 - 0.25 * dx]),
(1 + 1j),
)
self.assertAlmostEqual(
sin_data_complex_plus_one.evaluate_with_spline(
[2 * np.pi + 0.25 * dx]
),
(1 + 1j),
)
# Test __call__
self.assertAlmostEqual(
sin_data_complex([np.pi / 3]),
(1 + 1j) * np.sin(np.pi / 3),
)
# Test on a point of the grid
point = [sin_data.grid.coordinates_1d[0][2]]
self.assertAlmostEqual(
sin_data_complex(point),
(1 + 1j) * np.sin(point[0]),
)
# Test on a point outside the grid with the lookup table
with self.assertRaises(ValueError):
sin_data_complex._nearest_neighbor_interpolation(
np.array([1000]),
ext=2,
),
# Test on a point outside the grid with the lookup table and
# ext = 1
self.assertEqual(
sin_data_complex.evaluate_with_spline(
[1000], ext=1, piecewise_constant=True
),
0,
)
# Vector input
self.assertTrue(
np.allclose(
sin_data_complex.evaluate_with_spline(
[[np.pi / 3], [np.pi / 4]]
),
np.array(
[
(1 + 1j) * np.sin(np.pi / 3),
(1 + 1j) * np.sin(np.pi / 4),
]
),
)
)
# Vector input in, vector input out
self.assertEqual(sin_data_complex([[1]]).shape, (1,))
# Now 2d
def product(x, y):
return x * (y + 2)
prod_data = gdu.sample_function(product, [101, 101], [0, 0], [3, 3])
prod_data_complex = (1 + 1j) * prod_data
self.assertAlmostEqual(
prod_data_complex.evaluate_with_spline((2, 3)),
(1 + 1j) * 10,
)
# Vector input
self.assertTrue(
np.allclose(
prod_data_complex.evaluate_with_spline([(1, 0), (2, 3)]),
np.array([(1 + 1j) * 2, (1 + 1j) * 10]),
)
)
self.assertTrue(
np.allclose(
prod_data_complex.evaluate_with_spline(
[[(1, 0), (2, 3)], [(3, 1), (0, 0)]]
),
np.array([[(1 + 1j) * 2, (1 + 1j) * 10], [(1 + 1j) * 9, 0]]),
)
)
# Real data
self.assertAlmostEqual(
prod_data.evaluate_with_spline((2, 3)),
10,
)
# Extrapolate outside
self.assertAlmostEqual(
prod_data.evaluate_with_spline((20, 20), ext=1), 0
)
self.assertAlmostEqual(
prod_data_complex.evaluate_with_spline((20, 20), ext=1), 0
)
self.assertTrue(prod_data_complex.spline_real.bounds_error)
self.assertTrue(prod_data_complex.spline_imag.bounds_error)
# Test on a UniformGrid
sin_data = gdu.sample_function(np.sin, 12000, 0, 2 * np.pi)
linspace = gd.UniformGrid(101, x0=0, x1=3)
output = sin_data(linspace)
self.assertTrue(
np.allclose(output.data, np.sin(linspace.coordinates()))
)
# Incompatible dimensions
with self.assertRaises(ValueError):
sin_data(gd.UniformGrid([101, 201], x0=[0, 1], x1=[3, 4]))
# Test with grid that has a flat dimension
prod_data_flat = gdu.sample_function_from_uniformgrid(
product, gd.UniformGrid([101, 1], x0=[0, 0], dx=[1, 3])
)
# y = 1 is in the flat cell, where y = 0, and here we are using nearest
# interpolation
self.assertAlmostEqual(prod_data_flat((1, 1)), 2)
# Vector
self.assertCountEqual(prod_data_flat([(1, 1), (2, 1)]), [2, 4])