<filename>src/models/TensorGrid.py
# TO-DO: Memory efficiency
from __future__ import annotations

from collections import deque
from typing import Dict, List, Optional, Union

import numpy as np
import pyro
import pyro.distributions as distributions
import pyro.distributions.constraints as constraints
import torch
import torch.nn as nn
import torch.nn.functional as F
from pyro.nn import PyroModule, PyroParam
from torch.distributions import transform_to


class TensorGrid(PyroModule):
    def __init__(
        self,
        num_components: Union[int, List[int]],
        dists: Union[nn.ModuleList, List[distributions.Distribution]],
        component_logits: List[torch.Tensor] = None,
        dist_params: List[Dict[str, torch.Tensor]] = None,
        grid_length: int = None,
    ) -> None:
        """Tensor Grid is a generalization of Tensor Train and can be seen as a \
        train of Tensor Trains. This implementation uses a single shared `alpha_0` \
        and utilizes a grid structure for the data and associated mixture weights.

        Args:
            num_components (Union[int, List[int]]):
                Amount of mixture components in each mixture. Either a single shared \
                amount or list of specific amounts.
            dists (Union[nn.ModuleList, List[distributions.Distribution]]):
                List of distributions to describe each dimension.
            component_logits (List[torch.Tensor], optional):
                Mixture component weights as logits. \
                Defaults to equally weighted mixtures.
            dist_params (List[Dict[str, torch.Tensor]], optional):
                List of distribution parameters, formatted as if passed as **kwargs to \
                the underlying Distribution. \
                Defaults uniformly distributed random parameters, \
                transformed to fit parameter constraints.
            grid_length (int, optional):
                Length of each grid row. Defaults to `floor(sqrt(len(dists)))`.
        """
        super().__init__()
        # Easy access to .device
        self.__empty__ = nn.Parameter(torch.empty(0))

        self.M = len(dists)
        self.Ks = (
            num_components
            if isinstance(num_components, list)
            else [num_components] * (self.M + 1)
        )

        self.grid_length = (
            grid_length if grid_length is not None else int(np.sqrt(self.M))
        )

        self._mask: Optional[torch.Tensor] = None
        shape = [self.Ks[0]] + [1] * self.M
        self.logits = PyroParam(
            torch.ones(*shape)
            if component_logits is None
            else component_logits[0].view(*shape),
            constraint=constraints.real,
        )

        self.components = nn.ModuleList()
        for i, dist in enumerate(dists):
            comp = PyroModule(f"dim_{i}")
            comp.dist = dist

            shape = [1] * (self.M + 1)
            if i % self.grid_length == 0:  # first element in row
                shape[0] = self.Ks[0]
            else:
                shape[i] = self.Ks[i]
            if (i + 1) > self.grid_length:  # dependence on previous row
                shape[i + 1 - self.grid_length] = self.Ks[i + 1 - self.grid_length]
            shape[i + 1] = self.Ks[i + 1]

            # named _logits to avoid potential collision with distribution parameters
            comp._logits = PyroParam(
                torch.ones(*shape)
                if component_logits is None
                else component_logits[i + 1].view(*shape),
                constraint=constraints.real,
            )
            for param, constraint in dist.arg_constraints.items():
                setattr(
                    comp,
                    param,
                    PyroParam(
                        transform_to(constraint)(torch.empty(*shape).uniform_())
                        if dist_params is None
                        else dist_params[i][param].view(*shape),
                        constraint=constraint,
                    ),
                )
            self.components.append(comp)

    def mask(self, mask: Union[torch.BoolTensor, torch.Tensor] = None) -> TensorGrid:
        """Sets mask for model in place. Similar to Distribution.mask

        Args:
            mask (Union[torch.BoolTensor, torch.Tensor], optional): \
                None or masking tensor with shape (M,) or (N, M). Defaults to None.

        Returns:
            TensorGrid: self
        """
        assert (
            mask is None or 1 <= mask.dim() <= 2
        ), "Mask should be 1- or 2-dimensional or None."
        if mask is not None and mask.dim() == 1:
            mask = mask.unsqueeze(0)
        self._mask = mask
        return self

    def forward(
        self, data: torch.Tensor = None, n_samples: int = 1, impute: bool = False
    ) -> torch.Tensor:
        """Generate random samples if data is None. If data is defined, pass data through \
        model. Used for Pyro inference algorithms.

        Args:
            data (torch.Tensor, optional):
                Input data. Defaults to None.
            n_samples (int, optional):
                Amount of samples to generate if data is None. \
                Defaults to 1.
            impute (bool, optional):
                Should nan-values be imputed automatically. \
                Defaults to False

        Returns:
            torch.Tensor: Either observed samples or generated samples.
        """
        device = self.__empty__.device
        if data is not None and data.dim() == 1:
            data = data.unsqueeze(0)
        N, M = (n_samples, self.M) if data is None else data.shape
        assert M == self.M, f"Invalid data shape. Expected {self.M} features, got {M}."

        x = torch.empty((N, M)).to(device)
        with pyro.plate("data", N):
            previous_ks: deque = deque([], self.grid_length)
            k0 = pyro.sample(
                "k_0",
                distributions.Categorical(logits=self.logits.squeeze()),
                infer={"enumerate": "parallel"},
            )
            previous_ks.append(k0)
            for m in pyro.markov(range(M), history=self.grid_length):
                obs = data[:, m] if data is not None else None
                comp = self.components[m]
                if m < self.grid_length:
                    k_prev = previous_ks[-1]
                    k_m = pyro.sample(
                        f"k_{m+1}",
                        distributions.Categorical(
                            logits=comp._logits.squeeze()[k_prev]
                        ),
                        infer={"enumerate": "parallel"},
                    )
                    params = {
                        p: v.squeeze()[k_prev, k_m]
                        for (p, v) in comp.named_pyro_params()
                        if p != "_logits" and p != "probs"
                    }
                else:
                    k_prev_l = (
                        k0
                        if m % self.grid_length == 0
                        else previous_ks[-self.grid_length]
                    )
                    k_prev_u = previous_ks[-1]
                    k_m = pyro.sample(
                        f"k_{m+1}",
                        distributions.Categorical(
                            logits=comp._logits.squeeze()[k_prev_l, k_prev_u]
                        ),
                        infer={"enumerate": "parallel"},
                    )
                    params = {
                        p: v.squeeze()[k_prev_l, k_prev_u, k_m]
                        for (p, v) in comp.named_pyro_params()
                        if p != "_logits" and p != "probs"
                    }

                dist = comp.dist(**params)
                if obs is not None and impute:
                    mask = ~torch.isnan(obs)
                    x_imputed = pyro.sample(f"imputed_x_{m}", dist.mask(False))
                    obs = obs.clone()
                    obs[~mask] = x_imputed[~mask].to(dtype=obs.dtype)
                elif self._mask is not None:
                    dist = dist.mask(self._mask[:, m])
                x[:, m] = pyro.sample(f"x_{m+1}", dist, obs=obs)
                previous_ks.append(k_m)
        return x

    def log_prob(self, x: torch.Tensor) -> torch.Tensor:
        """Evaluates log probability densities for each of sample in a batch of samples

        Args:
            x (torch.Tensor): A single value (M,) or batch of values (N, M)

        Returns:
            torch.Tensor: Log probability of x
        """
        device = self.__empty__.device
        x = x.to(device)

        if x.dim() == 1:
            x = x.unsqueeze(0)
        N, M = x.shape
        assert M == self.M, f"Invalid data shape. Expected {self.M} features, got {M}."

        log_p = F.log_softmax(self.logits, 0).unsqueeze(0)  # add batch dimension
        idx: List[Optional[Union[slice, ellipsis, int]]] = [None] * (M + 3)  # noqa
        idx[0] = Ellipsis
        # consumed_dimensions = [False] * M

        for m in range(0, M, self.grid_length):
            idx[-1] = m
            comp = self.components[m]
            params = {p: v for (p, v) in comp.named_pyro_params() if p != "_logits" and p != "probs"}
            dist = comp.dist(**params)
            if self._mask is not None:
                dist = dist.mask(self._mask[idx])
            log_p = log_p + F.log_softmax(comp._logits, m + 1) + dist.log_prob(x[idx])
            # consumed_dimensions[m] = True
        log_p = torch.logsumexp(log_p, -(M + 1), keepdim=True)

        for c in range(self.grid_length):
            # doing column by column minimizes memory scaling
            for m in range(c, M, self.grid_length):
                if m + 1 < M and c < self.grid_length - 1:
                    idx[-1] = m + 1
                    comp = self.components[m + 1]
                    params = {
                        p: v for (p, v) in comp.named_pyro_params() if p != "_logits" and p != "probs"
                    }
                    dist = comp.dist(**params)
                    if self._mask is not None:
                        dist = dist.mask(self._mask[idx])
                    log_p = (
                        log_p
                        + F.log_softmax(comp._logits, m + 2)
                        + dist.log_prob(x[idx])
                    )
                log_p = torch.logsumexp(log_p, -(M - m), keepdim=True)
        return log_p.squeeze()

<filename>src/models/TensorGrid.py
# TO-DO: Memory efficiency
from __future__ import annotations

from collections import deque
from typing import Dict, List, Optional, Union

import numpy as np
import pyro
import pyro.distributions as distributions
import pyro.distributions.constraints as constraints
import torch
import torch.nn as nn
import torch.nn.functional as F
from pyro.nn import PyroModule, PyroParam
from torch.distributions import transform_to


class TensorGrid(PyroModule):
    def __init__(
        self,
        num_components: Union[int, List[int]],
        dists: Union[nn.ModuleList, List[distributions.Distribution]],
        component_logits: List[torch.Tensor] = None,
        dist_params: List[Dict[str, torch.Tensor]] = None,
        grid_length: int = None,
    ) -> None:
        """Tensor Grid is a generalization of Tensor Train and can be seen as a \
        train of Tensor Trains. This implementation uses a single shared `alpha_0` \
        and utilizes a grid structure for the data and associated mixture weights.

        Args:
            num_components (Union[int, List[int]]):
                Amount of mixture components in each mixture. Either a single shared \
                amount or list of specific amounts.
            dists (Union[nn.ModuleList, List[distributions.Distribution]]):
                List of distributions to describe each dimension.
            component_logits (List[torch.Tensor], optional):
                Mixture component weights as logits. \
                Defaults to equally weighted mixtures.
            dist_params (List[Dict[str, torch.Tensor]], optional):
                List of distribution parameters, formatted as if passed as **kwargs to \
                the underlying Distribution. \
                Defaults uniformly distributed random parameters, \
                transformed to fit parameter constraints.
            grid_length (int, optional):
                Length of each grid row. Defaults to `floor(sqrt(len(dists)))`.
        """
        super().__init__()
        # Easy access to .device
        self.__empty__ = nn.Parameter(torch.empty(0))

        self.M = len(dists)
        self.Ks = (
            num_components
            if isinstance(num_components, list)
            else [num_components] * (self.M + 1)
        )

        self.grid_length = (
            grid_length if grid_length is not None else int(np.sqrt(self.M))
        )

        self._mask: Optional[torch.Tensor] = None
        shape = [self.Ks[0]] + [1] * self.M
        self.logits = PyroParam(
            torch.ones(*shape)
            if component_logits is None
            else component_logits[0].view(*shape),
            constraint=constraints.real,
        )

        self.components = nn.ModuleList()
        for i, dist in enumerate(dists):
            comp = PyroModule(f"dim_{i}")
            comp.dist = dist

            shape = [1] * (self.M + 1)
            if i % self.grid_length == 0:  # first element in row
                shape[0] = self.Ks[0]
            else:
                shape[i] = self.Ks[i]
            if (i + 1) > self.grid_length:  # dependence on previous row
                shape[i + 1 - self.grid_length] = self.Ks[i + 1 - self.grid_length]
            shape[i + 1] = self.Ks[i + 1]

            # named _logits to avoid potential collision with distribution parameters
            comp._logits = PyroParam(
                torch.ones(*shape)
                if component_logits is None
                else component_logits[i + 1].view(*shape),
                constraint=constraints.real,
            )
            for param, constraint in dist.arg_constraints.items():
                setattr(
                    comp,
                    param,
                    PyroParam(
                        transform_to(constraint)(torch.empty(*shape).uniform_())
                        if dist_params is None
                        else dist_params[i][param].view(*shape),
                        constraint=constraint,
                    ),
                )
            self.components.append(comp)

    def mask(self, mask: Union[torch.BoolTensor, torch.Tensor] = None) -> TensorGrid:
        """Sets mask for model in place. Similar to Distribution.mask

        Args:
            mask (Union[torch.BoolTensor, torch.Tensor], optional): \
                None or masking tensor with shape (M,) or (N, M). Defaults to None.

        Returns:
            TensorGrid: self
        """
        assert (
            mask is None or 1 <= mask.dim() <= 2
        ), "Mask should be 1- or 2-dimensional or None."
        if mask is not None and mask.dim() == 1:
            mask = mask.unsqueeze(0)
        self._mask = mask
        return self

    def forward(
        self, data: torch.Tensor = None, n_samples: int = 1, impute: bool = False
    ) -> torch.Tensor:
        """Generate random samples if data is None. If data is defined, pass data through \
        model. Used for Pyro inference algorithms.

        Args:
            data (torch.Tensor, optional):
                Input data. Defaults to None.
            n_samples (int, optional):
                Amount of samples to generate if data is None. \
                Defaults to 1.
            impute (bool, optional):
                Should nan-values be imputed automatically. \
                Defaults to False

        Returns:
            torch.Tensor: Either observed samples or generated samples.
        """
        device = self.__empty__.device
        if data is not None and data.dim() == 1:
            data = data.unsqueeze(0)
        N, M = (n_samples, self.M) if data is None else data.shape
        assert M == self.M, f"Invalid data shape. Expected {self.M} features, got {M}."

        x = torch.empty((N, M)).to(device)
        with pyro.plate("data", N):
            previous_ks: deque = deque([], self.grid_length)
            k0 = pyro.sample(
                "k_0",
                distributions.Categorical(logits=self.logits.squeeze()),
                infer={"enumerate": "parallel"},
            )
            previous_ks.append(k0)
            for m in pyro.markov(range(M), history=self.grid_length):
                obs = data[:, m] if data is not None else None
                comp = self.components[m]
                if m < self.grid_length:
                    k_prev = previous_ks[-1]
                    k_m = pyro.sample(
                        f"k_{m+1}",
                        distributions.Categorical(
                            logits=comp._logits.squeeze()[k_prev]
                        ),
                        infer={"enumerate": "parallel"},
                    )
                    params = {
                        p: v.squeeze()[k_prev, k_m]
                        for (p, v) in comp.named_pyro_params()
                        if p != "_logits" and p != "probs"
                    }
                else:
                    k_prev_l = (
                        k0
                        if m % self.grid_length == 0
                        else previous_ks[-self.grid_length]
                    )
                    k_prev_u = previous_ks[-1]
                    k_m = pyro.sample(
                        f"k_{m+1}",
                        distributions.Categorical(
                            logits=comp._logits.squeeze()[k_prev_l, k_prev_u]
                        ),
                        infer={"enumerate": "parallel"},
                    )
                    params = {
                        p: v.squeeze()[k_prev_l, k_prev_u, k_m]
                        for (p, v) in comp.named_pyro_params()
                        if p != "_logits" and p != "probs"
                    }

                dist = comp.dist(**params)
                if obs is not None and impute:
                    mask = ~torch.isnan(obs)
                    x_imputed = pyro.sample(f"imputed_x_{m}", dist.mask(False))
                    obs = obs.clone()
                    obs[~mask] = x_imputed[~mask].to(dtype=obs.dtype)
                elif self._mask is not None:
                    dist = dist.mask(self._mask[:, m])
                x[:, m] = pyro.sample(f"x_{m+1}", dist, obs=obs)
                previous_ks.append(k_m)
        return x

    def log_prob(self, x: torch.Tensor) -> torch.Tensor:
        """Evaluates log probability densities for each of sample in a batch of samples

        Args:
            x (torch.Tensor): A single value (M,) or batch of values (N, M)

        Returns:
            torch.Tensor: Log probability of x
        """
        device = self.__empty__.device
        x = x.to(device)

        if x.dim() == 1:
            x = x.unsqueeze(0)
        N, M = x.shape
        assert M == self.M, f"Invalid data shape. Expected {self.M} features, got {M}."

        log_p = F.log_softmax(self.logits, 0).unsqueeze(0)  # add batch dimension
        idx: List[Optional[Union[slice, ellipsis, int]]] = [None] * (M + 3)  # noqa
        idx[0] = Ellipsis
        # consumed_dimensions = [False] * M

        for m in range(0, M, self.grid_length):
            idx[-1] = m
            comp = self.components[m]
            params = {p: v for (p, v) in comp.named_pyro_params() if p != "_logits" and p != "probs"}
            dist = comp.dist(**params)
            if self._mask is not None:
                dist = dist.mask(self._mask[idx])
            log_p = log_p + F.log_softmax(comp._logits, m + 1) + dist.log_prob(x[idx])
            # consumed_dimensions[m] = True
        log_p = torch.logsumexp(log_p, -(M + 1), keepdim=True)

        for c in range(self.grid_length):
            # doing column by column minimizes memory scaling
            for m in range(c, M, self.grid_length):
                if m + 1 < M and c < self.grid_length - 1:
                    idx[-1] = m + 1
                    comp = self.components[m + 1]
                    params = {
                        p: v for (p, v) in comp.named_pyro_params() if p != "_logits" and p != "probs"
                    }
                    dist = comp.dist(**params)
                    if self._mask is not None:
                        dist = dist.mask(self._mask[idx])
                    log_p = (
                        log_p
                        + F.log_softmax(comp._logits, m + 2)
                        + dist.log_prob(x[idx])
                    )
                log_p = torch.logsumexp(log_p, -(M - m), keepdim=True)
        return log_p.squeeze()

# TO-DO: Memory efficiency Tensor Grid is a generalization of Tensor Train and can be seen as a \         train of Tensor Trains. This implementation uses a single shared `alpha_0` \         and utilizes a grid structure for the data and associated mixture weights.          Args:             num_components (Union[int, List[int]]):                 Amount of mixture components in each mixture. Either a single shared \                 amount or list of specific amounts.             dists (Union[nn.ModuleList, List[distributions.Distribution]]):                 List of distributions to describe each dimension.             component_logits (List[torch.Tensor], optional):                 Mixture component weights as logits. \                 Defaults to equally weighted mixtures.             dist_params (List[Dict[str, torch.Tensor]], optional):                 List of distribution parameters, formatted as if passed as **kwargs to \                 the underlying Distribution. \                 Defaults uniformly distributed random parameters, \                 transformed to fit parameter constraints.             grid_length (int, optional):                 Length of each grid row. Defaults to `floor(sqrt(len(dists)))`. # Easy access to .device # first element in row # dependence on previous row # named _logits to avoid potential collision with distribution parameters Sets mask for model in place. Similar to Distribution.mask          Args:             mask (Union[torch.BoolTensor, torch.Tensor], optional): \                 None or masking tensor with shape (M,) or (N, M). Defaults to None.          Returns:             TensorGrid: self Generate random samples if data is None. If data is defined, pass data through \         model. Used for Pyro inference algorithms.          Args:             data (torch.Tensor, optional):                 Input data. Defaults to None.             n_samples (int, optional):                 Amount of samples to generate if data is None. \                 Defaults to 1.             impute (bool, optional):                 Should nan-values be imputed automatically. \                 Defaults to False          Returns:             torch.Tensor: Either observed samples or generated samples. Evaluates log probability densities for each of sample in a batch of samples          Args:             x (torch.Tensor): A single value (M,) or batch of values (N, M)          Returns:             torch.Tensor: Log probability of x # add batch dimension # noqa # consumed_dimensions = [False] * M # consumed_dimensions[m] = True # doing column by column minimizes memory scaling