"""
FRAGMENT-MNP Output
===================
Provides functionality for processing and visulalising model output data.
"""
import uuid
import numpy as np
import numpy.typing as npt
import matplotlib.pyplot as plt
[docs]class FMNPOutput():
"""
Class that holds output data from the model and provides plotting
functionalities.
Parameters
----------
t : np.ndarray, shape (n_timesteps,)
Time series over which the model was run
c: np.ndarray, shape (n_size_classes, n_timesteps)
Mass concentrations for each size class over the time
series
n : np.ndarray, shape (n_size_classes, n_timesteps)
Particle number concentrations for each size class over
the time series
c_diss : np.ndarray, shape (n_size_classes, n_timesteps)
Mass concentrations of dissolved organics
n_diss : np.ndarray, shape (n_size_classes, n_timesteps)
Particle number concentrations lost from size classes due
to dissolution
soln : Bunch object from scipy.integrate.solve_ivp return
The solution to the model ODE, passed directly from the
scipy.integrate.solve_ivp method
psd : np.ndarray, shape (n_size_classes, )
Particle size distribution - the average diameters of
each of the particle size classes
"""
__slots__ = ['t', 'c', 'n', 'c_diss', 'n_diss', 'n_timesteps',
'n_size_classes', 'soln', 'psd', 'id']
def __init__(self,
t: npt.NDArray,
c: npt.NDArray,
n: npt.NDArray,
c_diss: npt.NDArray,
n_diss: npt.NDArray,
soln, psd,
id=None) -> None:
"""
Initialise the output data object
"""
# Set the main output data
self.t = t
self.c = c
self.n = n
self.c_diss = c_diss
self.n_diss = n_diss
self.soln = soln
self.psd = psd
# Save the number of timesteps and size classes
self.n_timesteps = self.t.shape[0]
self.n_size_classes = self.c.shape[0]
# Set the ID based on what we've been given, or
# give a unique ID
if id is None:
self.id = uuid.uuid4()
else:
self.id = id
[docs] def plot(self,
type: str = 'mass_conc',
plot_dissolution: bool = False,
log_yaxis=False,
units=None,
cmap='viridis',
show_legend=True,
size_classes_to_plot=None):
"""
Plot the output data by choosing from a number of
pre-defined plot types.
Parameters
----------
type : str, default='mass_conc'
Tells the function what type of plot to produce.
Either `particle_number_conc` or `mass_conc`.
plot_dissolution : bool, default=False
Should dissolution be plotted on a separate y-axis
log_yaxis: bool or str, default=False
True and "log" plots the y axis on a log scale,
"symlog" plots the y axis on a symlog scale (useful
if the data contain zeros), and False plots the y
axis on a linear scale
units: dict or str, default=None
Units to be used in axis labels. Must either be "SI" to
use SI units, "dim" to use dimensional labels ("mass",
"volume", etc), or a dictionary containing elements for
`mass`, `volume` and `time`. E.g. `units={'mass': 'kg',
'volume': 'm3', 'time': 's', 'length': 'nm'}`. Note that
these units are used purely for axis labelling and are
not used to modify the model output data (which is unit
agnostic). If `None` (the default), no units are added to
labels.
cmap: str, default='viridis'
The colormap to use for the plot. Must be one of the
colormaps `available in matplotlib
<https://matplotlib.org/stable/gallery/color/colormap_reference.html>`.
Note that these are case-sensitive.
show_legend: bool, default=True
Should size classes be shown on a legend?
size_classes_to_plot: list of ints, default=None
Only plot specific size classes by providing a list of
size class indices to plot, where 0 is the index of the
smallest size class. By default, all size classes are
plotted.
Returns
-------
(matplotlib.figure.Figure, matplotlib.axes.Axes)
Matplotlib figure and axes objects for the plot
"""
unit_labels = self._construct_units(units)
# Are we plotting number or mass concentrations? Set the
# y values and labels accordingly
if type == 'particle_number_conc':
# Set the yaxis label
ylabel = 'Particle number concentration'
if unit_labels is not None:
ylabel += f' [{unit_labels["number_conc"]}]'
# Set the yaxis values, transposed to pass to mpl
yvals = self.n.T
elif type == 'mass_conc':
ylabel = 'Mass concentration'
if unit_labels is not None:
ylabel += f' [{unit_labels["mass_conc"]}]'
yvals = self.c.T
else:
# Raise an error if an invalid plot type is given
raise ValueError(f'Invalid option for plot `type`: {type}. '
'Should be `particle_number_conc` or'
'`mass_conc`.')
# Only plot the size classes we have been asked to
if size_classes_to_plot is not None:
yvals = yvals[:, size_classes_to_plot]
# Set the xaxis label
xlabel = 'Time'
if unit_labels is not None:
xlabel += f' [{unit_labels["time"]}]'
# Change the colourmap to something useful (Viridis)
cmap_ = plt.colormaps[cmap]
plt.rcParams['axes.prop_cycle'] = \
plt.cycler('color',
cmap_(np.linspace(0, 1, yvals.shape[1])))
# Create the figure and axes - we need to this because we need to
# add a second axis for the dissolution data
fig, ax1 = plt.subplots()
# Construct the xlabel with(out) units
xlabel = 'Time'
if unit_labels is not None:
xlabel += f' [{unit_labels["time"]}]'
# Add the x and y labels
ax1.set_xlabel(xlabel)
ax1.set_ylabel(ylabel)
# Should the y axis be logged?
if log_yaxis in [True, 'log']:
ax1.set_yscale('log')
elif log_yaxis == 'symlog':
ax1.set_yscale('symlog')
# Plot the concentrations
ax1.plot(self.t, yvals)
# Add a size class legend, if requested
if show_legend:
# Construct the legend with(out) units
legend = [f'{d:<1g}' for d in self.psd]
if unit_labels is not None:
legend = [f'{sc} {unit_labels["length"]}' for sc in legend]
legend = np.array(legend)
# Only show the requested size classes on the legend
if size_classes_to_plot is not None:
legend = legend[size_classes_to_plot]
ax1.legend(legend)
# Create and format the dissolution y axis
if plot_dissolution:
ax2 = ax1.twinx()
# Construct the ylabel
ylabel_diss = 'Dissolution mass concentration'
if unit_labels is not None:
ylabel_diss += f' [{unit_labels["mass_conc"]}]'
ax2.set_ylabel(ylabel_diss)
ax2.plot(self.t, np.sum(self.c_diss, axis=0), '--')
# Always show the dissolution legend to make it clear
# which line is dissolution
ax2.legend(['Dissolution'])
# Should the y axis be logged?
if log_yaxis in [True, 'log']:
ax2.set_yscale('log')
elif log_yaxis == 'symlog':
ax2.set_yscale('symlog')
return fig, (ax1, ax2)
else:
return fig, ax1
def _construct_units(self, units):
"""
Construct strings for axis labels based on the units
dictionary or string provided. Options given in `plot()`
method.
"""
units_out = {}
if (type(units) == dict
and {'mass', 'volume', 'time', 'length'} <= units.keys()):
# If we've been provided a dict of units, use these
units_out = units.copy()
elif type(units) == str and units.lower() == 'si':
# Use SI units
units_out['mass'] = 'kg'
units_out['volume'] = 'm3'
units_out['time'] = 's'
units_out['length'] = 'm'
elif type(units) == str and units.lower() == 'dim':
# Use dimensional labels as the units
units_out['mass'] = 'mass'
units_out['volume'] = 'volume'
units_out['time'] = 'time'
units_out['length'] = 'length'
elif units is None:
# Nothing in, nothing out...
return None
else:
raise ValueError('`units` parameter of `plot` function must be '
'a dictionary with keys {`mass`, `volume`, `time` '
'and `length`}, `SI` to denote use of SI '
'units, or `dim` to denote use of dimensions.')
# Construct the required compound units
units_out['mass_conc'] = f'{units_out["mass"]}/{units_out["volume"]}'
units_out['number_conc'] = f'/{units_out["volume"]}'
units_out['rate'] = f'/{units_out["time"]}'
# Return the unit strings
return units_out
