Functions to support GLM and ground network subsetting

pyxlma/plot/interactive.py

@@ -7,6 +7,7 @@
 
 from pyxlma.plot.xlma_plot_feature import color_by_time, plot_points, setup_hist, plot_3d_grid, subset
 from pyxlma.plot.xlma_base_plot import subplot_labels, inset_view, BlankPlot
+from pyxlma.xarray_util import generic_subset
 
 from ipywidgets import Output
 output = Output()
@@ -357,3 +358,51 @@ def make_plot(self):
 
 
 
+def get_glm_plot_subset(interactive_plot, glm):
+    """ Use the plot limits in the interactive plot to subset GLM data.
+
+        glm may be an xarray Dataset or a glmtools GLMDataset.
+
+        Returns in xarray Dataset subsetted to match the plot.
+    """
+
+    from glmtools.io.glm import GLMDataset
+
+    # We need the subsetting functionality attached to the GLMDataset class, which 
+    # can prune the flash-group-event hierarchy to a self-consistent sub-tree.
+    if isinstance(glm, xr.Dataset):
+        glm = GLMDataset(glm, check_area_units=False, change_energy_units=False)
+    else:
+        assert isinstance(glm, GLMDataset)
+
+    xlim = interactive_plot.bounds['x']
+    ylim = interactive_plot.bounds['y']    
+    tlim = interactive_plot.bounds['t']
+    start, end = np.datetime64(tlim[0]), np.datetime64(tlim[1])
+
+
+    # Find the groups in the time range.
+    # In some GLM datasets, perhaps all, the event times are incorrect.
+    # Probably missing some unsigned stuff.
+    # That is why above we use the group times only and select events by 
+    # parent ID through reduce_to_entities
+    # print(glm_sub.event_time_offset.min().data, glm_sub.event_time_offset.max().data)
+    # print(glm_sub.group_time_offset.min().data, glm_sub.group_time_offset.max().data)
+    # print(glm_sub)
+
+    glm_bounds = {'group_time_offset':slice(start,end),}
+    glm_sub = generic_subset(glm.dataset, glm.dataset.group_id.dims[0], glm_bounds)
+    glm_sub = glm.reduce_to_entities('group_id', glm_sub.group_id.data)
+
+    # Recreate the GLMDataset from the reduced dataset.
+    # There's probably some way to do this all in one step, perhaps by using the 
+    # common set of group_ids. But it may not be faster in the end anyway.
+    glm = GLMDataset(glm_sub, check_area_units=False, change_energy_units=False)
+
+    # Find the events that are in the view, and keep only their parent groups.
+    glm_bounds = {'event_lat':slice(ylim[0], ylim[1]),
+                  'event_lon':slice(xlim[0], xlim[1])}
+    glm_sub = generic_subset(glm_sub, glm_sub.event_id.dims[0], glm_bounds)
+    glm_sub = glm.reduce_to_entities('group_id', glm_sub.event_parent_group_id.data)
+
+    return glm_sub

pyxlma/plot/xlma_plot_feature.py

@@ -192,12 +192,16 @@ def plot_2d_network_points(bk_plot, netw_data, actual_height=None, fake_ic_heigh
     return art_out
 
 
-def plot_glm_events(glm, bk_plot, fake_alt=[0, 1], should_parallax_correct=True, poly_kwargs={}, vlines_kwargs={}):
+def plot_glm_events(glm, bk_plot, fake_alt=[0, 1], should_parallax_correct=True, poly_kwargs={}, vlines_kwargs={}, points_kwargs={}):
     """
     Plot event-level data from a glmtools dataset on a pyxlma.plot.xlma_base_plot.BlankPlot object.
     Events that occupy the same pixel have their energies summed and plotted on the planview axis, event locations
     are plotted on the lat/lon/time axes with an altitude specified as fake_alt.
     Requires glmtools to be installed.
+
+    The group latitudes and longitudes are plotted as in the LCFA data file, without parallax correction.
+
+    Does not do any subsetting of the GLM dataset; see pyxlma.plot.interactive.get_glm_plot_subset for that functionality.
 
     Parameters
     ----------
@@ -209,7 +213,8 @@ def plot_glm_events(glm, bk_plot, fake_alt=[0, 1], should_parallax_correct=True,
         the axes relative coordinates to plot the vertical lines for GLM events in the cross section, default [0, 1],
         the full height of the axes.
     should_parallax_correct : bool
-        whether to correct the GLM event locations for parallax effect. See https://doi.org/10.1029/2019JD030874 for more information.
+        whether to correct the GLM event locations for parallax effect. 
+        See https://doi.org/10.1029/2019JD030874 for more information.
     poly_kwargs : dict
         dictionary of additional keyword arguments to be passed to matplotlib Polygon
     vlines_kwargs : dict
@@ -277,10 +282,18 @@ def plot_glm_events(glm, bk_plot, fake_alt=[0, 1], should_parallax_correct=True,
     pc = PatchCollection(patches, **poly_kwargs)
     pc.set_array(evrad.data)
     bk_plot.ax_plan.add_collection(pc)
-    th_handle = bk_plot.ax_th.vlines(glm.event_time_offset.data, fake_alt[0], fake_alt[1], transform=bk_plot.ax_th.get_xaxis_transform(), **vlines_kwargs)
-    lon_handle = bk_plot.ax_lon.vlines(glm.event_lon, fake_alt[0], fake_alt[1], transform=bk_plot.ax_lon.get_xaxis_transform(), **vlines_kwargs)
-    lat_handle = bk_plot.ax_lat.hlines(glm.event_lat, fake_alt[0], fake_alt[1], transform=bk_plot.ax_lat.get_yaxis_transform(), **vlines_kwargs)
-    art_out = [pc, th_handle, lon_handle, lat_handle]
+    th_handle = bk_plot.ax_th.vlines(glm.group_time_offset.data, fake_alt[0], fake_alt[1], 
+        transform=bk_plot.ax_th.get_xaxis_transform(), **vlines_kwargs)
+    vert_proj_point_alt = np.ones_like(glm.group_lon) * (fake_alt[0]+fake_alt[1])/2.0
+    lon_handle = bk_plot.ax_lon.plot(glm.group_lon, vert_proj_point_alt, 
+        transform=bk_plot.ax_lon.get_xaxis_transform(), **points_kwargs)
+    lat_handle = bk_plot.ax_lat.plot(vert_proj_point_alt, glm.group_lat,
+        transform=bk_plot.ax_lat.get_yaxis_transform(), **points_kwargs)
+    plan_handle = bk_plot.ax_plan.plot(glm.group_lon, glm.group_lat, **points_kwargs)
+
+    # lon_handle = bk_plot.ax_lon.vlines(glm.event_lon, fake_alt[0], fake_alt[1], transform=bk_plot.ax_lon.get_xaxis_transform(), **vlines_kwargs)
+    # lat_handle = bk_plot.ax_lat.hlines(glm.event_lat, fake_alt[0], fake_alt[1], transform=bk_plot.ax_lat.get_yaxis_transform(), **vlines_kwargs)
+    art_out = [pc, th_handle] + lon_handle + lat_handle + plan_handle
     return art_out
 
 

pyxlma/xarray_util.py

@@ -2,6 +2,24 @@
 import xarray as xr
 import numpy as np
 
+def generic_subset(d, dim, slices):
+    """ slices is a dictionary of variable names to slice objects or boolean masks,
+        assumed to all apply to the same dimension dim within xarray dataset d"""
+    mask = None
+    for k, sl in slices.items():
+        # print(k, sl)
+        if hasattr(sl, 'start'):
+            this_mask = (d[k] >= sl.start) & (d[k] <= sl.stop)
+        else:
+            this_mask = sl
+        if mask is None:
+            mask = this_mask
+        else:
+            mask = mask & this_mask
+        # print(mask.sum())
+    return d[{dim:mask}]
+
+
 def get_1d_dims(d):
     """
     Find all dimensions in a dataset that are purely 1-dimensional,