alive_sieve.py

"""
Handle alive/dead masks
"""
import torch


class AliveSieve(object):
    def __init__(self, batch_size, enable_cuda):
        """
        assume all alive to start with, with given batch_size
        """
        self.batch_size = batch_size
        self.enable_cuda = enable_cuda
        self.type_constr = torch.cuda if enable_cuda else torch
        self.alive_mask = self.type_constr.ByteTensor(batch_size).fill_(1)
        self.alive_idxes = self.mask_to_idxes(self.alive_mask)
        """
        out_idxes are the indexes of the current members of alivesieve, in the original batch
        basically, we popualte it initially with 0...batch_size-1 , and then
        sieve this too each time we call sieve_self
        """
        self.out_idxes = self.alive_idxes.clone()

    @staticmethod
    def mask_to_idxes(mask):
        return mask.view(-1).nonzero().long().view(-1)

    def mark_dead(self, dead_mask):
        """
        sets the mask to 0 at relevant positions
        doesnt remove them yet, ie doesnt call 'sieve'
        """
        if dead_mask.max() == 0:
            return
        dead_idxes = self.mask_to_idxes(dead_mask)
        self.alive_mask[dead_idxes] = 0
        self.alive_idxes = self.mask_to_idxes(self.alive_mask)

    def get_dead_idxes(self):
        dead_mask = 1 - self.alive_mask
        return dead_mask.nonzero().long().view(-1)

    def any_alive(self):
        return self.alive_mask.max() == 1

    def all_dead(self):
        return self.alive_mask.max() == 0

    def set_dead_global(self, target, v):
        """
        What this does is assign v to target, masked to only modify items marked as dead
        It assumes that target is a global tensor, ie hasnt been sieved
        So we're going to use out_idxes to index into this
        """
        dead_idxes = self.get_dead_idxes()
        if len(dead_idxes) == 0:
            return
        target[self.out_idxes[dead_idxes]] = v

    def self_sieve_(self):
        """
        removes all dead from alive_mask
        basically equivalent to resetting the batch_size, and
        recreating an all-1s mask
        """
        self.out_idxes = self.out_idxes[self.alive_idxes]

        self.batch_size = self.alive_mask.int().sum()
        self.alive_mask = self.type_constr.ByteTensor(self.batch_size).fill_(1)
        self.alive_idxes = self.mask_to_idxes(self.alive_mask)

    def sieve_tensor(self, t):
        """
        returns t sieved by current alive mask (ie anything not alive is
        not returned)
        note: the returned tensor is NOT guaranteed to be a reference into
        the old one. In fact, I'm fairly sure it will never in fact be (but
        this isnt guaranteed either; I havent checked...)
        """
        return t[self.alive_idxes]

    def sieve_list(self, alist):
        """
        returns alist sieved by current alive mask (ie anything not alive is
        not returned)
        """
        return [alist[b] for b in self.alive_idxes]


class SievePlayback(object):
    """
    Given a list of alive_masks, will loop through these, providing a tuple of
    t, and the global idxes at each loop step
    """
    def __init__(self, alive_masks, enable_cuda):
        self.alive_masks = alive_masks
        self.type_constr = torch.cuda if enable_cuda else torch

    def __iter__(self):
        batch_size = self.alive_masks[0].size()[0]
        global_idxes = self.type_constr.ByteTensor(batch_size).fill_(1).nonzero().long().view(-1)
        T = len(self.alive_masks)
        for t in range(T):
            self.batch_size = len(global_idxes)
            yield t, global_idxes
            mask = self.alive_masks[t]
            if mask.max() == 0:
                return
            global_idxes = global_idxes[mask.nonzero().long().view(-1)]