.venv/lib/python3.12/site-packages/contourpy/__init__.py

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from contourpy._contourpy import (
    ContourGenerator,
    FillType,
    LineType,
    Mpl2005ContourGenerator,
    Mpl2014ContourGenerator,
    SerialContourGenerator,
    ThreadedContourGenerator,
    ZInterp,
    max_threads,
)
from contourpy._version import __version__
from contourpy.chunk import calc_chunk_sizes
from contourpy.convert import convert_filled, convert_lines
from contourpy.dechunk import dechunk_filled, dechunk_lines
from contourpy.enum_util import as_fill_type, as_line_type, as_z_interp

if TYPE_CHECKING:
    from typing import Any

    from numpy.typing import ArrayLike

    from ._contourpy import CoordinateArray, MaskArray

__all__ = [
    "__version__",
    "contour_generator",
    "convert_filled",
    "convert_lines",
    "dechunk_filled",
    "dechunk_lines",
    "max_threads",
    "FillType",
    "LineType",
    "ContourGenerator",
    "Mpl2005ContourGenerator",
    "Mpl2014ContourGenerator",
    "SerialContourGenerator",
    "ThreadedContourGenerator",
    "ZInterp",
]


# Simple mapping of algorithm name to class name.
_class_lookup: dict[str, type[ContourGenerator]] = {
    "mpl2005": Mpl2005ContourGenerator,
    "mpl2014": Mpl2014ContourGenerator,
    "serial": SerialContourGenerator,
    "threaded": ThreadedContourGenerator,
}


def _remove_z_mask(
    z: ArrayLike | np.ma.MaskedArray[Any, Any] | None,
) -> tuple[CoordinateArray, MaskArray | None]:
    # Preserve mask if present.
    z_array = np.ma.asarray(z, dtype=np.float64)  # type: ignore[no-untyped-call]
    z_masked = np.ma.masked_invalid(z_array, copy=False)  # type: ignore[no-untyped-call]

    if np.ma.is_masked(z_masked):  # type: ignore[no-untyped-call]
        mask = np.ma.getmask(z_masked)  # type: ignore[no-untyped-call]
    else:
        mask = None

    return np.ma.getdata(z_masked), mask  # type: ignore[no-untyped-call]


def contour_generator(
    x: ArrayLike | None = None,
    y: ArrayLike | None = None,
    z: ArrayLike | np.ma.MaskedArray[Any, Any] | None = None,
    *,
    name: str = "serial",
    corner_mask: bool | None = None,
    line_type: LineType | str | None = None,
    fill_type: FillType | str | None = None,
    chunk_size: int | tuple[int, int] | None = None,
    chunk_count: int | tuple[int, int] | None = None,
    total_chunk_count: int | None = None,
    quad_as_tri: bool = False,
    z_interp: ZInterp | str | None = ZInterp.Linear,
    thread_count: int = 0,
) -> ContourGenerator:
    """Create and return a :class:`~contourpy._contourpy.ContourGenerator` object.

    The class and properties of the returned :class:`~contourpy._contourpy.ContourGenerator` are
    determined by the function arguments, with sensible defaults.

    Args:
        x (array-like of shape (ny, nx) or (nx,), optional): The x-coordinates of the ``z`` values.
            May be 2D with the same shape as ``z.shape``, or 1D with length ``nx = z.shape[1]``.
            If not specified are assumed to be ``np.arange(nx)``. Must be ordered monotonically.
        y (array-like of shape (ny, nx) or (ny,), optional): The y-coordinates of the ``z`` values.
            May be 2D with the same shape as ``z.shape``, or 1D with length ``ny = z.shape[0]``.
            If not specified are assumed to be ``np.arange(ny)``. Must be ordered monotonically.
        z (array-like of shape (ny, nx), may be a masked array): The 2D gridded values to calculate
            the contours of.  May be a masked array, and any invalid values (``np.inf`` or
            ``np.nan``) will also be masked out.
        name (str): Algorithm name, one of ``"serial"``, ``"threaded"``, ``"mpl2005"`` or
            ``"mpl2014"``, default ``"serial"``.
        corner_mask (bool, optional): Enable/disable corner masking, which only has an effect if
            ``z`` is a masked array. If ``False``, any quad touching a masked point is masked out.
            If ``True``, only the triangular corners of quads nearest these points are always masked
            out, other triangular corners comprising three unmasked points are contoured as usual.
            If not specified, uses the default provided by the algorithm ``name``.
        line_type (LineType or str, optional): The format of contour line data returned from calls
            to :meth:`~contourpy.ContourGenerator.lines`, specified either as a
            :class:`~contourpy.LineType` or its string equivalent such as ``"SeparateCode"``.
            If not specified, uses the default provided by the algorithm ``name``.
        fill_type (FillType or str, optional): The format of filled contour data returned from calls
            to :meth:`~contourpy.ContourGenerator.filled`, specified either as a
            :class:`~contourpy.FillType` or its string equivalent such as ``"OuterOffset"``.
            If not specified, uses the default provided by the algorithm ``name``.
        chunk_size (int or tuple(int, int), optional): Chunk size in (y, x) directions, or the same
            size in both directions if only one value is specified.
        chunk_count (int or tuple(int, int), optional): Chunk count in (y, x) directions, or the
            same count in both directions if only one value is specified.
        total_chunk_count (int, optional): Total number of chunks.
        quad_as_tri (bool): Enable/disable treating quads as 4 triangles, default ``False``.
            If ``False``, a contour line within a quad is a straight line between points on two of
            its edges. If ``True``, each full quad is divided into 4 triangles using a virtual point
            at the centre (mean x, y of the corner points) and a contour line is piecewise linear
            within those triangles. Corner-masked triangles are not affected by this setting, only
            full unmasked quads.
        z_interp (ZInterp or str, optional): How to interpolate ``z`` values when determining where
            contour lines intersect the edges of quads and the ``z`` values of the central points of
            quads, specified either as a :class:`~contourpy.ZInterp` or its string equivalent such
            as ``"Log"``. Default is ``ZInterp.Linear``.
        thread_count (int): Number of threads to use for contour calculation, default 0. Threads can
            only be used with an algorithm ``name`` that supports threads (currently only
            ``name="threaded"``) and there must be at least the same number of chunks as threads.
            If ``thread_count=0`` and ``name="threaded"`` then it uses the maximum number of threads
            as determined by the C++11 call ``std::thread::hardware_concurrency()``. If ``name`` is
            something other than ``"threaded"`` then the ``thread_count`` will be set to ``1``.

    Return:
        :class:`~contourpy._contourpy.ContourGenerator`.

    Note:
        A maximum of one of ``chunk_size``, ``chunk_count`` and ``total_chunk_count`` may be
        specified.

    Warning:
        The ``name="mpl2005"`` algorithm does not implement chunking for contour lines.
    """
    x = np.asarray(x, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)
    z, mask = _remove_z_mask(z)

    # Check arguments: z.
    if z.ndim != 2:
        raise TypeError(f"Input z must be 2D, not {z.ndim}D")

    if z.shape[0] < 2 or z.shape[1] < 2:
        raise TypeError(f"Input z must be at least a (2, 2) shaped array, but has shape {z.shape}")

    ny, nx = z.shape

    # Check arguments: x and y.
    if x.ndim != y.ndim:
        raise TypeError(f"Number of dimensions of x ({x.ndim}) and y ({y.ndim}) do not match")

    if x.ndim == 0:
        x = np.arange(nx, dtype=np.float64)
        y = np.arange(ny, dtype=np.float64)
        x, y = np.meshgrid(x, y)
    elif x.ndim == 1:
        if len(x) != nx:
            raise TypeError(f"Length of x ({len(x)}) must match number of columns in z ({nx})")
        if len(y) != ny:
            raise TypeError(f"Length of y ({len(y)}) must match number of rows in z ({ny})")
        x, y = np.meshgrid(x, y)
    elif x.ndim == 2:
        if x.shape != z.shape:
            raise TypeError(f"Shapes of x {x.shape} and z {z.shape} do not match")
        if y.shape != z.shape:
            raise TypeError(f"Shapes of y {y.shape} and z {z.shape} do not match")
    else:
        raise TypeError(f"Inputs x and y must be None, 1D or 2D, not {x.ndim}D")

    # Check mask shape just in case.
    if mask is not None and mask.shape != z.shape:
        raise ValueError("If mask is set it must be a 2D array with the same shape as z")

    # Check arguments: name.
    if name not in _class_lookup:
        raise ValueError(f"Unrecognised contour generator name: {name}")

    # Check arguments: chunk_size, chunk_count and total_chunk_count.
    y_chunk_size, x_chunk_size = calc_chunk_sizes(
        chunk_size, chunk_count, total_chunk_count, ny, nx)

    cls = _class_lookup[name]

    # Check arguments: corner_mask.
    if corner_mask is None:
        # Set it to default, which is True if the algorithm supports it.
        corner_mask = cls.supports_corner_mask()
    elif corner_mask and not cls.supports_corner_mask():
        raise ValueError(f"{name} contour generator does not support corner_mask=True")

    # Check arguments: line_type.
    if line_type is None:
        line_type = cls.default_line_type
    else:
        line_type = as_line_type(line_type)

    if not cls.supports_line_type(line_type):
        raise ValueError(f"{name} contour generator does not support line_type {line_type}")

    # Check arguments: fill_type.
    if fill_type is None:
        fill_type = cls.default_fill_type
    else:
        fill_type = as_fill_type(fill_type)

    if not cls.supports_fill_type(fill_type):
        raise ValueError(f"{name} contour generator does not support fill_type {fill_type}")

    # Check arguments: quad_as_tri.
    if quad_as_tri and not cls.supports_quad_as_tri():
        raise ValueError(f"{name} contour generator does not support quad_as_tri=True")

    # Check arguments: z_interp.
    if z_interp is None:
        z_interp = ZInterp.Linear
    else:
        z_interp = as_z_interp(z_interp)

    if z_interp != ZInterp.Linear and not cls.supports_z_interp():
        raise ValueError(f"{name} contour generator does not support z_interp {z_interp}")

    # Check arguments: thread_count.
    if thread_count not in (0, 1) and not cls.supports_threads():
        raise ValueError(f"{name} contour generator does not support thread_count {thread_count}")

    # Prepare args and kwargs for contour generator constructor.
    args = [x, y, z, mask]
    kwargs: dict[str, int | bool | LineType | FillType | ZInterp] = {
        "x_chunk_size": x_chunk_size,
        "y_chunk_size": y_chunk_size,
    }

    if name not in ("mpl2005", "mpl2014"):
        kwargs["line_type"] = line_type
        kwargs["fill_type"] = fill_type

    if cls.supports_corner_mask():
        kwargs["corner_mask"] = corner_mask

    if cls.supports_quad_as_tri():
        kwargs["quad_as_tri"] = quad_as_tri

    if cls.supports_z_interp():
        kwargs["z_interp"] = z_interp

    if cls.supports_threads():
        kwargs["thread_count"] = thread_count

    # Create contour generator.
    return cls(*args, **kwargs)