import logging
import pandas as pd
model = "opendrift"
logger = logging.getLogger(model)
[docs]
def check_plots(which_plots, export_variables, drift_model):
"""Check that input plot options are valid,
particularly that the necessary export variables are included.
Parameters
----------
which_plots : dict
Dictionary of plot options.
export_variables : list
Variables to be saved to file.
drift_model : str
Which OpenDrift model is being used.
Raises
------
ValueError
If "all" is specified with other plot options.
ValueError
If necessary export variables are missing.
ValueError
If oil budget plot is requested with a drift model other than OpenOil.
"""
if (
len([k for k in which_plots.keys() if "oil" in k]) > 0
and drift_model != "OpenOil"
):
raise ValueError("Oil budget plot only available for OpenOil drift model")
# if "all", make sure it is the only plot option
# since it will create its own dictionary of plot options
if "all" in which_plots and len(which_plots) > 1:
raise ValueError(
"If 'all' is specified for plots, it must be the only plot option."
)
# check for cases that require an export variable
if export_variables is not None:
# flatten dict
which_plots_flat = pd.json_normalize(which_plots).to_dict(orient="records")
if len(which_plots_flat) == 0:
return
which_plots_flat = which_plots_flat[0]
plot_options_to_check = [
"linecolor",
"color",
"background",
"variable",
"markersize",
]
missing_variables = []
for plot_option in plot_options_to_check:
variables = [
v
for k, v in which_plots_flat.items()
if plot_option == k.split(".")[1] and v not in export_variables
]
missing_variables.extend(variables)
# handle oil budget separately
vars_needed = [
"x_wind",
"y_wind",
"x_sea_water_velocity",
"y_sea_water_velocity",
]
variables = [
list(set(vars_needed) - set(export_variables))
for k, v in which_plots_flat.items()
if "show_wind_and_current" in k
and v
and (len(set(vars_needed) - set(export_variables)) > 0)
]
if len(variables) > 0:
variables = variables[0]
missing_variables.extend(variables)
if len(missing_variables) > 0:
raise ValueError(
f"Missing export variables for the following plot options: {list(set(missing_variables))}"
)
[docs]
def make_filename_string(plot_name, filename, kwargs):
"""Create a filename string based on the kwargs."""
# modify filename with plot kwargs
if len(kwargs) > 0:
filename += "_" + "_".join(
[
f"{key}_{value}"
for key, value in kwargs.items()
if key not in ["filetype"]
]
)
# include filetype suffix
if "spaghetti" in plot_name or "oil" in plot_name or "property" in plot_name:
if "filetype" not in kwargs:
filename += ".png"
else:
filename += "." + kwargs["filetype"]
kwargs.pop("filetype")
elif "animation" in plot_name:
if "filetype" not in kwargs:
filename += ".mp4"
else:
filename += "." + kwargs["filetype"]
kwargs.pop("filetype")
else:
raise ValueError(f"Invalid plot name: {plot_name}")
return filename
[docs]
def plot(plot_name, input_kwargs, o, filename, drift_model):
"""Create a plot based on the plot name and input kwargs.
Parameters
----------
plot_name : str
Must contain key word to identify the plot type.
input_kwargs : dict
Optional plot kwargs.
o : _type_
OpenDrift simulation object.
filename : str
Filename substring to save the plot.
drift_model : str
Which OpenDrift model is being used.
"""
if "spaghetti" in plot_name or plot_name == "all":
filename += "_spaghetti"
# add input plot input_kwargs to the default kwargs
kwargs = {"fast": True}
kwargs.update(input_kwargs)
# modify filename with plot kwargs
filename = make_filename_string(plot_name, filename, kwargs)
o.plot(filename=filename, **kwargs)
elif (
"animation" in plot_name and "profile" not in plot_name
) or plot_name == "all":
# add input plot kwargs to the default kwargs
kwargs = {"fast": True, "fps": 4}
# selections for plotting landmask depend on whether the global
# landmask is being used ("auto landmask")
# if not using the global land mask, then plot the wetdry model
# landmask with these special inputs to look best.
# cannot plot land separate from simulation since then the reader
# is not available with the information. Checking for readers checks
# for that.
if (
o._config["general:use_auto_landmask"]["value"] == False
and len(o.env.readers.keys()) > 0
):
kwargs.update(
{
"background": "land_binary_mask",
"cmap": "cmo.deep",
"hide_landmask": True, # hides global land mask, which is not being used in this case
}
)
kwargs.update(input_kwargs)
# modify filename with plot kwargs
filename = make_filename_string(plot_name, filename, kwargs)
o.animation(filename=filename, **kwargs)
elif ("animation" in plot_name and "profile" in plot_name) or plot_name == "all":
filename += "_profile"
# add input plot kwargs to the default kwargs
kwargs = {"fps": 4}
kwargs.update(input_kwargs)
# modify filename with plot kwargs
filename = make_filename_string(plot_name, filename, kwargs)
o.animation_profile(filename=filename, **kwargs)
elif "oil" in plot_name or plot_name == "all":
filename += "_oil"
plot_name = "oil"
# add input plot kwargs to the default kwargs
kwargs = {}
kwargs = {
"show_wind_and_current": True,
"show_watercontent_and_viscosity": True,
}
kwargs.update(input_kwargs)
# modify filename with plot kwargs
filename = make_filename_string(plot_name, filename, kwargs)
if drift_model == "OpenOil":
o.plot_oil_budget(filename=filename, **kwargs)
else:
raise ValueError("Oil budget plot only available for OpenOil drift model")
elif "property" in plot_name: # or plot_name == "all":
if not "variable" in input_kwargs:
raise ValueError(
"Property plot must be specified as 'variable' in the plots dictionary."
)
filename += "_property"
plot_name = "property"
# add input plot kwargs to the default kwargs
kwargs = {}
kwargs.update(input_kwargs)
# modify filename with plot kwargs
filename = make_filename_string(plot_name, filename, kwargs)
o.plot_property(filename=filename, **kwargs)
logger.info(f"Saved plot to {filename}")
return filename
[docs]
def make_plots(which_plots, o, filename, drift_model):
"""Run through each plot key and make the plot.
This extra level of abstraction is necessary to allow for multiple
types of each plot to be made, for example, a spaghetti plot and a
spaghetti plot with tracks colored.
"""
if "all" in which_plots:
which_plots = {
"spaghetti": {},
"property_depth": {"variable": "z"},
"property_depth_mean": {"variable": "z", "mean": True},
"animation": {},
"animation_profile": {},
}
if drift_model == "OpenOil":
which_plots["oil"] = {}
for which_plot in which_plots:
filename_out = plot(
which_plot, which_plots[which_plot], o, filename, drift_model
)
# save filename for plot
which_plots[which_plot]["filename"] = filename_out
return which_plots
[docs]
def make_plots_after_simulation(output_filepath, plots="all"):
"""Make plots after a simulation has been run.
Parameters
----------
output_filepath : str
Path to the output file from the simulation.
which_plots : dict
Dictionary of plot options.
Returns
-------
dict
Dictionary of plot options with the filename of each plot.
"""
import opendrift as od
# load output file
o = od.open(str(output_filepath))
# want output_file to not include any suffix
output_filepath = (
str(output_filepath)
.replace(".nc", "")
.replace(".parquet", "")
.replace(".parq", "")
)
# figure out drift_model
drift_model = type(o).__name__
# make plots
plots = make_plots(plots, o, output_filepath, drift_model)
return plots