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)
