Merge pull request #31 from nvaytet/bug_in_plane_plots

Fix bug in plane plots

Author: nvaytet

setup.cfg

@@ -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

src/osyris/plot/plane.py

@@ -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
@@ -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):
@@ -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]))
@@ -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])
@@ -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]):
@@ -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:
@@ -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,
@@ -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,