from ...utils.integration import quad_table
from ...backend_obj import backend, ArrayLike
[docs]
def pixel_center_meshgrid(
extent: tuple[int, int, int, int], pad: int, upsample: int, dtype: type, device: str
) -> tuple:
i = backend.linspace(
extent[0] - 0.5 + 0.5 / upsample - pad / upsample,
extent[1] - 0.5 - 0.5 / upsample + pad / upsample,
(extent[1] - extent[0]) * upsample + 2 * pad,
dtype=dtype,
device=device,
)
j = backend.linspace(
extent[2] - 0.5 + 0.5 / upsample - pad / upsample,
extent[3] - 0.5 - 0.5 / upsample + pad / upsample,
(extent[3] - extent[2]) * upsample + 2 * pad,
dtype=dtype,
device=device,
)
return backend.meshgrid(i, j, indexing="ij")
[docs]
def cmos_pixel_center_meshgrid(
extent: tuple[int, int, int, int],
pad: int,
upsample: int,
loc: tuple[float, float],
dtype,
device,
) -> tuple:
if upsample > 1:
raise NotImplementedError("CMOS pixel meshgrid does not currently support upsampling.")
i, j = pixel_center_meshgrid(extent, pad, upsample, dtype, device)
return i + loc[0], j + loc[1]
[docs]
def pixel_corner_meshgrid(
extent: tuple[int, int, int, int], pad: int, upsample: int, dtype, device
) -> tuple:
i, j = pixel_center_meshgrid(
(extent[0], extent[1] + 1, extent[2], extent[3] + 1), pad, upsample, dtype, device
)
return i - 0.5 / upsample, j - 0.5 / upsample
[docs]
def pixel_simpsons_meshgrid(
extent: tuple[int, int, int, int], pad: int, upsample: int, dtype, device
) -> tuple:
return pixel_corner_meshgrid(extent, 2 * pad, 2 * upsample, dtype, device)
[docs]
def pixel_quad_meshgrid(
extent: tuple[int, int, int, int], pad: int, upsample: int, dtype, device, order=3
) -> tuple:
i, j = pixel_center_meshgrid(extent, pad, upsample, dtype, device)
di, dj, w = quad_table(order, dtype, device)
i = backend.repeat(i[..., None], order**2, -1) + di.flatten() / upsample
j = backend.repeat(j[..., None], order**2, -1) + dj.flatten() / upsample
return i, j, w.flatten()
[docs]
def rotate(theta: ArrayLike, x: ArrayLike, y: ArrayLike) -> tuple:
"""
Applies a rotation matrix to the X,Y coordinates
"""
s = backend.sin(theta)
c = backend.cos(theta)
return c * x - s * y, s * x + c * y