def pytest_configure(config):
try:
import funsor
except ImportError:
pass
else:
funsor.set_backend("torch")
def __init__(self, *args, **kwargs):
if collapse._coerce is None:
import funsor
from funsor.distribution import CoerceDistributionToFunsor
funsor.set_backend("jax")
collapse._coerce = CoerceDistributionToFunsor("jax")
super().__init__(*args, **kwargs)
def main(args):
funsor.set_backend("torch")
# XXX Temporary fix after https://github.com/pyro-ppl/pyro/pull/2701
import pyro
pyro.enable_validation(False)
encoder = Encoder()
decoder = Decoder()
encode = funsor.function(Reals[28, 28], (Reals[20], Reals[20]))(encoder)
decode = funsor.function(Reals[20], Reals[28, 28])(decoder)
@funsor.interpretation(funsor.montecarlo.MonteCarlo())
def loss_function(data, subsample_scale):
# Lazily sample from the guide.
loc, scale = encode(data)
q = funsor.Independent(dist.Normal(loc['i'], scale['i'], value='z_i'),
'z', 'i', 'z_i')
# Evaluate the model likelihood at the lazy value z.
probs = decode('z')
p = dist.Bernoulli(probs['x', 'y'], value=data['x', 'y'])
p = p.reduce(ops.add, {'x', 'y'})
# Construct an elbo. This is where sampling happens.
elbo = funsor.Integrate(q, p - q, 'z')
elbo = elbo.reduce(ops.add, 'batch') * subsample_scale
loss = -elbo
return loss
train_loader = torch.utils.data.DataLoader(datasets.MNIST(
DATA_PATH, train=True, download=True, transform=transforms.ToTensor()),
batch_size=args.batch_size,
shuffle=True)
encoder.train()
decoder.train()
optimizer = optim.Adam(list(encoder.parameters()) +
list(decoder.parameters()),
lr=1e-3)
for epoch in range(args.num_epochs):
train_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
subsample_scale = float(len(train_loader.dataset) / len(data))
data = data[:, 0, :, :]
data = funsor.Tensor(data, OrderedDict(batch=Bint[len(data)]))
optimizer.zero_grad()
loss = loss_function(data, subsample_scale)
assert isinstance(loss, funsor.Tensor), loss.pretty()
loss.data.backward()
train_loss += loss.item()
optimizer.step()
if batch_idx % 50 == 0:
print(' loss = {}'.format(loss.item()))
if batch_idx and args.smoke_test:
return
print('epoch {} train_loss = {}'.format(epoch, train_loss))
def __init__(self, *args, **kwargs):
if CollapseMessenger._coerce is None:
import funsor
from funsor.distribution import CoerceDistributionToFunsor
funsor.set_backend("torch")
CollapseMessenger._coerce = CoerceDistributionToFunsor("torch")
self._block = False
super().__init__(*args, **kwargs)
def _import_funsor():
try:
import funsor
except ImportError as e:
raise ImportError(
'AutoGaussian(..., backend="funsor") requires funsor. '
"Try installing via: pip install pyro-ppl[funsor]"
) from e
funsor.set_backend("torch")
return funsor
def main(args):
funsor.set_backend("torch")
# Declare parameters.
trans_probs = torch.tensor([[0.2, 0.8], [0.7, 0.3]], requires_grad=True)
emit_probs = torch.tensor([[0.4, 0.6], [0.1, 0.9]], requires_grad=True)
params = [trans_probs, emit_probs]
# A discrete HMM model.
def model(data):
log_prob = funsor.to_funsor(0.)
trans = dist.Categorical(probs=funsor.Tensor(
trans_probs,
inputs=OrderedDict([('prev', funsor.Bint[args.hidden_dim])]),
))
emit = dist.Categorical(probs=funsor.Tensor(
emit_probs,
inputs=OrderedDict([('latent', funsor.Bint[args.hidden_dim])]),
))
x_curr = funsor.Number(0, args.hidden_dim)
for t, y in enumerate(data):
x_prev = x_curr
# A delayed sample statement.
x_curr = funsor.Variable('x_{}'.format(t),
funsor.Bint[args.hidden_dim])
log_prob += trans(prev=x_prev, value=x_curr)
if not args.lazy and isinstance(x_prev, funsor.Variable):
log_prob = log_prob.reduce(ops.logaddexp, x_prev.name)
log_prob += emit(latent=x_curr, value=funsor.Tensor(y, dtype=2))
log_prob = log_prob.reduce(ops.logaddexp)
return log_prob
# Train model parameters.
data = torch.ones(args.time_steps, dtype=torch.long)
optim = torch.optim.Adam(params, lr=args.learning_rate)
for step in range(args.train_steps):
optim.zero_grad()
if args.lazy:
with interpretation(lazy):
log_prob = apply_optimizer(model(data))
log_prob = reinterpret(log_prob)
else:
log_prob = model(data)
assert not log_prob.inputs, 'free variables remain'
loss = -log_prob.data
loss.backward()
optim.step()
def main(args):
funsor.set_backend("torch")
# Define a basic model with a single Normal latent random variable `loc`
# and a batch of Normally distributed observations.
def model(data):
loc = pyro.sample("loc", dist.Normal(0., 1.))
with pyro.plate("data", len(data), dim=-1):
pyro.sample("obs", dist.Normal(loc, 1.), obs=data)
# Define a guide (i.e. variational distribution) with a Normal
# distribution over the latent random variable `loc`.
def guide(data):
guide_loc = pyro.param("guide_loc", torch.tensor(0.))
guide_scale = pyro.param("guide_scale", torch.tensor(1.),
constraint=constraints.positive)
pyro.sample("loc", dist.Normal(guide_loc, guide_scale))
# Generate some data.
torch.manual_seed(0)
data = torch.randn(100) + 3.0
# Because the API in minipyro matches that of Pyro proper,
# training code works with generic Pyro implementations.
with pyro_backend(args.backend), interpretation(MonteCarlo()):
# Construct an SVI object so we can do variational inference on our
# model/guide pair.
Elbo = infer.JitTrace_ELBO if args.jit else infer.Trace_ELBO
elbo = Elbo()
adam = optim.Adam({"lr": args.learning_rate})
svi = infer.SVI(model, guide, adam, elbo)
# Basic training loop
pyro.get_param_store().clear()
for step in range(args.num_steps):
loss = svi.step(data)
if args.verbose and step % 100 == 0:
print("step {} loss = {}".format(step, loss))
# Report the final values of the variational parameters
# in the guide after training.
if args.verbose:
for name in pyro.get_param_store():
value = pyro.param(name).data
print("{} = {}".format(name, value.detach().cpu().numpy()))
# For this simple (conjugate) model we know the exact posterior. In
# particular we know that the variational distribution should be
# centered near 3.0. So let's check this explicitly.
assert (pyro.param("guide_loc") - 3.0).abs() < 0.1
def main(args):
funsor.set_backend("torch")
# Declare parameters.
trans_noise = torch.tensor(0.1, requires_grad=True)
emit_noise = torch.tensor(0.5, requires_grad=True)
params = [trans_noise, emit_noise]
# A Gaussian HMM model.
def model(data):
log_prob = funsor.to_funsor(0.)
x_curr = funsor.Tensor(torch.tensor(0.))
for t, y in enumerate(data):
x_prev = x_curr
# A delayed sample statement.
x_curr = funsor.Variable('x_{}'.format(t), funsor.Real)
log_prob += dist.Normal(1 + x_prev / 2., trans_noise, value=x_curr)
# Optionally marginalize out the previous state.
if t > 0 and not args.lazy:
log_prob = log_prob.reduce(ops.logaddexp, x_prev.name)
# An observe statement.
log_prob += dist.Normal(0.5 + 3 * x_curr, emit_noise, value=y)
# Marginalize out all remaining delayed variables.
log_prob = log_prob.reduce(ops.logaddexp)
return log_prob
# Train model parameters.
torch.manual_seed(0)
data = torch.randn(args.time_steps)
optim = torch.optim.Adam(params, lr=args.learning_rate)
for step in range(args.train_steps):
optim.zero_grad()
if args.lazy:
with interpretation(lazy):
log_prob = apply_optimizer(model(data))
log_prob = reinterpret(log_prob)
else:
log_prob = model(data)
assert not log_prob.inputs, 'free variables remain'
loss = -log_prob.data
loss.backward()
optim.step()
if args.verbose and step % 10 == 0:
print('step {} loss = {}'.format(step, loss.item()))
def main(args):
funsor.set_backend("torch")
if args.force or not args.metrics_filename or not os.path.exists(
args.metrics_filename):
results = track(args)
else:
results = torch.load(args.metrics_filename)
if args.plot_filename:
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot
import numpy as np
seeds = set(seed for seed, _, _ in results)
X = args.num_frames
pyplot.figure(figsize=(5, 1.4), dpi=300)
pos_error = np.array(
[[results[s, 0, f]['final_pos_error'] for s in seeds]
for f in args.num_frames])
mse = (pos_error**2).mean(axis=1)
std = (pos_error**2).std(axis=1) / len(seeds)**0.5
pyplot.plot(X, mse**0.5, 'k--')
pyplot.fill_between(X, (mse - std)**0.5, (mse + std)**0.5,
color='black',
alpha=0.15,
lw=0)
pos_error = np.array(
[[results[s, 1, f]['final_pos_error'] for s in seeds]
for f in args.num_frames])
mse = (pos_error**2).mean(axis=1)
std = (pos_error**2).std(axis=1) / len(seeds)**0.5
pyplot.plot(X, mse**0.5, 'r-')
pyplot.fill_between(X, (mse - std)**0.5, (mse + std)**0.5,
color='red',
alpha=0.15,
lw=0)
pyplot.ylabel('Position RMSE')
pyplot.xlabel('Track Length')
pyplot.xticks((5, 10, 15, 20, 25, 30))
pyplot.xlim(5, 30)
pyplot.tight_layout(0)
pyplot.savefig(args.plot_filename)
def main(args):
funsor.set_backend("torch")
torch.manual_seed(args.seed)
print_ = print if args.verbose else lambda msg: None
print_('Data:')
data = torch.distributions.Categorical(torch.ones(2)).sample((args.size, ))
assert data.shape == (args.size, )
data = Tensor(data, OrderedDict(i=Bint[args.size]), dtype=2)
print_(data)
print_('Model:')
m = model(args.size)
print_(m.pretty())
print_('Eager log_prob:')
obs = {str(i): data(i) for i in range(args.size)}
log_prob = m(**obs)
print_(log_prob)
def test_gaussian_funsor(batch_shape):
# This tests sample distribution, rsample gradients, log_prob, and log_prob
# gradients for both Pyro's and Funsor's Gaussian.
import funsor
funsor.set_backend("torch")
num_samples = 100000
# Declare unconstrained parameters.
loc = torch.randn(batch_shape + (3, )).requires_grad_()
t = transform_to(constraints.positive_definite)
m = torch.randn(batch_shape + (3, 3))
precision_unconstrained = t.inv(m @ m.transpose(-1, -2)).requires_grad_()
# Transform to constrained space.
log_normalizer = torch.zeros(batch_shape)
precision = t(precision_unconstrained)
info_vec = (precision @ loc[..., None])[..., 0]
def check_equal(actual, expected, atol=0.01, rtol=0):
assert_close(actual.data, expected.data, atol=atol, rtol=rtol)
grads = torch.autograd.grad(
(actual - expected).abs().sum(),
[loc, precision_unconstrained],
retain_graph=True,
)
for grad in grads:
assert grad.abs().max() < atol
entropy = dist.MultivariateNormal(loc,
precision_matrix=precision).entropy()
# Monte carlo estimate entropy via pyro.
p_gaussian = Gaussian(log_normalizer, info_vec, precision)
p_log_Z = p_gaussian.event_logsumexp()
p_rsamples = p_gaussian.rsample((num_samples, ))
pp_entropy = (p_log_Z - p_gaussian.log_density(p_rsamples)).mean(0)
check_equal(pp_entropy, entropy)
# Monte carlo estimate entropy via funsor.
inputs = OrderedDict([(k, funsor.Bint[v])
for k, v in zip("ij", batch_shape)])
inputs["x"] = funsor.Reals[3]
f_gaussian = funsor.gaussian.Gaussian(mean=loc,
precision=precision,
inputs=inputs)
f_log_Z = f_gaussian.reduce(funsor.ops.logaddexp, "x")
sample_inputs = OrderedDict(particle=funsor.Bint[num_samples])
deltas = f_gaussian.sample("x", sample_inputs)
f_rsamples = funsor.montecarlo.extract_samples(deltas)["x"]
ff_entropy = (f_log_Z - f_gaussian(x=f_rsamples)).reduce(
funsor.ops.mean, "particle")
check_equal(ff_entropy.data, entropy)
# Check Funsor's .rsample against Pyro's .log_prob.
pf_entropy = (p_log_Z - p_gaussian.log_density(f_rsamples.data)).mean(0)
check_equal(pf_entropy, entropy)
# Check Pyro's .rsample against Funsor's .log_prob.
fp_rsamples = funsor.Tensor(p_rsamples)["particle"]
for i in "ij"[:len(batch_shape)]:
fp_rsamples = fp_rsamples[i]
fp_entropy = (f_log_Z - f_gaussian(x=fp_rsamples)).reduce(
funsor.ops.mean, "particle")
check_equal(fp_entropy.data, entropy)
def stats(
model: avail_models = typer.Option("cosmos",
help="Tapqir model",
prompt="Tapqir model"),
cuda: bool = typer.Option(
partial(get_default, "cuda"),
"--cuda/--cpu",
help="Run computations on GPU or CPU",
prompt="Run computations on GPU?",
show_default=False,
),
nbatch_size: int = typer.Option(
partial(get_default, "nbatch-size"),
"--nbatch-size",
"-n",
help="AOI batch size",
prompt="AOI batch size",
),
fbatch_size: int = typer.Option(
partial(get_default, "fbatch-size"),
"--fbatch-size",
"-f",
help="Frame batch size",
prompt="Frame batch size",
),
matlab: bool = typer.Option(
partial(get_default, "matlab"),
"--matlab",
help="Save parameters in matlab format",
prompt="Save parameters in matlab format?",
),
funsor: bool = typer.Option(False,
"--funsor/--pyro",
help="Use funsor or pyro backend"),
no_input: bool = typer.Option(
False,
"--no-input",
help="Use defaults values.",
is_eager=True,
callback=deactivate_prompts,
),
):
from pyroapi import pyro_backend
from tapqir.models import models
logger = logging.getLogger("tapqir")
global DEFAULTS
cd = DEFAULTS["cd"]
dtype = "double"
device = "cuda" if cuda else "cpu"
backend = "funsor" if funsor else "pyro"
settings = {}
settings["device"] = device
settings["dtype"] = dtype
# pyro backend
if backend == "pyro":
PYRO_BACKEND = "pyro"
elif backend == "funsor":
import funsor
import pyro.contrib.funsor # noqa: F401
funsor.set_backend("torch")
PYRO_BACKEND = "contrib.funsor"
else:
raise ValueError("Only pyro and funsor backends are supported.")
with pyro_backend(PYRO_BACKEND):
logger.info("Computing stats ...")
model = models[model](**settings)
try:
model.load(cd)
except TapqirFileNotFoundError:
logger.exception("Failed to load data file")
return 1
model.load_checkpoint(param_only=True)
model.nbatch_size = nbatch_size
model.fbatch_size = fbatch_size
try:
model.compute_stats(save_matlab=matlab)
except CudaOutOfMemoryError:
logger.exception("Failed to compute stats")
return 1
logger.info("Computing stats: Done")
return 0
def fit(
model: avail_models = typer.Option(
"cosmos", help="Tapqir model", prompt="Tapqir model"),
cuda: bool = typer.Option(
partial(get_default, "cuda"),
"--cuda/--cpu",
help="Run computations on GPU or CPU",
prompt="Run computations on GPU?",
show_default=False,
),
nbatch_size: int = typer.Option(
partial(get_default, "nbatch-size"),
"--nbatch-size",
"-n",
help="AOI batch size",
prompt="AOI batch size",
),
fbatch_size: int = typer.Option(
partial(get_default, "fbatch-size"),
"--fbatch-size",
"-f",
help="Frame batch size",
prompt="Frame batch size",
),
learning_rate: float = typer.Option(
partial(get_default, "learning-rate"),
"--learning-rate",
"-lr",
help="Learning rate",
prompt="Learning rate",
),
num_iter: int = typer.Option(
0,
"--num-iter",
"-it",
help="Number of iterations",
prompt="Number of iterations",
),
k_max: int = typer.Option(
2, "--k-max", "-k", help="Maximum number of spots per image"),
matlab: bool = typer.Option(
partial(get_default, "matlab"),
"--matlab",
help="Save parameters in matlab format",
prompt="Save parameters in matlab format?",
),
funsor: bool = typer.Option(
False, "--funsor/--pyro", help="Use funsor or pyro backend"),
pykeops: bool = typer.Option(
True,
"--pykeops/--no-pykeops",
help="Use pykeops backend for offset marginalization",
),
overwrite: bool = typer.Option(
True,
"--overwrite",
"-w",
help="Overwrite defaults values.",
prompt="Overwrite defaults values?",
),
no_input: bool = typer.Option(
False,
"--no-input",
help="Disable interactive prompt.",
is_eager=True,
callback=deactivate_prompts,
),
progress_bar=None,
):
"""
Fit the data to the selected model.
Available models:
* cosmos: single-color time-independent co-localization model.\n
"""
global DEFAULTS
cd = DEFAULTS["cd"]
if progress_bar is None:
progress_bar = tqdm
from pyroapi import pyro_backend
from tapqir.models import models
logger = logging.getLogger("tapqir")
settings = {}
settings["K"] = k_max
settings["device"] = "cuda" if cuda else "cpu"
settings["dtype"] = "double"
settings["use_pykeops"] = pykeops
# priors settings
settings["priors"] = DEFAULTS["priors"]
if overwrite:
DEFAULTS["cuda"] = cuda
DEFAULTS["nbatch-size"] = nbatch_size
DEFAULTS["fbatch-size"] = fbatch_size
DEFAULTS["learning-rate"] = learning_rate
DEFAULTS["matlab"] = matlab
with open(cd / ".tapqir" / "config.yaml", "w") as cfg_file:
yaml.dump(
{key: value
for key, value in DEFAULTS.items() if key != "cd"},
cfg_file,
sort_keys=False,
)
backend = "funsor" if funsor else "pyro"
if model == "cosmos+hmm":
backend = "funsor" # hmm requires funsor backend
if backend == "pyro":
PYRO_BACKEND = "pyro"
elif backend == "funsor":
import funsor
import pyro.contrib.funsor # noqa: F401
funsor.set_backend("torch")
PYRO_BACKEND = "contrib.funsor"
else:
raise ValueError("Only pyro and funsor backends are supported.")
with pyro_backend(PYRO_BACKEND):
logger.info("Fitting the data ...")
model = models[model](**settings)
try:
model.load(cd)
except TapqirFileNotFoundError as err:
logger.exception(f"Failed to load {err.name} file")
return 1
model.init(learning_rate, nbatch_size, fbatch_size)
try:
model.run(num_iter, progress_bar=progress_bar)
except CudaOutOfMemoryError:
logger.exception("Failed to fit the data")
return 1
logger.info("Fitting the data: Done")
logger.info("Computing stats ...")
try:
model.compute_stats(save_matlab=matlab)
except CudaOutOfMemoryError:
logger.exception("Failed to compute stats")
return 1
logger.info("Computing stats: Done")
return 0
def stats(
model: Model = typer.Option("cosmos", help="Tapqir model", prompt="Tapqir model"),
channels: List[int] = typer.Option(
[0],
help="Color-channel numbers to analyze",
prompt="Channel numbers (space separated if multiple)",
),
cuda: bool = typer.Option(
partial(get_default, "cuda"),
"--cuda/--cpu",
help="Run computations on GPU or CPU",
prompt="Run computations on GPU?",
show_default=False,
),
nbatch_size: int = typer.Option(
partial(get_default, "nbatch-size"),
"--nbatch-size",
"-n",
help="AOI batch size",
prompt="AOI batch size",
),
fbatch_size: int = typer.Option(
partial(get_default, "fbatch-size"),
"--fbatch-size",
"-f",
help="Frame batch size",
prompt="Frame batch size",
),
matlab: bool = typer.Option(
partial(get_default, "matlab"),
"--matlab",
help="Save parameters in matlab format",
prompt="Save parameters in matlab format?",
),
funsor: bool = typer.Option(
False, "--funsor/--pyro", help="Use funsor or pyro backend"
),
no_input: bool = typer.Option(
False,
"--no-input",
help="Use defaults values.",
is_eager=True,
callback=deactivate_prompts,
),
):
from pyroapi import pyro_backend
from tapqir.models import models
from tapqir.utils.stats import save_stats
global DEFAULTS
cd = DEFAULTS["cd"]
dtype = "double"
device = "cuda" if cuda else "cpu"
backend = "funsor" if funsor else "pyro"
# pyro backend
if backend == "pyro":
PYRO_BACKEND = "pyro"
elif backend == "funsor":
import funsor
import pyro.contrib.funsor # noqa: F401
funsor.set_backend("torch")
PYRO_BACKEND = "contrib.funsor"
else:
raise ValueError("Only pyro and funsor backends are supported.")
with pyro_backend(PYRO_BACKEND):
model = models[model](1, 2, channels, device, dtype)
model.load(cd)
model.load_checkpoint(param_only=True)
model.nbatch_size = nbatch_size
model.fbatch_size = fbatch_size
typer.echo("Computing stats ...")
save_stats(model, cd, save_matlab=matlab)
typer.echo("Computing stats: Done")
# Copyright Contributors to the Pyro project.
# SPDX-License-Identifier: Apache-2.0
try:
import funsor
except ImportError:
raise ImportError(
"Looking like you want to do inference for models with "
"discrete latent variables. This is an experimental feature. "
"You need to install `funsor` to be able to use this feature. "
"It can be installed with `pip install funsor`.")
from numpyro.contrib.funsor.enum_messenger import enum, infer_config, markov, plate, to_data, to_funsor, trace
from numpyro.contrib.funsor.infer_util import config_enumerate, log_density, plate_to_enum_plate
funsor.set_backend("jax")
__all__ = [
"config_enumerate",
"enum",
"infer_config",
"log_density",
"markov",
"plate",
"plate_to_enum_plate",
"to_data",
"to_funsor",
"trace",
]
def main(args):
funsor.set_backend("torch")
# download and pre-process EEG data if not in test mode
if not args.test:
download_data()
N_val, N_test = 149, 200
data = np.loadtxt('eeg.dat', delimiter=',', skiprows=19)
print("[raw data shape] {}".format(data.shape))
data = data[::20, :]
print("[data shape after thinning] {}".format(data.shape))
eye_state = [int(d) for d in data[:, -1].tolist()]
data = torch.tensor(data[:, :-1]).float()
# in test mode (for continuous integration on github) so create fake data
else:
data = torch.randn(10, 3)
N_val, N_test = 2, 2
T, obs_dim = data.shape
N_train = T - N_test - N_val
np.random.seed(0)
rand_perm = np.random.permutation(N_val + N_test)
val_indices = rand_perm[0:N_val]
test_indices = rand_perm[N_val:]
data_mean = data[0:N_train, :].mean(0)
data -= data_mean
data_std = data[0:N_train, :].std(0)
data /= data_std
print("Length of time series T: {} Observation dimension: {}".format(T, obs_dim))
print("N_train: {} N_val: {} N_test: {}".format(N_train, N_val, N_test))
torch.manual_seed(args.seed)
# set up model
slds = SLDS(num_components=args.num_components, hidden_dim=args.hidden_dim, obs_dim=obs_dim,
fine_observation_noise=args.fon, fine_transition_noise=args.ftn,
fine_observation_matrix=args.fom, fine_transition_matrix=args.ftm,
moment_matching_lag=args.moment_matching_lag)
# set up optimizer
adam = torch.optim.Adam(slds.parameters(), lr=args.learning_rate, betas=(args.beta1, 0.999), amsgrad=True)
scheduler = torch.optim.lr_scheduler.ExponentialLR(adam, gamma=args.gamma)
ts = [time.time()]
report_frequency = 1
# training loop
for step in range(args.num_steps):
nll = -slds.log_prob(data[0:N_train, :]) / N_train
nll.backward()
if step == 5:
scheduler.base_lrs[0] *= 0.20
adam.step()
scheduler.step()
adam.zero_grad()
if step % report_frequency == 0 or step == args.num_steps - 1:
step_dt = ts[-1] - ts[-2] if step > 0 else 0.0
pred_mse, pred_LLs = slds.filter_and_predict(data[0:N_train + N_val + N_test, :])
val_mse = pred_mse[val_indices].mean().item()
test_mse = pred_mse[test_indices].mean().item()
val_ll = pred_LLs[val_indices].mean().item()
test_ll = pred_LLs[test_indices].mean().item()
stats = "[step %03d] train_nll: %.5f val_mse: %.5f val_ll: %.5f test_mse: %.5f test_ll: %.5f\t(dt: %.2f)"
print(stats % (step, nll.item(), val_mse, val_ll, test_mse, test_ll, step_dt))
ts.append(time.time())
# plot predictions and smoothed means
if args.plot:
assert not args.test
predicted_mse, LLs, pred_means, pred_vars, smooth_means, smooth_probs = \
slds.filter_and_predict(data, smoothing=True)
pred_means = pred_means.data.numpy()
pred_stds = pred_vars.sqrt().data.numpy()
smooth_means = smooth_means.data.numpy()
smooth_probs = smooth_probs.data.numpy()
import matplotlib
matplotlib.use('Agg') # noqa: E402
import matplotlib.pyplot as plt
f, axes = plt.subplots(4, 1, figsize=(12, 8), sharex=True)
T = data.size(0)
N_valtest = N_val + N_test
to_seconds = 117.0 / T
for k, ax in enumerate(axes[:-1]):
which = [0, 4, 10][k]
ax.plot(to_seconds * np.arange(T), data[:, which], 'ko', markersize=2)
ax.plot(to_seconds * np.arange(N_train), smooth_means[:N_train, which], ls='solid', color='r')
ax.plot(to_seconds * (N_train + np.arange(N_valtest)),
pred_means[-N_valtest:, which], ls='solid', color='b')
ax.fill_between(to_seconds * (N_train + np.arange(N_valtest)),
pred_means[-N_valtest:, which] - 1.645 * pred_stds[-N_valtest:, which],
pred_means[-N_valtest:, which] + 1.645 * pred_stds[-N_valtest:, which],
color='lightblue')
ax.set_ylabel("$y_{%d}$" % (which + 1), fontsize=20)
ax.tick_params(axis='both', which='major', labelsize=14)
axes[-1].plot(to_seconds * np.arange(T), eye_state, 'k', ls='solid')
axes[-1].plot(to_seconds * np.arange(T), smooth_probs, 'r', ls='solid')
axes[-1].set_xlabel("Time (s)", fontsize=20)
axes[-1].set_ylabel("Eye state", fontsize=20)
axes[-1].tick_params(axis='both', which='major', labelsize=14)
plt.tight_layout(pad=0.7)
plt.savefig('eeg.pdf')
def run_expt(args):
funsor.set_backend("torch")
optim = args["optim"]
lr = args["learnrate"]
schedule = [] if not args["schedule"] else [
int(i) for i in args["schedule"].split(",")
]
# default these to "none" instead of None, which argparse does for some reason
args["group"] = "none" if args["group"] is None else args["group"]
args["individual"] = "none" if args["individual"] is None else args[
"individual"]
random_effects = {"group": args["group"], "individual": args["individual"]}
pyro.enable_validation(args["validation"])
pyro.set_rng_seed(
args["seed"]) # reproducible random effect parameter init
if args["cuda"]:
torch.set_default_tensor_type(torch.cuda.FloatTensor)
if args["dataset"] == "seal":
filename = os.path.join(args["folder"], "prep_seal_data.csv")
config = prepare_seal(filename, random_effects)
elif args["dataset"] == "fake":
fake_sizes = {
"state": args["size_state"],
"random": args["size_random"],
"group": args["size_group"],
"individual": args["size_individual"],
"timesteps": args["size_timesteps"],
}
config = prepare_fake(fake_sizes, random_effects)
else:
raise ValueError("Dataset {} not yet included".format(args["dataset"]))
if args["smoke"]:
args["timesteps"] = 2
config["sizes"]["timesteps"] = 3
if args["truncate"] > 0:
config["sizes"]["timesteps"] = args["truncate"]
config["zeroinflation"] = args["zeroinflation"]
model = Model(config)
guide = Guide(config)
loss_fn = parallel_loss_fn
if args["jit"]:
loss_fn = torch.jit.trace(
lambda: loss_fn(model, guide, args["parallel"]), ())
else:
loss_fn = functools.partial(loss_fn, model, guide, args["parallel"])
# count the number of parameters once
num_parameters = aic_num_parameters(model, guide)
losses = []
# TODO support continuous random effects with monte carlo
assert random_effects["group"] != "continuous"
assert random_effects["individual"] != "continuous"
with pyro.poutine.trace(param_only=True) as param_capture:
loss_fn()
params = [
site["value"].unconstrained()
for site in param_capture.trace.nodes.values()
]
if optim == "sgd":
optimizer = torch.optim.Adam(params, lr=lr)
elif optim == "lbfgs":
optimizer = torch.optim.LBFGS(params, lr=lr)
else:
raise ValueError("{} not supported optimizer".format(optim))
if schedule:
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=schedule,
gamma=0.5)
schedule_step_loss = False
else:
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, 'min')
schedule_step_loss = True
for t in range(args["timesteps"]):
def closure():
optimizer.zero_grad()
loss = loss_fn()
loss.backward()
return loss
loss = optimizer.step(closure)
scheduler.step(loss.item() if schedule_step_loss else t)
losses.append(loss.item())
print("Loss: {}, AIC[{}]: ".format(loss.item(), t),
2. * loss + 2. * num_parameters)
aic_final = 2. * losses[-1] + 2. * num_parameters
print("AIC final: {}".format(aic_final))
results = {}
results["args"] = args
results["sizes"] = config["sizes"]
results["likelihoods"] = losses
results["likelihood_final"] = losses[-1]
results["aic_final"] = aic_final
results["aic_num_parameters"] = num_parameters
if args["resultsdir"] is not None and os.path.exists(args["resultsdir"]):
re_str = "g" + ("n" if args["group"] is None else
"d" if args["group"] == "discrete" else "c")
re_str += "i" + ("n" if args["individual"] is None else
"d" if args["individual"] == "discrete" else "c")
results_filename = "expt_{}_{}_{}.json".format(
args["dataset"], re_str,
str(uuid.uuid4().hex)[0:5])
with open(os.path.join(args["resultsdir"], results_filename),
"w") as f:
json.dump(results, f)
return results
# SPDX-License-Identifier: Apache-2.0
import logging
import pytest
import torch
from pyro.ops.indexing import Vindex
from tests.common import xfail_param
# put all funsor-related imports here, so test collection works without funsor
try:
import funsor
import pyro.contrib.funsor
funsor.set_backend("torch")
from pyroapi import distributions as dist
from pyroapi import infer, pyro
from tests.contrib.funsor.test_valid_models_enum import assert_ok
except ImportError:
pytestmark = pytest.mark.skip(reason="funsor is not installed")
logger = logging.getLogger(__name__)
@pytest.mark.parametrize('enumerate_', [None, "parallel", "sequential"])
def test_enum_discrete_non_enumerated_plate_ok(enumerate_):
def model():
pyro.sample("w", dist.Bernoulli(0.5), infer={'enumerate': 'parallel'})
def fit(
model: Model = typer.Option("cosmos", help="Tapqir model", prompt="Tapqir model"),
channels: List[int] = typer.Option(
[0],
help="Color-channel numbers to analyze",
prompt="Channel numbers (space separated if multiple)",
),
cuda: bool = typer.Option(
partial(get_default, "cuda"),
"--cuda/--cpu",
help="Run computations on GPU or CPU",
prompt="Run computations on GPU?",
show_default=False,
),
nbatch_size: int = typer.Option(
partial(get_default, "nbatch-size"),
"--nbatch-size",
"-n",
help="AOI batch size",
prompt="AOI batch size",
),
fbatch_size: int = typer.Option(
partial(get_default, "fbatch-size"),
"--fbatch-size",
"-f",
help="Frame batch size",
prompt="Frame batch size",
),
learning_rate: float = typer.Option(
partial(get_default, "learning-rate"),
"--learning-rate",
"-lr",
help="Learning rate",
prompt="Learning rate",
),
num_iter: int = typer.Option(
0,
"--num-iter",
"-it",
help="Number of iterations",
prompt="Number of iterations",
),
k_max: int = typer.Option(
2, "--k-max", "-k", help="Maximum number of spots per image"
),
matlab: bool = typer.Option(
partial(get_default, "matlab"),
"--matlab",
help="Save parameters in matlab format",
prompt="Save parameters in matlab format?",
),
funsor: bool = typer.Option(
False, "--funsor/--pyro", help="Use funsor or pyro backend"
),
pykeops: bool = typer.Option(
True,
"--pykeops/--no-pykeops",
help="Use pykeops backend for offset marginalization",
),
overwrite: bool = typer.Option(
True,
"--overwrite",
"-w",
help="Overwrite defaults values.",
prompt="Overwrite defaults values?",
),
no_input: bool = typer.Option(
False,
"--no-input",
help="Disable interactive prompt.",
is_eager=True,
callback=deactivate_prompts,
),
progress_bar=None,
):
"""
Fit the data to the selected model.
Available models:
* cosmos: single-color time-independent co-localization model.\n
"""
global DEFAULTS
cd = DEFAULTS["cd"]
if progress_bar is None:
progress_bar = tqdm
from pyroapi import pyro_backend
from tapqir.models import models
from tapqir.utils.stats import save_stats
settings = {}
settings["S"] = 1
settings["K"] = k_max
settings["channels"] = channels
settings["device"] = "cuda" if cuda else "cpu"
settings["dtype"] = "double"
settings["use_pykeops"] = pykeops
# priors settings
settings["background_mean_std"] = DEFAULTS["background_mean_std"]
settings["background_std_std"] = DEFAULTS["background_std_std"]
settings["lamda_rate"] = DEFAULTS["lamda_rate"]
settings["height_std"] = DEFAULTS["height_std"]
settings["width_min"] = DEFAULTS["width_min"]
settings["width_max"] = DEFAULTS["width_max"]
settings["proximity_rate"] = DEFAULTS["proximity_rate"]
settings["gain_std"] = DEFAULTS["gain_std"]
if overwrite:
DEFAULTS["cuda"] = cuda
DEFAULTS["nbatch-size"] = nbatch_size
DEFAULTS["fbatch-size"] = fbatch_size
DEFAULTS["learning-rate"] = learning_rate
DEFAULTS["matlab"] = matlab
with open(cd / ".tapqir" / "config.yaml", "w") as cfg_file:
yaml.dump(
{key: value for key, value in DEFAULTS.items() if key != "cd"},
cfg_file,
sort_keys=False,
)
backend = "funsor" if funsor else "pyro"
if backend == "pyro":
PYRO_BACKEND = "pyro"
elif backend == "funsor":
import funsor
import pyro.contrib.funsor # noqa: F401
funsor.set_backend("torch")
PYRO_BACKEND = "contrib.funsor"
else:
raise ValueError("Only pyro and funsor backends are supported.")
with pyro_backend(PYRO_BACKEND):
model = models[model](**settings)
model.load(cd)
model.init(learning_rate, nbatch_size, fbatch_size)
typer.echo("Fitting the data ...")
exit_code = model.run(num_iter, progress_bar=progress_bar)
if exit_code == 0:
typer.echo("Fitting the data: Done")
elif exit_code == 1:
typer.echo("The model hasn't converged!")
typer.echo("Computing stats ...")
save_stats(model, cd, save_matlab=matlab)
typer.echo("Computing stats: Done")
