def __init__(self, n_components, assignments, dt=0.001, label=None, decoder_cache=NoDecoderCache(), save_file=""): if not h5py_available: raise Exception("h5py not available.") self.n_components = n_components self.assignments = assignments if not save_file and not mpi_sim_available: raise ValueError( "mpi_sim.so is unavailable, so nengo_mpi can only save " "network files (cannot run simulations). However, save_file " "argument was empty." ) # Only create a working simulator if our goal is not to simply # save the network to a file self.mpi_sim = MpiSimulator() if not save_file else None self.h5_compression = "gzip" self.op_strings = defaultdict(list) self.probe_strings = defaultdict(list) self.all_probe_strings = [] if not save_file: save_file = tempfile.mktemp() self.save_file_name = save_file # for each component, stores the keys of the signals that have # to be sent and received, respectively self.send_signals = defaultdict(list) self.recv_signals = defaultdict(list) # for each component, stores the signals that have # already been added to that component. self.signals = defaultdict(list) self.signal_key_set = defaultdict(set) self.total_signal_size = defaultdict(int) # component index (int) -> list of operators # stores the operators for each component self.component_ops = defaultdict(list) # probe -> C++ key (long int) # Used to query the C++ simulator for probe data self.probe_keys = {} self._object_context = [None] # high-level nengo object -> list of operators # stores the operators implementing each high-level object self.object_ops = defaultdict(list) self._mpi_tag = 0 self.pyfunc_args = [] self.probed_connections = set() super(MpiModel, self).__init__(dt, label, decoder_cache)
class MpiModel(builder.Model): """Output of the MpiBuilder, used by nengo_mpi.Simulator. MpiModel differs from Model in the reference implementation in that as the model is built, MpiModel keeps track of the high-level nengo object (e.g. Ensemble) currently being built. This permits it to track which operators implement which high-level objects, so that those operators can later be added to the correct MPI component (required since MPI components are specified in terms of the high-level objects). Parameters ---------- n_components: int Number of components that the network will be divided into. assignments: A dictionary mapping from high-level objects to component indices (ints). All high-level objects in the Network must appear as keys in this dictionary. dt: float Step length. decoder_cache: DecoderCache DecoderCache object to use. save_file: string Name of file to save the built network to. If a non-empty string is provided, then instead of creating a runnable simulator, the result of building the model is saved to a file which can be loaded by the executables bin/nengo_mpi and bin/nengo_cpp to run simulations. """ def __init__(self, n_components, assignments, dt=0.001, label=None, decoder_cache=NoDecoderCache(), save_file=""): if not h5py_available: raise Exception("h5py not available.") self.n_components = n_components self.assignments = assignments if not save_file and not mpi_sim_available: raise ValueError( "mpi_sim.so is unavailable, so nengo_mpi can only save " "network files (cannot run simulations). However, save_file " "argument was empty." ) # Only create a working simulator if our goal is not to simply # save the network to a file self.mpi_sim = MpiSimulator() if not save_file else None self.h5_compression = "gzip" self.op_strings = defaultdict(list) self.probe_strings = defaultdict(list) self.all_probe_strings = [] if not save_file: save_file = tempfile.mktemp() self.save_file_name = save_file # for each component, stores the keys of the signals that have # to be sent and received, respectively self.send_signals = defaultdict(list) self.recv_signals = defaultdict(list) # for each component, stores the signals that have # already been added to that component. self.signals = defaultdict(list) self.signal_key_set = defaultdict(set) self.total_signal_size = defaultdict(int) # component index (int) -> list of operators # stores the operators for each component self.component_ops = defaultdict(list) # probe -> C++ key (long int) # Used to query the C++ simulator for probe data self.probe_keys = {} self._object_context = [None] # high-level nengo object -> list of operators # stores the operators implementing each high-level object self.object_ops = defaultdict(list) self._mpi_tag = 0 self.pyfunc_args = [] self.probed_connections = set() super(MpiModel, self).__init__(dt, label, decoder_cache) @property def runnable(self): """ Return whether this MpiModel can immediately run a simulation. If save_file was not an empty string when __init__ was called, then this should always return False. Otherwise, returns True iff self.finalize_build has been called and completed successfully. """ return self.mpi_sim is not None def __str__(self): return "MpiModel: %s" % self.label def sanitize(self, s): s = re.sub("([0-9])L", lambda x: x.groups()[0], s) return s def build(self, obj, *args, **kwargs): """ Overrides Model.build """ return MpiBuilder.build(self, obj, *args, **kwargs) def _next_mpi_tag(self): """ Return the next mpi tag. Used to ensure that each Connection which straddles a component boundary uses a unique tag. """ mpi_tag = self._mpi_tag self._mpi_tag += 1 return mpi_tag def push_object(self, obj): """ Push high-level object onto context stack. So that we can record which operators implement the object. """ self._object_context.append(obj) def pop_object(self): """ Pop high-level object off of context stack. Once an object has been popped from the context stack, we know that it has finished building, and we can add the operators that implement the object to the MpiModel. We add the operators for the object to the component that the object is assigned to (which is stored in self.assignments). The only exceptions are Connections; when we pop a Connection whose pre object and post object are on different components, then some of operators implementing the connection go to one component, and there rest go to another. """ obj = self._object_context.pop() if not isinstance(obj, Connection): component = self.assignments[obj] self.assign_ops(component, self.object_ops[obj]) return conn = obj pre_component = self.assignments[conn.pre_obj] post_component = self.assignments[conn.post_obj] if pre_component == post_component: self.assign_ops(pre_component, self.object_ops[conn]) else: if conn.learning_rule_type: raise PartitionError("A Connection crossing a component boundary " "must not contain a learning rule.") if conn in self.probed_connections: raise PartitionError("A Connection crossing a component boundary " "must not be probed.") try: synapse_op = (op for op in self.object_ops[conn] if isinstance(op, builder.synapses.SimSynapse)).next() except StopIteration: raise PartitionError( "A Connection crossing a component boundary " "must have a synapse so that it contains an `update` " "operation." ) signal = synapse_op.output tag = self._next_mpi_tag() self.send_signals[pre_component].append((signal, tag, post_component)) self.recv_signals[post_component].append((signal, tag, pre_component)) pre_ops, post_ops = split_connection(self.object_ops[conn], signal) self.assign_ops(pre_component, pre_ops) self.assign_ops(post_component, post_ops) def assign_ops(self, component, ops): """ Assign a collection of operators to a component. Parameters ---------- component: int Component to add the operators to. ops: collection of nengo.builder.operator.Operator instances Operators to add the model. """ for op in ops: for signal in op.all_signals: key = make_key(signal) if key not in self.signal_key_set[component]: logger.debug("Component %d: Adding signal %s with key: %s", component, signal, make_key(signal)) self.signal_key_set[component].add(key) self.signals[component].append((key, signal)) self.total_signal_size[component] += signal.size self.component_ops[component].extend(ops) def add_op(self, op): """ Add operator to model. Overrides Model.add_op. Records that the operator was added as part of building the object that is on top of the _object_context stack. """ self.object_ops[self._object_context[-1]].append(op) def finalize_build(self): """ Finalize the build step. Called once the MpiBuilder has finished running. Finalizes operators and probes, converting them to strings. Then writes all relevant information (signals, ops and probes for each component) to an HDF5 file. Then, if self.mpi_sim is not None (so we want to create a runnable MPI simulator), calls self.mpi_sim.load_file, which tells the C++ code to load the HDF5 file we have just written and create a working simulator. """ all_ops = list(chain(*[self.component_ops[component] for component in range(self.n_components)])) dg = operator_depencency_graph(all_ops) global_ordering = [op for op in toposort(dg) if hasattr(op, "make_step")] self.global_ordering = {op: i for i, op in enumerate(global_ordering)} self._finalize_ops() self._finalize_probes() with h5.File(self.save_file_name, "w") as save_file: save_file.attrs["dt"] = self.dt save_file.attrs["n_components"] = self.n_components for component in range(self.n_components): component_group = save_file.create_group(str(component)) # signals signals = self.signals[component] signal_dset = component_group.create_dataset( "signals", (self.total_signal_size[component],), dtype="float64", compression=self.h5_compression ) offset = 0 for key, sig in signals: A = sig.base._value if A.ndim == 0: A = np.reshape(A, (1, 1)) if A.dtype != np.float64: A = A.astype(np.float64) signal_dset[offset : offset + A.size] = A.flatten() offset += A.size # signal keys component_group.create_dataset( "signal_keys", data=[long(key) for key, sig in signals], dtype="int64", compression=self.h5_compression, ) # signal shapes def pad(x): return (1, 1) if len(x) == 0 else ((x[0], 1) if len(x) == 1 else x) component_group.create_dataset( "signal_shapes", data=np.array([pad(sig.shape) for key, sig in signals]), dtype="u2", compression=self.h5_compression, ) # signal_labels signal_labels = [str(p[1]) for p in signals] store_string_list(component_group, "signal_labels", signal_labels, compression=self.h5_compression) # operators op_strings = self.op_strings[component] store_string_list(component_group, "operators", op_strings, compression=self.h5_compression) # probes probe_strings = self.probe_strings[component] store_string_list(component_group, "probes", probe_strings, compression=self.h5_compression) probe_strings = self.probe_strings[component] store_string_list(save_file, "probe_info", self.all_probe_strings, compression=self.h5_compression) if self.mpi_sim is not None: self.mpi_sim.load_network(self.save_file_name) os.remove(self.save_file_name) for args in self.pyfunc_args: f = { "N": self.mpi_sim.create_PyFunc, "I": self.mpi_sim.create_PyFuncI, "O": self.mpi_sim.create_PyFuncO, "IO": self.mpi_sim.create_PyFuncIO, }[args[0]] f(*args[1:]) self.mpi_sim.finalize_build() def _finalize_ops(self): """ Finalize operators. Main jobs are to create MpiSend and MpiRecv opreators based on send_signals and recv_signals, and to turn all ops belonging to the `component` into strings, which are then stored in self.op_strings. PyFunc ops are the only exception, as it is not generally possible to encode an arbitrary python function as a string. """ for component in range(self.n_components): send_signals = self.send_signals[component] recv_signals = self.recv_signals[component] component_ops = self.component_ops[component] # Store some info to make the next two loops faster for i, op in enumerate(component_ops): for sig in op.updates: if not hasattr(sig, "updated_by"): sig.updated_by = [] sig.updated_by.append(op) for sig in op.reads: if not hasattr(sig, "read_by"): sig.read_by = [] sig.read_by.append(op) for signal, tag, dst in send_signals: mpi_send = MpiSend(dst, tag, signal) assert len(signal.updated_by) == 1 # Put the send after the op that updates the signal. self.global_ordering[mpi_send] = self.global_ordering[signal.updated_by[0]] + 0.5 component_ops.append(mpi_send) for signal, tag, src in recv_signals: mpi_recv = MpiRecv(src, tag, signal) assert len(signal.read_by) > 0 # Put the recv in front of the first op that reads the signal. self.global_ordering[mpi_recv] = self.global_ordering[signal.read_by[0]] - 0.5 component_ops.append(mpi_recv) # Sort to make the ordering take effect. op_order = sorted(component_ops, key=self.global_ordering.__getitem__) self.component_ops[component] = op_order for op in op_order: op_type = type(op) if op_type == builder.node.SimPyFunc: if not self.runnable: raise Exception("Cannot create SimPyFunc operator " "when saving to file.") t_in = op.t_in fn = op.fn x = op.x if x is None: if op.output is None: pyfunc_args = ["N", fn, t_in] else: pyfunc_args = ["O", make_checked_func(fn, t_in, False), t_in, signal_to_string(op.output)] else: input_array = x.value if op.output is None: pyfunc_args = ["I", fn, t_in, signal_to_string(x), input_array] else: pyfunc_args = [ "IO", make_checked_func(fn, t_in, True), t_in, signal_to_string(x), input_array, signal_to_string(op.output), ] self.pyfunc_args.append(pyfunc_args + [self.global_ordering[op]]) else: op_string = self._op_to_string(op) if op_string: logger.debug("Component %d: Adding operator with string: %s", component, op_string) self.op_strings[component].append(op_string) def _op_to_string(self, op): """ Convert operator into a string. Such strings will eventually be used to construct operators in the C++ code. See the MpiSimulatorChunk::add_op for details on how these strings are used by the C++ code. We also prepend the operator's index in the global ordering of operators, which allows the C++ code to put the operators in the appropriate order. """ op_type = type(op) if op_type == builder.operator.Reset: op_args = ["Reset", signal_to_string(op.dst), op.value] elif op_type == builder.operator.Copy: op_args = ["Copy", signal_to_string(op.dst), signal_to_string(op.src)] elif op_type == builder.operator.SlicedCopy: try: seq_A = list(iter(op.a_slice)) start_A, stop_A, step_A = 0, 0, 0 except: seq_A = [] if op.a_slice == Ellipsis: start_A, stop_A, step_A = 0, op.a.size, 1 else: start_A, stop_A, step_A = op.a_slice.indices(op.a.size) try: seq_B = list(iter(op.b_slice)) start_B, stop_B, step_B = 0, 0, 0 except: seq_B = [] if op.b_slice == Ellipsis: start_B, stop_B, step_B = 0, op.b.size, 1 else: start_B, stop_B, step_B = op.b_slice.indices(op.b.size) op_args = [ "SlicedCopy", signal_to_string(op.b), signal_to_string(op.a), int(op.inc), start_A, stop_A, step_A, start_B, stop_B, step_B, str(seq_A), str(seq_B), ] elif op_type == builder.operator.DotInc: op_args = ["DotInc", signal_to_string(op.A), signal_to_string(op.X), signal_to_string(op.Y)] elif op_type == builder.operator.ElementwiseInc: op_args = ["ElementwiseInc", signal_to_string(op.A), signal_to_string(op.X), signal_to_string(op.Y)] elif op_type == builder.neurons.SimNeurons: n_neurons = op.J.size neuron_type = type(op.neurons) if neuron_type is LIF: tau_ref = op.neurons.tau_ref tau_rc = op.neurons.tau_rc min_voltage = op.neurons.min_voltage voltage_signal = signal_to_string(op.states[0]) ref_time_signal = signal_to_string(op.states[1]) op_args = [ "LIF", n_neurons, tau_rc, tau_ref, min_voltage, self.dt, signal_to_string(op.J), signal_to_string(op.output), voltage_signal, ref_time_signal, ] elif neuron_type is LIFRate: tau_ref = op.neurons.tau_ref tau_rc = op.neurons.tau_rc op_args = ["LIFRate", n_neurons, tau_rc, tau_ref, signal_to_string(op.J), signal_to_string(op.output)] elif neuron_type is AdaptiveLIF: tau_n = op.neurons.tau_n inc_n = op.neurons.inc_n tau_rc = op.neurons.tau_rc tau_ref = op.neurons.tau_ref min_voltage = op.neurons.min_voltage voltage_signal = signal_to_string(op.states[0]) ref_time_signal = signal_to_string(op.states[1]) adaptation = signal_to_string(op.states[2]) op_args = [ "AdaptiveLIF", n_neurons, tau_n, inc_n, tau_rc, tau_ref, min_voltage, self.dt, signal_to_string(op.J), signal_to_string(op.output), voltage_signal, ref_time_signal, adaptation, ] elif neuron_type is AdaptiveLIFRate: tau_n = op.neurons.tau_n inc_n = op.neurons.inc_n tau_rc = op.neurons.tau_rc tau_ref = op.neurons.tau_ref adaptation = signal_to_string(op.states[0]) op_args = [ "AdaptiveLIFRate", n_neurons, tau_n, inc_n, tau_rc, tau_ref, self.dt, signal_to_string(op.J), signal_to_string(op.output), adaptation, ] elif neuron_type is RectifiedLinear: op_args = ["RectifiedLinear", n_neurons, signal_to_string(op.J), signal_to_string(op.output)] elif neuron_type is Sigmoid: op_args = [ "Sigmoid", n_neurons, op.neurons.tau_ref, signal_to_string(op.J), signal_to_string(op.output), ] elif neuron_type is Izhikevich: tau_recovery = op.neurons.tau_recovery coupling = op.neurons.coupling reset_voltage = op.neurons.reset_voltage reset_recovery = op.neurons.reset_recovery voltage = signal_to_string(op.states[0]) recovery = signal_to_string(op.states[1]) op_args = [ "Izhikevich", n_neurons, tau_recovery, coupling, reset_voltage, reset_recovery, self.dt, signal_to_string(op.J), signal_to_string(op.output), voltage, recovery, ] else: raise NotImplementedError("nengo_mpi cannot handle neurons of type " + str(neuron_type)) elif op_type == builder.synapses.SimSynapse: if isinstance(op.synapse, LinearFilter): step = op.synapse.make_step(self.dt, []) den = step.den num = step.num if len(num) == 1 and len(den) == 0: op_args = ["NoDenSynapse", signal_to_string(op.input), signal_to_string(op.output), num[0]] elif len(num) == 1 and len(den) == 1: op_args = ["SimpleSynapse", signal_to_string(op.input), signal_to_string(op.output), den[0], num[0]] else: op_args = [ "Synapse", signal_to_string(op.input), signal_to_string(op.output), ",".join(map(str, num)), ",".join(map(str, den)), ] elif isinstance(op.synapse, Triangle): f = op.synapse.make_step(self.dt, op.output) closures = get_closures(f) n0 = closures["n0"] ndiff = closures["ndiff"] x = closures["x"] n_taps = x.maxlen op_args = [ "TriangleSynapse", signal_to_string(op.input), signal_to_string(op.output), n0, ndiff, n_taps, ] else: raise NotImplementedError("nengo_mpi cannot handle synapses of " "type %s" % type(op.synapse)) elif op_type == builder.processes.SimProcess: process_type = type(op.process) if process_type is WhiteNoise: assert type(op.process.dist) is nengo.dists.Gaussian mean = op.process.dist.mean std = op.process.dist.std do_scale = op.process.scale op_args = [ "WhiteNoise", signal_to_string(op.output), float(mean), float(std), int(do_scale), int(op.inc), self.dt, ] elif process_type is WhiteSignal: f = op.process.make_step(0, op.output.size, self.dt, np.random.RandomState()) closures = get_closures(f) assert closures["dt"] == self.dt coefs = closures["signal"] op_args = ["WhiteSignal", signal_to_string(op.output), ndarray_to_string(coefs)] elif process_type in [FilteredNoise, BrownNoise]: raise NotImplementedError("nengo_mpi cannot handle processes of " "type %s" % str(process_type)) else: raise NotImplementedError("Unrecognized process type: %s." % str(process_type)) elif op_type == builder.learning_rules.SimBCM: op_args = [ "BCM", signal_to_string(op.pre_filtered), signal_to_string(op.post_filtered), signal_to_string(op.theta), signal_to_string(op.delta), op.learning_rate, self.dt, ] elif op_type == builder.learning_rules.SimOja: op_args = [ "Oja", signal_to_string(op.pre_filtered), signal_to_string(op.post_filtered), signal_to_string(op.weights), signal_to_string(op.delta), op.learning_rate, self.dt, op.beta, ] elif op_type == builder.learning_rules.SimVoja: op_args = [ "Voja", signal_to_string(op.pre_decoded), signal_to_string(op.post_filtered), signal_to_string(op.scaled_encoders), signal_to_string(op.delta), signal_to_string(op.learning_signal), ",".join(map(str, op.scale)), op.learning_rate, self.dt, ] elif op_type == builder.operator.PreserveValue: logger.debug("Skipping PreserveValue, operator: %s, signal: %s", str(op.dst), signal_to_string(op.dst)) op_args = [] elif op_type == MpiSend: signal_key = make_key(op.signal) op_args = ["MpiSend", op.dst, op.tag, signal_key] elif op_type == MpiRecv: signal_key = make_key(op.signal) op_args = ["MpiRecv", op.src, op.tag, signal_key] elif op_type == SpaunStimulusOperator: output = signal_to_string(op.output) op_args = [ "SpaunStimulus", output, op.stimulus_sequence, op.present_interval, op.present_blanks, op.identifier, ] else: raise NotImplementedError("nengo_mpi cannot handle operator of " "type %s" % str(op_type)) if op_args: op_args = [self.global_ordering[op]] + op_args op_string = OP_DELIM.join(map(str, op_args)) op_string = op_string.replace(" ", "") op_string = op_string.replace("(", "") op_string = op_string.replace(")", "") return op_string def _finalize_probes(self): """ Finalize probes. Main job is to convert all probes in self.probes into strings, which then get stored in self.probe_strings. """ for probe in self.probes: period = 1 if probe.sample_every is None else probe.sample_every / self.dt probe_key = make_key(probe) self.probe_keys[probe] = probe_key signal = self.sig[probe]["in"] signal_string = signal_to_string(signal) component = self.assignments[probe] logger.debug( "Component: %d: Adding probe of signal %s.\n" "probe_key: %d, signal_string: %s, period: %d", component, str(signal), probe_key, signal_string, period, ) probe_string = PROBE_DELIM.join(str(i) for i in [component, probe_key, signal_string, period, str(probe)]) self.probe_strings[component].append(probe_string) self.all_probe_strings.append(probe_string)