def stochasticCorrelate(self, num_correlate, choices=None):
if num_correlate == 0:
return self
domshape = self.head.shape
batch_size = domshape[0]
num_pixs = h.product(domshape[1:])
num_correlate = min(num_correlate, num_pixs)
ucc_mask = h.ones([batch_size, num_pixs]).long()
cc_indx_batch_beta = h.cudify(
torch.multinomial(
ucc_mask.to_dtype(), num_correlate,
replacement=False)) if choices is None else choices
return self.correlate(cc_indx_batch_beta)
def stochasticDecorrelate(self,
num_decorrelate,
choices=None,
num_to_keep=False):
dummy = self.dummyDecorrelate(num_decorrelate)
if dummy is not None:
return dummy
num_error_terms = self.errors.shape[0]
batch_size = self.head.shape[0]
ucc_mask = h.ones([batch_size, self.errors.shape[0]]).long()
cc_indx_batch_err = h.cudify(
torch.multinomial(
ucc_mask.to_dtype(),
num_decorrelate if num_to_keep else num_error_terms -
num_decorrelate,
replacement=False)) if choices is None else choices
return self.decorrelate(cc_indx_batch_err)
