Merge pull request #31 from neu-spiral/develop

Develop

Author: esennesh

htfa_torch/dtfa.py

@@ -5,6 +5,7 @@
 
 import collections
 import numpy as np
+import scipy.spatial
 import torch
 from torch.autograd import Variable
 import probtorch
@@ -53,15 +54,12 @@ def __init__(self, hyper_means, num_times, num_blocks,
                 'sigma': torch.sqrt(torch.rand(self._num_blocks, self._num_factors)),
             },
             'weights': {
-                'mu': torch.randn(self._num_blocks, self._num_times,
-                                  self._num_factors),
+                'mu':  hyper_means['weights'].mean(0).unsqueeze(0).expand(
+                    self._num_blocks, self._num_times, self._num_factors
+                ),
                 'sigma': torch.ones(self._num_blocks, self._num_times,
                                     self._num_factors),
             },
-            'voxel_noise': {
-                'mu': torch.ones(self._num_blocks),
-                'sigma': torch.sqrt(torch.rand(self._num_blocks))
-            }
         })
 
         super(self.__class__, self).__init__(params, guide=True)
@@ -95,15 +93,6 @@ def forward(self, trace, params, times=None, blocks=None,
         # We only expand the parameters for which we're actually going to sample
         # values in this very method, and thus want to expand to get multiple
         # particles.
-        voxel_noise_params = params['block']['voxel_noise']
-        if num_particles and num_particles > 0:
-            voxel_noise_params = utils.unsqueeze_and_expand_vardict(
-                params['block']['voxel_noise'], 0, num_particles, True
-            )
-        voxel_noise = trace.normal(voxel_noise_params['mu'],
-                                   voxel_noise_params['sigma'],
-                                   name='voxel_noise')
-
         if blocks is None:
             blocks = list(range(self._num_blocks))
 
@@ -125,16 +114,17 @@ def forward(self, trace, params, times=None, blocks=None,
             factor_centers += [fc]
             factor_log_widths += [flw]
 
-        return weights, factor_centers, factor_log_widths, voxel_noise
+        return weights, factor_centers, factor_log_widths
 
 class HTFAGuide(nn.Module):
     """Variational guide for hierarchical topographic factor analysis"""
     def __init__(self, query, num_factors=tfa_models.NUM_FACTORS):
         super(self.__class__, self).__init__()
         self._num_blocks = len(query)
-        self._num_times = niidb.query_min_time(query)
+        self._num_times = niidb.query_max_time(query)
 
-        b = np.random.choice(self._num_blocks, 1)[0]
+        b = max(range(self._num_blocks), key=lambda b: query[b].end_time -
+                query[b].start_time)
         query[b].load()
         centers, widths, weights = utils.initial_hypermeans(
             query[b].activations.numpy().T, query[b].locations.numpy(),
@@ -161,7 +151,7 @@ def forward(self, trace, times=None, blocks=None,
 
 class HTFAGenerativeHyperParams(tfa_models.HyperParams):
     def __init__(self, brain_center, brain_center_std_dev, num_blocks,
-                 num_factors=tfa_models.NUM_FACTORS):
+                 num_factors=tfa_models.NUM_FACTORS, volume=None):
         self._num_factors = num_factors
         self._num_blocks = num_blocks
 
@@ -177,20 +167,22 @@ def __init__(self, brain_center, brain_center_std_dev, num_blocks,
         params['template']['factor_centers']['sigma'] =\
             brain_center_std_dev.expand(self._num_factors, 3)
 
+        coefficient = 1.0
+        if volume is not None:
+            coefficient = np.log(np.cbrt(volume / self._num_factors))
         params['template']['factor_log_widths']['mu'] =\
-            torch.ones(self._num_factors)
+            coefficient * torch.ones(self._num_factors)
         params['template']['factor_log_widths']['sigma'] =\
             tfa_models.SOURCE_LOG_WIDTH_STD_DEV * torch.ones(self._num_factors)
 
         params['block'] = {
             'factor_center_noise': torch.ones(self._num_blocks),
             'factor_log_width_noise': torch.ones(self._num_blocks),
             'weights': {
-                'mu': torch.rand(self._num_blocks, self._num_factors),
+                'mu': torch.zeros(self._num_blocks, self._num_factors),
                 'sigma': tfa_models.SOURCE_WEIGHT_STD_DEV *\
                          torch.ones(self._num_blocks, self._num_factors)
             },
-            'voxel_noise': utils.gaussian_populator(self._num_blocks)
         }
         super(self.__class__, self).__init__(params, guide=False)
 
@@ -216,19 +208,16 @@ def __init__(self, num_blocks, num_times):
                             for b in range(self._num_blocks)]
 
     def forward(self, trace, params, template, times=None, blocks=None,
-                guide=probtorch.Trace()):
-        voxel_noise = trace.normal(params['block']['voxel_noise']['mu'],
-                                   params['block']['voxel_noise']['sigma'],
-                                   value=guide['voxel_noise'],
-                                   name='voxel_noise')
-
+                guide=probtorch.Trace(), weights_params=None):
         if blocks is None:
             blocks = list(range(self._num_blocks))
+        if times is None:
+            times = (0, self._num_times)
 
         weights = []
         factor_centers = []
         factor_log_widths = []
-        for b in blocks:
+        for (i, b) in enumerate(blocks):
             sparams = utils.vardict({
                 'factor_centers': {
                     'mu': template['factor_centers'],
@@ -243,56 +232,58 @@ def forward(self, trace, params, template, times=None, blocks=None,
                     'sigma': params['block']['weights']['sigma'][b],
                 }
             })
+            if weights_params is not None:
+                sparams['weights'] = weights_params[i]
             w, fc, flw = self._tfa_priors[b](trace, sparams, times=times,
                                              guide=guide)
             weights += [w]
             factor_centers += [fc]
             factor_log_widths += [flw]
 
-        return weights, factor_centers, factor_log_widths, voxel_noise
+        return weights, factor_centers, factor_log_widths
 
 class HTFAModel(nn.Module):
     """Generative model for hierarchical topographic factor analysis"""
-    def __init__(self, query, num_blocks, num_times,
-                 num_factors=tfa_models.NUM_FACTORS):
+    def __init__(self, locations, num_blocks, num_times,
+                 num_factors=tfa_models.NUM_FACTORS, volume=None):
         super(self.__class__, self).__init__()
 
         self._num_factors = num_factors
         self._num_blocks = num_blocks
         self._num_times = num_times
 
-        b = np.random.choice(self._num_blocks, 1)[0]
-        query[b].load()
-        center, center_sigma = utils.brain_centroid(query[b].locations)
+        center, center_sigma = utils.brain_centroid(locations)
+        hull = scipy.spatial.ConvexHull(locations)
+        if volume is not None:
+            volume = hull.volume
 
         self._hyperparams = HTFAGenerativeHyperParams(center, center_sigma,
                                                       self._num_blocks,
-                                                      self._num_factors)
+                                                      self._num_factors,
+                                                      volume=volume)
         self._template_prior = HTFAGenerativeTemplatePrior()
         self._subject_prior = HTFAGenerativeSubjectPrior(
             self._num_blocks, self._num_times
         )
-        for block in query:
-            block.load()
         self.likelihoods = [tfa_models.TFAGenerativeLikelihood(
-            query[b].locations, self._num_times[b], tfa_models.VOXEL_NOISE,
+            locations, self._num_times[b], tfa_models.VOXEL_NOISE,
             block=b, register_locations=False
         ) for b in range(self._num_blocks)]
         for b, block_likelihood in enumerate(self.likelihoods):
             self.add_module('likelihood' + str(b), block_likelihood)
 
     def forward(self, trace, times=None, guide=probtorch.Trace(), blocks=None,
-                observations=[]):
+                observations=[], weights_params=None):
         if blocks is None:
             blocks = list(range(self._num_blocks))
         params = self._hyperparams.state_vardict()
 
         template = self._template_prior(trace, params, guide=guide)
-        weights, centers, log_widths, voxel_noise = self._subject_prior(
-            trace, params, template, times=times, blocks=blocks, guide=guide
+        weights, centers, log_widths = self._subject_prior(
+            trace, params, template, times=times, blocks=blocks, guide=guide,
+            weights_params=weights_params
         )
 
         return [self.likelihoods[b](trace, weights[i], centers[i], log_widths[i],
-                                    times=times, observations=observations[i],
-                                    voxel_noise=voxel_noise)
+                                    times=times, observations=observations[i])
                 for (i, b) in enumerate(blocks)]

htfa_torch/dtfa_models.py

@@ -114,6 +114,18 @@ def wrapped_table_method(*args, **kwargs):
             return wrapped_table_method
         return attr
 
+def query_max_time(qiter):
+    result = -1
+    for block in qiter:
+        load = block.activations is None
+        if load:
+            block.load()
+        if result < 0 or block.end_time > result:
+            result = block.end_time
+        if load:
+            block.unload()
+    return result
+
 def query_min_time(qiter):
     result = -1
     for block in qiter:

htfa_torch/htfa.py

@@ -44,26 +44,54 @@ def free_energy(q, p, num_particles=tfa_models.NUM_PARTICLES):
         sample_dim = None
     return -probtorch.objectives.montecarlo.elbo(q, p, sample_dim=sample_dim)
 
-def hierarchical_elbo(q, p, rv_weight=lambda x: 1.0,
+def hierarchical_elbo(q, p, rv_weight=lambda x, prior=True: 1.0,
                       num_particles=tfa_models.NUM_PARTICLES,
                       sample_dim=None, batch_dim=None):
     if num_particles and num_particles > 0:
         sample_dim = 0
     else:
         sample_dim = None
 
-    weight_rvs = utils.inverse(rv_weight, p)
-    weighted_elbo = 0.0
-    for weight, rvs in weight_rvs.items():
+    weighted_log_likelihood = 0.0
+    weighted_prior_kl = 0.0
+    for rv in p:
         local_elbo = p.log_joint(sample_dim=sample_dim, batch_dim=batch_dim,
-                                 nodes=rvs) -\
+                                 nodes=[rv]) -\
                      q.log_joint(sample_dim=sample_dim, batch_dim=batch_dim,
-                                 nodes=rvs)
-        weighted_elbo += weight * local_elbo.sum()
-    return weighted_elbo
+                                 nodes=[rv])
+        if sample_dim is not None:
+            local_elbo = local_elbo.mean(dim=sample_dim)
+        if p[rv].observed and rv not in q:
+            weighted_log_likelihood += rv_weight(rv, False) * local_elbo
+        else:
+            weighted_prior_kl -= rv_weight(rv, True) * local_elbo
+    weighted_elbo = weighted_log_likelihood - weighted_prior_kl
+    return weighted_elbo, weighted_log_likelihood, weighted_prior_kl
+
+def componentized_elbo(q, p, rv_weight=lambda x: 1.0,
+                       num_particles=tfa_models.NUM_PARTICLES, sample_dim=None,
+                       batch_dim=None):
+    if num_particles and num_particles > 0:
+        sample_dim = 0
+    else:
+        sample_dim = None
+
+    trace_elbos = {}
+    for rv in p:
+        local_elbo = p.log_joint(sample_dim=sample_dim, batch_dim=batch_dim,
+                                 nodes=[rv]) -\
+                     q.log_joint(sample_dim=sample_dim, batch_dim=batch_dim,
+                                 nodes=[rv])
+        if sample_dim is not None:
+            local_elbo = local_elbo.mean(dim=sample_dim)
+        trace_elbos[rv] = rv_weight(rv) * local_elbo
+
+    weighted_elbo = sum([trace_elbos[rv] for rv in trace_elbos])
+    return weighted_elbo, trace_elbos
 
 def hierarchical_free_energy(*args, **kwargs):
-    return -hierarchical_elbo(*args, **kwargs)
+    elbo, ll, kl = hierarchical_elbo(*args, **kwargs)
+    return -elbo, ll, kl
 
 def log_likelihood(q, p, num_particles=tfa_models.NUM_PARTICLES):
     """The expected log-likelihood of observed data under the proposal distribution"""
@@ -75,11 +103,12 @@ def log_likelihood(q, p, num_particles=tfa_models.NUM_PARTICLES):
 
 class TopographicalFactorAnalysis:
     """Overall container for a run of TFA"""
-    def __init__(self, data_file, num_factors=tfa_models.NUM_FACTORS):
+    def __init__(self, data_file, num_factors=tfa_models.NUM_FACTORS,
+                 zscore=True):
         self.num_factors = num_factors
 
         self.voxel_activations, self.voxel_locations, self._name,\
-            self._template = utils.load_dataset(data_file)
+            self._template = utils.load_dataset(data_file, zscore=zscore)
 
         # Pull out relevant dimensions: the number of times-of-recording, and
         # the number of voxels in each timewise "slice"

htfa_torch/htfa_models.py

@@ -131,7 +131,7 @@ def forward(self, trace, params, times=None, num_particles=NUM_PARTICLES):
 
         weights = trace.normal(weight_params['mu'],
                                weight_params['sigma'],
-                               name='Weights' + str(self.block))
+                               name='Weights%dt%d-%d' % (self.block, times[0], times[1]))
 
         centers = trace.normal(params['factor_centers']['mu'],
                                params['factor_centers']['sigma'],
@@ -196,8 +196,8 @@ def forward(self, trace, params, times=None, guide=probtorch.Trace()):
 
         weights = trace.normal(weight_params['mu'],
                                weight_params['sigma'],
-                               value=guide['Weights' + str(self.block)],
-                               name='Weights' + str(self.block))
+                               value=guide['Weights%dt%d-%d' % (self.block, times[0], times[1])],
+                               name='Weights%dt%d-%d' % (self.block, times[0], times[1]))
 
         factor_centers = trace.normal(params['factor_centers']['mu'],
                                       params['factor_centers']['sigma'],
@@ -224,18 +224,17 @@ def __init__(self, locations, num_times, voxel_noise=VOXEL_NOISE, block=0,
         self.block = block
 
     def forward(self, trace, weights, centers, log_widths, times=None,
-                observations=collections.defaultdict(), voxel_noise=None):
+                observations=collections.defaultdict()):
         if times is None:
             times = (0, self._num_times)
-        if voxel_noise is None:
-            voxel_noise = self._voxel_noise
 
-        factors = radial_basis(Variable(self.voxel_locations), centers,
-                               log_widths)
+        factors = radial_basis(Variable(self.voxel_locations,
+                                        requires_grad=True),
+                               centers, log_widths)
 
         activations = trace.normal(weights @ factors,
                                    self._voxel_noise, value=observations['Y'],
-                                   name='Y' + str(self.block))
+                                   name='Y%dt%d-%d' % (self.block, times[0], times[1]))
         return activations
 
 class TFAModel(nn.Module):