def test_3d(self):
data_files = [PltFile(self.DATA_FILE_3D), PltFile(self.DATA_FILE_3D_CHK)]
for pf in data_files:
porosity = pf.get_level_data("Enthalpy")
coords = porosity.coords
x = coords["x"]
y = coords["y"]
z = coords["y"]
self.assertEqual(pf.space_dim, 3)
self.assertEqual(len(z), 32)
self.assertEqual(len(x), 32)
self.assertEqual(len(y), 32)
x, y, z = pf.get_mesh_grid(extend_grid=True)
assert len(x) == 33
assert x[0, 0, 0] == 0.0
assert x[-1, -1, -1] == 1.0
x, y, z = pf.get_mesh_grid(extend_grid=False)
dx = 1.0 / 32
assert len(x) == 32
assert x[0, 0, 0] == dx / 2
assert x[-1, -1, -1] == 1.0 - dx / 2
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_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_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_get_box_comp_data(self):
pf = PltFile(self.DATA_FILE)
data = pf.get_box_comp_data(
np.array([]), 0, (0, 0), "A", (0, 0), {"i": [], "j": []}
)
assert len(data.coords["i"]) == 0
assert len(data) == 0
# Test when n_cells_dir is inconsistent with coords
data = pf.get_box_comp_data(
np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]),
0,
(0, 6),
"A",
n_cells_dir=(3, 2),
coords={"i": [0, 1], "j": [0, 1, 2]},
)
assert len(data.coords["i"]) == 2
assert data[1, 2] == 0.6
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()
def open_dataset(filename, inputs_file="inputs"):
"""
Get dataset with each variable and coordinates labelled differently on each level
Parameters
----------
filename :
inputs_file :
Returns
-------
"""
pf = PltFile(filename, inputs_file=inputs_file)
# Save as:
# comp_i, with coords x_i, y_i for each level i
ds_levs = []
for level in pf.get_levels():
ds_lev = pf.ds_levels[level]
if pf.space_dim == 3:
ds_lev = ds_lev.transpose("z", "y", "x")
elif pf.space_dim == 2:
ds_lev = ds_lev.transpose("y", "x")
# Rename appropriately
keys = list(ds_lev.keys())
coords = list(ds_lev.coords)
keys_lev = [k + get_lev_str(level) for k in keys]
coords_lev = [c + get_lev_str(level) for c in coords]
ds_lev = ds_lev.rename(dict(zip(keys, keys_lev)))
ds_lev = ds_lev.rename(dict(zip(coords, coords_lev)))
ds_lev = ds_lev.drop("level" + get_lev_str(level))
ds_levs.append(ds_lev)
ds = xr.merge(ds_levs)
ds.attrs["num_levels"] = pf.num_levels
ds.attrs["time"] = pf.time
ds.attrs["iteration"] = pf.iteration
ds.attrs["max_level"] = pf.max_level
ds.attrs["num_components"] = pf.num_comps
ds.attrs["SpaceDim"] = pf.space_dim
ds.attrs["components"] = pf.comp_names
ds.attrs["dims"] = ("x", "y", "z")[0:pf.space_dim]
ds.attrs["dx"] = []
ds.attrs["dt"] = []
ds.attrs["ref_ratio"] = []
ds.attrs["boxes"] = []
for level in pf.get_levels():
ds.attrs["boxes"].append(pf.levels[level][PltFile.BOXES])
ds.attrs["dx"].append(pf.levels[level][PltFile.DX])
ds.attrs["dt"].append(pf.levels[level][PltFile.DT])
ds.attrs["ref_ratio"].append(pf.levels[level][PltFile.REF_RATIO])
return ds
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_get_data(self):
pf = PltFile(self.DATA_FILE)
assert float(pf.get_data("Porosity")[13, 14]) == 1.0
pf = PltFile("no data")
assert pf.get_data("Porosity") is None
def convert_chombo(input_file, output_path):
data = PltFile(input_file)
data.python_index_ordering = False
if not data.space_dim == 3:
print("Error - input file is not 3D")
return
output_path = Path(output_path)
# fields_to_include = ["Porosity", "Bulk concentration"]
fields_to_include = [
"Temperature", "Porosity", "Pressure", "zAdvection velocity"
]
print(data.comp_names)
amr_dataset = get_amr_dataset(data)
point_data = {f: np.array(amr_dataset[f]) for f in fields_to_include}
x, y = data.get_mesh_grid_for_level(data.get_levels()[-1])
dx = x[1] - x[0]
num_cells = 1.0 / dx
create_http_dataset(output_path,
"fields",
cell_data=None,
point_data=point_data,
dx=dx)
vtk_boxes = []
colours = [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]]
dirs = ("x", "y", "z")
# print(data.domain_size)
# top_corner = np.array(data.domain_size[data.space_dim:]) - dx / 2
# Don't do level 0 box
for level in data.get_levels():
chombo_boxes = data.levels[level][PltFile.BOXES]
colour = colours[level]
x, y = data.get_mesh_grid_for_level(level)
dx = x[1] - x[0]
top_corner = np.array(data.domain_size[data.space_dim:]) - dx / 2
print(top_corner)
print(chombo_boxes)
for box in chombo_boxes:
vtk_box = {}
for i in range(data.space_dim):
# Note this length goes from the center of one outside to the center of the next
# If it was face to face, it would be one cell bigger
vtk_box[dirs[i] +
"Length"] = dx * (box[i + data.space_dim] + 1 - box[i])
center = np.array([
dx * (box[i + data.space_dim] + 1 + box[i]) / 2
for i in range(data.space_dim)
])
# center[2] = top_corner[2] - center[2]
vtk_box["center"] = list(center)
vtk_box["colour"] = colour
vtk_boxes.append(vtk_box)
print(vtk_boxes)
metadata = {"boxes": vtk_boxes, "time": data.time}
with open(output_path.joinpath("metadata.json"), 'w') as outfile:
json.dump(metadata, outfile)
def get_plt_file(self, filename):
return PltFile(os.path.join(self.test_folder, filename), load_data=True)