Skip to content
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

Adding an enhancement feature #8

Open
wants to merge 1 commit into
base: main
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
22 changes: 21 additions & 1 deletion enlilviz/enlil.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
"""Main enlil data module."""
import functools
import numpy as np


def _validate_satellite(func):
Expand Down Expand Up @@ -106,7 +107,7 @@ def get_satellite_data(self, satellite, var, coord=None):
varname += '_' + coord
return self.ds.loc[{'satellite': satellite}][varname]

def get_slice(self, var, slice_plane, time=None):
def get_slice(self, var, slice_plane, time=None, enhance=None):
"""Get a 2D slice of data.

Parameters
Expand All @@ -118,6 +119,10 @@ def get_slice(self, var, slice_plane, time=None):
Note that a slice in lon, will return data with r/lat coordinates.
time : datetime-like, optional
Time of interest. If left out, all times will be returned.
enhance : int, optional
Enhanced resolution factor to increase the field by through
interpolation. enhance = 4 will increase each dimension by a
factor of 4, creating a 16x bigger data array.

Returns
-------
Expand All @@ -129,6 +134,21 @@ def get_slice(self, var, slice_plane, time=None):
if time is not None:
da = da.sel({'t': time}, method='nearest')

# Interpolating the data to a higher resolution
if enhance:
new_r = np.linspace(self.ds.r[0], self.ds.r[-1],
len(self.ds.r)*enhance)
new_lon = np.linspace(self.ds.lon[0], self.ds.lon[-1],
len(self.ds.lon)*enhance)
new_lat = np.linspace(self.ds.lat[0], self.ds.lat[-1],
len(self.ds.lat)*enhance)
if slice_plane == 'r':
da = da.interp(lat=new_lat, lon=new_lon)
elif slice_plane == 'lat':
da = da.interp(r=new_r, lon=new_lon)
elif slice_plane == 'lon':
da = da.interp(r=new_r, lat=new_lat)

return da


Expand Down
41 changes: 20 additions & 21 deletions enlilviz/plotting/plots.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,8 +180,9 @@ class LatitudeSlice(_BasePlot):
Variable to plot (den, vel)
"""

def __init__(self, enlil_run, var, ax=None):
def __init__(self, enlil_run, var, enhance=None, ax=None):
self.var = var
self.enhance = enhance
self.ax = plt.subplot(projection='polar') if ax is None else ax
super().__init__(enlil_run, self.ax)

Expand All @@ -191,17 +192,17 @@ def _init_plot(self):
r = run.r
lon = np.deg2rad(run.lon)

# Make the longitude-radial mesh
r_lon, lon_r = _mesh_grid(r, lon)

# Get the polarity data
inner_pol, outer_pol = self._get_polarity_data('lat')

ax = self.ax
ax.axis('off')

cmap, norm = CMAP_LOOKUP[self.var]
data = run.get_slice(self.var, 'lat', self.time)
data = run.get_slice(self.var, 'lat', self.time, self.enhance)
# Make the longitude-radial mesh
r_lon, lon_r = _mesh_grid(data.r, np.deg2rad(data.lon))

mesh = ax.pcolormesh(lon_r, r_lon, data,
cmap=cmap, norm=norm,
shading='flat')
Expand Down Expand Up @@ -247,7 +248,7 @@ def update(self):
"""Update all variable quantities within the plot."""
run = self.enlil_run
lon = np.deg2rad(run.lon)
data = run.get_slice(self.var, 'lat', self.time)
data = run.get_slice(self.var, 'lat', self.time, self.enhance)
self.plot_data['mesh'].set_array(data.values.flatten())

inner_pol, outer_pol = self._get_polarity_data('lat')
Expand All @@ -268,8 +269,9 @@ class LongitudeSlice(_BasePlot):
Variable to plot (den, vel)
"""

def __init__(self, enlil_run, var, ax=None):
def __init__(self, enlil_run, var, enhance=None, ax=None):
self.var = var
self.enhance = enhance
self.ax = plt.subplot(projection='polar') if ax is None else ax
super().__init__(enlil_run, self.ax)

Expand All @@ -279,17 +281,16 @@ def _init_plot(self):
r = run.r
lat = np.deg2rad(run.lat)

# Make the longitude-radial mesh
r_lat, lat_r = _mesh_grid(r, lat)

# Get the polarity data
inner_pol, outer_pol = self._get_polarity_data('lon')

ax = self.ax
ax.axis('off')

cmap, norm = CMAP_LOOKUP[self.var]
data = run.get_slice(self.var, 'lon', self.time)
data = run.get_slice(self.var, 'lon', self.time, self.enhance)
# Make the longitude-radial mesh
r_lat, lat_r = _mesh_grid(data.r, np.deg2rad(data.lat))
mesh = ax.pcolormesh(lat_r, r_lat, data,
cmap=cmap, norm=norm, shading='flat')
self.plot_data['mesh'] = mesh
Expand Down Expand Up @@ -332,7 +333,7 @@ def update(self):
"""Update all variable quantities within the plot."""
run = self.enlil_run
lat = np.deg2rad(run.lat)
data = run.get_slice(self.var, 'lon', self.time)
data = run.get_slice(self.var, 'lon', self.time, self.enhance)
self.plot_data['mesh'].set_array(data.values.flatten())

inner_pol, outer_pol = self._get_polarity_data('lon')
Expand All @@ -354,26 +355,24 @@ class RadialSlice(_BasePlot):
Variable to plot (den, vel)
"""

def __init__(self, enlil_run, var, ax=None):
def __init__(self, enlil_run, var, enhance=None, ax=None):
self.var = var
self.enhance = enhance
self.ax = plt.subplot(projection='mollweide') if ax is None else ax
super().__init__(enlil_run, self.ax)

def _init_plot(self):
"""Initialize the axis with all of the plot data."""
run = self.enlil_run
lon = np.deg2rad(run.lon)
lat = np.deg2rad(run.lat)

# Make the longitude-latitude mesh
lon_lat, lat_lon = _mesh_grid(lon, lat)

ax = self.ax
ax.axis('off')

cmap, norm = CMAP_LOOKUP[self.var]
# Need to transpose the data for plotting
data = run.get_slice(self.var, 'r', self.time).T
data = run.get_slice(self.var, 'r', self.time, enhance=self.enhance).T
# Make the longitude-latitude mesh
lon_lat, lat_lon = _mesh_grid(np.deg2rad(data.lon),
np.deg2rad(data.lat))
mesh = ax.pcolormesh(lon_lat, lat_lon, data,
cmap=cmap, norm=norm,
shading='flat')
Expand All @@ -400,7 +399,7 @@ def update(self):
"""Update all variable quantities within the plot."""
run = self.enlil_run
# transpose the data to match the coordinates
data = run.get_slice(self.var, 'r', self.time).T
data = run.get_slice(self.var, 'r', self.time, enhance=self.enhance).T
self.plot_data['mesh'].set_array(data.values.flatten())

for sat in ['Earth', 'STEREO_A', 'STEREO_B']:
Expand Down