from __future__ import annotations
from enum import Enum
import numpy as np
from numpy.typing import ArrayLike, NDArray
from cmap import Color
[docs]def as_array_1d(x: ArrayLike, dtype=None) -> NDArray[np.number]:
x = np.asarray(x, dtype=dtype)
if x.ndim != 1:
raise ValueError(f"Expected 1D array, got {x.ndim}D array")
if x.dtype.kind not in "iuf":
raise ValueError(f"Input {x!r} did not return a numeric array")
return x
[docs]def normalize_xy(*args) -> tuple[NDArray[np.number], NDArray[np.number]]:
if len(args) == 1:
arr = np.asarray(args[0])
if arr.dtype.kind not in "iuf":
raise ValueError(f"Input {args[0]!r} did not return a numeric array")
if arr.ndim == 1:
ydata = arr
xdata = np.arange(ydata.size)
elif arr.ndim == 2 and arr.shape[1] == 2:
xdata = arr[:, 0]
ydata = arr[:, 1]
else:
raise ValueError(
"Expected 1D array or 2D array with shape (N, 2), " f"got {arr.shape}"
)
elif len(args) == 2:
if np.isscalar(args[0]) and np.isscalar(args[1]):
return np.array(args[0]), np.array(args[1])
xdata = as_array_1d(args[0])
if not hasattr(args[1], "__array__") and callable(args[1]):
ydata = args[1](xdata)
else:
ydata = as_array_1d(args[1])
if xdata.size != ydata.size:
raise ValueError(
"Expected xdata and ydata to have the same size, "
f"got {xdata.size} and {ydata.size}"
)
else:
raise TypeError(f"Expected 1 or 2 positional arguments, got {len(args)}")
return xdata, ydata
[docs]def arr_color(color) -> np.ndarray:
"""Normalize a color input to a 4-element float array."""
try:
c = Color(color)
except Exception as e:
raise ValueError(f"Invalid input for a color: {color!r}") from None
return np.array(c.rgba, dtype=np.float32)
[docs]def hex_color(color) -> str:
"""Normalize a color input to a #RRGGBBAA string."""
return Color(color).hex
[docs]def rgba_str_color(color) -> str:
"""Normalize a color input to a rgba(r, g, b, a) string."""
return Color(color).rgba_string
[docs]def as_any_1d_array(x: float, size: int, dtype=None) -> np.ndarray:
if np.isscalar(x) or isinstance(x, Enum):
out = np.full((size,), x, dtype=dtype)
else:
out = np.asarray(x, dtype=dtype)
if out.shape != (size,):
raise ValueError(f"Expected shape ({size},), got {out.shape}")
return out
[docs]def as_color_array(color, size: int) -> NDArray[np.float32]:
if isinstance(color, str): # e.g. color = "black"
col = arr_color(color)
return np.repeat(col[np.newaxis, :], size, axis=0)
if isinstance(color, np.ndarray):
if color.dtype.kind in "OU":
if color.shape != (size,):
raise ValueError(
f"Expected color array of shape ({size},), got {color.shape}"
)
return np.stack([arr_color(each) for each in color], axis=0)
elif color.shape in [(3,), (4,)]:
col = arr_color(color)
return np.repeat(col[np.newaxis, :], size, axis=0)
elif color.shape in [(size, 3), (size, 4)]:
return color
else:
raise ValueError(
f"Color array must have shape (3,), (4,), (N={size}, 3), or (N={size}, 4) "
f"but got {color.shape}"
)
arr = np.array(color)
return as_color_array(arr, size)