From 2cc193321c2306e64a69401fa2d9f25bfc54d01b Mon Sep 17 00:00:00 2001 From: Eric Bruning Date: Sun, 5 Jan 2025 14:19:03 -0600 Subject: [PATCH 1/2] Functions to support GLM subsetting --- pyxlma/plot/interactive.py | 49 ++++++++++++++++++++++++++++++++++++++ pyxlma/xarray_util.py | 18 ++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/pyxlma/plot/interactive.py b/pyxlma/plot/interactive.py index c8038d0..42816da 100644 --- a/pyxlma/plot/interactive.py +++ b/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 diff --git a/pyxlma/xarray_util.py b/pyxlma/xarray_util.py index 79b67cc..742dcc8 100644 --- a/pyxlma/xarray_util.py +++ b/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, From 61e2d492d0588b2d04bb10bc81b2df9d37d3c068 Mon Sep 17 00:00:00 2001 From: Eric Bruning Date: Sun, 5 Jan 2025 15:37:59 -0600 Subject: [PATCH 2/2] Update GLM plot style, show only groups --- pyxlma/plot/xlma_plot_feature.py | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/pyxlma/plot/xlma_plot_feature.py b/pyxlma/plot/xlma_plot_feature.py index 1f01800..bc7cdae 100644 --- a/pyxlma/plot/xlma_plot_feature.py +++ b/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