From 2cc193321c2306e64a69401fa2d9f25bfc54d01b Mon Sep 17 00:00:00 2001 From: Eric Bruning Date: Sun, 5 Jan 2025 14:19:03 -0600 Subject: [PATCH 1/5] 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/5] 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 From 2d85171bf7505e54b871c0fc4db850f308f37285 Mon Sep 17 00:00:00 2001 From: Eric Bruning Date: Wed, 25 Jun 2025 19:11:57 -0500 Subject: [PATCH 3/5] reduce_to_entities doesn't like empty GLM datasets --- pyxlma/plot/interactive.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyxlma/plot/interactive.py b/pyxlma/plot/interactive.py index 42816da..de077de 100644 --- a/pyxlma/plot/interactive.py +++ b/pyxlma/plot/interactive.py @@ -392,6 +392,9 @@ def get_glm_plot_subset(interactive_plot, glm): glm_bounds = {'group_time_offset':slice(start,end),} glm_sub = generic_subset(glm.dataset, glm.dataset.group_id.dims[0], glm_bounds) + if glm_sub.group_id.data.shape[0] < 1: + # No data, so just empty everything + return glm.dataset[{'number_of_events':[], 'number_of_groups':[], 'number_of_flashes':[]}] glm_sub = glm.reduce_to_entities('group_id', glm_sub.group_id.data) # Recreate the GLMDataset from the reduced dataset. From 2a6a162ca865ac15b2f7f0b94d93b54af66cd603 Mon Sep 17 00:00:00 2001 From: Eric Bruning Date: Wed, 25 Jun 2025 19:16:55 -0500 Subject: [PATCH 4/5] Add 2D network subsetting --- pyxlma/plot/interactive.py | 51 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/pyxlma/plot/interactive.py b/pyxlma/plot/interactive.py index de077de..cb16b4f 100644 --- a/pyxlma/plot/interactive.py +++ b/pyxlma/plot/interactive.py @@ -409,3 +409,54 @@ def get_glm_plot_subset(interactive_plot, glm): glm_sub = glm.reduce_to_entities('group_id', glm_sub.event_parent_group_id.data) return glm_sub + + +def get_2d_network_subset(interactive_plot, netw_data, adjust_alt=True, pad_factor=0.01): + """ + Subset 2D network point data to the current view of an interactive plot. + + Parameters + ---------- + interactive_plot : `pyxlma.plot.interactive.InteractiveLMAPlot` + The XLMA plot from which the bounds of the current view are extracted + netw_data : `pandas.DataFrame` + 2d network point location data of the type expected by + `pyxlma.plot.xlma_plot_feature.plot_2d_network_points` + adjust_alt : bool + if True, adjust the `icheight` column of netw_sub to fit within the plot range + if False, reject any points whose `icheight` is outside the plot range. + pad_factor : float + if adjust_alt==True, factor of the plot's total altitude height range to use + as padding on the minimum and maximum altitude. + + Returns + ------- + netw_sub : `pandas.DataFrame` + subset netw_data to match the current plot + """ + xlim = interactive_plot.bounds['x'] + ylim = interactive_plot.bounds['y'] + zlim = [z*1000.0 for z in interactive_plot.bounds['z']] + tlim = interactive_plot.bounds['t'] + start, end = np.datetime64(tlim[0]), np.datetime64(tlim[1]) + z_span = zlim[1] - zlim[0] + + in_lat = (netw_data['latitude'] >= ylim[0]) & (netw_data['latitude'] < ylim[1]) + in_lon = (netw_data['longitude'] >= xlim[0]) & (netw_data['longitude'] < xlim[1]) + in_time = (netw_data['datetime'] >= start) & (netw_data['datetime'] < end) + + # Must copy, or it gives a view of the original dataframe/ + netw_sub = netw_data[in_lat & in_lon & in_time].copy() + + # if altitude greater or less than top or bottom of plot, move to top or bottom. + if adjust_alt: + min_z = zlim[0]+pad_factor*z_span + max_z = zlim[1]-pad_factor*z_span + below = (netw_sub['icheight'] < min_z) + above = (netw_sub['icheight'] >= max_z) + netw_sub['icheight'] = np.where(below, min_z, netw_sub['icheight']) + netw_sub['icheight'] = np.where(above, max_z, netw_sub['icheight']) + else: + in_alt = (netw_sub['icheight'] >= zlim[0]) & (netw_sub['icheight'] < zlim[1]) + netw_sub = netw_sub[in_alt] + return netw_sub \ No newline at end of file From c349af4086dbe7a040200307ca1f8667e5762305 Mon Sep 17 00:00:00 2001 From: Eric Bruning Date: Wed, 25 Jun 2025 21:10:40 -0500 Subject: [PATCH 5/5] Enforce numpy arrays --- pyxlma/plot/leader_speed.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pyxlma/plot/leader_speed.py b/pyxlma/plot/leader_speed.py index 8bc853d..513376b 100644 --- a/pyxlma/plot/leader_speed.py +++ b/pyxlma/plot/leader_speed.py @@ -32,10 +32,10 @@ def get_time_distance(lat, lon, time, lat0, lon0, time0): def time_distance_plot_interactive(interactive_lma, ax): - lat = interactive_lma.this_lma_lat - lon = interactive_lma.this_lma_lon - alt = interactive_lma.this_lma_alt - time = interactive_lma.this_lma_time + lat = np.asarray(interactive_lma.this_lma_lat) + lon = np.asarray(interactive_lma.this_lma_lon) + alt = np.asarray(interactive_lma.this_lma_alt) + time = np.asarray(interactive_lma.this_lma_time) first = np.nanargmin(time) distance_from_origin, time_from_origin = get_time_distance(lat, lon, time,