Skip to content

Functions to support GLM and ground network subsetting #55

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 49 additions & 0 deletions pyxlma/plot/interactive.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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
25 changes: 19 additions & 6 deletions pyxlma/plot/xlma_plot_feature.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
----------
Expand All @@ -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
Expand Down Expand Up @@ -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)
Comment on lines +288 to +292
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See test result below, should this maybe be a call to scatter instead of plot (as the default/"no points_kwargs specified" for scatter is probably more in line with what's expected?
image


# 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


Expand Down
18 changes: 18 additions & 0 deletions pyxlma/xarray_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading