def model_3(sequences, lengths, args, batch_size=None, include_prior=True): with ignore_jit_warnings(): num_sequences, max_length, data_dim = map(int, sequences.shape) assert lengths.shape == (num_sequences, ) assert lengths.max() <= max_length hidden_dim = int(args.hidden_dim**0.5) # split between w and x with handlers.mask(mask=include_prior): probs_w = pyro.sample( "probs_w", dist.Dirichlet(0.9 * torch.eye(hidden_dim) + 0.1).to_event(1)) probs_x = pyro.sample( "probs_x", dist.Dirichlet(0.9 * torch.eye(hidden_dim) + 0.1).to_event(1)) probs_y = pyro.sample( "probs_y", dist.Beta(0.1, 0.9).expand([hidden_dim, hidden_dim, data_dim]).to_event(3)) tones_plate = pyro.plate("tones", data_dim, dim=-1) with pyro.plate("sequences", num_sequences, batch_size, dim=-2) as batch: lengths = lengths[batch] w, x = 0, 0 for t in pyro.markov(range(max_length if args.jit else lengths.max())): with handlers.mask(mask=(t < lengths).unsqueeze(-1)): w = pyro.sample("w_{}".format(t), dist.Categorical(probs_w[w]), infer={"enumerate": "parallel"}) x = pyro.sample("x_{}".format(t), dist.Categorical(probs_x[x]), infer={"enumerate": "parallel"}) with tones_plate as tones: pyro.sample("y_{}".format(t), dist.Bernoulli(probs_y[w, x, tones]), obs=sequences[batch, t])
def model_5(sequences, lengths, args, batch_size=None, include_prior=True): with ignore_jit_warnings(): num_sequences, max_length, data_dim = map(int, sequences.shape) assert lengths.shape == (num_sequences, ) assert lengths.max() <= max_length # Initialize a global module instance if needed. global tones_generator if tones_generator is None: tones_generator = TonesGenerator(args, data_dim) pyro.module("tones_generator", tones_generator) with handlers.mask(mask=include_prior): probs_x = pyro.sample( "probs_x", dist.Dirichlet(0.9 * torch.eye(args.hidden_dim) + 0.1).to_event(1)) with pyro.plate("sequences", num_sequences, batch_size, dim=-2) as batch: lengths = lengths[batch] x = 0 y = torch.zeros(data_dim) for t in pyro.markov(range(max_length if args.jit else lengths.max())): with handlers.mask(mask=(t < lengths).unsqueeze(-1)): x = pyro.sample("x_{}".format(t), dist.Categorical(probs_x[x]), infer={"enumerate": "parallel"}) # Note that since each tone depends on all tones at a previous time step # the tones at different time steps now need to live in separate plates. with pyro.plate("tones_{}".format(t), data_dim, dim=-1): y = pyro.sample( "y_{}".format(t), dist.Bernoulli(logits=tones_generator(x, y)), obs=sequences[batch, t])
def model_2(sequences, lengths, args, batch_size=None, include_prior=True): with ignore_jit_warnings(): num_sequences, max_length, data_dim = map(int, sequences.shape) assert lengths.shape == (num_sequences, ) assert lengths.max() <= max_length with handlers.mask(mask=include_prior): probs_x = pyro.sample( "probs_x", dist.Dirichlet(0.9 * torch.eye(args.hidden_dim) + 0.1).to_event(1)) probs_y = pyro.sample( "probs_y", dist.Beta(0.1, 0.9).expand([args.hidden_dim, 2, data_dim]).to_event(3)) tones_plate = pyro.plate("tones", data_dim, dim=-1) with pyro.plate("sequences", num_sequences, batch_size, dim=-2) as batch: lengths = lengths[batch] x, y = 0, 0 for t in pyro.markov(range(max_length if args.jit else lengths.max())): with handlers.mask(mask=(t < lengths).unsqueeze(-1)): x = pyro.sample("x_{}".format(t), dist.Categorical(probs_x[x]), infer={"enumerate": "parallel"}) # Note the broadcasting tricks here: to index probs_y on tensors x and y, # we also need a final tensor for the tones dimension. This is conveniently # provided by the plate associated with that dimension. with tones_plate as tones: y = pyro.sample("y_{}".format(t), dist.Bernoulli(probs_y[x, y, tones]), obs=sequences[batch, t]).long()
def model_0(sequences, lengths, args, batch_size=None, include_prior=True): assert not torch._C._get_tracing_state() num_sequences, max_length, data_dim = sequences.shape with handlers.mask(mask=include_prior): # Our prior on transition probabilities will be: # stay in the same state with 90% probability; uniformly jump to another # state with 10% probability. probs_x = pyro.sample( "probs_x", dist.Dirichlet(0.9 * torch.eye(args.hidden_dim) + 0.1).to_event(1)) # We put a weak prior on the conditional probability of a tone sounding. # We know that on average about 4 of 88 tones are active, so we'll set a # rough weak prior of 10% of the notes being active at any one time. probs_y = pyro.sample( "probs_y", dist.Beta(0.1, 0.9).expand([args.hidden_dim, data_dim]).to_event(2)) # In this first model we'll sequentially iterate over sequences in a # minibatch; this will make it easy to reason about tensor shapes. tones_plate = pyro.plate("tones", data_dim, dim=-1) for i in pyro.plate("sequences", len(sequences), batch_size): length = lengths[i] sequence = sequences[i, :length] x = 0 for t in pyro.markov(range(length)): # On the next line, we'll overwrite the value of x with an updated # value. If we wanted to record all x values, we could instead # write x[t] = pyro.sample(...x[t-1]...). x = pyro.sample("x_{}_{}".format(i, t), dist.Categorical(probs_x[x]), infer={"enumerate": "parallel"}) with tones_plate: pyro.sample("y_{}_{}".format(i, t), dist.Bernoulli(probs_y[x.squeeze(-1)]), obs=sequence[t])
def model_4(sequences, lengths, args, batch_size=None, include_prior=True): with ignore_jit_warnings(): num_sequences, max_length, data_dim = map(int, sequences.shape) assert lengths.shape == (num_sequences,) assert lengths.max() <= max_length hidden_dim = int(args.hidden_dim**0.5) # split between w and x with handlers.mask(mask=include_prior): probs_w = pyro.sample( "probs_w", dist.Dirichlet(0.9 * torch.eye(hidden_dim) + 0.1).to_event(1) ) probs_x = pyro.sample( "probs_x", dist.Dirichlet(0.9 * torch.eye(hidden_dim) + 0.1) .expand_by([hidden_dim]) .to_event(2), ) probs_y = pyro.sample( "probs_y", dist.Beta(0.1, 0.9).expand([hidden_dim, hidden_dim, data_dim]).to_event(3), ) tones_plate = pyro.plate("tones", data_dim, dim=-1) with pyro.plate("sequences", num_sequences, batch_size, dim=-2) as batch: lengths = lengths[batch] # Note the broadcasting tricks here: we declare a hidden torch.arange and # ensure that w and x are always tensors so we can unsqueeze them below, # thus ensuring that the x sample sites have correct distribution shape. w = x = torch.tensor(0, dtype=torch.long) for t in pyro.markov(range(max_length if args.jit else lengths.max())): with handlers.mask(mask=(t < lengths).unsqueeze(-1)): w = pyro.sample( "w_{}".format(t), dist.Categorical(probs_w[w]), infer={"enumerate": "parallel"}, ) x = pyro.sample( "x_{}".format(t), dist.Categorical(Vindex(probs_x)[w, x]), infer={"enumerate": "parallel"}, ) with tones_plate as tones: pyro.sample( "y_{}".format(t), dist.Bernoulli(probs_y[w, x, tones]), obs=sequences[batch, t], )
def model_1(sequences, lengths, args, batch_size=None, include_prior=True): # Sometimes it is safe to ignore jit warnings. Here we use the # pyro.util.ignore_jit_warnings context manager to silence warnings about # conversion to integer, since we know all three numbers will be the same # across all invocations to the model. with ignore_jit_warnings(): num_sequences, max_length, data_dim = map(int, sequences.shape) assert lengths.shape == (num_sequences,) assert lengths.max() <= max_length with handlers.mask(mask=include_prior): probs_x = pyro.sample( "probs_x", dist.Dirichlet(0.9 * torch.eye(args.hidden_dim) + 0.1).to_event(1), ) probs_y = pyro.sample( "probs_y", dist.Beta(0.1, 0.9).expand([args.hidden_dim, data_dim]).to_event(2), ) tones_plate = pyro.plate("tones", data_dim, dim=-1) # We subsample batch_size items out of num_sequences items. Note that since # we're using dim=-1 for the notes plate, we need to batch over a different # dimension, here dim=-2. with pyro.plate("sequences", num_sequences, batch_size, dim=-2) as batch: lengths = lengths[batch] x = 0 # If we are not using the jit, then we can vary the program structure # each call by running for a dynamically determined number of time # steps, lengths.max(). However if we are using the jit, then we try to # keep a single program structure for all minibatches; the fixed # structure ends up being faster since each program structure would # need to trigger a new jit compile stage. for t in pyro.markov(range(max_length if args.jit else lengths.max())): with handlers.mask(mask=(t < lengths).unsqueeze(-1)): x = pyro.sample( "x_{}".format(t), dist.Categorical(probs_x[x]), infer={"enumerate": "parallel"}, ) with tones_plate: pyro.sample( "y_{}".format(t), dist.Bernoulli(probs_y[x.squeeze(-1)]), obs=sequences[batch, t], )
def model_6(sequences, lengths, args, batch_size=None, include_prior=False): num_sequences, max_length, data_dim = sequences.shape assert lengths.shape == (num_sequences, ) assert lengths.max() <= max_length hidden_dim = args.hidden_dim if not args.raftery_parameterization: # Explicitly parameterize the full tensor of transition probabilities, which # has hidden_dim cubed entries. probs_x = pyro.param("probs_x", torch.rand(hidden_dim, hidden_dim, hidden_dim), constraint=constraints.simplex) else: # Use the more parsimonious "Raftery" parameterization of # the tensor of transition probabilities. See reference: # Raftery, A. E. A model for high-order markov chains. # Journal of the Royal Statistical Society. 1985. probs_x1 = pyro.param("probs_x1", torch.rand(hidden_dim, hidden_dim), constraint=constraints.simplex) probs_x2 = pyro.param("probs_x2", torch.rand(hidden_dim, hidden_dim), constraint=constraints.simplex) mix_lambda = pyro.param("mix_lambda", torch.tensor(0.5), constraint=constraints.unit_interval) # we use broadcasting to combine two tensors of shape (hidden_dim, hidden_dim) and # (hidden_dim, 1, hidden_dim) to obtain a tensor of shape (hidden_dim, hidden_dim, hidden_dim) probs_x = mix_lambda * probs_x1 + (1.0 - mix_lambda) * probs_x2.unsqueeze(-2) probs_y = pyro.param("probs_y", torch.rand(hidden_dim, data_dim), constraint=constraints.unit_interval) tones_plate = pyro.plate("tones", data_dim, dim=-1) with pyro.plate("sequences", num_sequences, batch_size, dim=-2) as batch: lengths = lengths[batch] x_curr, x_prev = torch.tensor(0), torch.tensor(0) # we need to pass the argument `history=2' to `pyro.markov()` # since our model is now 2-markov for t in pyro.markov(range(lengths.max()), history=2): with handlers.mask(mask=(t < lengths).unsqueeze(-1)): probs_x_t = Vindex(probs_x)[x_prev, x_curr] x_prev, x_curr = x_curr, pyro.sample( "x_{}".format(t), dist.Categorical(probs_x_t), infer={"enumerate": "parallel"}) with tones_plate: probs_y_t = probs_y[x_curr.squeeze(-1)] pyro.sample("y_{}".format(t), dist.Bernoulli(probs_y_t), obs=sequences[batch, t])
def model_7(sequences, lengths, args, batch_size=None, include_prior=True): with ignore_jit_warnings(): num_sequences, max_length, data_dim = map(int, sequences.shape) assert lengths.shape == (num_sequences,) assert lengths.max() <= max_length with handlers.mask(mask=include_prior): probs_x = pyro.sample( "probs_x", dist.Dirichlet(0.9 * torch.eye(args.hidden_dim) + 0.1).to_event(1), ) probs_y = pyro.sample( "probs_y", dist.Beta(0.1, 0.9).expand([args.hidden_dim, data_dim]).to_event(2), ) tones_plate = pyro.plate("tones", data_dim, dim=-1) # Note that since we're using dim=-2 for the time dimension, we need # to batch sequences over a different dimension, here dim=-3. with pyro.plate("sequences", num_sequences, batch_size, dim=-3) as batch: lengths = lengths[batch] batch = batch[:, None] x_prev = 0 # To vectorize time dimension we use pyro.vectorized_markov(name=...). # With the help of Vindex and additional unsqueezes we can ensure that # dimensions line up properly. for t in pyro.vectorized_markov( name="time", size=int(max_length if args.jit else lengths.max()), dim=-2 ): with handlers.mask(mask=(t < lengths.unsqueeze(-1)).unsqueeze(-1)): x_curr = pyro.sample( "x_{}".format(t), dist.Categorical(probs_x[x_prev]), infer={"enumerate": "parallel"}, ) with tones_plate: pyro.sample( "y_{}".format(t), dist.Bernoulli(probs_y[x_curr.squeeze(-1)]), obs=Vindex(sequences)[batch, t], )
def model(z=None): p = pyro.param("p", torch.tensor([0.75, 0.25])) z = pyro.sample("z", dist.Categorical(p), obs=z) logger.info("z.shape = {}".format(z.shape)) with pyro.plate("data", 3), handlers.mask(mask=mask): pyro.sample("x", dist.Normal(z.type_as(data), 1.0), obs=data)
def guide(self): r""" Variational Distribution """ # global parameters pyro.sample( "gain", dist.Gamma( pyro.param("gain_loc") * pyro.param("gain_beta"), pyro.param("gain_beta"), ), ) pyro.sample( "alpha", dist.Dirichlet( pyro.param("alpha_mean") * pyro.param("alpha_size")).to_event(1), ) pyro.sample( "pi", dist.Dirichlet(pyro.param("pi_mean") * pyro.param("pi_size")).to_event(1), ) pyro.sample( "lamda", dist.Gamma( pyro.param("lamda_loc") * pyro.param("lamda_beta"), pyro.param("lamda_beta"), ).to_event(1), ) pyro.sample( "proximity", AffineBeta( pyro.param("proximity_loc"), pyro.param("proximity_size"), 0, (self.data.P + 1) / math.sqrt(12), ), ) # aoi sites aois = pyro.plate( "aois", self.data.Nt, subsample=self.n, subsample_size=self.nbatch_size, dim=-2, ) # time frames frames = pyro.plate( "frames", self.data.F, subsample=self.f, subsample_size=self.fbatch_size, dim=-1, ) with aois as ndx: ndx = ndx[:, None] mask = Vindex(self.data.mask)[ndx].to(self.device) with handlers.mask(mask=mask): pyro.sample( "background_mean", dist.Delta( Vindex( pyro.param("background_mean_loc"))[ndx, 0]).to_event(1), ) pyro.sample( "background_std", dist.Delta( Vindex( pyro.param("background_std_loc"))[ndx, 0]).to_event(1), ) with frames as fdx: # sample background intensity pyro.sample( "background", dist.Gamma( Vindex(pyro.param("b_loc"))[ndx, fdx] * Vindex(pyro.param("b_beta"))[ndx, fdx], Vindex(pyro.param("b_beta"))[ndx, fdx], ).to_event(1), ) for qdx in range(self.Q): for kdx in range(self.K): # sample spot presence m m = pyro.sample( f"m_k{kdx}_q{qdx}", dist.Bernoulli( Vindex(pyro.param("m_probs"))[kdx, ndx, fdx, qdx]), infer={"enumerate": "parallel"}, ) with handlers.mask(mask=m > 0): # sample spot variables pyro.sample( f"height_k{kdx}_q{qdx}", dist.Gamma( Vindex(pyro.param("h_loc"))[kdx, ndx, fdx, qdx] * Vindex(pyro.param("h_beta"))[kdx, ndx, fdx, qdx], Vindex(pyro.param("h_beta"))[kdx, ndx, fdx, qdx], ), ) pyro.sample( f"width_k{kdx}_q{qdx}", AffineBeta( Vindex(pyro.param("w_mean"))[kdx, ndx, fdx, qdx], Vindex(pyro.param("w_size"))[kdx, ndx, fdx, qdx], self.priors["width_min"], self.priors["width_max"], ), ) pyro.sample( f"x_k{kdx}_q{qdx}", AffineBeta( Vindex(pyro.param("x_mean"))[kdx, ndx, fdx, qdx], Vindex(pyro.param("size"))[kdx, ndx, fdx, qdx], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), ) pyro.sample( f"y_k{kdx}_q{qdx}", AffineBeta( Vindex(pyro.param("y_mean"))[kdx, ndx, fdx, qdx], Vindex(pyro.param("size"))[kdx, ndx, fdx, qdx], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), )
def model(self): """ **Generative Model** """ # global parameters gain = pyro.sample("gain", dist.HalfNormal(self.priors["gain_std"])) init = pyro.sample( "init", dist.Dirichlet(torch.ones(self.Q, self.S + 1) / (self.S + 1)).to_event(1), ) init = expand_offtarget(init) trans = pyro.sample( "trans", dist.Dirichlet( torch.ones(self.Q, self.S + 1, self.S + 1) / (self.S + 1)).to_event(2), ) trans = expand_offtarget(trans) lamda = pyro.sample( "lamda", dist.Exponential(torch.full( (self.Q, ), self.priors["lamda_rate"])).to_event(1), ) proximity = pyro.sample( "proximity", dist.Exponential(self.priors["proximity_rate"])) size = torch.stack( ( torch.full_like(proximity, 2.0), (((self.data.P + 1) / (2 * proximity))**2 - 1), ), dim=-1, ) # spots spots = pyro.plate("spots", self.K) # aoi sites aois = pyro.plate( "aois", self.data.Nt, subsample=self.n, subsample_size=self.nbatch_size, dim=-3, ) # time frames frames = (pyro.vectorized_markov( name="frames", size=self.data.F, dim=-2) if self.vectorized else pyro.markov(range(self.data.F))) # color channels channels = pyro.plate( "channels", self.data.C, dim=-1, ) with channels as cdx, aois as ndx: ndx = ndx[:, None, None] mask = Vindex(self.data.mask)[ndx].to(self.device) with handlers.mask(mask=mask): # background mean and std background_mean = pyro.sample( "background_mean", dist.HalfNormal(self.priors["background_mean_std"]), ) background_std = pyro.sample( "background_std", dist.HalfNormal(self.priors["background_std_std"])) z_prev = None for fdx in frames: if self.vectorized: fsx, fdx = fdx fdx = torch.as_tensor(fdx) fdx = fdx.unsqueeze(-1) else: fsx = fdx # fetch data obs, target_locs, is_ontarget = self.data.fetch( ndx, fdx, cdx) # sample background intensity background = pyro.sample( f"background_f{fsx}", dist.Gamma( (background_mean / background_std)**2, background_mean / background_std**2, ), ) # sample hidden model state (1+S,) z_probs = (Vindex(init)[..., cdx, :, is_ontarget.long()] if z_prev is None else Vindex(trans)[..., cdx, z_prev, :, is_ontarget.long()]) z_curr = pyro.sample(f"z_f{fsx}", dist.Categorical(z_probs)) theta = pyro.sample( f"theta_f{fsx}", dist.Categorical( Vindex(probs_theta( self.K, self.device))[torch.clamp(z_curr, min=0, max=1)]), infer={"enumerate": "parallel"}, ) onehot_theta = one_hot(theta, num_classes=1 + self.K) ms, heights, widths, xs, ys = [], [], [], [], [] for kdx in spots: specific = onehot_theta[..., 1 + kdx] # spot presence m_probs = Vindex(probs_m(lamda, self.K))[..., cdx, theta, kdx] m = pyro.sample( f"m_k{kdx}_f{fsx}", dist.Categorical( torch.stack((1 - m_probs, m_probs), -1)), ) with handlers.mask(mask=m > 0): # sample spot variables height = pyro.sample( f"height_k{kdx}_f{fsx}", dist.HalfNormal(self.priors["height_std"]), ) width = pyro.sample( f"width_k{kdx}_f{fsx}", AffineBeta( 1.5, 2, self.priors["width_min"], self.priors["width_max"], ), ) x = pyro.sample( f"x_k{kdx}_f{fsx}", AffineBeta( 0, Vindex(size)[..., specific], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), ) y = pyro.sample( f"y_k{kdx}_f{fsx}", AffineBeta( 0, Vindex(size)[..., specific], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), ) # append ms.append(m) heights.append(height) widths.append(width) xs.append(x) ys.append(y) # observed data pyro.sample( f"data_f{fsx}", KSMOGN( torch.stack(heights, -1), torch.stack(widths, -1), torch.stack(xs, -1), torch.stack(ys, -1), target_locs, background, gain, self.data.offset.samples, self.data.offset.logits.to(self.dtype), self.data.P, torch.stack(torch.broadcast_tensors(*ms), -1), use_pykeops=self.use_pykeops, ), obs=obs, ) z_prev = z_curr
def guide(self): """ **Variational Distribution** """ # global parameters pyro.sample( "gain", dist.Gamma( pyro.param("gain_loc") * pyro.param("gain_beta"), pyro.param("gain_beta"), ), ) pyro.sample( "init", dist.Dirichlet(pyro.param("init_mean") * pyro.param("init_size")).to_event(1), ) pyro.sample( "trans", dist.Dirichlet( pyro.param("trans_mean") * pyro.param("trans_size")).to_event(2), ) pyro.sample( "lamda", dist.Gamma( pyro.param("lamda_loc") * pyro.param("lamda_beta"), pyro.param("lamda_beta"), ).to_event(1), ) pyro.sample( "proximity", AffineBeta( pyro.param("proximity_loc"), pyro.param("proximity_size"), 0, (self.data.P + 1) / math.sqrt(12), ), ) # spots spots = pyro.plate("spots", self.K) # aoi sites aois = pyro.plate( "aois", self.data.Nt, subsample=self.n, subsample_size=self.nbatch_size, dim=-3, ) # time frames frames = (pyro.vectorized_markov( name="frames", size=self.data.F, dim=-2) if self.vectorized else pyro.markov(range(self.data.F))) # color channels channels = pyro.plate( "channels", self.data.C, dim=-1, ) with channels as cdx, aois as ndx: ndx = ndx[:, None, None] mask = Vindex(self.data.mask)[ndx].to(self.device) with handlers.mask(mask=mask): pyro.sample( "background_mean", dist.Delta( Vindex(pyro.param("background_mean_loc"))[ndx, 0, cdx]), ) pyro.sample( "background_std", dist.Delta( Vindex(pyro.param("background_std_loc"))[ndx, 0, cdx]), ) z_prev = None for fdx in frames: if self.vectorized: fsx, fdx = fdx fdx = torch.as_tensor(fdx) fdx = fdx.unsqueeze(-1) else: fsx = fdx # sample background intensity pyro.sample( f"background_f{fsx}", dist.Gamma( Vindex(pyro.param("b_loc"))[ndx, fdx, cdx] * Vindex(pyro.param("b_beta"))[ndx, fdx, cdx], Vindex(pyro.param("b_beta"))[ndx, fdx, cdx], ), ) # sample hidden model state z_probs = (Vindex(pyro.param("z_trans"))[ndx, fdx, cdx, 0] if z_prev is None else Vindex( pyro.param("z_trans"))[ndx, fdx, cdx, z_prev]) z_curr = pyro.sample( f"z_f{fsx}", dist.Categorical(z_probs), infer={"enumerate": "parallel"}, ) for kdx in spots: # spot presence m_probs = Vindex(pyro.param("m_probs"))[z_curr, kdx, ndx, fdx, cdx] m = pyro.sample( f"m_k{kdx}_f{fsx}", dist.Categorical( torch.stack((1 - m_probs, m_probs), -1)), infer={"enumerate": "parallel"}, ) with handlers.mask(mask=m > 0): # sample spot variables pyro.sample( f"height_k{kdx}_f{fsx}", dist.Gamma( Vindex(pyro.param("h_loc"))[kdx, ndx, fdx, cdx] * Vindex(pyro.param("h_beta"))[kdx, ndx, fdx, cdx], Vindex(pyro.param("h_beta"))[kdx, ndx, fdx, cdx], ), ) pyro.sample( f"width_k{kdx}_f{fsx}", AffineBeta( Vindex(pyro.param("w_mean"))[kdx, ndx, fdx, cdx], Vindex(pyro.param("w_size"))[kdx, ndx, fdx, cdx], self.priors["width_min"], self.priors["width_max"], ), ) pyro.sample( f"x_k{kdx}_f{fsx}", AffineBeta( Vindex(pyro.param("x_mean"))[kdx, ndx, fdx, cdx], Vindex(pyro.param("size"))[kdx, ndx, fdx, cdx], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), ) pyro.sample( f"y_k{kdx}_f{fsx}", AffineBeta( Vindex(pyro.param("y_mean"))[kdx, ndx, fdx, cdx], Vindex(pyro.param("size"))[kdx, ndx, fdx, cdx], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), ) z_prev = z_curr
def model(self): r""" **Generative Model** Model parameters: +-----------------+-----------+-------------------------------------+ | Parameter | Shape | Description | +=================+===========+=====================================+ | |g| - :math:`g` | (1,) | camera gain | +-----------------+-----------+-------------------------------------+ | |sigma| - |prox|| (1,) | proximity | +-----------------+-----------+-------------------------------------+ | ``lamda`` - |ld|| (1,) | average rate of target-nonspecific | | | | binding | +-----------------+-----------+-------------------------------------+ | ``pi`` - |pi| | (1,) | average binding probability of | | | | target-specific binding | +-----------------+-----------+-------------------------------------+ | |bg| - |b| | (N, F) | background intensity | +-----------------+-----------+-------------------------------------+ | |z| - :math:`z` | (N, F) | target-specific spot presence | +-----------------+-----------+-------------------------------------+ | |t| - |theta| | (N, F) | target-specific spot index | +-----------------+-----------+-------------------------------------+ | |m| - :math:`m` | (K, N, F) | spot presence indicator | +-----------------+-----------+-------------------------------------+ | |h| - :math:`h` | (K, N, F) | spot intensity | +-----------------+-----------+-------------------------------------+ | |w| - :math:`w` | (K, N, F) | spot width | +-----------------+-----------+-------------------------------------+ | |x| - :math:`x` | (K, N, F) | spot position on x-axis | +-----------------+-----------+-------------------------------------+ | |y| - :math:`y` | (K, N, F) | spot position on y-axis | +-----------------+-----------+-------------------------------------+ | |D| - :math:`D` | |shape| | observed images | +-----------------+-----------+-------------------------------------+ .. |ps| replace:: :math:`p(\mathsf{specific})` .. |theta| replace:: :math:`\theta` .. |prox| replace:: :math:`\sigma^{xy}` .. |ld| replace:: :math:`\lambda` .. |b| replace:: :math:`b` .. |shape| replace:: (N, F, P, P) .. |sigma| replace:: ``proximity`` .. |bg| replace:: ``background`` .. |h| replace:: ``height`` .. |w| replace:: ``width`` .. |D| replace:: ``data`` .. |m| replace:: ``m`` .. |z| replace:: ``z`` .. |t| replace:: ``theta`` .. |x| replace:: ``x`` .. |y| replace:: ``y`` .. |pi| replace:: :math:`\pi` .. |g| replace:: ``gain`` Full joint distribution: .. math:: \begin{aligned} p(D, \phi) =~&p(g) p(\sigma^{xy}) p(\pi) p(\lambda) \prod_{\mathsf{AOI}} \left[ p(\mu^b) p(\sigma^b) \prod_{\mathsf{frame}} \left[ \vphantom{\prod_{F}} p(b | \mu^b, \sigma^b) p(z | \pi) p(\theta | z) \vphantom{\prod_{\substack{\mathsf{pixelX} \\ \mathsf{pixelY}}}} \cdot \right. \right. \\ &\prod_{\mathsf{spot}} \left[ \vphantom{\prod_{F}} p(m | \theta, \lambda) p(h) p(w) p(x | \sigma^{xy}, \theta) p(y | \sigma^{xy}, \theta) \right] \left. \left. \prod_{\substack{\mathsf{pixelX} \\ \mathsf{pixelY}}} \sum_{\delta} p(\delta) p(D | \mu^I, g, \delta) \right] \right] \end{aligned} :math:`z` and :math:`\theta` marginalized joint distribution: .. math:: \begin{aligned} \sum_{z, \theta} p(D, \phi) =~&p(g) p(\sigma^{xy}) p(\pi) p(\lambda) \prod_{\mathsf{AOI}} \left[ p(\mu^b) p(\sigma^b) \prod_{\mathsf{frame}} \left[ \vphantom{\prod_{F}} p(b | \mu^b, \sigma^b) \sum_{z} p(z | \pi) \sum_{\theta} p(\theta | z) \vphantom{\prod_{\substack{\mathsf{pixelX} \\ \mathsf{pixelY}}}} \cdot \right. \right. \\ &\prod_{\mathsf{spot}} \left[ \vphantom{\prod_{F}} p(m | \theta, \lambda) p(h) p(w) p(x | \sigma^{xy}, \theta) p(y | \sigma^{xy}, \theta) \right] \left. \left. \prod_{\substack{\mathsf{pixelX} \\ \mathsf{pixelY}}} \sum_{\delta} p(\delta) p(D | \mu^I, g, \delta) \right] \right] \end{aligned} """ # global parameters gain = pyro.sample("gain", dist.HalfNormal(self.gain_std)) pi = pyro.sample("pi", dist.Dirichlet(torch.ones(self.S + 1) / (self.S + 1))) pi = expand_offtarget(pi) lamda = pyro.sample("lamda", dist.Exponential(self.lamda_rate)) proximity = pyro.sample("proximity", dist.Exponential(self.proximity_rate)) size = torch.stack( ( torch.full_like(proximity, 2.0), (((self.data.P + 1) / (2 * proximity))**2 - 1), ), dim=-1, ) # spots spots = pyro.plate("spots", self.K) # aoi sites aois = pyro.plate( "aois", self.data.Nt, subsample=self.n, subsample_size=self.nbatch_size, dim=-2, ) # time frames frames = pyro.plate( "frames", self.data.F, subsample=self.f, subsample_size=self.fbatch_size, dim=-1, ) with aois as ndx: ndx = ndx[:, None] # background mean and std background_mean = pyro.sample( "background_mean", dist.HalfNormal(self.background_mean_std)) background_std = pyro.sample( "background_std", dist.HalfNormal(self.background_std_std)) with frames as fdx: # fetch data obs, target_locs, is_ontarget = self.data.fetch( ndx, fdx, self.cdx) # sample background intensity background = pyro.sample( "background", dist.Gamma( (background_mean / background_std)**2, background_mean / background_std**2, ), ) # sample hidden model state (1+S,) z = pyro.sample( "z", dist.Categorical(Vindex(pi)[..., :, is_ontarget.long()]), infer={"enumerate": "parallel"}, ) theta = pyro.sample( "theta", dist.Categorical( Vindex(probs_theta(self.K, self.device))[torch.clamp(z, min=0, max=1)]), infer={"enumerate": "parallel"}, ) onehot_theta = one_hot(theta, num_classes=1 + self.K) ms, heights, widths, xs, ys = [], [], [], [], [] for kdx in spots: specific = onehot_theta[..., 1 + kdx] # spot presence m = pyro.sample( f"m_{kdx}", dist.Bernoulli( Vindex(probs_m(lamda, self.K))[..., theta, kdx]), ) with handlers.mask(mask=m > 0): # sample spot variables height = pyro.sample( f"height_{kdx}", dist.HalfNormal(self.height_std), ) width = pyro.sample( f"width_{kdx}", AffineBeta( 1.5, 2, self.width_min, self.width_max, ), ) x = pyro.sample( f"x_{kdx}", AffineBeta( 0, Vindex(size)[..., specific], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), ) y = pyro.sample( f"y_{kdx}", AffineBeta( 0, Vindex(size)[..., specific], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), ) # append ms.append(m) heights.append(height) widths.append(width) xs.append(x) ys.append(y) # observed data pyro.sample( "data", KSMOGN( torch.stack(heights, -1), torch.stack(widths, -1), torch.stack(xs, -1), torch.stack(ys, -1), target_locs, background, gain, self.data.offset.samples, self.data.offset.logits.to(self.dtype), self.data.P, torch.stack(torch.broadcast_tensors(*ms), -1), self.use_pykeops, ), obs=obs, )
def guide(self): r""" **Variational Distribution** .. math:: \begin{aligned} q(\phi \setminus \{z, \theta\}) =~&q(g) q(\sigma^{xy}) q(\pi) q(\lambda) \cdot \\ &\prod_{\mathsf{AOI}} \left[ q(\mu^b) q(\sigma^b) \prod_{\mathsf{frame}} \left[ \vphantom{\prod_{F}} q(b) \prod_{\mathsf{spot}} q(m) q(h | m) q(w | m) q(x | m) q(y | m) \right] \right] \end{aligned} """ # global parameters pyro.sample( "gain", dist.Gamma( pyro.param("gain_loc") * pyro.param("gain_beta"), pyro.param("gain_beta"), ), ) pyro.sample( "pi", dist.Dirichlet(pyro.param("pi_mean") * pyro.param("pi_size"))) pyro.sample( "lamda", dist.Gamma( pyro.param("lamda_loc") * pyro.param("lamda_beta"), pyro.param("lamda_beta"), ), ) pyro.sample( "proximity", AffineBeta( pyro.param("proximity_loc"), pyro.param("proximity_size"), 0, (self.data.P + 1) / math.sqrt(12), ), ) # spots spots = pyro.plate("spots", self.K) # aoi sites aois = pyro.plate( "aois", self.data.Nt, subsample=self.n, subsample_size=self.nbatch_size, dim=-2, ) # time frames frames = pyro.plate( "frames", self.data.F, subsample=self.f, subsample_size=self.fbatch_size, dim=-1, ) with aois as ndx: ndx = ndx[:, None] pyro.sample( "background_mean", dist.Delta(Vindex(pyro.param("background_mean_loc"))[ndx, 0]), ) pyro.sample( "background_std", dist.Delta(Vindex(pyro.param("background_std_loc"))[ndx, 0]), ) with frames as fdx: # sample background intensity pyro.sample( "background", dist.Gamma( Vindex(pyro.param("b_loc"))[ndx, fdx] * Vindex(pyro.param("b_beta"))[ndx, fdx], Vindex(pyro.param("b_beta"))[ndx, fdx], ), ) for kdx in spots: # sample spot presence m m = pyro.sample( f"m_{kdx}", dist.Bernoulli( Vindex(pyro.param("m_probs"))[kdx, ndx, fdx]), infer={"enumerate": "parallel"}, ) with handlers.mask(mask=m > 0): # sample spot variables pyro.sample( f"height_{kdx}", dist.Gamma( Vindex(pyro.param("h_loc"))[kdx, ndx, fdx] * Vindex(pyro.param("h_beta"))[kdx, ndx, fdx], Vindex(pyro.param("h_beta"))[kdx, ndx, fdx], ), ) pyro.sample( f"width_{kdx}", AffineBeta( Vindex(pyro.param("w_mean"))[kdx, ndx, fdx], Vindex(pyro.param("w_size"))[kdx, ndx, fdx], 0.75, 2.25, ), ) pyro.sample( f"x_{kdx}", AffineBeta( Vindex(pyro.param("x_mean"))[kdx, ndx, fdx], Vindex(pyro.param("size"))[kdx, ndx, fdx], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), ) pyro.sample( f"y_{kdx}", AffineBeta( Vindex(pyro.param("y_mean"))[kdx, ndx, fdx], Vindex(pyro.param("size"))[kdx, ndx, fdx], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), )
def guide(self): """ **Variational Distribution** """ # global parameters pyro.sample( "gain", dist.Gamma( pyro.param("gain_loc") * pyro.param("gain_beta"), pyro.param("gain_beta"), ), ) pyro.sample( "init", dist.Dirichlet(pyro.param("init_mean") * pyro.param("init_size"))) pyro.sample( "trans", dist.Dirichlet( pyro.param("trans_mean") * pyro.param("trans_size")).to_event(1), ) pyro.sample( "lamda", dist.Gamma( pyro.param("lamda_loc") * pyro.param("lamda_beta"), pyro.param("lamda_beta"), ), ) pyro.sample( "proximity", AffineBeta( pyro.param("proximity_loc"), pyro.param("proximity_size"), 0, (self.data.P + 1) / math.sqrt(12), ), ) # spots spots = pyro.plate("spots", self.K) # aoi sites aois = pyro.plate( "aois", self.data.Nt, subsample=self.n, subsample_size=self.nbatch_size, dim=-2, ) # time frames frames = (pyro.vectorized_markov( name="frames", size=self.data.F, dim=-1) if self.vectorized else pyro.markov(range(self.data.F))) with aois as ndx: ndx = ndx[:, None] pyro.sample( "background_mean", dist.Delta(Vindex(pyro.param("background_mean_loc"))[ndx, 0]), ) pyro.sample( "background_std", dist.Delta(Vindex(pyro.param("background_std_loc"))[ndx, 0]), ) z_prev = None for fdx in frames: if self.vectorized: fsx, fdx = fdx else: fsx = fdx # sample background intensity pyro.sample( f"background_{fsx}", dist.Gamma( Vindex(pyro.param("b_loc"))[ndx, fdx] * Vindex(pyro.param("b_beta"))[ndx, fdx], Vindex(pyro.param("b_beta"))[ndx, fdx], ), ) # sample hidden model state z_probs = (Vindex(pyro.param("z_trans"))[ndx, fdx, 0] if isinstance(fdx, int) and fdx < 1 else Vindex( pyro.param("z_trans"))[ndx, fdx, z_prev]) z_curr = pyro.sample( f"z_{fsx}", dist.Categorical(z_probs), infer={"enumerate": "parallel"}, ) for kdx in spots: # spot presence m_probs = Vindex(pyro.param("m_probs"))[z_curr, kdx, ndx, fdx] m = pyro.sample( f"m_{kdx}_{fsx}", dist.Categorical( torch.stack((1 - m_probs, m_probs), -1)), infer={"enumerate": "parallel"}, ) with handlers.mask(mask=m > 0): # sample spot variables pyro.sample( f"height_{kdx}_{fsx}", dist.Gamma( Vindex(pyro.param("h_loc"))[kdx, ndx, fdx] * Vindex(pyro.param("h_beta"))[kdx, ndx, fdx], Vindex(pyro.param("h_beta"))[kdx, ndx, fdx], ), ) pyro.sample( f"width_{kdx}_{fsx}", AffineBeta( Vindex(pyro.param("w_mean"))[kdx, ndx, fdx], Vindex(pyro.param("w_size"))[kdx, ndx, fdx], 0.75, 2.25, ), ) pyro.sample( f"x_{kdx}_{fsx}", AffineBeta( Vindex(pyro.param("x_mean"))[kdx, ndx, fdx], Vindex(pyro.param("size"))[kdx, ndx, fdx], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), ) pyro.sample( f"y_{kdx}_{fsx}", AffineBeta( Vindex(pyro.param("y_mean"))[kdx, ndx, fdx], Vindex(pyro.param("size"))[kdx, ndx, fdx], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), ) z_prev = z_curr
def model(self): r""" Generative Model """ # global parameters gain = pyro.sample("gain", dist.HalfNormal(self.priors["gain_std"])) alpha = pyro.sample( "alpha", dist.Dirichlet( torch.ones((self.Q, self.data.C)) + torch.eye(self.Q) * 9).to_event(1), ) pi = pyro.sample( "pi", dist.Dirichlet(torch.ones( (self.Q, self.S + 1)) / (self.S + 1)).to_event(1), ) pi = expand_offtarget(pi) lamda = pyro.sample( "lamda", dist.Exponential(torch.full( (self.Q, ), self.priors["lamda_rate"])).to_event(1), ) proximity = pyro.sample( "proximity", dist.Exponential(self.priors["proximity_rate"])) size = torch.stack( ( torch.full_like(proximity, 2.0), (((self.data.P + 1) / (2 * proximity))**2 - 1), ), dim=-1, ) # aoi sites aois = pyro.plate( "aois", self.data.Nt, subsample=self.n, subsample_size=self.nbatch_size, dim=-2, ) # time frames frames = pyro.plate( "frames", self.data.F, subsample=self.f, subsample_size=self.fbatch_size, dim=-1, ) with aois as ndx: ndx = ndx[:, None] mask = Vindex(self.data.mask)[ndx].to(self.device) with handlers.mask(mask=mask): # background mean and std background_mean = pyro.sample( "background_mean", dist.HalfNormal(self.priors["background_mean_std"]).expand( (self.data.C, )).to_event(1), ) background_std = pyro.sample( "background_std", dist.HalfNormal(self.priors["background_std_std"]).expand( (self.data.C, )).to_event(1), ) with frames as fdx: # fetch data obs, target_locs, is_ontarget = self.data.fetch( ndx.unsqueeze(-1), fdx.unsqueeze(-1), torch.arange(self.data.C)) # sample background intensity background = pyro.sample( "background", dist.Gamma( (background_mean / background_std)**2, background_mean / background_std**2, ).to_event(1), ) ms, heights, widths, xs, ys = [], [], [], [], [] is_ontarget = is_ontarget.squeeze(-1) for qdx in range(self.Q): # sample hidden model state (1+S,) z_probs = Vindex(pi)[..., qdx, :, is_ontarget.long()] z = pyro.sample( f"z_q{qdx}", dist.Categorical(z_probs), infer={"enumerate": "parallel"}, ) theta = pyro.sample( f"theta_q{qdx}", dist.Categorical( Vindex(probs_theta( self.K, self.device))[torch.clamp(z, min=0, max=1)]), infer={"enumerate": "parallel"}, ) onehot_theta = one_hot(theta, num_classes=1 + self.K) for kdx in range(self.K): specific = onehot_theta[..., 1 + kdx] # spot presence m = pyro.sample( f"m_k{kdx}_q{qdx}", dist.Bernoulli( Vindex(probs_m(lamda, self.K))[..., qdx, theta, kdx]), ) with handlers.mask(mask=m > 0): # sample spot variables height = pyro.sample( f"height_k{kdx}_q{qdx}", dist.HalfNormal(self.priors["height_std"]), ) width = pyro.sample( f"width_k{kdx}_q{qdx}", AffineBeta( 1.5, 2, self.priors["width_min"], self.priors["width_max"], ), ) x = pyro.sample( f"x_k{kdx}_q{qdx}", AffineBeta( 0, Vindex(size)[..., specific], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), ) y = pyro.sample( f"y_k{kdx}_q{qdx}", AffineBeta( 0, Vindex(size)[..., specific], -(self.data.P + 1) / 2, (self.data.P + 1) / 2, ), ) # append ms.append(m) heights.append(height) widths.append(width) xs.append(x) ys.append(y) heights = torch.stack( [ torch.stack(heights[q * self.K:(1 + q) * self.K], -1) for q in range(self.Q) ], -2, ) widths = torch.stack( [ torch.stack(widths[q * self.K:(1 + q) * self.K], -1) for q in range(self.Q) ], -2, ) xs = torch.stack( [ torch.stack(xs[q * self.K:(1 + q) * self.K], -1) for q in range(self.Q) ], -2, ) ys = torch.stack( [ torch.stack(ys[q * self.Q:(1 + q) * self.K], -1) for q in range(self.Q) ], -2, ) ms = torch.broadcast_tensors(*ms) ms = torch.stack( [ torch.stack(ms[q * self.Q:(1 + q) * self.K], -1) for q in range(self.Q) ], -2, ) # observed data pyro.sample( "data", KSMOGN( heights, widths, xs, ys, target_locs, background, gain, self.data.offset.samples, self.data.offset.logits.to(self.dtype), self.data.P, ms, alpha, use_pykeops=self.use_pykeops, ), obs=obs, )
def ttfb_model(data, control, Tmax): r""" Eq. 4 and Eq. 7 in:: @article{friedman2015multi, title={Multi-wavelength single-molecule fluorescence analysis of transcription mechanisms}, author={Friedman, Larry J and Gelles, Jeff}, journal={Methods}, volume={86}, pages={27--36}, year={2015}, publisher={Elsevier} } :param data: time prior to the first binding at the target location :param control: time prior to the first binding at the control location :param Tmax: entire observation interval """ ka = pyro.param( "ka", lambda: torch.full((data.shape[0], 1), 0.001), constraint=constraints.positive, ) kns = pyro.param( "kns", lambda: torch.full((data.shape[0], 1), 0.001), constraint=constraints.positive, ) Af = pyro.param( "Af", lambda: torch.full((data.shape[0], 1), 0.9), constraint=constraints.unit_interval, ) k = torch.stack([kns, ka + kns]) # on-target data mask = (data < Tmax) & (data > 0) tau = data.masked_fill(~mask, 1.0) with pyro.plate("bootstrap", data.shape[0], dim=-2) as bdx: with pyro.plate("N", data.shape[1], dim=-1): active = pyro.sample( "active", dist.Bernoulli(Af), infer={"enumerate": "parallel"} ) with handlers.mask(mask=(data == Tmax)): pyro.factor("Tmax", -Vindex(k)[active.long().squeeze(-1), bdx] * Tmax) # pyro.factor("Tmax", -k * Tmax) with handlers.mask(mask=mask): pyro.sample( "tau", dist.Exponential(Vindex(k)[active.long().squeeze(-1), bdx]), obs=tau, ) # pyro.sample("tau", dist.Exponential(k), obs=tau) # negative control data if control is not None: mask = (control < Tmax) & (control > 0) tauc = control.masked_fill(~mask, 1.0) with pyro.plate("bootstrapc", control.shape[0], dim=-2): with pyro.plate("Nc", control.shape[1], dim=-1): with handlers.mask(mask=(control == Tmax)): pyro.factor("Tmaxc", -kns * Tmax) with handlers.mask(mask=mask): pyro.sample("tauc", dist.Exponential(kns), obs=tauc)