def __init__(self, in_features, a=None, trainable=True):
'''
Initialization.
Args:
in_features: shape of the input
a: trainable parameter
trainable: sets `a` as a trainable parameter
`a` is initialized to 1 by default, higher values = higher-frequency,
5-50 is a good starting point if you already think your data is periodic,
consider starting lower e.g. 0.5 if you think not, but don't worry,
`a` will be trained along with the rest of your model
'''
super(Snake, self).__init__()
self.in_features = in_features if isinstance(in_features,
list) else [in_features]
# Initialize `a`
if a is not None:
self.a = Parameter(torch.ones(self.in_features) *
a) # create a tensor out of alpha
else:
m = Exponential(torch.tensor([0.1]))
self.a = Parameter(
(m.rsample(self.in_features)
).squeeze()) # random init = mix of frequencies
self.a.requiresGrad = trainable # set the training of `a` to true
class EmitterSamplerBlinking(EmitterSamplerFrameIndependent):
def __init__(self, *, structure: structure_prior.StructurePrior, intensity_mu_sig: tuple, lifetime: float,
frame_range: tuple, xy_unit: str, px_size: tuple, density=None, em_avg=None, intensity_th=None):
"""
Args:
structure:
intensity_mu_sig:
lifetime:
xy_unit:
px_size:
frame_range: specifies the frame range
density:
em_avg:
intensity_th:
"""
super().__init__(structure=structure,
photon_range=None,
xy_unit=xy_unit,
px_size=px_size,
density=density,
em_avg=em_avg)
self.n_sampler = np.random.poisson
self.frame_range = frame_range
self.intensity_mu_sig = intensity_mu_sig
self.intensity_dist = torch.distributions.normal.Normal(self.intensity_mu_sig[0],
self.intensity_mu_sig[1])
self.intensity_th = intensity_th if intensity_th is not None else 1e-8
self.lifetime_avg = lifetime
self.lifetime_dist = Exponential(1 / self.lifetime_avg) # parse the rate not the scale ...
self.t0_dist = torch.distributions.uniform.Uniform(*self._frame_range_plus)
"""
Determine the total number of emitters. Depends on lifetime, frames and emitters.
(lifetime + 1) because of binning effect.
"""
self._emitter_av_total = self._em_avg * self._num_frames_plus / (self.lifetime_avg + 1)
@property
def _frame_range_plus(self):
"""
Frame range including buffer in front and end to account for build up effects.
"""
return self.frame_range[0] - 3 * self.lifetime_avg, self.frame_range[1] + 3 * self.lifetime_avg
@property
def num_frames(self):
return self.frame_range[1] - self.frame_range[0] + 1
@property
def _num_frames_plus(self):
return self._frame_range_plus[1] - self._frame_range_plus[0] + 1
def sample(self):
"""
Return sampled EmitterSet in the specified frame range.
Returns:
EmitterSet
"""
n = self.n_sampler(self._emitter_av_total)
loose_em = self.sample_loose_emitter(n=n)
em = loose_em.return_emitterset()
em = em.get_subset_frame(*self.frame_range) # because the simulated frame range is larger
return em
def sample_n(self, *args, **kwargs):
raise NotImplementedError
def sample_loose_emitter(self, n) -> decode.generic.emitter.LooseEmitterSet:
"""
Generate loose EmitterSet. Loose emitters are emitters that are not yet binned to frames.
Args:
n: number of 'loose' emitters
Returns:
LooseEmitterSet
"""
xyz = self.structure.sample(n)
"""Draw from intensity distribution but clamp the value so as not to fall below 0."""
intensity = torch.clamp(self.intensity_dist.sample((n,)), self.intensity_th)
"""Distribute emitters in time. Increase the range a bit."""
t0 = self.t0_dist.sample((n,))
ontime = self.lifetime_dist.rsample((n,))
return decode.generic.emitter.LooseEmitterSet(xyz, intensity, ontime, t0, id=torch.arange(n).long(),
xy_unit=self.xy_unit, px_size=self.px_size)
@classmethod
def parse(cls, param, structure, frames: tuple):
return cls(structure=structure,
intensity_mu_sig=param.Simulation.intensity_mu_sig,
lifetime=param.Simulation.lifetime_avg,
xy_unit=param.Simulation.xy_unit,
px_size=param.Camera.px_size,
frame_range=frames,
density=param.Simulation.density,
em_avg=param.Simulation.emitter_av,
intensity_th=param.Simulation.intensity_th)