stack.py

import torch
from caskade import Module, forward, Param
from torch.nn.functional import avg_pool2d
from utils import meshgrid, Pixelated

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.float64


class Cutout(Module):
    def __init__(
        self,
        data,
        mask,
        variance,
        upsample,
        psf,
        starf_mode="param",
        center=None,
        starf=None,
        sky=None,
        variance_scale=1.0,
        anneal=0.0,
        psf_upsample=4,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.data = torch.as_tensor(data, device=DEVICE, dtype=DTYPE)
        self.mask = torch.as_tensor(mask, device=DEVICE, dtype=torch.bool)
        self.variance = torch.as_tensor(variance, device=DEVICE, dtype=DTYPE)
        self.upsample = upsample
        self.psf = Pixelated(
            name="PSF",
            x0=0.0,
            y0=0.0,
            image=psf,
            pixelscale=1.0 / psf_upsample,
        )

        self.center = Param(f"center", center, shape=(2,))
        self.starf = Param(f"starf", starf)
        self.sky = Param("sky", sky)

        self.thx, self.thy = meshgrid(
            1 / upsample,
            data.shape[1] * upsample,
            data.shape[0] * upsample,
            device=DEVICE,
            dtype=DTYPE,
        )

        self.starf_mode = starf_mode
        self.variance_scale = variance_scale
        self.anneal = anneal
        self.center_mask = False
        self.meta.bad_obj_mask = torch.zeros_like(self._mask)

    def to(self, device, dtype):
        super().to(device, dtype)
        self.data = self.data.to(device, dtype)
        self._mask = self._mask.to(device, torch.bool)
        self.meta.bad_obj_mask = self.meta.bad_obj_mask.to(device, torch.bool)
        self.variance = self.variance.to(device, dtype)
        self.thx = self.thx.to(device, dtype)
        self.thy = self.thy.to(device, dtype)
        return self

    @property
    def mask(self):
        mask = self._mask | self.meta.bad_obj_mask
        if self.center_mask:
            center_mask = torch.zeros_like(self._mask)
            center_mask[: int(center_mask.shape[0] // 4)] = True
            center_mask[-int(center_mask.shape[0] // 4) :] = True
            center_mask[:, : int(center_mask.shape[0] // 4)] = True
            center_mask[:, -int(center_mask.shape[0] // 4) :] = True
            return mask | center_mask
        return mask

    @mask.setter
    def mask(self, value):
        self._mask = value

    @property
    def starf_mode(self):
        return self._starf_mode

    @starf_mode.setter
    def starf_mode(self, value):
        self._starf_mode = value
        if value == "exact":
            self.starf = torch.ones((), device=self.data.device, dtype=self.data.dtype)

    @forward
    def _sample(self, center=None, sky=None, model_perturb=None):
        sample = avg_pool2d(
            self.psf.brightness(
                self.thx + center[0],
                self.thy + center[1],
                padding_mode="clamp",
            )[
                None
            ][None],
            self.upsample,
        ).squeeze()
        return sample

    @forward
    def max_starf(self, sky=None):
        return ((~self.mask) * (self.data - sky)).amax()

    @forward
    def exact_starf(self, sample=None, sky=None):
        if sample is None:
            sample = self._sample()
        data = self.data - sky
        starf = ((~self.mask) * sample * data / self.variance).sum() / (
            (~self.mask) * sample**2 / self.variance
        ).sum()
        return starf

    @forward
    def sample(self, sample=None, starf=None, sky=None):
        if sample is None:
            sample = self._sample()

        if self.starf_mode == "exact":
            starf = self.exact_starf(sample=sample, sky=sky)

        return sample * starf + sky

    @forward
    def residuals(self, sample=None):
        if sample is None:
            sample = self.sample()
        return (self.data - sample) / self.variance.sqrt()

    @forward
    def set_prior_bounds(self, center=None, size=3):
        self.prior_center = center.detach()
        self.prior_size = size

    @forward
    def log_prior(self, center=None):
        ob = torch.tensor(-torch.inf, device=center.device)
        ib = torch.tensor(0.0, device=center.device)
        return torch.where(
            torch.any((center - self.prior_center).abs() > self.prior_size),
            ob,
            ib,
        )

    @forward
    def log_likelihood(self, sample=None, anneal=0.0):
        if sample is None:
            sample = self.sample()
        if not isinstance(anneal, (float, int)):
            assert sample.shape == anneal.shape
        return (
            -0.5
            * (
                (~self.mask)
                * (sample - self.data) ** 2
                / (self.variance * self.variance_scale + self.anneal + anneal)
            ).sum()
        )

    @forward
    def log_posterior(self, sample=None, anneal=0.0):
        lp = self.log_prior()
        ll = self.log_likelihood(sample=sample, anneal=anneal)
        return lp + ll