Skip to content

Commit

Permalink
Merge pull request #89 from wtbarnes/viz-bug-fixes
Browse files Browse the repository at this point in the history 
Minor fixes in viz module
  • Loading branch information
@wtbarnes
wtbarnes authored Feb 6, 2020
2 parents d7748d2 + e5c9aa4 commit 962a886
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 30 deletions.
23 changes: 12 additions & 11 deletions pydrad/visualize/animate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,26 +20,27 @@ def animate_strand(strand, **kwargs):
limits = kwargs.get('limits', {})
if 'limits' in kwargs:
del kwargs['limits']
# Initialize figure
fig, axes = _setup_figure(strand[0], limits, **kwargs)
# Make sure the lines stay the same color
if 'color' not in plot_kwargs:
plot_kwargs['color'] = 'C0'
l1a, l1b, l2a, l2b, l3a, l3b, l4 = _plot_profile(strand[0],
axes,
**plot_kwargs)
# Update function

def update_plot(i):
profile = strand[i]
l1a.set_data(profile.coordinate.to(u.cm), profile.electron_temperature.to(u.MK))
l1b.set_data(profile.coordinate.to(u.cm), profile.ion_temperature.to(u.MK))
l2a.set_data(profile.coordinate.to(u.cm), profile.electron_density)
l2b.set_data(profile.coordinate.to(u.cm), profile.ion_density)
l3a.set_data(profile.coordinate.to(u.cm), profile.electron_pressure)
l3b.set_data(profile.coordinate.to(u.cm), profile.ion_pressure)
l4.set_data(profile.coordinate.to(u.cm), profile.velocity)
fig.suptitle(r'$t={:4.0f}$ {}'.format(
p = strand[i]
l1a.set_data(p.coordinate.to(u.Mm), p.electron_temperature.to(u.MK))
l1b.set_data(p.coordinate.to(u.Mm), p.ion_temperature.to(u.MK))
l2a.set_data(p.coordinate.to(u.Mm), p.electron_density.to(u.cm**(-3)))
l2b.set_data(p.coordinate.to(u.Mm), p.ion_density.to(u.cm**(-3)))
l3a.set_data(p.coordinate.to(u.Mm), p.electron_pressure.to(u.dyne/(u.cm**2)))
l3b.set_data(p.coordinate.to(u.Mm), p.ion_pressure.to(u.dyne/(u.cm**2)))
l4.set_data(p.coordinate.to(u.Mm), p.velocity.to(u.km/u.s))
fig.suptitle(r'$t={:.0f}$ {}'.format(
strand.time[i].value, strand.time[i].unit), y=0.905)
return l1a, l1b, l2a, l2b, l3a, l3b, l4

return FuncAnimation(
fig, update_plot, blit=kwargs.get('blit', True), frames=len(strand),
interval=kwargs.get('interval', 10), repeat=kwargs.get('repeat', True))
45 changes: 26 additions & 19 deletions pydrad/visualize/plot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ def plot_strand(strand, limits=None, cmap='viridis', **kwargs):
# Parameters
strand (#pydrad.parse.Strand): Loop strand object
limits (`dict`): Set axes limits for hydrodynamic quantities, optional
cmap (`str`): The colormap to map the timestep index to
plot_kwargs (`dict`): Any keyword arguments used matplotlib.plot, optional
figsize (`tuple`): Width and height of figure, optional
"""
Expand All @@ -68,17 +69,23 @@ def plot_strand(strand, limits=None, cmap='viridis', **kwargs):
fig, axes = _setup_figure(strand[0], limits, **kwargs)
colors = matplotlib.colors.LinearSegmentedColormap.from_list(
'', plt.get_cmap(cmap).colors, N=len(strand))
# NOTE: once strand indexing is fixed, we can index it directly
for i, p in enumerate(strand):
plot_kwargs['color'] = colors(i)
_ = _plot_profile(p, axes, **plot_kwargs)
plt.show()


def plot_profile(profile, **kwargs):
limits = kwargs.get('limits', {})
if 'limits' in kwargs:
del kwargs['limits']
"""
Plot hydrodynamic quantites at a single timestep
# Parameters
profile (#pydrad.parse.Strand): Loop profile object
limits (`dict`): Set axes limits for hydrodynamic quantities, optional
plot_kwargs (`dict`): Any keyword arguments used matplotlib.plot, optional
figsize (`tuple`): Width and height of figure, optional
"""
limits = kwargs.pop('limits', {})
plot_kwargs = kwargs.get('plot_kwargs', {})
fig, axes = _setup_figure(profile, limits, **kwargs)
_plot_profile(profile, axes, **plot_kwargs)
Expand All @@ -95,13 +102,13 @@ def _setup_figure(profile, limits, **kwargs):
axes[0, 1].set_yscale('log')
axes[1, 0].set_ylim(limits.get('pressure', (0.1, 1e2)))
axes[1, 0].set_yscale('log')
axes[1, 1].set_ylim(limits.get('velocity', (-5e7, 5e7)))
axes[1, 1].set_xlim(0, profile.coordinate[-1].to(u.cm).value)
axes[1, 1].set_ylim(limits.get('velocity', (-1e2, 1e2)))
axes[1, 1].set_xlim(profile.coordinate[[0, -1]].to(u.Mm).value)
# Labels
axes[0, 0].set_ylabel(r'$T$ [MK]')
axes[0, 1].set_ylabel(r'$n$ [cm$^{-3}$]')
axes[1, 0].set_ylabel(r'$P$ [dyne cm$^{-2}$ s$^{-1}$]')
axes[1, 1].set_ylabel(r'$v$ [cm s$^{-1}$]')
axes[1, 1].set_ylabel(r'$v$ [km s$^{-1}$]')
axes[1, 0].set_xlabel(r'$s$ [Mm]')
axes[1, 1].set_xlabel(r'$s$ [Mm]')

Expand All @@ -110,44 +117,44 @@ def _setup_figure(profile, limits, **kwargs):

def _plot_profile(profile, axes, **kwargs):
line1a, = axes[0, 0].plot(
profile.coordinate.to(u.cm),
profile.coordinate.to(u.Mm),
profile.electron_temperature.to(u.MK),
**kwargs,
ls='-'
)
line1b, = axes[0, 0].plot(
profile.coordinate.to(u.cm),
profile.coordinate.to(u.Mm),
profile.ion_temperature.to(u.MK),
**kwargs,
ls='--'
)
line2a, = axes[0, 1].plot(
profile.coordinate.to(u.cm),
profile.electron_density,
profile.coordinate.to(u.Mm),
profile.electron_density.to(u.cm**(-3)),
**kwargs,
ls='-'
)
line2b, = axes[0, 1].plot(
profile.coordinate.to(u.cm),
profile.ion_density,
profile.coordinate.to(u.Mm),
profile.ion_density.to(u.cm**(-3)),
**kwargs,
ls='--'
)
line3a, = axes[1, 0].plot(
profile.coordinate.to(u.cm),
profile.electron_pressure,
profile.coordinate.to(u.Mm),
profile.electron_pressure.to(u.dyne / (u.cm**2)),
**kwargs,
ls='-'
)
line3b, = axes[1, 0].plot(
profile.coordinate.to(u.cm),
profile.ion_pressure,
profile.coordinate.to(u.Mm),
profile.ion_pressure.to(u.dyne / (u.cm**2)),
**kwargs,
ls='--'
)
line4, = axes[1, 1].plot(
profile.coordinate.to(u.cm),
profile.velocity,
profile.coordinate.to(u.Mm),
profile.velocity.to(u.km/u.s),
**kwargs
)
return line1a, line1b, line2a, line2b, line3a, line3b, line4

0 comments on commit 962a886

Please sign in to comment.