Skip to content

Commit

Permalink
Merge pull request #31 from nvaytet/bug_in_plane_plots
Browse files Browse the repository at this point in the history 
Fix bug in plane plots
  • Loading branch information
@nvaytet
nvaytet authored Oct 4, 2021
2 parents d1b6670 + 83075b1 commit 446d687
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 63 deletions.
2 changes: 1 addition & 1 deletion setup.cfg
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[metadata]
name = osyris
version = 2.2.1
version = 2.2.2
author = Neil Vaytet
author_email = [email protected]
description = A package to visualize AMR data from the RAMSES code
Expand Down
129 changes: 67 additions & 62 deletions src/osyris/plot/plane.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,14 @@
from scipy.stats import binned_statistic_2d


def _add_scatter(datax, datay, to_scatter, origin, datadx, dir_vecs, dx, dy, ax):
def _add_scatter(to_scatter, origin, dir_vecs, dx, dy, ax):
xyz = to_scatter[0]["data"] - origin
viewport = max(dx.magnitude, dy.magnitude)
radius = None
if "s" in to_scatter[0]["params"]:
size = to_scatter[0]["params"]["s"]
if isinstance(size, Array) or isinstance(size, Quantity):
radius = size.to(datax.unit.units)
radius = size.to(dx.units)
to_scatter[0]["params"]["s"] = radius
if radius is None:
# Fudge factor to select sinks close to the plane
Expand Down Expand Up @@ -115,65 +115,64 @@ def plane(*layers,

# Distance to the plane
xyz = dataset["amr"]["xyz"] - origin
diagonal = dataset["amr"]["dx"] * np.sqrt(dataset.meta["ndim"]) * 0.5
dist1 = np.sum(xyz * dir_vecs[0], axis=1)
diagonal_close = dataset["amr"]["dx"] * 0.5 * np.sqrt(dataset.meta["ndim"])
dist_close = np.sum(xyz * dir_vecs[0], axis=1)
# Create an array of indices to allow further narrowing of the selection below
global_selection = np.arange(len(dataset["amr"]["dx"]))

# Select cells in contact with plane
select = np.ravel(np.where(np.abs(dist1) <= diagonal))
global_selection = global_selection[select]
global_indices = np.arange(len(dataset["amr"]["dx"]))
# Select cells in close to the plane, including factor of sqrt(ndim)
close_to_plane = np.ravel(np.where(np.abs(dist_close) <= diagonal_close))
indices_close_to_plane = global_indices[close_to_plane]

if len(select) == 0:
if len(indices_close_to_plane) == 0:
raise RuntimeError("No cells were selected to construct the plane. "
"The resulting figure would be empty.")

# Project coordinates onto the plane by taking dot product with axes vectors
coords = xyz[select]
datax = np.inner(coords, dir_vecs[1])
datay = np.inner(coords, dir_vecs[2])
datadx = 0.5 * dataset["amr"]["dx"][select]

# Get limits
limits = {
'xmin': np.amin(datax - datadx).values,
'xmax': np.amax(datax + datadx).values,
'ymin': np.amin(datay - datadx).values,
'ymax': np.amax(datay + datadx).values
}

# Define slice extent
if dx is None:
xmin = limits['xmin']
xmax = limits['xmax']
ymin = limits['ymin']
ymax = limits['ymax']
else:
xmin = None
if dx is not None:
xmin = -0.5 * dx.magnitude
xmax = xmin + dx.magnitude
ymin = -0.5 * dy.magnitude
ymax = ymin + dy.magnitude
# Limit selection further by using distance from center
dist2 = coords - datadx * np.sqrt(dataset.meta["ndim"])
select2 = np.ravel(
radial_distance = xyz[indices_close_to_plane] - 0.5 * dataset["amr"]["dx"][
indices_close_to_plane] * np.sqrt(dataset.meta["ndim"])
radial_selection = np.ravel(
np.where(
np.abs(dist2.norm.values) <= max(dx.magnitude, dy.magnitude) * 0.6 *
np.sqrt(2.0)))
coords = coords[select2]
datax = datax[select2]
datay = datay[select2]
datadx = datadx[select2]
global_selection = global_selection[select2]
np.abs(radial_distance.norm.values) <= max(dx.magnitude, dy.magnitude) *
0.6 * np.sqrt(2.0)))
indices_close_to_plane = indices_close_to_plane[radial_selection]

# Select cells touching the plane, excluding factor of sqrt(ndim)
dist_touching = np.sum(xyz[indices_close_to_plane] * dir_vecs[0], axis=1)
diagonal_touching = dataset["amr"]["dx"][indices_close_to_plane] * 0.5
touching_plane = np.ravel(np.where(np.abs(dist_touching) <= diagonal_touching))

# Project coordinates onto the plane by taking dot product with axes vectors
coords_close = xyz[indices_close_to_plane]
datax_close = np.inner(coords_close, dir_vecs[1])
datay_close = np.inner(coords_close, dir_vecs[2])
datadx_close = diagonal_touching

if xmin is None:
xmin = (datax_close - datadx_close).min().values
xmax = (datax_close + datadx_close).max().values
ymin = (datay_close - datadx_close).min().values
ymax = (datay_close + datadx_close).max().values

datax_touching = datax_close[touching_plane]
datay_touching = datay_close[touching_plane]

scalar_layer = []
to_binning = []
cell_variables = [] # contains the variables in cells close to the plane
to_binning = [] # a subset of cell_variables for only cells actually touching plane
for ind in range(len(to_process)):
if to_render[ind]["mode"] in ["vec", "stream"]:
if to_process[ind].ndim < 3:
uv = to_process[ind].array[global_selection]
uv = to_process[ind].array[indices_close_to_plane]
else:
uv = np.inner(to_process[ind].array.take(global_selection, axis=0),
dir_vecs[1:])
uv = np.inner(
to_process[ind].array.take(indices_close_to_plane, axis=0),
dir_vecs[1:])
w = None
if "color" in to_render[ind]["params"]:
if isinstance(to_render[ind]["params"]["color"], Array):
Expand All @@ -183,14 +182,22 @@ def plane(*layers,
if w is None:
w = np.linalg.norm(uv, axis=1)
else:
w = w.take(global_selection, axis=0)
to_binning.append(apply_mask(uv[:, 0]))
to_binning.append(apply_mask(uv[:, 1]))
to_binning.append(apply_mask(w))
w = w.take(indices_close_to_plane, axis=0)
vec_u = apply_mask(uv[:, 0])
vec_v = apply_mask(uv[:, 1])
vec_w = apply_mask(w)
cell_variables.append(vec_u)
cell_variables.append(vec_v)
cell_variables.append(vec_w)
scalar_layer.append(False)
to_binning.append(vec_u[touching_plane])
to_binning.append(vec_v[touching_plane])
to_binning.append(vec_w[touching_plane])
else:
to_binning.append(apply_mask(to_process[ind].norm.values[global_selection]))
var = apply_mask(to_process[ind].norm.values[indices_close_to_plane])
cell_variables.append(var)
scalar_layer.append(True)
to_binning.append(var[touching_plane])

# Buffer for counts
to_binning.append(np.ones_like(to_binning[0]))
Expand All @@ -206,8 +213,8 @@ def plane(*layers,
ycenters = to_bin_centers(yedges)

# First histogram the cell centers into the grid bins
binned, _, _, _ = binned_statistic_2d(x=apply_mask(datay.array),
y=apply_mask(datax.array),
binned, _, _, _ = binned_statistic_2d(x=apply_mask(datay_touching.array),
y=apply_mask(datax_touching.array),
values=to_binning,
statistic="mean",
bins=[yedges, xedges])
Expand All @@ -224,10 +231,11 @@ def plane(*layers,
ygrid.shape + (1, )) * dir_vecs[2]
# We only need to search in the cells above a certain size
large_cells = np.ravel(
np.where(datadx >= 0.25 * (min(xedges[1] - xedges[0], yedges[1] - yedges[0]))))
coords = coords[large_cells]
large_cells_dx = datadx.array[large_cells]
global_indices = np.arange(len(datadx))[large_cells]
np.where(datadx_close >= 0.25 *
(min(xedges[1] - xedges[0], yedges[1] - yedges[0]))))
coords = coords_close[large_cells]
large_cells_dx = datadx_close.array[large_cells]
large_cells_indices = np.arange(len(datadx_close))[large_cells]

# To keep memory usage down to a minimum, we process the image one column at a time
for i in range(indices.shape[-1]):
Expand Down Expand Up @@ -255,7 +263,7 @@ def plane(*layers,
inds = np.logical_and.reduce(
[np.abs(d) <= large_cells_dx[column] for d in distance_to_cell])
index_found = inds.max(axis=-1)
index_value = global_indices[column][inds.argmax(axis=-1)]
index_value = large_cells_indices[column][inds.argmax(axis=-1)]
indices[:, i][index_found] = index_value[index_found]
mask[:, i][np.logical_and(~index_found, condition[:, i])] = True
else:
Expand All @@ -266,13 +274,13 @@ def plane(*layers,
# Now we fill the arrays to be sent to the renderer, also constructing vectors
counter = 0
for ind in range(len(to_render)):
binned[counter][condition] = to_binning[counter][indices][condition]
binned[counter][condition] = cell_variables[counter][indices][condition]
if scalar_layer[ind]:
to_render[ind]["data"] = ma.masked_where(mask, binned[counter], copy=False)
counter += 1
else:
for j in range(counter + 1, counter + 3):
binned[j][condition] = to_binning[j][indices][condition]
binned[j][condition] = cell_variables[j][indices][condition]
to_render[ind]["data"] = ma.masked_where(mask_vec,
np.array([
binned[counter].T,
Expand All @@ -293,11 +301,8 @@ def plane(*layers,

# Add scatter layer
if len(to_scatter) > 0:
_add_scatter(datax=datax,
datay=datay,
to_scatter=to_scatter,
_add_scatter(to_scatter=to_scatter,
origin=origin,
datadx=datadx,
dir_vecs=dir_vecs,
dx=dx,
dy=dy,
Expand Down

0 comments on commit 446d687

Please sign in to comment.