Source code for snnib.blender.network

"""base class and function to create a `SNNIB` network

- exposed (and called) via blender UI elements
- controls the look of the generated network

Exceptions

Classes
    - `Network` -- container of the `SNNIB` network

Functions
    - `generate_template_neuron()` -- generates a template neuron and adds it to the scene

Other Objects
"""

#%%imports
import bpy
import bmesh

import importlib
import json
import logging
import numpy as np
from typing import List

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

from snnib import scaling
from . import utils
from . import spiketrain

importlib.reload(utils)
importlib.reload(spiketrain)


#%%definitions
[docs] class Network: """container representing a spiking neural network Attributes - `network_container` -- container defining the boundaries of the `SNNIB` network inside [blender](https://www.blender.org/) - `template_neuron` -- object serving as template for all neurons in the network - `network_file` -- `SNNIB` compatible file representing a network output from a simulation Inferred Attributes - `axon_length` - `float` - length of the axon root - axon only starts to branch out after that point - obtained from [blender](https://www.blender.org/) UI - `axon_objects` - `List[bpy.types.Object]` - objects of all axons that are part of the network - `n_neurons` - `int` - number of neurons contained in the network - obtained from [blender](https://www.blender.org/) UI - only relevant for randomly generated network - `neuron_objects` - `List[bpy.types.Object]` - objects of all neurons that are part of the network - `p_synapses` - `float` - probability of a synapse forming between any two neurons - obtained from [blender](https://www.blender.org/) UI - only relevant for randomly generated network - `p_spike` - `float` - probability of a neuron spiking at any time - i.e., on for every neuron for every frame `p_spike` a spike is emitted with probability `p_spike` - obtained from [blender](https://www.blender.org/) UI - only relevant for randomly generated network - `Rng` - `np.random.Generator` - random number generator to use for network generation - `seed` - `int` - random seed - obtained from [blender](https://www.blender.org/) UI - only relevant for randomly generated network Methods - `_get_mean_outconnection()` -- returns mean direction of outgoing connections of some neuron - `generate_network()` -- generates a random network - `read_network()` -- loads a network from a file - `setup_container()` -- sets up the network container - `draw_neurons()` -- creates neurons and adds them to the scene - `draw_synapses()` -- creates outgoing connections and ads them to the scene Dependencies - `bpy` - `bmesh` - `json` - `logging` - `numpy` - `typing` """
[docs] def __init__(self, network_container:bpy.types.Object, template_neuron:bpy.types.Object, network_file:str=None, ): """constructor - generates a random network if no `network_file` was found - otherwise reads the network from `network_file` Parameters - `network_container` - `bpy.types.Object` - object representing the network container in the scene - bounding box of `network_container` defines the region the network will occupy - neuron coordinates will be remapped to lie within the bounding box - `template_neuron` - `bpy.types.Object` - object representing a template to use for every neuron contained in the network - neurons in the network are instanced from `template_neuron` - you only need to modify `template_neuron` in order to procedurally apply the changes to all neurons in the network - `network_file` - `str`, optional - file containing a save network to load into `SNNIB` - the default is `None` - use random generation instead of reading from a file Raises Returns """ #user input self.network_container = network_container self.template_neuron = template_neuron if network_file in [None,""]: self.network_file = None else: self.network_file = network_file #scene attributes self.axon_length = bpy.context.scene.snnib_props.axon_length self.n_neurons = bpy.context.scene.snnib_props.n_neurons self.p_spike = bpy.context.scene.snnib_props.p_spike self.p_synapses = bpy.context.scene.snnib_props.p_synapses self.seed = bpy.context.scene.snnib_props.seed #get RNG self.Rng = np.random.default_rng(seed=self.seed) if self.network_file is None: #generate random network logger.info(f"generating random network") self.generate_network() else: #read from file logger.info(f"reading {self.network_file}") self.read_network() #inferred attributes self.neuron_objects = [] self.axon_objects = [] return
[docs] def _get_mean_outconnection(self, pre_idx:int, ) -> np.ndarray: """returns vector describing the direction of mean outgoing connection from neuron `pre_idx` - returns random direction if no outgoing connections Parameters - `pre_idx` - `int` - index of the neuron to compute the direction for Raises Returns -`direction` - `np.ndarray` - mean direction of outgoing connections """ #get postsynaptic neurons post_idxs = [int(s_out["post"]) for s_out in filter(lambda s: s["pre"]==pre_idx, self.synapses)] if len(post_idxs) == 0: #random direction direction = self.Rng.random(3) - 0.5 #random direction else: #compute mean connection direction direction = np.array([(self.neurons[pn]["coords"]-self.neurons[pre_idx]["coords"]) for pn in post_idxs]).mean(axis=0) return direction
[docs] def generate_network(self,): """generates random network based on user input - will - generate random coordinates within `self.network_container`s bounding box - generate random spiketrains for every neuron - for every frame there is a probability `self.p_spike` to emit a spike - generate random connections between neurons - for every pair of neurons a connections exists with probability `self.p_synapses` - generate metadata corresponding to the render settings Parameters Raises Returns """ #generate neurons at random locations coords = utils.random.random_points_bbox(self.network_container, self.Rng, self.n_neurons) #TODO (actual inside): coords = utils.random.random_points_raycast(self.network_container, self.Rng, self.n_neurons) #generate random spikes n_frames = bpy.context.scene.frame_end - bpy.context.scene.frame_start spiketrains = [[ frame for frame in range(n_frames) if (self.Rng.random() < self.p_spike) #check if spike occurred at current frame ] for n in range(len(coords))] #generate random synapses synapses = self.Rng.random((self.n_neurons,self.n_neurons)) #synapse weight synapses *= (self.Rng.random(synapses.shape) < self.p_synapses) #connection probability synapses *= (1-np.eye(self.n_neurons,self.n_neurons)) #prevent self-connection (for now) connected = (synapses > 0) synapses = np.append(np.transpose(np.where(connected)), synapses[connected].reshape(-1,1), axis=1) #set attributes self.meta = dict( t_sim=n_frames, t_sim_unit="frame", dt=1, dt_unit="frame", steps=n_frames, ) self.neurons = [dict(coords=coords[n], spiketrain=spiketrains[n]) for n in range(self.n_neurons)] self.synapses = [dict(pre=s[0], post=s[1], w=s[2]) for s in synapses] self.n_synapses = len(synapses) logger.debug(self.neurons[0]) logger.debug(self.synapses[0]) return
[docs] def read_network(self): """reads a network from `self.network_file` - will - read the network elements and parameters - map the contained neuron coordinates to `self.network_container`s bounding box Parameters Raises Returns """ #get context locations bb_min, bb_max = utils.mesh_utils.get_bbox(self.network_container) with open(self.network_file, "r") as f: data = json.load(f) #read and adjust coordinates to bbox coords = np.array([[n[0],n[1],n[2]] for n in data["neurons"]]) coords = scaling.minmaxscale(coords, bb_min, bb_max, axis=0) self.neurons = [dict( coords=coords[nidx], spiketrain=set(np.round(n[3],0).astype(int))) for nidx, n in enumerate(data["neurons"])] self.meta = data["meta"] self.synapses = [dict(pre=s[0], post=s[1], w=s[2]) for s in data["synapses"]] self.n_neurons = len(self.neurons) self.n_synapses = len(self.synapses) self.p_synapses = self.n_synapses / self.n_neurons**2 logger.debug(self.neurons[0]) logger.debug(self.synapses[0]) return
[docs] def setup_container(self): """sets up the network container - will - clear all children of `self.container` - add the respective geo nodes node tree Parameters Raises Returns """ logger.info(f"setting up network container") #initial cleanup logger.warning(f"clearing children of network container") for child in self.network_container.children_recursive: bpy.data.objects.remove(child, do_unlink=True) #add geonodes (if none existent) if len([mod for mod in self.network_container.modifiers if mod.type=='NODES']) == 0: gn = self.network_container.modifiers.new(name="Network.Container", type='NODES') gn.node_group = bpy.data.node_groups["SnnibNetworkContainer"] #other settings self.network_container.hide_render = True return
[docs] def draw_neurons(self,): """draws neurons into the scene - will - instantiate neurons based on `self.template_neuron` - initialize axons (axon roots) - apply geo nodes inputs - adjust geo nodes mappings (to roughly match actual simulation) Parameters Raises Returns """ logger.info(f"drawing {self.n_neurons} neurons") #generating instances for n in range(self.n_neurons): #create new neuron based on `self.template_neuron` neuron_idx = f"Neuron" axon_idx = f"Axon" ######## #NEURON# ######## #instantiation neuron_obj = bpy.data.objects.new(neuron_idx, self.template_neuron.data) neuron_obj.location = self.neurons[n]["coords"] # neuron_obj.rotation_euler = self.Rng.uniform(0, 2*np.pi, 3) #random orientation # utils.mesh_utils.apply_rotation(neuron_obj) #add to network neurons self.neuron_objects.append(neuron_obj) ###### #AXON# ###### #get axon direction # axon_direction = self.Rng.random(3) - 0.5 #random direction # axon_direction = np.array([0,0,1]) #in z axon_direction = self._get_mean_outconnection(n) #mean outgoing connection axon_direction /= np.linalg.norm(axon_direction) #normalize #init neurons axon neuron_verts = neuron_obj.data.vertices axon_data = bpy.data.curves.new(name=axon_idx, type='CURVE') axon_data.dimensions = '3D' _ = utils.mesh_utils.add_spline2data( axon_data, coords=[ (0,0,0), #because child of neuron axon_direction * 0.6 * self.axon_length, #always in z (neuron rotation controls actual orientation) axon_direction * 1.0 * self.axon_length, #always in z (neuron rotation controls actual orientation) ], scale=0.1, #small scale to make sure other outgoing connections have the same root handle_type='ALIGNED', ) ##create object axon_obj = bpy.data.objects.new(axon_idx, axon_data) self.axon_objects.append(axon_obj) #copy geonodes neuron_gn = utils.geo_nodes_utils.copy_geonodes(self.template_neuron, neuron_obj) neuron_gn = [mod for mod in neuron_obj.modifiers if mod.type == 'NODES'][0] #first geonodes modifier of the template object socket_mapping = {item.name:item.identifier for item in neuron_gn.node_group.interface.items_tree} #map sockets to name for easy access ##set geonodes inputs neuron_gn[socket_mapping["Axon Curve"]] = axon_obj # neuron_gn[socket_mapping["Spiketrain"]] = bpy.data.images["SpikeTrain.Main.001"] #TODO: adjust neuron_gn[socket_mapping["Spiketrain"]] = spiketrain.make_spike_texture( spike_steps=self.neurons[n]["spiketrain"], steps=self.meta["steps"], img_name=f"Spiketrain.{neuron_obj.name}", override=False, ) ##adjust geonodes mapping neuron_gn[socket_mapping["Seed"]] = int(self.Rng.integers(0,10000)) #make sure every set of dendrites in unique gn_steps_node = utils.geo_nodes_utils.get_node_by_label(neuron_gn.node_group, "Number of Simulation Steps") gn_steps_node.outputs["Value"].default_value = self.meta["steps"] gn_st_stretch = utils.geo_nodes_utils.get_node_by_label(neuron_gn.node_group, "SpikeTrain.Stretch") gn_st_stretch.outputs["Value"].default_value = self.meta["t_sim"]/self.meta["steps"] #equivalent to `self.meta["dt"]` #parenting axon_obj.parent = neuron_obj neuron_obj.parent = self.network_container for obj in [neuron_obj, axon_obj]: #add to scene (link to all collections that the parent is in) for col in self.network_container.users_collection: col.objects.link(obj) return
[docs] def draw_synapses(self,): """draws synapses into the scene - will - add spline to respective axon object for every existing synapse Parameters Raises Returns """ logger.info(f"drawing {self.n_synapses} synapses") #create synapses (additional splines appended to axon that connect to postsynaptic neuron) for s in range(self.n_synapses): #synapse parameters synapse = self.synapses[s] pre_neuron = self.neuron_objects[int(synapse["pre"])] pre_axon = self.axon_objects[int(synapse["pre"])] post_neuron = self.neuron_objects[int(synapse["post"])] offset = post_neuron.location - pre_neuron.location #destination relative to `pre_neuron` pre_axon_data = pre_axon.data axon_root_points = [bp.co for bp in pre_axon_data.splines[0].bezier_points] _ = utils.mesh_utils.add_spline2data( pre_axon_data, #add another outgoing connection to the axon (new, detached spline) coords=[ *axon_root_points, #start the same as the axon root offset, #target point ], scale=[ *[0.1]*len(axon_root_points), 0.5, ], handle_type=[ *['ALIGNED']*len(axon_root_points), 'ALIGNED', ], ) return
[docs] def generate_template_neuron( name:str, ) -> bpy.types.Object: """returns generated template neuron object - will - generate a template neuron - create a cube - convert to a sphere - add respective geometry nodes (as a single user copy) - add the object (with name `name`) to the scene Parameters - `name` - `str` - name to use for the template neuron Raises Returns - `neuron_obj` - `bpy.types.Object` - generated neuron object """ #create cube bpy.ops.mesh.primitive_cube_add( size=1, location=(0, 0, 0), ) neuron_obj = bpy.context.active_object neuron_obj.name = name neuron_obj.data.name = name #convert to sphere radius = .5 bm = bmesh.new() bm.from_mesh(neuron_obj.data) bmesh.ops.subdivide_edges( #subdivide bm, edges=bm.edges, cuts=5, use_grid_fill=True, ) for v in bm.verts: #convert to sphere (all vertices at constant radius from object origin) v.co = v.co.normalized() * radius bm.to_mesh(neuron_obj.data) bm.free() #add geonodes gn = neuron_obj.modifiers.new(name="Neuron.Axon", type='NODES') gn.node_group = bpy.data.node_groups["SnnibNeuronNeurites"].copy() #make single user copy of node group return neuron_obj
#%%registering def register(): pass def unregister(): pass