Image Resampling and Combination

from typing import Optional, Union, Tuple
import numpy as np

class Drizzle:
    """
    A class for managing resampling and co-adding of multiple images onto a
    common output grid using the drizzle algorithm.

    The main method is add_image(). Outputs are accessible as properties:
    out_img, out_wht, out_ctx, and total_exptime.
    """

    def __init__(
        self,
        kernel: str = "square",
        fillval: Optional[Union[str, float]] = None,
        out_shape: Optional[Tuple[int, int]] = None,
        out_img: Optional[np.ndarray] = None,
        out_wht: Optional[np.ndarray] = None,
        out_ctx: Optional[np.ndarray] = None,
        exptime: float = 0.0,
        begin_ctx_id: int = 0,
        max_ctx_id: Optional[int] = None,
        disable_ctx: bool = False
    ):
        """
        Parameters:
        - kernel: Kernel for flux distribution ("square", "gaussian", "point", "turbo", "lanczos2", "lanczos3")
        - fillval: Value for output pixels without input contributions (None, "INDEF", or numeric string)
        - out_shape: Shape of output images (Ny, Nx). Required if no output arrays provided
        - out_img: Pre-allocated output image array (float32). If None, will be created
        - out_wht: Pre-allocated weight array (float32). If None, will be created
        - out_ctx: Pre-allocated context array (int32). If None, will be created (unless disable_ctx=True)
        - exptime: Exposure time of previously resampled images (if providing existing arrays)
        - begin_ctx_id: Starting context ID number (0-based) for first image
        - max_ctx_id: Maximum expected context ID for pre-allocation
        - disable_ctx: If True, disables context image creation
        """

    def add_image(
        self,
        data: np.ndarray,
        exptime: float,
        pixmap: np.ndarray,
        scale: float = 1.0,
        weight_map: Optional[np.ndarray] = None,
        wht_scale: float = 1.0,
        pixfrac: float = 1.0,
        in_units: str = 'cps',
        xmin: Optional[int] = None,
        xmax: Optional[int] = None,
        ymin: Optional[int] = None,
        ymax: Optional[int] = None
    ) -> Tuple[float, float]:
        """
        Resample and add an image to the cumulative output image.

        Parameters:
        - data: Input image to be drizzled (2D numpy array)
        - exptime: Exposure time of input image (positive number)
        - pixmap: Coordinate mapping array of shape (Ny, Nx, 2)
                  pixmap[..., 0] = X-coordinates, pixmap[..., 1] = Y-coordinates
        - scale: Pixel scale of input image (conceptually linear dimension of input pixel)
        - weight_map: Pixel-by-pixel weighting array. Must match data shape. If None, uniform weights used
        - wht_scale: Scaling factor applied to pixel weighting
        - pixfrac: Fraction of pixel flux is confined to (0-1). 1=full pixel, 0.5=half pixel
        - in_units: Units of input image ("cps" for counts per second, "counts" for total counts)
        - xmin, xmax, ymin, ymax: Bounding box on input image. Only pixels inside contribute to output

        Returns:
        Tuple of (nmiss, nskip):
        - nmiss: Number of pixels that were ignored and did not contribute to output
        - nskip: Number of lines that were ignored and did not contribute to output
        """

    # Properties
    @property
    def fillval(self) -> str:
        """Fill value for output pixels without contributions from input images."""

    @property
    def kernel(self) -> str:
        """Resampling kernel."""

    @property
    def ctx_id(self) -> Optional[int]:
        """Context image ID (0-based) of the next image to be resampled."""

    @property
    def out_img(self) -> Optional[np.ndarray]:
        """Output resampled image (float32 array)."""

    @property
    def out_wht(self) -> Optional[np.ndarray]:
        """Output weight image (float32 array)."""

    @property
    def out_ctx(self) -> Optional[np.ndarray]:
        """Output context image (int32 array, 2D or 3D)."""

    @property
    def total_exptime(self) -> float:
        """Total exposure time of all resampled images."""

def blot_image(
    data: np.ndarray,
    pixmap: np.ndarray,
    pix_ratio: float,
    exptime: float,
    output_pixel_shape: Tuple[int, int],
    interp: str = 'poly5',
    sinscl: float = 1.0
) -> np.ndarray:
    """
    Resample input image onto output grid using interpolation (not flux-conserving).

    Typically used to resample drizzled output back to input coordinate grids.
    Works best with well-sampled images like drizzled outputs.

    Parameters:
    - data: Input image in units of 'cps' (2D array)
    - pixmap: Coordinate mapping array of shape (Ny, Nx, 2)
              pixmap[..., 0] = X-coordinates, pixmap[..., 1] = Y-coordinates
    - pix_ratio: Ratio of input image pixel scale to output image pixel scale
    - exptime: Exposure time of input image
    - output_pixel_shape: Dimensions of output image (Nx, Ny)
    - interp: Interpolation method ("nearest", "linear", "poly3", "poly5", "sinc", "lan3", "lan5")
    - sinscl: Scaling factor for "sinc" interpolation

    Returns:
    2D numpy array containing the resampled image data (float32)
    """

import numpy as np
from drizzle.resample import Drizzle

# Create sample data
data = np.random.random((100, 100)).astype(np.float32)

# Create identity pixel mapping (no transformation)
y, x = np.indices((100, 100), dtype=np.float64)
pixmap = np.dstack([x, y])

# Initialize drizzler
drizzler = Drizzle(out_shape=(100, 100), kernel="square")

# Add the image
nmiss, nskip = drizzler.add_image(
    data=data,
    exptime=30.0,
    pixmap=pixmap,
    pixfrac=1.0
)

# Get results
output = drizzler.out_img
weights = drizzler.out_wht
context = drizzler.out_ctx

print(f"Missed pixels: {nmiss}, Skipped lines: {nskip}")
print(f"Total exposure time: {drizzler.total_exptime}")

import numpy as np
from drizzle.resample import Drizzle

# Initialize drizzler for combining multiple images
drizzler = Drizzle(out_shape=(200, 200), kernel="gaussian")

# Process multiple input images
for i in range(3):
    # Simulate input data and weights
    data = np.random.random((180, 180)).astype(np.float32)
    weights = np.ones_like(data) * (0.8 + 0.1 * i)  # Different weights per image

    # Simulate slight shifts in pixel mapping
    y, x = np.indices((180, 180), dtype=np.float64)
    x += i * 0.5  # Small shift
    y += i * 0.3
    pixmap = np.dstack([x, y])

    # Add each image
    nmiss, nskip = drizzler.add_image(
        data=data,
        exptime=15.0,
        pixmap=pixmap,
        weight_map=weights,
        pixfrac=0.8,
        scale=1.0
    )

    print(f"Image {i}: missed={nmiss}, skipped={nskip}")

# Final combined result
combined_image = drizzler.out_img
total_weights = drizzler.out_wht
print(f"Combined {drizzler.ctx_id} images, total exposure: {drizzler.total_exptime}s")

import numpy as np
from drizzle.resample import blot_image

# Simulate a well-sampled drizzled image
drizzled_data = np.random.random((150, 150)).astype(np.float32)

# Create pixel mapping from drizzled grid to original input grid
y, x = np.indices((120, 120), dtype=np.float64)
# Transform coordinates (reverse of original drizzling transformation)
x_mapped = x * 1.25  # Scale factor
y_mapped = y * 1.25
pixmap = np.dstack([x_mapped, y_mapped])

# Blot the drizzled image back to input grid
blotted = blot_image(
    data=drizzled_data,
    pixmap=pixmap,
    pix_ratio=1.25,
    exptime=30.0,
    output_pixel_shape=(120, 120),
    interp='poly5'
)

print(f"Blotted image shape: {blotted.shape}")

from drizzle.resample import Drizzle

try:
    # Invalid kernel
    drizzler = Drizzle(kernel="invalid_kernel")
except ValueError as e:
    print(f"Kernel error: {e}")

try:
    # Invalid exposure time
    drizzler = Drizzle(out_shape=(100, 100))
    drizzler.add_image(data, exptime=-1.0, pixmap=pixmap)
except ValueError as e:
    print(f"Exposure time error: {e}")

try:
    # Inconsistent array shapes
    data = np.ones((100, 100))
    pixmap = np.ones((50, 50, 2))  # Wrong shape
    drizzler.add_image(data, exptime=1.0, pixmap=pixmap)
except ValueError as e:
    print(f"Shape error: {e}")