Merge pull request #176 from yashrsharma44/extra_coords_plot

Support plotting to use pixel_edges instead of pixel_values

changelog/176.bugfix.rst

@@ -0,0 +1 @@
+Fixed the bug of using `pixel_edges` instead of `pixel_values` in plotting

ndcube/mixins/plotting.py

@@ -12,6 +12,7 @@
     from sunpy.visualization.imageanimator import ImageAnimator, ImageAnimatorWCS, LineAnimator
 
 from ndcube import utils
+from ndcube.utils.cube import _get_extra_coord_edges
 from ndcube.mixins import sequence_plotting
 
 __all__ = ['NDCubePlotMixin']
@@ -63,7 +64,7 @@ def plot(self, axes=None, plot_axis_indices=None, axes_coordinates=None,
 
         """
         # If old API is used, convert to new API.
-        plot_axis_indices, axes_coordiantes, axes_units, data_unit, kwargs = _support_101_plot_API(
+        plot_axis_indices, axes_coordinates, axes_units, data_unit, kwargs = _support_101_plot_API(
             plot_axis_indices, axes_coordinates, axes_units, data_unit, kwargs)
         # Check kwargs are in consistent formats and set default values if not done so by user.
         naxis = len(self.dimensions)
@@ -350,10 +351,10 @@ def _animate_cube_1D(self, plot_axis_index=-1, axes_coordinates=None,
         # Get real world axis values along axis to be plotted and enter into axes_ranges kwarg.
         if axes_coordinates[plot_axis_index] is None:
             xname = self.world_axis_physical_types[plot_axis_index]
-            xdata = self.axis_world_coords(plot_axis_index)
+            xdata = self.axis_world_coords(plot_axis_index, edges=True)
         elif isinstance(axes_coordinates[plot_axis_index], str):
             xname = axes_coordinates[plot_axis_index]
-            xdata = self.extra_coords[xname]["value"]
+            xdata = _get_extra_coord_edges(self.extra_coords[xname]["value"])
         else:
             xname = ""
             xdata = axes_coordinates[plot_axis_index]
@@ -370,6 +371,18 @@ def _animate_cube_1D(self, plot_axis_index=-1, axes_coordinates=None,
             else:
                 unit_x_axis = None
         # Put xdata back into axes_coordinates as a masked array.
+
+        if len(xdata.shape) > 1:
+
+            # Since LineAnimator currently only accepts 1-D arrays for the x-axis, collapse xdata
+            # to single dimension by taking mean along non-plotting axes.
+            index = utils.wcs.get_dependent_data_axes(self.wcs, plot_axis_index, self.missing_axes)
+            reduce_axis = np.where(index == np.array(plot_axis_index))[0]
+
+            index = np.delete(index, reduce_axis)
+            # Reduce the data by taking mean
+            xdata = np.mean(xdata, axis=tuple(index))
+
         axes_coordinates[plot_axis_index] = xdata
         # Set default x label
         default_xlabel = "{0} [{1}]".format(xname, unit_x_axis)

ndcube/mixins/sequence_plotting.py

@@ -10,6 +10,7 @@
     from sunpy.visualization.imageanimator import ImageAnimatorWCS, LineAnimator
 
 from ndcube import utils
+from ndcube.utils.cube import _get_extra_coord_edges
 
 __all__ = ['NDCubeSequencePlotMixin']
 
@@ -956,7 +957,7 @@ def __init__(self, seq, plot_axis_index=None, axis_ranges=None, unit_x_axis=None
         if axis_ranges is None:
             axis_ranges = [None] * len(seq.dimensions)
             if plot_axis_index == 0:
-                axis_ranges[plot_axis_index] = np.arange(len(seq.data))
+                axis_ranges[plot_axis_index] = _get_extra_coord_edges(np.arange(len(seq.data)), axis=plot_axis_index)
             else:
                 cube_plot_axis_index = plot_axis_index - 1
                 # Define unit of x-axis if not supplied by user.
@@ -966,8 +967,8 @@ def __init__(self, seq, plot_axis_index=None, axis_ranges=None, unit_x_axis=None
                     unit_x_axis = np.asarray(seq[0].wcs.wcs.cunit)[wcs_plot_axis_index]
                 # Get x-axis values from each cube and combine into a single
                 # array for axis_ranges kwargs.
-                x_axis_coords = _get_non_common_axis_x_axis_coords(seq.data, cube_plot_axis_index,
-                                                                   unit_x_axis)
+                x_axis_coords = _get_extra_coord_edges(_get_non_common_axis_x_axis_coords(seq.data, cube_plot_axis_index,
+                                                                   unit_x_axis), axis=plot_axis_index)
                 axis_ranges[plot_axis_index] = np.stack(x_axis_coords)
             # Set x-axis label.
             if xlabel is None:
@@ -1034,6 +1035,7 @@ def __init__(self, seq, plot_axis_index=None, axis_ranges=None, unit_x_axis=None
                         # the final x-axis coord array is the same shape as the data array.
                         # This will be used in determining the correct x-axis coords for
                         # each frame of the animation.
+
                         if len(extra_coord_axes) != data_concat.ndim:
                             x_axis_coords_copy = copy.deepcopy(x_axis_coords)
                             x_axis_coords = []
@@ -1045,6 +1047,8 @@ def __init__(self, seq, plot_axis_index=None, axis_ranges=None, unit_x_axis=None
                                 # same as the cube's data array.
                                 # First, create shape of pre-np.tiled x-coord array for the cube.
                                 coords_reshape = np.array([1] * seq[i].data.ndim)
+                                # Convert x_axis_cube_coords to a numpy array
+                                x_axis_cube_coords = np.array(x_axis_cube_coords)
                                 coords_reshape[extra_coord_axes] = x_axis_cube_coords.shape
                                 # Then reshape x-axis array to give it the dummy dimensions.
                                 x_axis_cube_coords = x_axis_cube_coords.reshape(
@@ -1067,7 +1071,7 @@ def __init__(self, seq, plot_axis_index=None, axis_ranges=None, unit_x_axis=None
                 if xlabel is None:
                     xlabel = "{0} [{1}]".format(axis_extra_coord, unit_x_axis)
                 # Re-enter x-axis values into axis_ranges
-                axis_ranges[plot_axis_index] = x_axis_coords
+                axis_ranges[plot_axis_index] = _get_extra_coord_edges(x_axis_coords, axis=plot_axis_index)
             # Else coordinate must have been defined manually.
             else:
                 if isinstance(axis_ranges[plot_axis_index], u.Quantity):
@@ -1085,7 +1089,6 @@ def __init__(self, seq, plot_axis_index=None, axis_ranges=None, unit_x_axis=None
         # Make label for y-axis.
         if ylabel is None:
             ylabel = "Data [{0}]".format(data_unit)
-
         super(LineAnimatorNDCubeSequence, self).__init__(
             data_concat, plot_axis_index=plot_axis_index, axis_ranges=axis_ranges,
             xlabel=xlabel, ylabel=ylabel, xlim=xlim, ylim=ylim, **kwargs)
@@ -1217,13 +1220,13 @@ def __init__(self, seq, plot_axis_index=None, axis_ranges=None, unit_x_axis=None
                     # The repeats is the inverse of dummy_reshape.
                     tile_shape = copy.deepcopy(cube_like_shape)
                     tile_shape[np.array(dependent_axes)] = 1
-                    x_axis_coords = np.tile(x_axis_cube_coords, tile_shape)
+                    x_axis_coords = _get_extra_coord_edges(np.tile(x_axis_cube_coords, tile_shape), axis=plot_axis_index)
             else:
                 # Get x-axis values from each cube and combine into a single
                 # array for axis_ranges kwargs.
                 x_axis_coords = _get_non_common_axis_x_axis_coords(seq.data, plot_axis_index,
                                                                    unit_x_axis)
-                axis_ranges[plot_axis_index] = np.concatenate(x_axis_coords, axis=seq._common_axis)
+                axis_ranges[plot_axis_index] = _get_extra_coord_edges(np.concatenate(x_axis_coords, axis=seq._common_axis), axis=plot_axis_index)
             # Set axis labels and limits, etc.
             if xlabel is None:
                 xlabel = "{0} [{1}]".format(

ndcube/tests/test_plotting.py

@@ -206,7 +206,7 @@ def test_cube_plot_1D_errors(test_input, test_kwargs, expected_error):
 @pytest.mark.parametrize("test_input, test_kwargs, expected_values", [
     (cube[0], {},
      (np.ma.masked_array(cube[0].data, cube[0].mask), "time [min]", "em.wl [m]",
-      (0.4, 1.6, 2e-11, 6e-11))),
+      (-0.5, 3.5, 2.5, -0.5))),
 
     (cube[0], {"axes_coordinates": ["bye", None], "axes_units": [None, u.cm]},
      (np.ma.masked_array(cube[0].data, cube[0].mask), "bye [m]", "em.wl [cm]",
@@ -327,7 +327,7 @@ def test_support_101_plot_API_errors(input_values):
     (cube_unit, {"plot_axis_indices": -1, "axes_coordinates": "bye"},
      (cube_data, cube_none_axis_ranges_axis2_bye, "bye [m]", "Data [J]")),
 
-    (cube, {"plot_axis_indices": -1, "axes_coordinates": np.arange(10, 10+cube.data.shape[-1])},
+    (cube, {"plot_axis_indices": -1, "axes_coordinates": np.arange(10 - 0.5, 0.5+10 + cube.data.shape[-1])},
      (cube_data, cube_none_axis_ranges_axis2_array, " [None]", "Data [None]"))
     ])
 def test_cube_plot_ND_as_1DAnimation(test_input, test_kwargs, expected_values):

ndcube/tests/test_sequence_plotting.py

@@ -7,10 +7,14 @@
 import astropy.units as u
 import matplotlib
 
+from sunpy.visualization.animator.base import edges_to_centers_nd
+
 from ndcube import NDCube, NDCubeSequence
+from ndcube.utils.cube import _get_extra_coord_edges
 from ndcube.utils.wcs import WCS
 import ndcube.mixins.sequence_plotting
 
+
 # sample data for tests
 # TODO: use a fixture reading from a test file. file TBD.
 data = np.array([[[1, 2, 3, 4], [2, 4, 5, 3], [0, -1, 2, 3]],
@@ -207,29 +211,29 @@
 x_axis_coords3 = np.array([0.4, 0.8, 1.2, 1.6]).reshape((1, 1, 4))
 new_x_axis_coords3_shape = u.Quantity(seq.dimensions, unit=u.pix).value.astype(int)
 new_x_axis_coords3_shape[-1] = 1
-none_axis_ranges_axis3 = [np.arange(0, len(seq.data)+1),
-                          np.array([0., 1., 2.]), np.arange(0, 4),
+none_axis_ranges_axis3 = [np.arange(0, len(seq.data)),
+                          np.array([0., 1.]), np.arange(0, 3),
                           np.tile(np.array(x_axis_coords3), new_x_axis_coords3_shape)]
 none_axis_ranges_axis0 = [np.arange(len(seq.data)),
-                          np.array([0., 1., 2.]), np.arange(0, 4),
-                          np.arange(0, int(seq.dimensions[-1].value)+1)]
-distance0_none_axis_ranges_axis0 = [seq.sequence_axis_extra_coords["distance"].value,
-                                    np.array([0., 1., 2.]), np.arange(0, 4),
-                                    np.arange(0, int(seq.dimensions[-1].value)+1)]
-distance0_none_axis_ranges_axis0_mm = [seq.sequence_axis_extra_coords["distance"].to("mm").value,
-                                       np.array([0., 1., 2.]), np.arange(0, 4),
-                                       np.arange(0, int(seq.dimensions[-1].value)+1)]
+                          np.array([0., 1.]), np.arange(0, 3),
+                          np.arange(0, int(seq.dimensions[-1].value))]
+distance0_none_axis_ranges_axis0 = [edges_to_centers_nd(_get_extra_coord_edges(seq.sequence_axis_extra_coords["distance"].value),0),
+                                    np.array([0., 1.]), np.arange(0, 3),
+                                    np.arange(0, int(seq.dimensions[-1].value))]
+distance0_none_axis_ranges_axis0_mm = [edges_to_centers_nd(_get_extra_coord_edges(seq.sequence_axis_extra_coords["distance"].to("mm").value),0),
+                                       np.array([0., 1.]), np.arange(0, 3),
+                                       np.arange(0, int(seq.dimensions[-1].value))]
 userrangequantity_none_axis_ranges_axis0 = [
-    np.arange(int(seq.dimensions[0].value)), np.array([0., 1., 2.]), np.arange(0, 4),
-    np.arange(0, int(seq.dimensions[-1].value)+1)]
+    np.arange(int(seq.dimensions[0].value)), np.array([0., 1.]), np.arange(0, 3),
+    np.arange(0, int(seq.dimensions[-1].value))]
 
 userrangequantity_none_axis_ranges_axis0_1e7 = [
-    (np.arange(int(seq.dimensions[0].value)) * u.J).to(u.erg).value, np.array([0., 1., 2.]),
-    np.arange(0, 4), np.arange(0, int(seq.dimensions[-1].value)+1)]
+    (np.arange(int(seq.dimensions[0].value)) * u.J).to(u.erg).value, np.array([0., 1.]),
+    np.arange(0, 3), np.arange(0, int(seq.dimensions[-1].value))]
 
 hi2_none_axis_ranges_axis2 = [
-    np.arange(0, len(seq.data)+1), np.array([0., 1., 2.]),
-    np.arange(int(seq.dimensions[2].value)), np.arange(0, int(seq.dimensions[-1].value)+1)]
+    np.arange(0, len(seq.data)), np.array([0., 1.]),
+    np.array([0, 1, 2]), np.arange(0, int(seq.dimensions[-1].value))]
 
 x_axis_coords1 = np.zeros(tuple([int(s.value) for s in seq.dimensions]))
 x_axis_coords1[0, 1] = 1.
@@ -239,8 +243,8 @@
 x_axis_coords1[3, 0] = 2.
 x_axis_coords1[3, 1] = 3.
 pix1_none_axis_ranges_axis1 = [
-    np.arange(0, len(seq.data)+1), x_axis_coords1, np.arange(0, 4),
-    np.arange(0, int(seq.dimensions[-1].value)+1)]
+    np.arange(0, len(seq.data)), x_axis_coords1, np.arange(0, 3),
+    np.arange(0, int(seq.dimensions[-1].value))]
 
 # Derive expected extents
 seq_axis1_lim_deg = [0.49998731, 0.99989848]
@@ -253,14 +257,14 @@
     seq.cube_like_dimensions, unit=u.pix).value.astype(int)
 cube_like_new_x_axis_coords2_shape[-1] = 1
 cubelike_none_axis_ranges_axis2 = [
-    np.arange(0, int(seq.cube_like_dimensions[0].value)+1), np.arange(0, 4),
+    np.arange(0, int(seq.cube_like_dimensions[0].value)), np.arange(0, 3),
     np.tile(x_axis_coords3, cube_like_new_x_axis_coords2_shape)]
 
 cubelike_none_axis_ranges_axis2_s = copy.deepcopy(cubelike_none_axis_ranges_axis2)
 cubelike_none_axis_ranges_axis2_s[2] = cubelike_none_axis_ranges_axis2_s[2] * 60.
 
-cubelike_none_axis_ranges_axis0 = [[0, 8], np.arange(0, 4),
-                                   np.arange(0, int(seq.cube_like_dimensions[-1].value)+1)]
+cubelike_none_axis_ranges_axis0 = [[-0.5, 7.5], np.arange(0, 3),
+                                   np.arange(0, int(seq.cube_like_dimensions[-1].value))]
 
 
 @pytest.mark.parametrize("test_input, test_kwargs, expected_values", [
@@ -684,14 +688,14 @@ def test_prep_axes_kwargs_errors(test_input, expected_error):
       (seq_stack.data.min(), seq_stack.data.max()))),
 
     (seq, {"plot_axis_indices": 0,
-           "axes_coordinates": userrangequantity_none_axis_ranges_axis0[0]*u.J},
+           "axes_coordinates": _get_extra_coord_edges(userrangequantity_none_axis_ranges_axis0[0])*u.J},
      (seq_stack.data, userrangequantity_none_axis_ranges_axis0, " [J]", "Data [None]",
       (userrangequantity_none_axis_ranges_axis0[0].min(),
        userrangequantity_none_axis_ranges_axis0[0].max()),
       (seq_stack.data.min(), seq_stack.data.max()))),
 
     (seq, {"plot_axis_indices": 0, "axes_units": u.erg,
-           "axes_coordinates": userrangequantity_none_axis_ranges_axis0[0]*u.J},
+           "axes_coordinates": _get_extra_coord_edges(userrangequantity_none_axis_ranges_axis0[0])*u.J},
      (seq_stack.data, userrangequantity_none_axis_ranges_axis0_1e7, " [erg]", "Data [None]",
       (userrangequantity_none_axis_ranges_axis0_1e7[0].min(),
        userrangequantity_none_axis_ranges_axis0_1e7[0].max()),

ndcube/utils/cube.py

@@ -5,6 +5,7 @@
 """
 
 import numpy as np
+from astropy.units import Quantity
 
 __all__ = ['wcs_axis_to_data_axis', 'data_axis_to_wcs_axis', 'select_order',
            'convert_extra_coords_dict_to_input_format', 'get_axis_number_from_axis_name']
@@ -195,3 +196,58 @@ def _get_dimension_for_pixel(axis_length, edges):
         False stands for pixel_value, while True stands for pixel_edge
     """
     return axis_length+1 if edges else axis_length
+
+def _get_extra_coord_edges(value, axis=-1):
+    """Gets the pixel_edges from the pixel_values
+     Parameters
+    ----------
+    value : `astropy.units.Quantity` or array-like
+        The Quantity object containing the values for a given `extra_coords`
+     axis : `int`
+        The axis about which pixel_edges needs to be calculated
+        Default value is -1, which is the last axis for a ndarray
+    """
+
+     # Checks for corner cases
+
+    if not isinstance(value, np.ndarray):
+        value = np.array(value)
+
+    # Get the shape of the Quantity object
+    shape = value.shape
+    if len(shape) == 1:
+
+        shape = len(value)
+        if isinstance(value, Quantity):
+            edges = np.zeros(shape+1) * value.unit
+        else:
+            edges = np.zeros(shape+1)
+
+         # Calculate the pixel_edges from the given pixel_values
+        edges[1:-1] = value[:-1] + (value[1:] - value[:-1]) / 2
+        edges[0] = value[0] - (value[1] - value[0]) / 2
+        edges[-1] = value[-1] + (value[-1] - value[-2]) / 2
+
+    else:
+        # Edit the shape of the new ndarray to increase the length
+        # by one for a given axis
+        shape = list(shape)
+        shape[axis] += 1
+        shape = tuple(shape)
+
+        if isinstance(value, Quantity):
+            edges = np.zeros(shape) * value.unit
+        else:
+            edges = np.zeros(shape)
+        # Shift the axis which is point of interest to last axis
+        value = np.moveaxis(value, axis, -1)
+        edges = np.moveaxis(edges, axis, -1)
+
+        # Calculate the pixel_edges from the given pixel_values
+        edges[...,1:-1] = value[...,:-1] + (value[...,1:] - value[...,:-1]) / 2
+        edges[...,0] = value[...,0] - (value[...,1] - value[...,0]) / 2
+        edges[...,-1] = value[...,-1] + (value[...,-1] - value[...,-2]) / 2
+
+        # Revert the shape of the edges array
+        edges = np.moveaxis(edges, -1, axis)
+    return edges