Conditionally disable probes

I think that adding an attribute to Probe to conditionally disable it would be a very useful feature.

Something like: nengo.Probe(connection,disable=True), which could be used as nengo.Probe(connection,disable=lambda t: True if t>threshold) to disable them after a certain time has passed.

As it is now, I have to set probe=None if I want to disable collection but then that messes with the rest of my code (for ex. sim.data[ probe ]) that expects probe to exist.

You could do the first part of your request with sample_every=np.inf on the probe. Here’s an example showing how you can apply this to all of the probes in your model with one line of code:

import nengo
import numpy as np

with nengo.Network() as model:
    model.config[nengo.Probe].sample_every = np.inf  # 'disable' all probes

    x = nengo.Ensemble(100, 1)
    p = nengo.Probe(x)  #, sample_every=np.inf)
    
with nengo.Simulator(model) as sim:
    sim.run(1)
    
print(sim.data[p].size)  # => 0

The second part of your request is a bit trickier. I think the easiest way to do something conditional right now would be with an approach like the one here: Spikes analysis.

To expand on @arvoelke’s post, feature requests can be added on the issues page for the respective code repositories. As an example, Nengo (core) feature requests can be added here!

As for how to implement conditional probes in Nengo currently, the best way to do it is by using Nengo nodes (see the spike analysis link that @arvoelke posted)

An update. I was playing around with it, and it looks like you can directly modify the probe data of the Nengo simulation, but it requires quite a bit of hackyness to do so (so, the Node approach is still a better implementation).

If you want to modify the Nengo probe data directly, here’s an example of how you can do so:

import matplotlib.pyplot as plt
import numpy as np
import nengo


# Probe data management class. Needed to keep a reference of the simulator object and
# probes to manage.
class ProbeDataManager:
    def __init__(self):
        self.sim_obj = None
        self.probe_list = []

        self.probe_t_max = 0.005
        self.sim_step_max = 0

    def add_probes(self, probe_list):
        self.probe_list = list(probe_list)

    def set_sim(self, sim):
        self.sim_obj = sim
        self.sim_step_max = int(self.probe_t_max / self.sim_obj.dt)

    def limit_probe_data(self, probe_obj):
        self.sim_obj._sim_data[probe_obj] = \
            self.sim_obj._sim_data[probe_obj][:self.sim_step_max]

    def __call__(self, t):
        if self.sim_obj is not None:
            for p in self.probe_list:
                self.limit_probe_data(p)

pdm = ProbeDataManager()
pdm.probe_t_max = 0.1  # Maximum 100 steps in the simulation

with nengo.Network() as model:
    node1 = nengo.Node(lambda t: np.sin(t * np.pi * 4))
    node2 = nengo.Node(lambda t: np.cos(t * np.pi * 2))

    # Get Nengo to run the probe data manager
    probe_manager = nengo.Node(pdm)

    p1 = nengo.Probe(node1)
    p2 = nengo.Probe(node2)

    # Add probes to probe data manager
    pdm.add_probes([p1, p2])

with nengo.Simulator(model) as sim:
    # Set simulator object to probe data manager
    pdm.set_sim(sim)

    # Run the simulation for 200 steps
    sim.run_steps(200)

# Run the limit_probe_data function once more to take care of out-of-order operations
pdm(0)

# Plot figures
plt.figure()
plt.plot(sim.data[p1])
plt.plot(sim.data[p2])
plt.show()

I think a better solution if you want to do something hacky using the simulator object would be to do it outside of the simulation; i.e., something like this

with nengo.Simulator(model) as sim:
    # Run the simulation for 200 steps
    for _ in range(200):
        sim.step()
        limit_probe_data(sim, probe_obj)

Objects in the model shouldn’t have a handle to a Simulator instance. Even if that might sort of work at times, the code is needlessly complicated and is hard to reason about.

1 Like

Thanks everyone!

I’m sharing my code for the conditional probe; it seems to work as expected:

class ConditionalProbe:
    def __init__( self, ens, probe_from=0 ):
        self.dimensions = ens.dimensions
        self.time = probe_from
        self.probed_data = [ [ ] for _ in range( self.dimensions ) ]
    
    def __call__( self, t, x ):
        if x.shape != (self.dimensions,):
            raise RuntimeError(
                    "Expected dimensions=%d; got shape: %s"
                    % (self.dimensions, x.shape)
                    )
        if t > 0 and t > self.time:
            for i, k in enumerate( x ):
                self.probed_data[ i ].append( k )
    
    @classmethod
    def setup( cls, ens, probe_from=0 ):
        obj = ConditionalProbe( ens, probe_from )
        output = nengo.Node( obj, size_in=ens.dimensions )
        nengo.Connection( ens, output, synapse=0.01 )
        
        return obj
    
    def get_conditional_probe( self ):
        return np.array( self.probed_data ).T

It can be used as:

with nengo.Network() as model:
    x = nengo.Ensemble(n_neurons=10, dimensions=1)

    # probe second half of simulation (with dt=0.001)
    cond_probe = ConditionalProbe.setup( x, probe_from=500 )


with nengo.Simulator(model) as sim:
    sim.run(1)


print(k = cond_probe.get_conditional_probe())
1 Like