def test_load(self):
# Test loading a file that doesn't exist
pf = PltFile("file that does not exist")
self.assertEqual(pf.defined, False)
# Test loading a file that does exist
pf = PltFile(self.DATA_FILE, load_data=True)
self.assertEqual(pf.num_levels, 3)
self.assertEqual(pf.plot_prefix, "plt")
self.assertEqual(pf.frame, 100)
# Test pretty printing object
logging.info(pf)
# Test loading after data already loaded (should do nothing)
pf.load_data()
# Check data is correct
for pf in [PltFile(self.CHK_DATA_FILE), PltFile(self.DATA_FILE)]:
pf.load_data()
max_enthalpy = float(pf.get_level_data("Enthalpy", level=0).max())
self.assertAlmostEqual(max_enthalpy, 6.302214, 6)
max_enthalpy = float(pf.get_level_data("Enthalpy", level=2).max())
self.assertAlmostEqual(max_enthalpy, 6.307355035152367, 6)
# Test removing data
pf.unload_data()
self.assertEqual(pf.data_loaded, False)
pf.load_data(zero_x=True)
assert pf.get_level_data("Enthalpy").coords["x"][0] == 0
assert np.array_equal(pf.get_levels(), [0, 1, 2])
outline = pf.level_outlines[2]
assert np.array_equal(outline.total_bounds, [0, 0.375, 1, 1])
assert outline.area[0] == 0.53125
assert pf.get_norm("Enthalpy").vmax == pytest.approx(6.30735504, 5)
assert pf.get_norm("Enthalpy").vmin == pytest.approx(1.49226642, 5)
x, y = pf.get_mesh_grid_for_level(1)
assert len(x) == 32
assert len(y) == 24
pf_no_name = PltFile(self.DATA_FILE_NO_FRAME)
self.assertEqual(pf_no_name.frame, -1)
pf_no_inputs = PltFile(self.DATA_FILE, inputs_file="does not exist")
self.assertIsNone(pf_no_inputs.inputs)
assert str(pf_no_inputs) == "<PltFile object for %s>" % self.DATA_FILE
def test_static_methods(self):
pf = PltFile(self.DATA_FILE, load_data=True)
x, y = pf.get_mesh_grid()
self.assertEqual(len(x), 16)
x, y = pf.get_mesh_grid(extend_grid=False)
self.assertEqual(len(x), 15)
porosity = pf.get_level_data("Porosity")
x, y = PltFile.get_mesh_grid_n(porosity)
self.assertEqual(len(x), 16)
x, y = PltFile.get_mesh_grid_xarray(porosity)
dx = 0.0625
assert len(x) == 16
assert len(y) == 16
assert x[0] == dx / 2
assert y[0] == dx / 2
assert x[-1] == 1 - dx / 2
assert y[-1] == 1 - dx / 2
x, y = PltFile.get_mesh_grid_xarray(porosity, grow=True)
assert len(x) == 17
assert x[0] == 0.0
assert y[0] == 0.0
assert x[-1] == 1.0
assert y[-1] == 1.0
def test_get_level_data(self):
pf = PltFile(self.DATA_FILE)
assert pf.num_levels == 3
assert pf.get_level_data("does not exist") is None
valid_only = pf.get_level_data("Porosity", valid_only=True)
all_data = pf.get_level_data("Porosity", valid_only=False)
assert len(valid_only.coords["x"]) == 16
assert len(valid_only.coords["y"]) == 16
assert np.isnan(valid_only[5, 5])
assert not all_data.equals(valid_only)
valid_only = pf.get_level_data("Porosity", level=1, valid_only=True)
all_data = pf.get_level_data("Porosity", level=1, valid_only=False)
assert len(valid_only.coords["x"]) == 32
assert len(valid_only.coords["y"]) == 24
assert np.isnan(valid_only[20, 20])
assert not all_data.equals(valid_only)
valid_only = pf.get_level_data("Porosity", level=2, valid_only=True)
all_data = pf.get_level_data("Porosity", level=2, valid_only=False)
assert len(valid_only.coords["x"]) == 64
assert len(valid_only.coords["y"]) == 40
assert valid_only.equals(all_data)
def test_child_methods(self):
pf = PltFile(self.DATA_FILE)
child = ChildPltFileForTesting(self.DATA_FILE)
pf.reflect = True
child.reflect = True
streamfunction_unreflected = pf.get_level_data("streamfunction")
streamfunction_reflected = child.get_level_data("streamfunction")
assert float(streamfunction_reflected[4, 4]) == -float(
streamfunction_unreflected[4, 4]
)
def test_plotting(self):
pf = PltFile(self.DATA_FILE)
pf.load_data(zero_x=True)
porosity = pf.get_level_data("Porosity")
fig = plt.figure()
ax = fig.gca()
ax.pcolormesh(porosity.x, porosity.y, porosity)
# Can only execute these tests if latex is installed
if is_installed("latex"):
pf.plot_outlines(ax)
pf.plot_field("Porosity")
def test_scale_slice_transform(self):
pf = PltFile(self.DATA_FILE)
indices = dict(y=slice(0, 3), x=slice(3, 5))
pf.set_scale_slice_transform(indices, True)
assert pf.indices == indices
assert pf.reflect == True
porosity = pf.get_level_data("Porosity")
assert len(porosity.coords["x"]) == 2
assert len(porosity.coords["y"]) == 3
assert float(porosity[0, 0]) == 1.0
pf.reset_scale_slice_transform()
assert pf.indices is None
assert pf.reflect is None
import matplotlib.pyplot as plt
from chombopy.plotting import PltFile, setup_mpl_latex
import matplotlib.cm as cm
pf = PltFile("../tests/data/plt000100.2d.hdf5")
setup_mpl_latex(14)
fig = plt.figure()
ax = plt.gca()
cmap = "viridis"
field = "Temperature"
# Get data for the temperature variable on level 2
for level in pf.get_levels():
temperature = pf.get_level_data(field, level)
# temperature is an xarray.DataSet object, which can be plotted using matplotlib
x, y = pf.get_mesh_grid_for_level(level=level, grow=True)
ax.pcolormesh(x, y, temperature, cmap=cmap)
# Or you can do some analysis using the xarray/numpy functionality
print(temperature.mean())
pf.plot_outlines(ax)
cbar = fig.colorbar(cm.ScalarMappable(norm=pf.get_norm(field), cmap=cmap),
ax=ax)
cbar.ax.set_ylabel(field)
plt.savefig("../docs/images/plt000100.jpg")
plt.show()