from itertools import count
class Data:
"""
Base class for all data quantities in the pipeline.
:param key: Unique identifier for the field.
:type key: str
:param symbol: Symbol for labels (e.g., r"$\rho$").
:type symbol: str
:param plot_coords: [row, col] position in the subplot grid, defaults to [0, 0].
:type plot_coords: list[int], optional
:param vmin: Minimum value for manual scaling, defaults to None.
:type vmin: float, optional
:param vmax: Maximum value for manual scaling, defaults to None.
:type vmax: float, optional
:param kwargs:
* **title** (str): Custom title for the plot. Defaults to `symbol`.
* **id** (str): Unique ID to distinguish instances of the same field nature.
* **xmin** (float): Minimum x-axis bound.
* **xmax** (float): Maximum x-axis bound.
* **ymin** (float): Minimum y-axis bound.
* **ymax** (float): Maximum y-axis bound.
* **xscale** (str): X-axis scaling type, e.g., 'linear' or 'log'.
* **yscale** (str): Y-axis scaling type, e.g., 'linear' or 'log'.
* **xlabel** (str): X-axis label.
* **ylabel** (str): Y-axis label.
* **style_kwargs** (dict): Style options forwarded to plotting calls.
* **parts_kwargs** (dict): Style options forwarded to particles plotting calls.
* **parts_color** (callable): Optional callable `parts_color(commonvtk)` returning a sequence of colors.
* **ref_function** (callable): Analytical function for comparison.
* **compute** (callable): Custom function to calculate new fields on the fly.
* **customize** (callable): If really what you want is not implemented, customize(ax, vtk) will do whatver you want on the corresponding ax.
"""
def __init__(self, key, symbol="", plot_coords=None, bounds=None, **kwargs):
self.key = key
self.symbol = symbol
self.plot_coords = plot_coords if plot_coords else [0, 0]
self.bounds = bounds if bounds else [None, None]
if self.bounds[0] is not None or self.bounds[1] is not None:
self.bounds_set = True
else:
self.bounds_set = False
self.title = kwargs.get(
"title", None
) # if None, will be replaced by symbol in ax
self.id = kwargs.get(
"id", None
) # some custom id, to distinguish different instances of the same field nature (for example tau)
self.xmin = kwargs.get("xmin", None)
self.xmax = kwargs.get("xmax", None)
self.ymin = kwargs.get("ymin", None)
self.ymax = kwargs.get("ymax", None)
self.xscale = kwargs.get("xscale", "linear")
self.yscale = kwargs.get("yscale", "linear")
# heatmaps have a `norm` attribute
self.xlabel = kwargs.get("xlabel", None)
self.ylabel = kwargs.get("ylabel", None)
self.style_kwargs = kwargs.get("style_kwargs", {})
default_parts_kwargs = {"marker": "x", "markersize": 0.2}
self.parts_kwargs = merge_default_to_dict(
default_parts_kwargs, kwargs.get("parts_kwargs", {})
)
self.parts_color = kwargs.get("parts_color", None)
self.points = []
self.values = []
self.ref_function = kwargs.get("ref_function", None)
self.pointsRef = []
self.valuesRef = []
default_ref_plot_kwargs = {
"zorder": 3,
"ls": "--",
"lw": 1,
"alpha": 0.8,
"color": "limegreen",
"label": "Reference",
}
if self.ref_function is not None:
if not hasattr(self.ref_function, "plot_kwargs"):
self.ref_function.plot_kwargs = {}
self.ref_function.plot_kwargs = merge_default_to_dict(
default_ref_plot_kwargs, self.ref_function.plot_kwargs
)
self.compute = kwargs.get("compute", None)
self.customize = kwargs.get("customize", None)
self.label_func = kwargs.get("label_func", None)
def set_bounds(self, bounds):
self.bounds = bounds
def set_ref_data(self, points, values):
self.pointsRef = points
self.valuesRef = values
def set_data(self, points, values):
self.points = points
self.values = values
def set_norm(self, norm):
self.norm = norm
supported_norms = ["linear", "log", "TwoSlopeNorm"]
if norm not in supported_norms:
raise Exception(
f"{norm} not implemented. Supported norms: {supported_norms}"
)
def set_default_xlabel(self, xlabel):
if self.xlabel is None:
self.xlabel = xlabel
def set_default_ylabel(self, ylabel):
if self.ylabel is None:
self.ylabel = ylabel
def __str__(self):
return self.key
[docs]
class MapMovie2D(Data):
r"""
2D spatial field :math:`f(x, z, t)` rendered as a heatmap (pcolormesh) animation.
"""
def __init__(
self,
key,
symbol="",
plot_coords=None,
norm="linear",
streamlines=None,
uids=None,
**kwargs,
):
r"""
Initializes a 2D movie field.
(Refer to :class:`Data` for base parameters)
:param norm: Colorbar scaling. Options usually include 'linear', 'log', or 'TwoSlopeNorm'.
Defaults to "linear".
:type norm: str, optional
:param streamlines: A list of two Idefix field keys used to show vector streamlines,
e.g., ``["VX1", "VX2"]``. Defaults to None.
:type streamlines: list[str], optional
:param uids: List of the particles uid. Their trajectories will be showed over the maps. To show every particle, set it to "all".
e.g., ``[1,2]``. Defaults to None.
:type uids: list[int] | Literal["all"] | None, optional
:param \**kwargs: Additional rendering options.
:keyword streamline_kwargs (dict): kwargs that will be passed to streamplot.
:keyword contours (Sequence[float] | None): Contour levels used to draw contour lines over the pcolormesh for this field. Defaults to None.
:keyword contour_color (str): Color of the contour lines. Defaults to "green".
"""
# streamlines should be a list like ["VX1", "VX2"]
super().__init__(key, symbol, plot_coords, **kwargs)
self.set_norm(norm)
self.streamlines = streamlines
if streamlines is not None:
if not isinstance(streamlines, (list, tuple)) or not len(streamlines) == 2:
raise Exception(
f"Invalid streamline configuration: {streamlines}. Expected a list/tuple of length 2."
)
default_streamline_kwargs = {
"linewidth": 0.2,
"arrowstyle": "->",
"color": "#d3d3d3",
# "color": (1, 1, 1, 0.5),
"density": 2,
}
self.streamline_kwargs = merge_default_to_dict(
default_streamline_kwargs, kwargs.get("streamline_kwargs", {})
)
self.contours = kwargs.get("contours", None)
self.contour_color = kwargs.get("contour_color", "green")
self.uids = uids
self.is_movie = True
self.is_timeline = False
def set_XYgrid(self, X, Y):
"""
Assign the spatial cartesian grid used for rendering the 2D pcolormesh.
:param X: 2D array of horizontal coordinates.
:type X: numpy.ndarray
:param Y: 2D array of vertical coordinates.
:type Y: numpy.ndarray
"""
self.X, self.Y = X, Y
class Field1D(Data):
"""
Base class for 1D fields :math:`f(x, t)`.
Increments a global counter for indexing in results arrays.
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.is_timeline = False
self.is_movie = True
[docs]
class LineMovie1D(Field1D):
"""
For :math:`f(x, t)` fields, renders as a line plot :math:`f(x, t)` that updates every frame.
"""
def __init__(
self,
key,
symbol="",
plot_coords=None,
bounds=None,
uids=None,
**kwargs,
):
super().__init__(key, symbol, plot_coords, bounds, **kwargs)
self.uids = uids
self.is_movie = True
self.is_timeline = False
def set_localqty(self, localqty):
self.localqty = localqty
[docs]
class SpaceTimeHeatmap(Field1D):
"""
For :math:`f(x, t)` fields, renders a space-time heatmap.
:keyword cmap: Colormap for the heatmap.
:keyword uids: List of particles' uids which trajectories will be displayed.
"""
instances = count(1)
def __init__(
self,
key,
symbol="",
plot_coords=None,
bounds=None,
norm="linear",
uids=None,
**kwargs,
):
super().__init__(key, symbol, plot_coords, bounds, **kwargs)
self.set_norm(norm)
self.uids = uids
self.is_timeline = True
self.is_movie = False
class OneComponentOneVariable(Data):
"""
A y(x) value where x can be any variable. If xqty is None, that means x is time and the quantity will be treated as a timeline.
Otherwise, it will be treated as a LineMovie1D.
"""
def __init__(
self,
key,
symbol="",
plot_coords=None,
bounds=None,
xqty=None,
**kwargs,
):
super().__init__(key, symbol, plot_coords, bounds, **kwargs)
if kwargs.get("uids", None) is not None:
raise Exception(
"For uid specific 1C1V quantity, please use PartQuantity instead."
)
self.uids = None
self.xqty = xqty # if None, it will be time.
self.is_timeline = xqty is None
self.is_movie = xqty is not None
[docs]
class PartQuantity(Data):
"""
Particular case of OneComponentOneVariable when the variable is time and that there is one value per particle (so not really one component but rather Npart components...)
Tracks Lagrangian particle properties over time.
:keyword: uids (optional) the ids of the particles wanted.
Defaults to "all" (all particles)
"""
def __init__(
self,
key,
symbol="",
plot_coords=None,
bounds=None,
xqty=None,
uids="all",
**kwargs,
):
super().__init__(key, symbol, plot_coords, bounds, **kwargs)
self.uids = uids
self.is_global = kwargs.get("is_global", False)
self.xqty = xqty # if None, it will be time.
self.is_timeline = xqty is None
self.is_movie = not self.is_timeline
self.colors = kwargs.get("colors", [])
class LocalQuantity(Data):
"""
Particle case of PartQuantity.
To access the quantity of the cell on which the particle is (No interpolation atm)
:keyword: uids (optional) the ids of the particles wanted.
Defaults to "all" (all particles)
"""
def __init__(
self,
key,
localkey,
symbol="",
plot_coords=None,
bounds=None,
uids="all",
**kwargs,
):
if kwargs.get("compute"):
raise ValueError("compute is not allwed for LocalQuantity")
super().__init__(key, symbol, plot_coords, bounds, **kwargs)
self.localkey = localkey
self.uids = uids
self.is_global = kwargs.get("is_global", False)
self.is_timeline = True
self.is_movie = False
self.colors = kwargs.get("colors", [])
def merge_default_to_dict(default_dict, final_dict):
for key, value in default_dict.items():
if key not in final_dict:
final_dict[key] = value
return final_dict