kaiserpister's picture
Upload folder using huggingface_hub
6eefbd7
from __future__ import annotations
from typing import TYPE_CHECKING, Any
import numpy as np
if TYPE_CHECKING:
from contourpy._contourpy import CoordinateArray
def simple(
shape: tuple[int, int],
want_mask: bool = False,
) -> tuple[
CoordinateArray, CoordinateArray, CoordinateArray | np.ma.MaskedArray[Any, Any]
]:
"""Return simple test data consisting of the sum of two gaussians.
Args:
shape (tuple(int, int)): 2D shape of data to return.
want_mask (bool, optional): Whether test data should be masked or not, default ``False``.
Return:
Tuple of 3 arrays: ``x``, ``y``, ``z`` test data, ``z`` will be masked if
``want_mask=True``.
"""
ny, nx = shape
x = np.arange(nx, dtype=np.float64)
y = np.arange(ny, dtype=np.float64)
x, y = np.meshgrid(x, y)
xscale = nx - 1.0
yscale = ny - 1.0
# z is sum of 2D gaussians.
amp = np.asarray([1.0, -1.0, 0.8, -0.9, 0.7])
mid = np.asarray([[0.4, 0.2], [0.3, 0.8], [0.9, 0.75], [0.7, 0.3], [0.05, 0.7]])
width = np.asarray([0.4, 0.2, 0.2, 0.2, 0.1])
z = np.zeros_like(x)
for i in range(len(amp)):
z += amp[i] * np.exp(
-((x / xscale - mid[i, 0]) ** 2 + (y / yscale - mid[i, 1]) ** 2)
/ width[i] ** 2
)
if want_mask:
mask = np.logical_or(
((x / xscale - 1.0) ** 2 / 0.2 + (y / yscale - 0.0) ** 2 / 0.1) < 1.0,
((x / xscale - 0.2) ** 2 / 0.02 + (y / yscale - 0.45) ** 2 / 0.08) < 1.0,
)
z = np.ma.array(z, mask=mask) # type: ignore[no-untyped-call]
return x, y, z
def random(
shape: tuple[int, int],
seed: int = 2187,
mask_fraction: float = 0.0,
) -> tuple[
CoordinateArray, CoordinateArray, CoordinateArray | np.ma.MaskedArray[Any, Any]
]:
"""Return random test data..
Args:
shape (tuple(int, int)): 2D shape of data to return.
seed (int, optional): Seed for random number generator, default 2187.
mask_fraction (float, optional): Fraction of elements to mask, default 0.
Return:
Tuple of 3 arrays: ``x``, ``y``, ``z`` test data, ``z`` will be masked if
``mask_fraction`` is greater than zero.
"""
ny, nx = shape
x = np.arange(nx, dtype=np.float64)
y = np.arange(ny, dtype=np.float64)
x, y = np.meshgrid(x, y)
rng = np.random.default_rng(seed)
z = rng.uniform(size=shape)
if mask_fraction > 0.0:
mask_fraction = min(mask_fraction, 0.99)
mask = rng.uniform(size=shape) < mask_fraction
z = np.ma.array(z, mask=mask) # type: ignore[no-untyped-call]
return x, y, z