Integrated lowpass noise

nengo_loihi/builder/nengo_dl.py

@@ -10,6 +10,7 @@
     discretize_tau_ref,
     LoihiLIF,
     LoihiSpikingRectifiedLinear,
+    LowpassIntegratedNoise,
     LowpassRCNoise,
 )
 
@@ -136,6 +137,67 @@ def generate(self, period, tau_rc=None):
         return (1.0 / d) * (q_rc / r_rc1 - q_s / r_s1)
 
 
+@NoiseBuilder.register(LowpassIntegratedNoise)
+class LowpassIntegratedNoiseBuilder(NoiseBuilder):
+    """nengo_dl builder for the LowpassIntegratedNoise model."""
+
+    def __init__(self, ops, signals, *args, **kwargs):
+        super(LowpassIntegratedNoiseBuilder, self).__init__(
+            ops, signals, *args, **kwargs
+        )
+
+        # tau_s is the time constant of the synaptic filter
+        tau_s = np.concatenate(
+            [
+                model.tau_s * np.ones((op.J.shape[0], 1), dtype=self.np_dtype)
+                for model, op in zip(self.noise_models, ops)
+            ]
+        )
+        self.tau_s = signals.constant(tau_s, dtype=self.dtype)
+
+    def generate(self, period, tau_rc=None):
+        r"""Get Tensorflow code for generating this noise.
+
+        This noise model combines the transfer function for a perfect integrator neuron
+        membrane (e.g. ``SpikingRectifiedLinear``) with that of a ``Lowpass`` filter::
+
+            H(s) = 1/s * 1/(tau*s + 1) = 1/s - tau/(tau*s + 1)
+            h(t) = 1 - exp(-t/tau)
+
+        where ``tau`` is the synapse time constant.
+
+        We wish to look at how this filter applies to a regular spike train of period
+        ``p`` as the length of the spike train approaches infinity. Since this filter
+        convolved with an infinite spike train will result in an infinite series, we
+        subtract the "ideal" response of the neuron, given by ``x(t) = t/p``.
+
+            s(t) = sum_{i=0}^{N=\infty} [1 - e^{-(ip + t)/\tau}] - x((N - 1)p + t)
+                 = N - e^{-t/\tau} / (1 - e^{-p/\tau}) - ((N - 1)p + t) / p
+                 = 1 - e^{-t/\tau} / (1 - e^{-p/\tau}) - t/p
+
+        Because the transfer function has a pure integrator in it, this series is
+        sensitive to offsets that happen at the beginning of the spike train, even
+        when looking at infinite trains. To accommodate this, we compute and subtract
+        the mean of the series.
+
+            E[s(t)] = 1/p \int_0^p s(t) dt = 1/2 - \tau/p
+        """
+        if tau_rc is not None:
+            raise ValueError(
+                "This noise model only works for neurons with pure "
+                "integrator membranes (e.g. `SpikingRectifiedLinear`)"
+            )
+
+        u01 = tf.random.uniform(tf.shape(period))
+        t = u01 * period
+        q_s = tf.exp(-t / self.tau_s)
+        r_s1 = -(tf.math.expm1(-period / self.tau_s))  # 1 - exp(-period/tau_s)
+        eta = (0.5 + self.tau_s / period) - q_s / r_s1 - t / period
+
+        x = tf.math.reciprocal(period)
+        return x * (1 + eta)
+
+
 @NoiseBuilder.register(AlphaRCNoise)
 class AlphaRCNoiseBuilder(NoiseBuilder):
     """nengo_dl builder for the AlphaRCNoise model."""
@@ -287,6 +349,7 @@ class LoihiSpikingRectifiedLinearBuilder(
 
     def __init__(self, ops, signals, config):
         super(LoihiSpikingRectifiedLinearBuilder, self).__init__(ops, signals, config)
+        self.spike_noise = NoiseBuilder.build(ops, signals, config)
 
         self.amplitude = signals.op_constant(
             [op.neurons for op in ops],
@@ -313,8 +376,9 @@ def _rate_step(self, J, dt):
 
         # --- compute Loihi rates (for forward pass)
         period = tf.math.reciprocal(tf.maximum(J, self.epsilon))
-        loihi_rates = self.alpha / tf.math.ceil(period / dt)
-        loihi_rates = tf.where(J > self.zero, loihi_rates, self.zeros)
+        period = dt * tf.math.ceil(period / dt)
+        loihi_rates = self.spike_noise.generate(period)
+        loihi_rates = tf.where(J > self.zero, self.amplitude * loihi_rates, self.zeros)
 
         # --- compute RectifiedLinear rates (for backward pass)
         rates = self.amplitude / (

nengo_loihi/neurons.py

@@ -148,6 +148,17 @@ class LoihiSpikingRectifiedLinear(SpikingRectifiedLinear):
     in e.g. ``nengo`` or ``nengo_dl`` to reproduce these unique Loihi effects.
     """
 
+    def __init__(self, amplitude=1, nengo_dl_noise=None):
+        super(LoihiSpikingRectifiedLinear, self).__init__(amplitude=amplitude)
+        self.nengo_dl_noise = nengo_dl_noise
+
+    @property
+    def _argreprs(self):
+        args = super(LoihiSpikingRectifiedLinear, self)._argreprs
+        if self.nengo_dl_noise is not None:
+            args.append("nengo_dl_noise=%s" % self.nengo_dl_noise)
+        return args
+
     def rates(self, x, gain, bias, dt=0.001):
         return loihi_spikingrectifiedlinear_rates(self, x, gain, bias, dt)
 
@@ -172,6 +183,30 @@ class NeuronOutputNoise:
     pass
 
 
+class LowpassIntegratedNoise(NeuronOutputNoise):
+    """Noise model combining Lowpass synapse and neuron membrane integration.
+
+    Samples "noise" (i.e. variability) from a regular spike train filtered
+    by the following transfer function, where :math:`\\tau_s` is the synapse time
+    constant, and the neuron membrane is a perfect integrator:
+
+    .. math::
+
+        H(s) = [(\\tau_s s + 1) (s)]^{-1}
+
+    Attributes
+    ----------
+    tau_s : float
+        Time constant for Lowpass synaptic filter.
+    """
+
+    def __init__(self, tau_s):
+        self.tau_s = tau_s
+
+    def __repr__(self):
+        return "%s(tau_s=%s)" % (type(self).__name__, self.tau_s)
+
+
 class LowpassRCNoise(NeuronOutputNoise):
     """Noise model combining Lowpass synapse and neuron membrane filters.
 

nengo_loihi/tests/test_neurons.py

@@ -13,6 +13,7 @@
     loihi_rates,
     LoihiLIF,
     LoihiSpikingRectifiedLinear,
+    LowpassIntegratedNoise,
     LowpassRCNoise,
     nengo_rates,
 )
@@ -365,6 +366,9 @@ def test_nengo_dl_neuron_grads(neuron_type, plt):
     [
         LoihiLIF(amplitude=0.3, nengo_dl_noise=LowpassRCNoise(0.001)),
         LoihiLIF(amplitude=0.3, nengo_dl_noise=AlphaRCNoise(0.001)),
+        LoihiSpikingRectifiedLinear(
+            amplitude=0.3, nengo_dl_noise=LowpassIntegratedNoise(0.001)
+        ),
     ],
 )
 def test_nengo_dl_noise(neuron_type, seed, plt):
@@ -395,12 +399,17 @@ def test_nengo_dl_noise(neuron_type, seed, plt):
     if isinstance(neuron_type.nengo_dl_noise, AlphaRCNoise):
         exp_model = 0.7 + 2.8 * np.exp(-0.22 * (x1 - 1))
         atol = 0.12 * exp_model.max()
+        mu_atol = 0.6  # depends on n_noise and variance of noise
     elif isinstance(neuron_type.nengo_dl_noise, LowpassRCNoise):
         exp_model = 1.5 + 2.2 * np.exp(-0.22 * (x1 - 1))
         atol = 0.2 * exp_model.max()
+        mu_atol = 0.6  # depends on n_noise and variance of noise
+    elif isinstance(neuron_type.nengo_dl_noise, LowpassIntegratedNoise):
+        exp_model = 10 * (1 - np.exp(-0.010 * (x1 - 1)))
+        atol = 0.2 * exp_model.max()
+        mu_atol = 2.0
 
     rtol = 0.2
-    mu_atol = 0.6  # depends on n_noise and variance of noise
 
     # --- plots
     plt.subplot(211)