From 1ac6c96aa6c7b507e924feccb0c73c0fc1b1441a Mon Sep 17 00:00:00 2001 From: Knut-Frode Dagestad Date: Thu, 14 Mar 2024 17:47:19 +0100 Subject: [PATCH] The size and color of particles of animation and animation_profile methods can now be scaled with any element or environment property by specifying marker=. Sizes can eventually be scaled by spcifying markersize_scaling. Transparency (alpha) can also be provided. Some examples are updated. --- examples/example_biodegradation.py | 25 ++-- examples/example_chemicaldrift.py | 8 +- examples/example_oilspill_seafloor.py | 4 +- opendrift/models/basemodel/__init__.py | 157 +++++++++++++++---------- 4 files changed, 115 insertions(+), 79 deletions(-) diff --git a/examples/example_biodegradation.py b/examples/example_biodegradation.py index 303e21cfa..daca5a864 100755 --- a/examples/example_biodegradation.py +++ b/examples/example_biodegradation.py @@ -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 diff --git a/examples/example_chemicaldrift.py b/examples/example_chemicaldrift.py index 9a52f8f48..e350d1c30 100755 --- a/examples/example_chemicaldrift.py +++ b/examples/example_chemicaldrift.py @@ -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, @@ -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, diff --git a/examples/example_oilspill_seafloor.py b/examples/example_oilspill_seafloor.py index 52a3c4424..c9ec8307d 100755 --- a/examples/example_oilspill_seafloor.py +++ b/examples/example_oilspill_seafloor.py @@ -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 diff --git a/opendrift/models/basemodel/__init__.py b/opendrift/models/basemodel/__init__.py index 2a00bcdc3..53d3ee19d 100644 --- a/opendrift/models/basemodel/__init__.py +++ b/opendrift/models/basemodel/__init__.py @@ -2508,6 +2508,7 @@ def animation(self, compare=None, compare_marker='o', background=None, + alpha=1, bgalpha=.5, vmin=None, vmax=None, @@ -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', @@ -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]) @@ -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], @@ -2796,6 +2803,7 @@ def plot_timestep(i): zorder=10, edgecolor=[], cmap=cmap, + alpha=alpha, s=markersize, vmin=vmin, vmax=vmax, @@ -3005,13 +3013,15 @@ 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.""" @@ -3019,42 +3029,46 @@ def animation_profile(self, 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 = \ @@ -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 @@ -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) @@ -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(), @@ -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()