from __future__ import annotations
import numpy as np
from matplotlib.cm import ScalarMappable
from pylbo.visualisation.modes.cartesian_3d import CartesianSlicePlot3D
from pylbo.visualisation.modes.mode_data import ModeVisualisationData
[docs]
class CylindricalSlicePlot3D(CartesianSlicePlot3D):
"""
Class for handling cylindrical 3D plots of the eigenmode solution.
Parameters
----------
data : ModeVisualisationData
The data for the visualisation.
u2 : np.ndarray
The :math:`\\theta` coordinate of the eigenmode solution.
u3 : np.ndarray
The :math:`z` coordinate of the eigenmode solution.
time : float
The time at which the eigenmode solution is calculated.
slicing_axis : str
The axis along which the eigenmode solution is sliced.
figsize : tuple[int, int]
The size of the figure.
vmin : float
The minimum value of the colourbar. If None, the minimum value of the
solution is used.
vmax : float
The maximum value of the colourbar. If None, the maximum value of the
solution is used.
"""
def __init__(
self,
data: ModeVisualisationData,
u2: np.ndarray,
u3: np.ndarray,
time: float,
slicing_axis: str,
figsize: tuple[int, int],
vmin: float = None,
vmax: float = None,
**kwargs,
) -> None:
if figsize is None:
figsize = (8, 8)
if slicing_axis == "theta":
raise NotImplementedError(
"3D slicing is not implemented for a theta slice."
)
super().__init__(data, u2, u3, time, slicing_axis, figsize, **kwargs)
self.vmin = np.min(self._solutions) if vmin is None else vmin
self.vmax = np.max(self._solutions) if vmax is None else vmax
self.set_contours(levels=25, fill=True)
[docs]
def set_plot_arrays(self) -> None:
self.solution_shape = (len(self._u1), len(self._u2))
for ef, omega in zip(self.data.eigenfunction, self.data.omega):
data = np.broadcast_to(ef, shape=reversed(self.solution_shape)).transpose()
self.ef_data.append({"ef": data, "omega": omega})
r_2d, theta_2d = np.meshgrid(self.data.ds.ef_grid, self._u2, indexing="ij")
self.u1_data = r_2d
self.u2_data = theta_2d
self.u3_data = self._u3
self.time_data = self._time
[docs]
def draw_solution(self) -> None:
level_kwargs = {}
if self._contour_levels is not None:
level_kwargs["levels"] = self._contour_levels
for i, z in enumerate(self._u3):
self._view[i] = self._contour_recipe(
self.u1_data * np.cos(self.u2_data),
self.u1_data * np.sin(self.u2_data),
self.solutions[..., i],
zdir="z",
offset=z,
alpha=max(0.4, 1 - i * 0.1),
vmin=self.vmin,
vmax=self.vmax,
**level_kwargs,
**self._kwargs,
)
self.cbar = self.fig.colorbar(
ScalarMappable(norm=self._view[0].norm, cmap=self._view[0].cmap),
cax=self.cbar_ax,
orientation="horizontal",
)
xmax = np.max(self._u1)
self.ax.set_xlim(-xmax, xmax)
self.ax.set_ylim(-xmax, xmax)
self.ax.set_zlim(np.min(self._u3), np.max(self._u3))
[docs]
def get_view_xlabel(self) -> str:
return "x"
[docs]
def get_view_ylabel(self) -> str:
return "y"