Source code for deepinv.physics.generator.noise

import torch
from deepinv.physics.generator import PhysicsGenerator


[docs] class SigmaGenerator(PhysicsGenerator): r""" Generator for the noise level :math:`\sigma` in the Gaussian noise model. The noise level is sampled uniformly from the interval :math:`[\text{sigma_min}, \text{sigma_max}]`. :param float sigma_min: minimum noise level :param float sigma_max: maximum noise level :param torch.Generator rng: random number generator :param str device: device where the tensor is stored :param torch.dtype dtype: data type of the generated tensor |sep| :Examples: >>> from deepinv.physics.generator import SigmaGenerator >>> generator = SigmaGenerator() >>> sigma_dict = generator.step(seed=0) # dict_keys(['sigma']) >>> print(sigma_dict['sigma']) tensor([0.2532]) """ def __init__( self, sigma_min=0.01, sigma_max=0.5, rng: torch.Generator = None, device: str = "cpu", dtype: torch.dtype = torch.float32, ): super().__init__(shape=(1,), rng=rng, device=device, dtype=dtype) self.sigma_min = sigma_min self.sigma_max = sigma_max
[docs] def step(self, batch_size=1, seed: int = None, **kwargs): r""" Generates a batch of noise levels. :param int batch_size: batch size :param int seed: the seed for the random number generator. :return: dictionary with key **'sigma'**: tensor of size (batch_size,). :rtype: dict """ self.rng_manual_seed(seed) sigma = ( torch.rand(batch_size, generator=self.rng, **self.factory_kwargs) * (self.sigma_max - self.sigma_min) + self.sigma_min ) return {"sigma": sigma}