--- a/pylearn/algorithms/mcRBM.py	Tue Aug 24 15:50:14 2010 -0400
+++ b/pylearn/algorithms/mcRBM.py	Tue Aug 24 16:00:08 2010 -0400
@@ -232,19 +232,6 @@
     return d
 
 
-
-
-# this is a little hack, probably should be removed
-# The logic about casting things to shared vars is busted anyway (wrt pickling)
-def as_shared(x, name=None, dtype=floatX):
-    if hasattr(x, 'type'):
-        return x
-    else:
-        if 'float' in str(x.dtype):
-            return shared(x.astype(floatX), name=name)
-        else:
-            return shared(x, name=name)
-
 def hidden_cov_units_preactivation_given_v(rbm, v, small=0.5):
     """Return argument to the sigmoid that would give mean of covariance hid units
 
@@ -321,7 +308,7 @@
         return W.shape[0]
 
 
-def sampler(rbm, n_particles, rng=7823748):
+def sampler(rbm, n_particles, n_visible=None, rng=7823748):
     """Return an `HMC_sampler` that will draw samples from the distribution over visible
     units specified by this RBM.
 
@@ -330,26 +317,31 @@
     """
     if not hasattr(rng, 'randn'):
         rng = np.random.RandomState(rng)
-
+    if n_visible is None:
+        n_visible = n_visible_units(rbm)
     rval = HMC_sampler(
-        positions = [as_shared(
+        positions = [shared(
             rng.randn(
                 n_particles,
-                n_visible_units(rbm)))],
+                n_visible).astype(floatX),
+            name='particles')],
         energy_fn = lambda p : free_energy_given_v(rbm, p[0]),
         seed=int(rng.randint(2**30)))
     return rval
 
 class MeanCovRBM(object):
-    """Container for mcRBM parameters that gives more convenient access to mcRBM methods.
+    """Container for mcRBM parameters
+
+    It provides parameter lookup by name, as well as a heuristic for initializing the
+    parameters for effective learning.
 
     Attributes:
 
-      - U - the covariance filters
-      - W - the mean filters
-      - a - the visible bias
-      - b - the covariance bias
-      - c - the mean bias
+      - U - the covariance filters (theano shared variable)
+      - W - the mean filters (theano shared variable)
+      - a - the visible bias (theano shared variable)
+      - b - the covariance bias (theano shared variable)
+      - c - the mean bias (theano shared variable)
     
     """
 
@@ -361,6 +353,10 @@
         self.__dict__.update(locals())
         del self.__dict__['self']
 
+    def __getitem__(self, idx):
+        # support unpacking of this container as if it were a tuple
+        return self.params[idx]
+
     @classmethod
     def new_from_dims(cls, n_I, n_K, n_J, rng = 8923402190):
         """
@@ -400,28 +396,6 @@
             d[key] = shared(d[key], name=key)
         self.__init__(**d)
 
-    def hmc_sampler(self, n_particles, rng=7823748):
-        """Return an HMC_sampler that will draw samples from the distribution over visible
-        units specified by this RBM.
-
-        :param n_particles: this many parallel chains will be simulated.
-        :param rng: seed or numpy RandomState object to initialize particles, and to drive the simulation.
-        """
-        return sampler(self.params, n_particles, rng)
-
-    def free_energy_given_v(self, v):
-        """Return expressions for F.E. of visible configuration `v`"""
-        return free_energy_given_v(self.params, v)
-
-    def contrastive_gradient(self, *args, **kwargs):
-        """Return a list of gradient expressions for self.params
-
-        :param pos_v: positive-phase sample of visible units
-        :param neg_v: negative-phase sample of visible units
-        """
-        return contrastive_gradient(self.params, *args, **kwargs)
-
-
 if __name__ == '__main__':
 
     print >> sys.stderr, "TODO: use P matrix (aka FH matrix)"
@@ -454,7 +428,7 @@
 
     rbm = MeanCovRBM.new_from_dims(n_I=R*C, n_K=n_K, n_J=n_J) 
 
-    sampler = rbm.hmc_sampler(n_particles=batchsize)
+    sampler = sampler(rbm, n_particles=batchsize)
 
     def l2(X):
         return numpy.sqrt((X**2).sum())
@@ -505,7 +479,7 @@
 
     imgs_fn = function([batch_idx], outputs=train_batch)
 
-    grads = rbm.contrastive_gradient(
+    grads = contrastive_gradient(rbm,
             pos_v=train_batch, 
             neg_v=sampler.positions[0],
             U_l1_penalty=s_l1_penalty,
@@ -522,11 +496,6 @@
                 (sgd_ups[1][1] - sgd_ups[1][0]).norm(2),
                 ],
             updates = sgd_ups)
-                #rbm.free_energy_given_v(train_batch).sum(),
-                #rbm.free_energy_given_v(train_batch,extra=1)[1][0].sum(),
-                #rbm.free_energy_given_v(train_batch,extra=1)[1][1].sum(),
-                #rbm.free_energy_given_v(train_batch,extra=1)[1][2].sum(),
-                #rbm.free_energy_given_v(train_batch,extra=1)[1][3].sum(),
     theano.printing.pydotprint(function([batch_idx, s_l1_penalty], grads[0]), 'grads0.png')
 
     print "Learning..."