--- a/pylearn/shared/layers/rust2005.py	Tue Nov 03 15:25:08 2009 -0500
+++ b/pylearn/shared/layers/rust2005.py	Tue Nov 03 15:26:01 2009 -0500
@@ -4,6 +4,10 @@
 
 This layer implements a model of simple and complex cell firing rate responses.
 
+
+:TODO: implement full model with variable exponents.  The current implementation fixes internal
+exponents to 2 and the external exponent to 1/2.
+
 """
 
 import numpy
@@ -18,25 +22,31 @@
 from theano.compile.sandbox import shared
 from theano.sandbox.softsign import softsign
 from theano.tensor.nnet import softplus
+from theano.sandbox.conv import ConvOp
 
 from .util import update_locals, add_logging
 
-def rust2005_act_from_filters(linpart, E_quad, S_quad):
+def rust2005_act_from_filters(linpart, E_quad, S_quad, eps):
+    """Return rust2005 activation from linear filter responses, as well as E and S terms
+
+    :param linpart: a single tensor of linear filter responses
+    :param E_quad: a list of tensors of linear filter responses
+    :param S_quad: a list of tensors of linear filter responses
+    :param eps: a scalar to add to the sum of squares before the sqrt
+
+    """
+    if isinstance(E_quad, theano.Variable):
+        raise TypeError('E_quad should be a list of Variables, not a Variable itself')
+    if isinstance(S_quad, theano.Variable):
+        raise TypeError('E_quad should be a list of Variables, not a Variable itself')
     sqrt = theano.tensor.sqrt
     softlin = theano.tensor.nnet.softplus(linpart)
-    E = sqrt(sum([E_quad_i**2 for E_quad_i in E_quad] + [1e-8, softlin**2]))
-    S = sqrt(sum([S_quad_i**2 for S_quad_i in S_quad] + [1e-8]))
-    return (E-S) / (1+E+S)
+    E = sqrt(sum([E_quad_i**2 for E_quad_i in E_quad] + [softlin**2], eps))
+    S = sqrt(sum([S_quad_i**2 for S_quad_i in S_quad], eps))
+    return (E-S) / (1+E+S), E, S
 
 class Rust2005(object):
-    """
-    shared variable version.
-
-    w is 3-tensor n_in x n_out x (1+n_E_quadratic + n_S_quadratic)
-
-    w 
-
-    """
+    """Energy-like complex cell activation function described in Rust et al. 2005 """
     #logging methods come from the add_logging() call below
     # _info, _debug, _warn, _error, _fatal
 
@@ -99,7 +109,7 @@
         b = shared(numpy.zeros((n_out,), dtype=dtype))
         return cls(input, w, b, n_out, n_E, n_S, epsilon, filter_shape, [w,b])
 
-    def img_from_weights(self, rows=12, cols=24, row_gap=1, col_gap=1, eps=1e-4):
+    def img_from_weights(self, rows=12, cols=25, row_gap=1, col_gap=1, eps=1e-4, triplegap=0):
         """ Return an image that visualizes all the weights in the layer.
 
         The current implentation returns a tiling in which every triple of columns is a logical
@@ -110,9 +120,12 @@
         """
         if cols % 3: #because there are three kinds of filters: linear, excitatory, inhibitory
             raise ValueError("cols must be multiple of 3")
+
+        n_triples = cols / 3
+
         filter_shape = self.filter_shape
         height = rows * (row_gap + filter_shape[0]) - row_gap
-        width = cols * (col_gap + filter_shape[1]) - col_gap
+        width = cols * (col_gap + filter_shape[1]) - col_gap + (n_triples-1) * triplegap
 
         out_array = numpy.zeros((height, width, 3), dtype='uint8')
 
@@ -123,10 +136,14 @@
         for r in xrange(rows):
             out_r_low = r*(row_gap + filter_shape[0])
             out_r_high = out_r_low + filter_shape[0]
+            extra_col_gap = 0 # a counter we'll use for the triplegap
             for c in xrange(cols):
-                out_c_low = c*(col_gap + filter_shape[1])
+                if c and (c%3==0):
+                    extra_col_gap += triplegap
+                out_c_low = c*(col_gap + filter_shape[1]) + extra_col_gap
                 out_c_high = out_c_low + filter_shape[1]
                 out_tile = out_array[out_r_low:out_r_high, out_c_low:out_c_high,:]
+                assert out_tile.shape[1] == filter_shape[1]
 
                 if c % 3 == 0: # linear filter
                     if w_col < w.shape[1]:
@@ -148,3 +165,119 @@
                         w_col += self.n_S_quadratic
         return Image.fromarray(out_array, 'RGB')
 add_logging(Rust2005)
+
+class Rust2005Conv(object):
+    """Convolutional version of `Rust2005`
+
+    :note:
+    The layer doesn't contain an option for downsampling. It makes sense to downsample the output
+    using DownsampleMaxPool, but downsampling is orthogonal to the behaviour of this
+    layer so it is not included.
+    
+    """
+
+    l1 = 0.0
+    l2_sqr = 0.0
+
+    eps=1e-6 # default epsilon to prevent sqrt(0) in quadratic filters
+
+    image_shape = None
+    filter_shape = None
+    output_shape = None
+    output_channels = None
+    output_examples = None
+
+    def __init__(self, linpart, Es, Ss, params, eps=eps):
+        """
+        """
+        eps = numpy.asarray(self.eps, linpart.dtype)
+        output, E, S = rust2005_act_from_filters(linpart, Es, Ss, eps)
+
+        update_locals(self, locals())
+
+
+    @classmethod
+    def new(cls, rng, input, image_shape, filter_shape, n_examples, n_filters, n_E, n_S,
+            n_channels=1,
+            eps=1.0e-6, dtype=None, conv_mode='valid',
+            w_range=None,
+            q_range=None
+            ):
+        """ Return Rust2005Conv layer
+
+        layer.output will be 4D tensor with shape (n_examples, n_filters, R, C) where R and C
+        depend on the image_shape, the filter_shape and the convolution mode.
+
+
+        :param rng: generator for randomized initial filters
+        :param input: symbolic input (4D tensor)
+        :type input: 4D tensor with shape (n_examples, n_channels, image_shape[0], image_shape[1])
+
+        :param image_shape:  rows, cols of every channel of every image
+        :param filter_shape: rows, cols of every filter
+        :param bsize: number of images to be treated
+        :param n_filters: number of filters (output will have this many channels)
+        :param n_channels: number of channels in each image and filter
+        :param n_E: number of squared exciting terms
+        :param n_S: number of squared inhibition terms
+        :param eps: epsilon to add to sum-of-squares in sqrt
+        :param dtype: dtype to use for new variables (Default: input.dtype)
+        :param conv_mode: convolution mode
+        :param w_range: linear weights will be drawn uniformly from this range
+        :type w_range: pair (lower_bound, upper_bound
+        :param q_range: quadratic weights will be drawn uniformly from this range
+        :type q_range: pair (lower_bound, upper_bound
+        """
+        if dtype is None:
+            dtype = input.dtype
+
+        irows, icols = image_shape
+        krows, kcols = filter_shape
+
+        conv = ConvOp((n_channels,irows, icols), (krows, kcols), n_filters, n_examples,
+                dx=1, dy=1, output_mode=conv_mode)
+
+        w_shp = (n_filters, n_channels, krows, kcols)
+        b_shp = (n_filters,)
+
+        if w_range is None:
+            w_low = -2.0/numpy.sqrt(image_shape[0] * image_shape[1])
+            w_high = 2.0/numpy.sqrt(image_shape[0] * image_shape[1])
+        else:
+            w_low, w_high = w_range
+
+        if q_range is None:
+            q_low, q_high = w_low, w_high
+        else:
+            q_low, q_high = w_range
+
+        w = shared(numpy.asarray(rng.uniform(low=w_low, high=w_high, size=w_shp), dtype=dtype))
+        b = shared(numpy.asarray(rng.uniform(low=w_low, high=w_low, size=b_shp), dtype=dtype))
+
+        E_w = [
+                shared(numpy.asarray(rng.uniform(low=q_low, high=q_high, size=w_shp), dtype=dtype))
+                for i in xrange(n_E)
+                ]
+        S_w = [
+                shared(numpy.asarray(rng.uniform(low=q_low, high=q_high, size=w_shp), dtype=dtype))
+                for i in xrange(n_S)
+                ]
+
+        rval = cls(
+                linpart=conv(input, w) + b.dimshuffle(0,'x','x'),
+                Es=[conv(input, e) for e in E_w],
+                Ss=[conv(input, s) for s in S_w],
+                params=[w,b]+E_w + S_w,
+                eps=eps)
+
+        # ignore bias in l1 (Yoshua's habit)
+        rval.l1 = sum(abs(p) for p in ([w]+E_w+S_w))
+        rval.l2_sqr = sum(p**2 for p in ([w]+E_w+S_w))
+        rval.image_shape = image_shape
+        rval.filter_shape = filter_shape
+        rval.output_shape = conv.outshp
+        rval.output_channels = n_filters # how many channels of *output*
+        rval.output_examples = n_examples
+
+        return rval
+add_logging(Rust2005Conv) # _debug, _info, _warn, _error, _fatal