Skip to content

Commit

Permalink
The size and color of particles of animation and animation_profile me…
Browse files Browse the repository at this point in the history 
…thods can now be scaled with any element or environment property by specifying marker=<property>. Sizes can eventually be scaled by spcifying markersize_scaling. Transparency (alpha) can also be provided. Some examples are updated.
  • Loading branch information
@knutfrode
knutfrode committed Mar 14, 2024
1 parent 271123d commit 1ac6c96
Show file tree
Hide file tree
Showing 4 changed files with 115 additions and 79 deletions.
25 changes: 15 additions & 10 deletions examples/example_biodegradation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,29 +15,34 @@
# No motion is needed for this test
o.set_config('environment:constant', {k: 0 for k in
['x_wind', 'y_wind', 'x_sea_water_velocity', 'y_sea_water_velocity']})
o.set_config('drift', {'current_uncertainty': 0, 'wind_uncertainty': 0})
o.set_config('drift', {'current_uncertainty': 0, 'wind_uncertainty': 0, 'horizontal_diffusivity': 100})

#%%
# Seeding some particles
o.set_config('drift:vertical_mixing', False) # Keeping oil at fixed depths
o.set_config('drift:vertical_mixing', True)
o.set_config('processes:biodegradation', True)
o.set_config('biodegradation:method', 'decay_rate')

#%%
# Fast decay for droplets, and slow decay for slick
decay = {'biodegradation_decay_rate_slick': np.log(2)/timedelta(days=100).total_seconds(),
'biodegradation_decay_rate_droplet': np.log(2)/timedelta(days=3).total_seconds(),
'oil_type': 'GENERIC MEDIUM CRUDE', 'm3_per_hour': .5}
decay = {'biodegradation_decay_rate_slick': np.log(2)/timedelta(days=5).total_seconds(),
'biodegradation_decay_rate_droplet': np.log(2)/timedelta(days=1).total_seconds(),
'oil_type': 'GENERIC MEDIUM CRUDE', 'm3_per_hour': .5, 'diameter': 1e-9} # small droplets

#%%
# Seed 5000 oil elements at surface, and 5000 elements at 100m depth
o.seed_elements(lon=4, lat=60.0, z=0, number=5000, time=datetime.now(), **decay)
o.seed_elements(lon=4, lat=60.0, z=-100, number=5000, time=datetime.now(), **decay)
# Seed 500 oil elements at surface, and 500 elements at 50m depth
o.seed_elements(lon=4, lat=60.0, z=0, number=500, time=datetime.now(), **decay)
o.seed_elements(lon=4, lat=60.0, z=-50, number=500, time=datetime.now(), **decay)

#%%
# Running model
o.run(duration=timedelta(hours=24), time_step=900)
o.run(duration=timedelta(hours=72), time_step=3600)

#%%
# Print and plot results
# Plot results
o.plot_oil_budget(show_watercontent_and_viscosity=False, show_wind_and_current=False)

o.animation_profile(markersize='mass_oil', markersize_scaling=50, color='z', alpha=.5)

#%%
# .. image:: /gallery/animations/example_biodegradation_0.gif
8 changes: 5 additions & 3 deletions examples/example_chemicaldrift.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,11 +85,12 @@
print('mass degraded : {:.3f}'.format(m_deg * 1e-6),' g {:.3f}'.format(m_deg/m_tot*100),'%')
print('mass volatilized : {:.3f}'.format(m_vol * 1e-6),' g {:.3f}'.format(m_vol/m_tot*100),'%')


legend=['dissolved', '', 'SPM', 'sediment', '']

o.animation_profile(color='specie',
markersize=5,
markersize='mass',
markersize_scaling=30,
alpha=.5,
vmin=0,vmax=o.nspecies-1,
legend = legend,
legend_loc = 3,
Expand All @@ -100,7 +101,8 @@

o.animation(color='specie',
markersize='mass',
markersize_scaling=100,
markersize_scaling=30,
alpha=.5,
vmin=0,vmax=o.nspecies-1,
colorbar=False,
fast = True,
Expand Down
4 changes: 2 additions & 2 deletions examples/example_oilspill_seafloor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,13 +39,13 @@

#%%
# Running model with a small timestep to resolve the boyant rising
o.run(duration=timedelta(hours=2), time_step=60, time_step_output=60)
o.run(duration=timedelta(hours=6), time_step=60, time_step_output=60)

#%%
# Print and plot results
print(o)

o.animation_profile()
o.animation_profile(markersize='z', color='z')
#%%
# .. image:: /gallery/animations/example_oilspill_seafloor_0.gif

Expand Down
157 changes: 93 additions & 64 deletions opendrift/models/basemodel/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2508,6 +2508,7 @@ def animation(self,
compare=None,
compare_marker='o',
background=None,
alpha=1,
bgalpha=.5,
vmin=None,
vmax=None,
Expand All @@ -2529,7 +2530,7 @@ def animation(self,
lcs=None,
surface_only=False,
markersize=20,
markersize_scaling=1,
markersize_scaling=None,
origin_marker=None,
legend=None,
legend_loc='best',
Expand Down Expand Up @@ -2652,8 +2653,9 @@ def plot_timestep(i):
y_deactive[index_of_last_deactivated < i]])

if isinstance(markersize, str):
points.set_sizes(markersize_scaling *
(self.history[markersize][:, i] / self.history[markersize].compressed()[0]))
points.set_sizes(markersize_scaling * np.abs(self.history[markersize][:, i]))
#points.set_sizes(markersize_scaling *
# np.abs((self.history[markersize][:, i] / self.history[markersize].compressed()[0])))

if color is not False: # Update colors
points.set_array(colorarray[:, i])
Expand Down Expand Up @@ -2779,13 +2781,18 @@ def plot_timestep(i):
else:
c = []

if isinstance(markersize, str):
if markersize_scaling is None:
markersize_scaling = 20
markersize_scaling = markersize_scaling / np.abs(self.history[markersize]).max()

if isinstance(markersize, str):
points = ax.scatter([], [],
c=c,
zorder=10,
edgecolor=[],
cmap=cmap,
alpha=.4,
alpha=alpha,
vmin=vmin,
vmax=vmax,
label=legend[0],
Expand All @@ -2796,6 +2803,7 @@ def plot_timestep(i):
zorder=10,
edgecolor=[],
cmap=cmap,
alpha=alpha,
s=markersize,
vmin=vmin,
vmax=vmax,
Expand Down Expand Up @@ -3005,56 +3013,62 @@ def __save_or_plot_animation__(self, figure, plot_timestep, filename,
def animation_profile(self,
filename=None,
compare=None,
legend=['', ''],
markersize=5,
markersize=20,
markersize_scaling=None,
alpha=1,
fps=20,
color=None,
cmap=None,
vmin=None,
vmax=None,
legend=None,
legend_loc=None):
"""Animate vertical profile of the last run."""

start_time = datetime.now()

def plot_timestep(i):
"""Sub function needed for matplotlib animation."""
#plt.gcf().gca().set_title(str(i))
ax.set_title('%s UTC' % times[i])
if PlotColors:
points.set_offsets(
np.array(
[x[range(x.shape[0]), i].T, z[range(x.shape[0]),
i].T]).T)
points.set_offsets(np.c_[x[range(x.shape[0]), i], z[range(x.shape[0]), i]])
if color is not None and compare is None:
points.set_array(colorarray[:, i])
else:
points.set_data(x[range(x.shape[0]), i], z[range(x.shape[0]),
i])

points_deactivated.set_data(
points_deactivated.set_offsets(np.c_[
x_deactive[index_of_last_deactivated < i],
z_deactive[index_of_last_deactivated < i])
z_deactive[index_of_last_deactivated < i]])

if isinstance(markersize, str):
points.set_sizes(np.abs(markersize_scaling * self.history[markersize][:, i]))
#points.set_sizes(np.abs(markersize_scaling *
# (self.history[markersize][:, i] / self.history[markersize].compressed()[0])))

if compare is not None:
points_other.set_data(x_other[range(x_other.shape[0]), i],
z_other[range(x_other.shape[0]), i])
points_other_deactivated.set_data(
x_other_deactive[index_of_last_deactivated_other < i],
z_other_deactive[index_of_last_deactivated_other < i])
points_other.set_offsets(np.c_[x_other[range(x_other.shape[0]), i],
z_other[range(x_other.shape[0]), i]])
points_other_deactivated.set_offsets(np.c_[
x_other_deactive[index_of_last_deactivated_other < i],
z_other_deactive[index_of_last_deactivated_other < i]])
return points, points_deactivated, points_other,
else:
return points, points_deactivated,

PlotColors = (compare is None) and (legend != ['', ''])
if PlotColors:
if cmap is None:
cmap = 'jet'
if isinstance(cmap, str):
cmap = matplotlib.cm.get_cmap(cmap)
if cmap is None:
cmap = 'jet'
if isinstance(cmap, str):
cmap = matplotlib.cm.get_cmap(cmap)
if color is not False:
if isinstance(color, str):
colorarray = self.get_property(color)[0].T
if vmin is None:
vmin = colorarray.min()
vmax = colorarray.max()

if color is not False:
if isinstance(color, str):
colorarray = self.get_property(color)[0].T
markercolor = self.plot_comparison_colors[0]
if color is None:
c = markercolor
else:
c = []

# Set up plot
index_of_first, index_of_last = \
Expand All @@ -3071,40 +3085,50 @@ def plot_timestep(i):
times = self.get_time_array()[0]
index_of_last_deactivated = \
index_of_last[self.elements_deactivated.ID-1]
if PlotColors:
points = ax.scatter([], [],
c=[],
zorder=10,
edgecolor=[],
cmap=cmap,
s=markersize,
vmin=vmin,
vmax=vmax)

markers = []
if legend is None:
if compare is None:
legs = ['', '']
else:
legs = legend
if isinstance(markersize, str):
ms = None
else:
ms = markersize
if isinstance(markersize, str):
if markersize_scaling is None:
markersize_scaling = 20
markersize_scaling = markersize_scaling / np.abs(self.history[markersize]).max()

points = ax.scatter([], [],
c=c,
zorder=10,
edgecolor=[],
alpha=alpha,
cmap=cmap,
s=ms,
label=legs[0],
vmin=vmin,
vmax=vmax)

markers = []
if compare is None and legend is not None: # Empty points to get legend
for legend_index in np.arange(len(legend)):
if legend[legend_index] != '':
markers.append(
matplotlib.lines.Line2D(
[0], [0],
marker='o',
linewidth=0,
markeredgewidth=0,
markerfacecolor=cmap(legend_index /
(len(legend) - 1)),
markersize=10,
label=legend[legend_index]))
[0], [0], marker='o', linewidth=0, markeredgewidth=0,
markerfacecolor=cmap(legend_index /
(len(legend) - 1)),
markersize=10,
label=legend[legend_index]))

legend = list(filter(None, legend))
leg = ax.legend(markers, legend, loc=legend_loc)
leg.set_zorder(20)
else:
points = plt.plot([], [],
'.k',
label=legend[0],
markersize=markersize)[0]

# Plot deactivated elements, with transparency
points_deactivated = plt.plot([], [], '.k', alpha=.3)[0]
points_deactivated = ax.scatter([], [], c=[], zorder=10, edgecolor=[],
cmap=cmap, s=ms, vmin=vmin, vmax=vmax)
x_deactive = self.elements_deactivated.lon
z_deactive = self.elements_deactivated.z

Expand All @@ -3118,13 +3142,19 @@ def plot_timestep(i):
other = compare
z_other = other.get_property('z')[0].T
x_other = self.get_property('lon')[0].T
points_other = plt.plot(x_other[0, 0],
z_other[0, 0],
'.r',
label=legend[1],
markersize=markersize)[0]
points_other = ax.scatter([], [],
c='r',
zorder=10,
alpha=alpha,
edgecolor=[],
cmap=cmap,
s=markersize,
label=legs[1],
vmin=vmin,
vmax=vmax)

# Plot deactivated elements, with transparency
points_other_deactivated = plt.plot([], [], '.r', alpha=.3)[0]
points_other_deactivated = ax.scatter([], [], c='r', cmap=cmap, s=markersize, alpha=.3)
x_other_deactive = other.elements_deactivated.lon
z_other_deactive = other.elements_deactivated.z
firstlast = np.ma.notmasked_edges(x_other, axis=1)
Expand Down Expand Up @@ -3153,7 +3183,7 @@ def plot_timestep(i):
-zmin,
color='cornflowerblue'))

if legend != ['', ''] and PlotColors is False:
if legend is not None and compare is not None:
plt.legend(loc=4)

self.__save_or_plot_animation__(plt.gcf(),
Expand Down Expand Up @@ -4534,7 +4564,6 @@ def __save_animation__(self, fig, plot_timestep, filename, frames, fps,

if writer is not None:
with writer.saving(fig, filename, None):
print(frames)
for i in frames if isinstance(frames, (list, range)) else range(frames):
plot_timestep(i)
writer.grab_frame()
Expand Down

0 comments on commit 1ac6c96

Please sign in to comment.