from __future__ import annotations from abc import abstractmethod from enum import Enum import io from typing import TYPE_CHECKING, Any, Literal, cast import matplotlib.pyplot as plt import numpy as np from contourpy import FillType, LineType, contour_generator if TYPE_CHECKING: from collections.abc import Sequence from matplotlib.axes import Axes from matplotlib.figure import Figure from contourpy._contourpy import ( CoordinateArray, FillReturn_OuterCode, LineReturn_SeparateCode, MaskArray, ) class PointType(Enum): CORNER = 0 LOWER = 1 UPPER = 2 class Corner(Enum): NW = 1 NE = 0 SW = 3 SE = 2 class Config: corner_mask: bool name: str quad_as_tri: bool show_text: bool arrowsize: float axes: Sequence[Axes] # Set in derived classes. axes_index: int fill_alpha: float fig: Figure | None # Set in derived classes. fontsize: float gap: float grid_kwargs: dict[str, Any] mask: MaskArray | Literal[False] marker_size: float text_gap: float title_fontsize: float x: CoordinateArray y: CoordinateArray def __init__( self, name: str, corner_mask: bool, quad_as_tri: bool, show_text: bool, ) -> None: self.name = name self.corner_mask = corner_mask self.quad_as_tri = quad_as_tri self.show_text = show_text self.arrowsize = 0.1 self.fontsize = 8 self.title_fontsize = 8 self.marker_size = 3.5 self.gap = 0.1 self.text_gap = self.gap/2 self.grid_kwargs = {"color": "gray", "linewidth": 0.35} self.fill_alpha = 0.2 self.x, self.y = np.meshgrid([0.0, 1.0], [0.0, 1.0]) self.mask = False import matplotlib as mpl mpl.use("Agg") def __del__(self) -> None: self.clear() def _arrow( self, ax: Axes, line_start: CoordinateArray, line_end: CoordinateArray, color: str, ) -> None: mid = 0.5*(line_start + line_end) along = line_end - line_start along /= np.sqrt(np.dot(along, along)) # Unit vector. right = np.asarray((along[1], -along[0])) arrow = np.stack(( mid - (along*0.5 - right)*self.arrowsize, mid + along*0.5*self.arrowsize, mid - (along*0.5 + right)*self.arrowsize)) ax.plot(arrow[:, 0], arrow[:, 1], "-", c=color) @abstractmethod def _decode_config(self, config: int, corner: Corner | None = None) -> tuple[int | None, ...]: pass @abstractmethod def _quad_lines( self, ax: Axes, z: np.ma.MaskedArray[Any, Any], zlower: float, zupper: float, corner: Corner | None, ) -> None: pass def _next_quad( self, config: int, corner: Corner | None = None, suffix: str = "", set_0: float = 0, set_1: float = 1, set_2: float = 2, ) -> None: zlower = 0.5 zupper = 1.5 ax = self.axes[self.axes_index] nw, ne, sw, se = self._decode_config(config, corner) # Quad edge. ax.axis("off") if corner is None: ax.plot([0, 1, 1, 0, 0], [0, 0, 1, 1, 0], **self.grid_kwargs) if self.quad_as_tri and not nw == ne == sw == se: ax.plot([0, 1], [0, 1], [0, 1], [1, 0], **self.grid_kwargs) title = f"{nw}{ne}{sw}{se}={config}{suffix}" elif corner == Corner.NW: ax.plot([0, 1, 0, 0], [0, 1, 1, 0], **self.grid_kwargs) title = f"{nw}{ne}{sw}={config}{suffix}" elif corner == Corner.NE: ax.plot([1, 1, 0, 1], [0, 1, 1, 0], **self.grid_kwargs) title = f"{nw}{ne}{se}={config}{suffix}" elif corner == Corner.SW: ax.plot([0, 1, 0, 0], [0, 0, 1, 0], **self.grid_kwargs) title = f"{nw}{sw}{se}={config}{suffix}" elif corner == Corner.SE: ax.plot([0, 1, 1, 0], [0, 0, 1, 0], **self.grid_kwargs) title = f"{ne}{sw}{se}={config}{suffix}" else: raise RuntimeError("Invalid corner") if self.show_text: ax.set_title(title, size=self.title_fontsize) # Axes bounds include gap for arrows and corner points. ax.set_xlim(-self.gap, 1.0+self.gap) ax.set_ylim(-self.gap, 1.0+self.gap) if self.show_text: # Text for corner and optional middle z-levels. fontsize = self.fontsize if corner != Corner.NW: ax.text(1+self.text_gap, 0, str(se), va="center", size=fontsize, ha="left") if corner != Corner.SW: ax.text(1+self.text_gap, 1, str(ne), va="center", size=fontsize, ha="left") if corner != Corner.NE: ax.text(-self.text_gap, 0, str(sw), va="center", size=fontsize, ha="right") if corner != Corner.SE: ax.text(-self.text_gap, 1, str(nw), va="center", size=fontsize, ha="right") if suffix: # Suffix (without brackets) is the z-level of quad middle. ax.text(0.5, 0.5, suffix[1:-1], va="center", size=fontsize, ha="center") lookup = {0: set_0, 1: set_1, 2: set_2, None: None} z_nw = lookup[nw] z_ne = lookup[ne] z_sw = lookup[sw] z_se = lookup[se] z = np.asarray([[z_sw, z_se], [z_nw, z_ne]], dtype=np.float64) z_masked = np.ma.array(z, mask=self.mask) # type: ignore[no-untyped-call] if suffix: zmean = z.mean() if suffix == "(0)": if zmean > zlower: raise RuntimeError( f"Inconsistent zlower {zlower} for config {config}{suffix} with middle_z " f"{zmean}") elif suffix in ["(1)", "(>0)", "(<2)"]: if zmean < zlower: raise RuntimeError( f"Inconsistent zlower {zlower} for config {config}{suffix} with middle_z " f"{zmean}") elif zmean > zupper: raise RuntimeError( f"Inconsistent zupper {zupper} for config {config}{suffix} with middle_z " f"{zmean}") elif suffix == "(2)": if zmean < zupper: raise RuntimeError( f"Inconsistent zupper {zupper} for config {config}{suffix} with middle_z " f"{zmean}") else: raise RuntimeError(f"Unexpected suffix {suffix} for config {config}") self._quad_lines(ax, z_masked, zlower, zupper, corner) self.axes_index += 1 def _set_mask(self, corner: Corner | None) -> None: if corner is None: self.mask = False else: self.mask = np.zeros(4, dtype=bool) self.mask[corner.value] = True self.mask.reshape((2, 2)) def clear(self) -> None: if self.fig: self.axes = [] plt.close(self.fig) self.fig = None def save_to_buffer(self) -> io.BytesIO: if not self.fig: raise RuntimeError("self.fig not set") buf = io.BytesIO() self.fig.savefig(buf, format="png") buf.seek(0) return buf def save_to_file(self, filename: str) -> None: if not self.fig: raise RuntimeError("self.fig not set") self.fig.savefig(filename) class ConfigFilledCommon(Config): def __init__(self, name: str, corner_mask: bool, quad_as_tri: bool, show_text: bool) -> None: super().__init__(name, corner_mask, quad_as_tri, show_text) def _decode_config(self, config: int, corner: Corner | None = None) -> tuple[int | None, ...]: if corner is None: nw = (config >> 6) & 0x3 ne = (config >> 4) & 0x3 sw = (config >> 2) & 0x3 se = config & 0x3 else: a = (config >> 4) & 0x3 b = (config >> 2) & 0x3 c = config & 0x3 if corner == Corner.NW: [nw, ne, sw, se] = [a, b, c, None] elif corner == Corner.NE: [nw, ne, sw, se] = [a, b, None, c] elif corner == Corner.SW: [nw, ne, sw, se] = [a, None, b, c] else: # Corner.SE [nw, ne, sw, se] = [None, a, b, c] if nw == 3 or ne == 3 or sw == 3 or se == 3: raise ValueError(f"Invalid config {config}") return nw, ne, sw, se def _quad_lines( self, ax: Axes, z: np.ma.MaskedArray[Any, Any], zlower: float, zupper: float, corner: Corner | None, ) -> None: cont_gen = contour_generator( self.x, self.y, z, name=self.name, corner_mask=self.corner_mask, quad_as_tri=self.quad_as_tri, fill_type=FillType.OuterCode) filled = cont_gen.filled(zlower, zupper) assert cont_gen.fill_type == FillType.OuterCode if TYPE_CHECKING: filled = cast(FillReturn_OuterCode, filled) lines = filled[0] # May be 0..2 polygons, and there cannot be any holes. for points in lines: n = len(points) ax.fill(points[:, 0], points[:, 1], c="C2", alpha=self.fill_alpha, ec=None) # Classify points, either corner, lower level or upper level. point_types = np.empty(n, dtype=PointType) for i in range(n-1): xy = points[i] x_on_quad_edge = (xy[0] in (0.0, 1.0)) y_on_quad_edge = (xy[1] in (0.0, 1.0)) if x_on_quad_edge and y_on_quad_edge: point_types[i] = PointType.CORNER elif corner is not None or x_on_quad_edge or y_on_quad_edge: if (self._interp(z, xy[0], xy[1], corner) < 0.5*(zlower + zupper)): point_types[i] = PointType.LOWER else: point_types[i] = PointType.UPPER else: # Point on quad diagonal from middle point, so same type as previous point. point_types[i] = point_types[i-1] # End point is the same as the start point. point_types[-1] = point_types[0] # Draw lines. for i in range(n-1): s = points[i] e = points[i+1] c = "C2" if point_types[i] == point_types[i+1] and point_types[i] != PointType.CORNER: c = "C3" if point_types[i] == PointType.LOWER else "C0" ax.plot([s[0], e[0]], [s[1], e[1]], "-", c=c) # Arrows on boundary and contour levels (the latter only if not quad_as_tri). if point_types[i] == point_types[i+1] and ( point_types[i] == PointType.CORNER or not self.quad_as_tri): self._arrow(ax, s, e, c) # Arrows on contour levels (lower and upper) for quad_as_tri. Only want a single # arrow on each line strip, so find and use the longest line segment. Care needed # identifying the longest due to rounding error, so find all segments within a small # length of the longest and take the max index of these. if self.quad_as_tri: for point_type in [PointType.LOWER, PointType.UPPER]: mask = np.equal(point_types, point_type) # type: ignore[call-overload] start_indices = np.nonzero(np.logical_and(mask[1:], ~mask[:-1]))[0] for start_index in start_indices: start_index = (start_index + 1) % (n-1) # Roll point_types array so start segment of line strip is at front. rolled = np.roll(point_types[:-1], -start_index) end_offset = np.nonzero(rolled != point_type)[0][0] pts = points[start_index:start_index+end_offset] # Line strip points. diff = np.diff(pts, axis=0) lengths = np.sum(np.square(diff), axis=1) # Segment lengths. offset = np.nonzero(lengths > lengths.max() - 1e-6)[0].max() i = (start_index + offset) % (n-1) c = "C3" if point_type == PointType.LOWER else "C0" self._arrow(ax, points[i], points[i+1], c) # Draw markers. for i in range(n-1): if point_types[i] != PointType.CORNER: c = "C3" if point_types[i] == PointType.LOWER else "C0" ax.plot(points[i][0], points[i][1], "o", c=c, ms=self.marker_size) def _interp( self, zquad: CoordinateArray, x: CoordinateArray, y: CoordinateArray, corner: Corner | None = None, ) -> float: # Interpolate zquad to determine value of z at (x, y). Could use # scipy.interp2d for this, but do not want scipy as a dependency just # for this. # (x, y) must lie on one of the edges of the quad. # zquad[0,0] = sw, zquad[0,1] = se, zquad[1,0] = nw, zquad[1,1] = ne ret: float if x == 0.0: ret = zquad[0, 0]*(1.0-y) + zquad[1, 0]*y elif x == 1.0: ret = zquad[0, 1]*(1.0-y) + zquad[1, 1]*y elif y == 0.0: ret = zquad[0, 0]*(1.0-x) + zquad[0, 1]*x elif y == 1.0: ret = zquad[1, 0]*(1.0-x) + zquad[1, 1]*x else: # corner is not None if corner in (Corner.NW, Corner.SE): ret = zquad[0, 0]*(1.0-y) + zquad[1, 1]*y else: ret = zquad[0, 1]*(1.0-y) + zquad[1, 0]*y return ret class ConfigFilled(ConfigFilledCommon): def __init__(self, name: str, quad_as_tri: bool = False, show_text: bool = True) -> None: super().__init__(name, False, quad_as_tri, show_text) subplot_kw = {"aspect": "equal"} if self.quad_as_tri: self.fig, axes = plt.subplots(18, 12, figsize=(10.4, 17.1), subplot_kw=subplot_kw) else: self.fig, axes = plt.subplots(9, 11, figsize=(9.4, 8.4), subplot_kw=subplot_kw) self.axes = axes.flatten() self.axes_index = 0 for config in range(171): try: (nw, ne, sw, se) = self._decode_config(config) except ValueError: continue assert nw is not None and ne is not None and sw is not None and se is not None all = np.asarray((nw, ne, sw, se)) # A quad is degenerate if all 4 quad edges include either the lower # or upper contour levels. degenerate_lower = (se == nw == 0 and sw > 0 and ne > 0) or \ (sw == ne == 0 and se > 0 and nw > 0) degenerate_upper = (se == nw == 2 and sw < 2 and ne < 2) or \ (sw == ne == 2 and se < 2 and nw < 2) if self.quad_as_tri: if not nw == ne == sw == se: zmax = all.max() zmin = all.min() zsum = all.sum() if zmin == 0 and zmax == 1: lookup_0 = {1: 0, 2: -1, 3: -1.8} self._next_quad(config, suffix="(0)", set_0=lookup_0[zsum]) lookup_1 = {1: 0.4, 2: 0.35, 3: 0} self._next_quad(config, suffix="(1)", set_0=lookup_1[zsum]) elif zmin == 0 and zmax == 2: count_1 = np.count_nonzero(all == 1) lookup_0 = {2: -0.3, 3: -1, 4: -2-count_1/2, 5: -5, 6: -6} # type: ignore[dict-item] self._next_quad(config, suffix="(0)", set_0=lookup_0[zsum]) if zsum <= 3: self._next_quad(config, suffix="(1)", set_2=3.01) elif zsum <= 5: self._next_quad(config, suffix="(1)") elif zsum == 6: self._next_quad(config, suffix="(1)", set_0=-1) lookup_2 = {2: 8.01, 3: 7.01, 4: 4+count_1/2, 5: 3.5, 6: 2.5} self._next_quad(config, suffix="(2)", set_2=lookup_2[zsum]) # type: ignore[arg-type] elif zmin == 1 and zmax == 2: lookup_1 = {5: 2, 6: 1.65, 7: 1.6} self._next_quad(config, suffix="(1)", set_2=lookup_1[zsum]) lookup_2 = {5: 4, 6: 3, 7: 2} self._next_quad(config, suffix="(2)", set_2=lookup_2[zsum]) # type: ignore[arg-type] else: raise RuntimeError(f"Invalid combination of zmin {zmin} and zmax {zmax}") else: self._next_quad(config) else: # !quad_as_tri if degenerate_lower and degenerate_upper: # middle needs to be 0, 1 and 2. self._next_quad(config, suffix="(0)", set_0=-1.01) self._next_quad(config, suffix="(1)") self._next_quad(config, suffix="(2)", set_2=3.01) elif degenerate_lower: # middle needs to be 0 and >0. self._next_quad(config, suffix="(0)", set_0=-0.01 if all.max() == 1 else -0.51) self._next_quad(config, suffix="(>0)", set_1=1.01) elif degenerate_upper: # middle needs to be <2 and 2. self._next_quad(config, suffix="(<2)") self._next_quad(config, suffix="(2)", set_2=2.01 if all.min() == 1 else 2.51) else: # Not degenerate. self._next_quad(config) # Hide unwanted axes. for i in range(self.axes_index, len(self.axes)): self.axes[i].axis("off") self.fig.tight_layout() class ConfigFilledCorner(ConfigFilledCommon): def __init__(self, name: str, show_text: bool = True) -> None: super().__init__(name, corner_mask=True, quad_as_tri=False, show_text=show_text) self.fig, axes = plt.subplots(8, 14, figsize=(12.1, 7.6), subplot_kw={"aspect": "equal"}) self.axes = axes.flatten() corners = [Corner.NW, Corner.NE, Corner.SW, Corner.SE] self.axes_index = 0 for icorner in range(4): corner = corners[icorner] self._set_mask(corner) for config in range(43): try: _ = self._decode_config(config, corner) except ValueError: continue self._next_quad(config, corner=corner) self.axes[self.axes_index].axis("off") self.axes_index += 1 self.fig.tight_layout() class ConfigLinesCommon(Config): def __init__(self, name: str, corner_mask: bool, quad_as_tri: bool, show_text: bool) -> None: super().__init__(name, corner_mask, quad_as_tri, show_text) def _decode_config(self, config: int, corner: Corner | None = None) -> tuple[int | None, ...]: if corner is None: nw = (config >> 3) & 0x1 ne = (config >> 2) & 0x1 sw = (config >> 1) & 0x1 se = config & 0x1 else: a = (config >> 2) & 0x1 b = (config >> 1) & 0x1 c = config & 0x1 if corner == Corner.NW: [nw, ne, sw, se] = [a, b, c, None] elif corner == Corner.NE: [nw, ne, sw, se] = [a, b, None, c] elif corner == Corner.SW: [nw, ne, sw, se] = [a, None, b, c] else: # Corner.SE [nw, ne, sw, se] = [None, a, b, c] return nw, ne, sw, se def _quad_lines( self, ax: Axes, z: np.ma.MaskedArray[Any, Any], zlower: float, zupper: float, corner: Corner | None, ) -> None: cont_gen = contour_generator( self.x, self.y, z, name=self.name, corner_mask=self.corner_mask, quad_as_tri=self.quad_as_tri, line_type=LineType.SeparateCode) lines_and_codes = cont_gen.lines(zlower) assert cont_gen.line_type == LineType.SeparateCode if TYPE_CHECKING: lines_and_codes = cast(LineReturn_SeparateCode, lines_and_codes) lines = lines_and_codes[0] for points in lines: ax.plot(points[:, 0], points[:, 1], "o-", c="C3", ms=self.marker_size) # Single arrow in middle segment (if even number of points) or the following # segment (if odd number of points). n = len(points) i = (n-1) // 2 self._arrow(ax, points[i], points[i+1], "C3") class ConfigLines(ConfigLinesCommon): def __init__(self, name: str, quad_as_tri: bool = False, show_text: bool = True) -> None: super().__init__(name, False, quad_as_tri, show_text) subplot_kw = {"aspect": "equal"} if self.quad_as_tri: self.fig, axes = plt.subplots(5, 6, figsize=(5.2, 4.75), subplot_kw=subplot_kw) else: self.fig, axes = plt.subplots(3, 6, figsize=(5.2, 2.85), subplot_kw=subplot_kw) self.axes = axes.flatten() self.axes_index = 0 for config in range(16): (nw, ne, sw, se) = self._decode_config(config) assert nw is not None and ne is not None and sw is not None and se is not None # A quad is degenerate if all 4 quad edges include the z level. degenerate = (se == nw == 0 and sw > 0 and ne > 0) or \ (sw == ne == 0 and se > 0 and nw > 0) if degenerate or (self.quad_as_tri and not nw == ne == sw == se): if self.quad_as_tri: count_1 = np.count_nonzero((nw, ne, sw, se)) count_0 = 4 - count_1 self._next_quad(config, suffix="(0)", set_0=1.01 - count_1) self._next_quad(config, suffix="(1)", set_1=0.01 + count_0) else: self._next_quad(config, suffix="(0)", set_1=0.99) self._next_quad(config, suffix="(1)", set_1=1.01) else: self._next_quad(config) self.fig.tight_layout() class ConfigLinesCorner(ConfigLinesCommon): # All 4 corners plotted together. def __init__(self, name: str, show_text: bool = True) -> None: super().__init__(name, corner_mask=True, quad_as_tri=False, show_text=show_text) self.fig, axes = plt.subplots(4, 8, figsize=(6.93, 3.8), subplot_kw={"aspect": "equal"}) self.axes = axes.flatten() self.axes_index = 0 corners = [Corner.NW, Corner.NE, Corner.SW, Corner.SE] for icorner in range(4): corner = corners[icorner] self._set_mask(corner) for config in range(8): self._next_quad(config, corner) self.fig.tight_layout()