def read_weights(weights_path,
weights_namespace,
synapse_name,
destination,
selection,
comm,
io_size,
weights_by_syn_id_dict,
logger=None):
if logger is not None:
logger.info(f"reading weights from namespace {weights_namespace}...")
if weights_path is not None:
weights_iter = \
scatter_read_cell_attribute_selection(weights_path, destination,
namespace=weights_namespace,
selection=selection,
comm=comm, io_size=io_size)
weights_gid_count = 0
weights_syn_count = 0
for this_gid, syn_weight_attr_dict in weights_iter:
syn_ids = syn_weight_attr_dict['syn_id']
weights = syn_weight_attr_dict[synapse_name]
for (syn_id, weight) in zip(syn_ids, weights):
weights_by_syn_id_dict[this_gid][int(syn_id)] = float(weight)
weights_gid_count += 1
weights_syn_count += len(syn_ids)
if logger is not None:
logger.info(
f'destination {destination}; gid {this_gid} has {len(syn_ids)} synaptic weights; '
f'min/max/mean is {(np.min(weights), np.max(weights), np.mean(weights))}'
)
if logger is not None:
logger.info(
f'destination {destination}: '
f'read initial synaptic weights for {weights_gid_count} gids and {weights_syn_count} syns; '
)
return weights_by_syn_id_dict
def read_weights(weights_path,
weights_namespace,
synapse_name,
destination,
selection,
comm,
io_size,
weights_by_syn_id_dict,
logger=None):
if logger is not None:
logger.info("reading weights from namespace %s..." % weights_namespace)
if weights_path is not None:
weights_iter = \
scatter_read_cell_attribute_selection(weights_path, destination,
namespace=weights_namespace,
selection=selection,
comm=comm, io_size=io_size)
weights_gid_count = 0
weights_syn_count = 0
for this_gid, syn_weight_attr_dict in weights_iter:
syn_ids = syn_weight_attr_dict['syn_id']
weights = syn_weight_attr_dict[synapse_name]
for (syn_id, weight) in zip(syn_ids, weights):
weights_by_syn_id_dict[this_gid][int(syn_id)] = float(weight)
weights_gid_count += 1
weights_syn_count += len(syn_ids)
if logger is not None:
logger.info(
'destination: %s; read initial synaptic weights for %i gids and %i syns'
% (destination, weights_gid_count, weights_syn_count))
return weights_by_syn_id_dict
def rate_maps_from_features(env,
pop_name,
input_features_path,
input_features_namespace,
cell_index_set,
time_range=None,
n_trials=1):
"""Initializes presynaptic spike sources from a file with input selectivity features represented as firing rates."""
if time_range is not None:
if time_range[0] is None:
time_range[0] = 0.0
spatial_resolution = float(env.stimulus_config['Spatial Resolution'])
temporal_resolution = float(env.stimulus_config['Temporal Resolution'])
this_input_features_namespace = '%s %s' % (input_features_namespace,
env.arena_id)
input_features_attr_names = [
'Selectivity Type', 'Num Fields', 'Field Width', 'Peak Rate',
'Module ID', 'Grid Spacing', 'Grid Orientation',
'Field Width Concentration Factor', 'X Offset', 'Y Offset'
]
selectivity_type_names = {
i: n
for n, i in viewitems(env.selectivity_types)
}
arena = env.stimulus_config['Arena'][env.arena_id]
arena_x, arena_y = stimulus.get_2D_arena_spatial_mesh(
arena=arena, spatial_resolution=spatial_resolution)
trajectory = arena.trajectories[env.trajectory_id]
t, x, y, d = stimulus.generate_linear_trajectory(
trajectory, temporal_resolution=temporal_resolution)
if time_range is not None:
t_range_inds = np.where((t < time_range[1]) & (t >= time_range[0]))[0]
t = t[t_range_inds]
x = x[t_range_inds]
y = y[t_range_inds]
d = d[t_range_inds]
input_rate_map_dict = {}
pop_index = int(env.Populations[pop_name])
input_features_iter = scatter_read_cell_attribute_selection(
input_features_path,
pop_name,
selection=list(cell_index_set),
namespace=this_input_features_namespace,
mask=set(input_features_attr_names),
comm=env.comm,
io_size=env.io_size)
for gid, selectivity_attr_dict in input_features_iter:
this_selectivity_type = selectivity_attr_dict['Selectivity Type'][0]
this_selectivity_type_name = selectivity_type_names[
this_selectivity_type]
input_cell_config = stimulus.get_input_cell_config(
selectivity_type=this_selectivity_type,
selectivity_type_names=selectivity_type_names,
selectivity_attr_dict=selectivity_attr_dict)
if input_cell_config.num_fields > 0:
rate_map = input_cell_config.get_rate_map(x=x, y=y)
input_rate_map_dict[gid] = rate_map
return input_rate_map_dict
def main(config, coordinates, field_width, gid, input_features_path,
input_features_namespaces, initial_weights_path,
output_features_namespace, output_features_path, output_weights_path,
reference_weights_path, h5types_path, synapse_name,
initial_weights_namespace, output_weights_namespace,
reference_weights_namespace, connections_path, destination, sources,
non_structured_sources, non_structured_weights_namespace,
non_structured_weights_path, arena_id, field_width_scale,
max_opt_iter, max_weight_decay_fraction, optimize_tol, peak_rate,
reference_weights_are_delta, arena_margin, target_amplitude, io_size,
chunk_size, value_chunk_size, cache_size, write_size, verbose,
dry_run, plot, show_fig, save_fig, debug):
"""
:param config: str (path to .yaml file)
:param input_features_path: str (path to .h5 file)
:param initial_weights_path: str (path to .h5 file)
:param initial_weights_namespace: str
:param output_weights_namespace: str
:param connections_path: str (path to .h5 file)
:param destination: str
:param sources: list of str
:param io_size:
:param chunk_size:
:param value_chunk_size:
:param write_size:
:param verbose:
:param dry_run:
:return:
"""
utils.config_logging(verbose)
script_name = __file__
logger = utils.get_script_logger(script_name)
comm = MPI.COMM_WORLD
rank = comm.rank
nranks = comm.size
if io_size == -1:
io_size = comm.size
if rank == 0:
logger.info(f'{comm.size} ranks have been allocated')
env = Env(comm=comm, config_file=config, io_size=io_size)
env.comm.barrier()
if plot and (not save_fig) and (not show_fig):
show_fig = True
if (not dry_run) and (rank == 0):
if not os.path.isfile(output_weights_path):
if initial_weights_path is not None:
input_file = h5py.File(initial_weights_path, 'r')
elif h5types_path is not None:
input_file = h5py.File(h5types_path, 'r')
else:
raise RuntimeError(
'h5types input path must be specified when weights path is not specified.'
)
output_file = h5py.File(output_weights_path, 'w')
input_file.copy('/H5Types', output_file)
input_file.close()
output_file.close()
env.comm.barrier()
LTD_output_weights_namespace = f'LTD {output_weights_namespace} {arena_id}'
LTP_output_weights_namespace = f'LTP {output_weights_namespace} {arena_id}'
this_input_features_namespaces = [
f'{input_features_namespace} {arena_id}'
for input_features_namespace in input_features_namespaces
]
selectivity_type_index = {
i: n
for n, i in viewitems(env.selectivity_types)
}
target_selectivity_type_name = 'place'
target_selectivity_type = env.selectivity_types[
target_selectivity_type_name]
features_attrs = defaultdict(dict)
source_features_attr_names = [
'Selectivity Type', 'Num Fields', 'Field Width', 'Peak Rate',
'Module ID', 'Grid Spacing', 'Grid Orientation',
'Field Width Concentration Factor', 'X Offset', 'Y Offset'
]
target_features_attr_names = [
'Selectivity Type', 'Num Fields', 'Field Width', 'Peak Rate',
'X Offset', 'Y Offset'
]
seed_offset = int(
env.model_config['Random Seeds']['GC Structured Weights'])
spatial_resolution = env.stimulus_config['Spatial Resolution'] # cm
arena = env.stimulus_config['Arena'][arena_id]
default_run_vel = arena.properties['default run velocity'] # cm/s
gid_count = 0
start_time = time.time()
target_gid_set = None
if len(gid) > 0:
target_gid_set = set(gid)
projections = [(source, destination) for source in sources]
graph_info = read_graph_info(connections_path,
namespaces=['Connections', 'Synapses'],
read_node_index=True)
for projection in projections:
if projection not in graph_info:
raise RuntimeError(
f'Projection {projection[0]} -> {projection[1]} is not present in connections file.'
)
if target_gid_set is None:
target_gid_set = set(graph_info[projection][1])
all_sources = sources + non_structured_sources
src_input_features_attr_dict = {source: {} for source in all_sources}
for source in sorted(all_sources):
this_src_input_features_attr_dict = {}
for this_input_features_namespace in this_input_features_namespaces:
logger.info(
f'Rank {rank}: Reading {this_input_features_namespace} feature data for cells in population {source}'
)
input_features_dict = scatter_read_cell_attributes(
input_features_path,
source,
namespaces=[this_input_features_namespace],
mask=set(source_features_attr_names),
comm=env.comm,
io_size=env.io_size)
for gid, attr_dict in input_features_dict[
this_input_features_namespace]:
this_src_input_features_attr_dict[gid] = attr_dict
src_input_features_attr_dict[
source] = this_src_input_features_attr_dict
source_gid_count = env.comm.reduce(
len(this_src_input_features_attr_dict), op=MPI.SUM, root=0)
if rank == 0:
logger.info(
f'Rank {rank}: Read feature data for {source_gid_count} cells in population {source}'
)
dst_gids = []
if target_gid_set is not None:
for i, gid in enumerate(target_gid_set):
if i % nranks == rank:
dst_gids.append(gid)
dst_input_features_attr_dict = {}
for this_input_features_namespace in this_input_features_namespaces:
feature_count = 0
gid_count = 0
logger.info(
f'Rank {rank}: reading {this_input_features_namespace} feature data for {len(dst_gids)} cells in population {destination}'
)
input_features_iter = scatter_read_cell_attribute_selection(
input_features_path,
destination,
namespace=this_input_features_namespace,
mask=set(target_features_attr_names),
selection=dst_gids,
io_size=env.io_size,
comm=env.comm)
for gid, attr_dict in input_features_iter:
gid_count += 1
if (len(coordinates) > 0) or (attr_dict['Num Fields'][0] > 0):
dst_input_features_attr_dict[gid] = attr_dict
feature_count += 1
logger.info(
f'Rank {rank}: read {this_input_features_namespace} feature data for '
f'{gid_count} / {feature_count} cells in population {destination}')
feature_count = env.comm.reduce(feature_count, op=MPI.SUM, root=0)
env.comm.barrier()
if rank == 0:
logger.info(
f'Read {this_input_features_namespace} feature data for {feature_count} cells in population {destination}'
)
feature_dst_gids = list(dst_input_features_attr_dict.keys())
all_feature_gids_per_rank = comm.allgather(feature_dst_gids)
all_feature_gids = sorted(
[item for sublist in all_feature_gids_per_rank for item in sublist])
request_dst_gids = []
for i, gid in enumerate(all_feature_gids):
if i % nranks == rank:
request_dst_gids.append(gid)
dst_input_features_attr_dict = exchange_input_features(
env.comm, request_dst_gids, dst_input_features_attr_dict)
dst_gids = list(dst_input_features_attr_dict.keys())
if rank == 0:
logger.info(
f"Rank {rank} feature dict is {dst_input_features_attr_dict}")
dst_count = env.comm.reduce(len(dst_gids), op=MPI.SUM, root=0)
logger.info(f"Rank {rank} has {len(dst_gids)} feature gids")
if rank == 0:
logger.info(f'Total {dst_count} feature gids')
max_dst_count = env.comm.allreduce(len(dst_gids), op=MPI.MAX)
env.comm.barrier()
max_iter_count = max_dst_count
output_features_dict = {}
LTP_output_weights_dict = {}
LTD_output_weights_dict = {}
non_structured_output_weights_dict = {}
for iter_count in range(max_iter_count):
gc.collect()
local_time = time.time()
selection = []
if len(dst_gids) > 0:
dst_gid = dst_gids.pop()
selection.append(dst_gid)
logger.info(f'Rank {rank} received gid {dst_gid}')
env.comm.barrier()
arena_margin_size = 0.
arena_margin = max(arena_margin, 0.)
target_selectivity_features_dict = {}
target_selectivity_config_dict = {}
target_field_width_dict = {}
for destination_gid in selection:
arena_margin_size = init_selectivity_config(
destination_gid,
spatial_resolution,
arena,
arena_margin,
arena_margin_size,
coordinates,
field_width,
field_width_scale,
peak_rate,
target_selectivity_type,
selectivity_type_index,
dst_input_features_attr_dict,
target_selectivity_features_dict,
target_selectivity_config_dict,
target_field_width_dict,
logger=logger)
arena_x, arena_y = stimulus.get_2D_arena_spatial_mesh(
arena, spatial_resolution, margin=arena_margin_size)
selection = list(target_selectivity_features_dict.keys())
initial_weights_by_source_gid_dict = defaultdict(lambda: dict())
initial_weights_by_syn_id_dict = \
read_weights(initial_weights_path, initial_weights_namespace, synapse_name,
destination, selection, env.comm, env.io_size, defaultdict(lambda: dict()),
logger=logger if rank == 0 else None)
non_structured_weights_by_source_gid_dict = defaultdict(lambda: dict())
non_structured_weights_by_syn_id_dict = None
if len(non_structured_sources) > 0:
non_structured_weights_by_syn_id_dict = \
read_weights(non_structured_weights_path, non_structured_weights_namespace, synapse_name,
destination, selection, env.comm, env.io_size, defaultdict(lambda: dict()),
logger=logger if rank == 0 else None)
reference_weights_by_syn_id_dict = None
reference_weights_by_source_gid_dict = defaultdict(lambda: dict())
if reference_weights_path is not None:
reference_weights_by_syn_id_dict = \
read_weights(reference_weights_path, reference_weights_namespace, synapse_name,
destination, selection, env.comm, env.io_size, defaultdict(lambda: dict()),
logger=logger if rank == 0 else None)
source_gid_set_dict = defaultdict(set)
syn_count_by_source_gid_dict = defaultdict(lambda: defaultdict(int))
syn_ids_by_source_gid_dict = defaultdict(lambda: defaultdict(list))
structured_syn_id_count = defaultdict(int)
non_structured_syn_id_count = defaultdict(int)
projections = [(source, destination) for source in all_sources]
edge_iter_dict, edge_attr_info = scatter_read_graph_selection(
connections_path,
selection=selection,
namespaces=['Synapses'],
projections=projections,
comm=env.comm,
io_size=env.io_size)
syn_counts_by_source = init_syn_weight_dicts(
destination, non_structured_sources, edge_iter_dict,
edge_attr_info, initial_weights_by_syn_id_dict,
initial_weights_by_source_gid_dict,
non_structured_weights_by_syn_id_dict,
non_structured_weights_by_source_gid_dict,
reference_weights_by_syn_id_dict,
reference_weights_by_source_gid_dict, source_gid_set_dict,
syn_count_by_source_gid_dict, syn_ids_by_source_gid_dict,
structured_syn_id_count, non_structured_syn_id_count)
for source in syn_counts_by_source:
for this_gid in syn_counts_by_source[source]:
count = syn_counts_by_source[source][this_gid]
logger.info(
f'Rank {rank}: destination: {destination}; gid {this_gid}; '
f'{count} edges from source population {source}')
input_rate_maps_by_source_gid_dict = {}
if len(non_structured_sources) > 0:
non_structured_input_rate_maps_by_source_gid_dict = {}
else:
non_structured_input_rate_maps_by_source_gid_dict = None
for source in all_sources:
source_gids = list(source_gid_set_dict[source])
if rank == 0:
logger.info(
f'Rank {rank}: getting feature data for {len(source_gids)} cells in population {source}'
)
this_src_input_features = exchange_input_features(
env.comm, source_gids, src_input_features_attr_dict[source])
count = 0
for this_gid in source_gids:
attr_dict = this_src_input_features[this_gid]
this_selectivity_type = attr_dict['Selectivity Type'][0]
this_selectivity_type_name = selectivity_type_index[
this_selectivity_type]
input_cell_config = stimulus.get_input_cell_config(
this_selectivity_type,
selectivity_type_index,
selectivity_attr_dict=attr_dict)
this_arena_rate_map = np.asarray(
input_cell_config.get_rate_map(arena_x, arena_y),
dtype=np.float32)
if source in non_structured_sources:
non_structured_input_rate_maps_by_source_gid_dict[
this_gid] = this_arena_rate_map
else:
input_rate_maps_by_source_gid_dict[
this_gid] = this_arena_rate_map
count += 1
for destination_gid in selection:
if is_interactive:
context.update(locals())
save_fig_path = None
if save_fig is not None:
save_fig_path = f'{save_fig}/Structured Weights {destination} {destination_gid}.png'
reference_weight_dict = None
if reference_weights_path is not None:
reference_weight_dict = reference_weights_by_source_gid_dict[
destination_gid]
LTP_delta_weights_dict, LTD_delta_weights_dict, arena_structured_map = \
synapses.generate_structured_weights(destination_gid,
target_map=target_selectivity_features_dict[destination_gid]['Arena Rate Map'],
initial_weight_dict=initial_weights_by_source_gid_dict[destination_gid],
#reference_weight_dict=reference_weight_dict,
#reference_weights_are_delta=reference_weights_are_delta,
#reference_weights_namespace=reference_weights_namespace,
input_rate_map_dict=input_rate_maps_by_source_gid_dict,
non_structured_input_rate_map_dict=non_structured_input_rate_maps_by_source_gid_dict,
non_structured_weights_dict=non_structured_weights_by_source_gid_dict[destination_gid],
syn_count_dict=syn_count_by_source_gid_dict[destination_gid],
max_opt_iter=max_opt_iter,
max_weight_decay_fraction=max_weight_decay_fraction,
target_amplitude=target_amplitude,
arena_x=arena_x, arena_y=arena_y,
optimize_tol=optimize_tol,
verbose=verbose if rank == 0 else False,
plot=plot, show_fig=show_fig,
save_fig=save_fig_path,
fig_kwargs={'gid': destination_gid,
'field_width': target_field_width_dict[destination_gid]})
input_rate_maps_by_source_gid_dict.clear()
target_map_flat = target_selectivity_features_dict[
destination_gid]['Arena Rate Map'].flat
arena_map_residual_mae = np.mean(
np.abs(arena_structured_map - target_map_flat))
output_features_dict[destination_gid] = \
{ fld: target_selectivity_features_dict[destination_gid][fld]
for fld in ['Selectivity Type',
'Num Fields',
'Field Width',
'Peak Rate',
'X Offset',
'Y Offset',]}
output_features_dict[destination_gid][
'Rate Map Residual Mean Error'] = np.asarray(
[arena_map_residual_mae], dtype=np.float32)
this_structured_syn_id_count = structured_syn_id_count[
destination_gid]
output_syn_ids = np.empty(this_structured_syn_id_count,
dtype='uint32')
LTD_output_weights = np.empty(this_structured_syn_id_count,
dtype='float32')
LTP_output_weights = np.empty(this_structured_syn_id_count,
dtype='float32')
i = 0
for source_gid in LTP_delta_weights_dict:
for syn_id in syn_ids_by_source_gid_dict[destination_gid][
source_gid]:
output_syn_ids[i] = syn_id
LTP_output_weights[i] = LTP_delta_weights_dict[source_gid]
LTD_output_weights[i] = LTD_delta_weights_dict[source_gid]
i += 1
LTP_output_weights_dict[destination_gid] = {
'syn_id': output_syn_ids,
synapse_name: LTP_output_weights
}
LTD_output_weights_dict[destination_gid] = {
'syn_id': output_syn_ids,
synapse_name: LTD_output_weights
}
this_non_structured_syn_id_count = non_structured_syn_id_count[
destination_gid]
i = 0
logger.info(
f'Rank {rank}; destination: {destination}; gid {destination_gid}; '
f'generated structured weights for {len(output_syn_ids)} inputs in {time.time() - local_time:.2f} s; '
f'residual error is {arena_map_residual_mae:.2f}')
gid_count += 1
gc.collect()
env.comm.barrier()
if (write_size > 0) and (iter_count % write_size == 0):
if not dry_run:
append_cell_attributes(output_weights_path,
destination,
LTD_output_weights_dict,
namespace=LTD_output_weights_namespace,
comm=env.comm,
io_size=env.io_size,
chunk_size=chunk_size,
value_chunk_size=value_chunk_size)
append_cell_attributes(output_weights_path,
destination,
LTP_output_weights_dict,
namespace=LTP_output_weights_namespace,
comm=env.comm,
io_size=env.io_size,
chunk_size=chunk_size,
value_chunk_size=value_chunk_size)
count = env.comm.reduce(len(LTP_output_weights_dict),
op=MPI.SUM,
root=0)
env.comm.barrier()
if rank == 0:
logger.info(
f'Destination: {destination}; appended weights for {count} cells'
)
if output_features_path is not None:
if output_features_namespace is None:
output_features_namespace = f'{target_selectivity_type_name.title()} Selectivity'
this_output_features_namespace = f'{output_features_namespace} {arena_id}'
append_cell_attributes(
output_features_path,
destination,
output_features_dict,
namespace=this_output_features_namespace)
count = env.comm.reduce(len(output_features_dict),
op=MPI.SUM,
root=0)
env.comm.barrier()
if rank == 0:
logger.info(
f'Destination: {destination}; appended selectivity features for {count} cells'
)
LTP_output_weights_dict.clear()
LTD_output_weights_dict.clear()
output_features_dict.clear()
gc.collect()
env.comm.barrier()
if (iter_count >= 10) and debug:
break
env.comm.barrier()
if not dry_run:
append_cell_attributes(output_weights_path,
destination,
LTD_output_weights_dict,
namespace=LTD_output_weights_namespace,
comm=env.comm,
io_size=env.io_size,
chunk_size=chunk_size,
value_chunk_size=value_chunk_size)
append_cell_attributes(output_weights_path,
destination,
LTP_output_weights_dict,
namespace=LTP_output_weights_namespace,
comm=env.comm,
io_size=env.io_size,
chunk_size=chunk_size,
value_chunk_size=value_chunk_size)
count = comm.reduce(len(LTP_output_weights_dict), op=MPI.SUM, root=0)
env.comm.barrier()
if rank == 0:
logger.info(
f'Destination: {destination}; appended weights for {count} cells'
)
if output_features_path is not None:
if output_features_namespace is None:
output_features_namespace = 'Selectivity Features'
this_output_features_namespace = f'{output_features_namespace} {arena_id}'
append_cell_attributes(output_features_path,
destination,
output_features_dict,
namespace=this_output_features_namespace)
count = env.comm.reduce(len(output_features_dict),
op=MPI.SUM,
root=0)
env.comm.barrier()
if rank == 0:
logger.info(
f'Destination: {destination}; appended selectivity features for {count} cells'
)
env.comm.barrier()
global_count = env.comm.gather(gid_count, root=0)
env.comm.barrier()
if rank == 0:
total_count = np.sum(global_count)
total_time = time.time() - start_time
logger.info(
f'Destination: {destination}; '
f'{env.comm.size} ranks assigned structured weights to {total_count} cells in {total_time:.2f} s'
)
def write_input_cell_selection(env,
input_sources,
write_selection_file_path,
populations=None,
write_kwds={}):
"""
Writes out predefined spike trains when only a subset of the network is instantiated.
:param env: an instance of the `dentate.Env` class
:param input_sources: a dictionary of the form { pop_name, gid_sources }
"""
if 'comm' not in write_kwds:
write_kwds['comm'] = env.comm
if 'io_size' not in write_kwds:
write_kwds['io_size'] = env.io_size
rank = int(env.comm.Get_rank())
nhosts = int(env.comm.Get_size())
dataset_path = env.dataset_path
input_file_path = env.data_file_path
if populations is None:
pop_names = sorted(env.celltypes.keys())
else:
pop_names = populations
for pop_name, gid_range in sorted(viewitems(input_sources)):
gc.collect()
if pop_name not in pop_names:
continue
spikes_output_dict = {}
if (env.cell_selection is not None) and (pop_name
in env.cell_selection):
local_gid_range = gid_range.difference(
set(env.cell_selection[pop_name]))
else:
local_gid_range = gid_range
gid_range = env.comm.allreduce(local_gid_range, op=mpi_op_set_union)
this_gid_range = set([])
for i, gid in enumerate(gid_range):
if i % nhosts == rank:
this_gid_range.add(gid)
has_spike_train = False
spike_input_source_loc = []
if (env.spike_input_attribute_info is not None) and (env.spike_input_ns
is not None):
if (pop_name in env.spike_input_attribute_info) and \
(env.spike_input_ns in env.spike_input_attribute_info[pop_name]):
has_spike_train = True
spike_input_source_loc.append(
(env.spike_input_path, env.spike_input_ns))
if (env.cell_attribute_info is not None) and (env.spike_input_ns
is not None):
if (pop_name in env.cell_attribute_info) and \
(env.spike_input_ns in env.cell_attribute_info[pop_name]):
has_spike_train = True
spike_input_source_loc.append(
(input_file_path, env.spike_input_ns))
if rank == 0:
logger.info(
'*** Reading spike trains for population %s: %d cells: has_spike_train = %s'
% (pop_name, len(this_gid_range), str(has_spike_train)))
if has_spike_train:
vecstim_attr_set = set(['t'])
if env.spike_input_attr is not None:
vecstim_attr_set.add(env.spike_input_attr)
if 'spike train' in env.celltypes[pop_name]:
vecstim_attr_set.add(
env.celltypes[pop_name]['spike train']['attribute'])
cell_spikes_iters = [ scatter_read_cell_attribute_selection(input_path, pop_name, \
list(this_gid_range), \
namespace=input_ns, \
mask=vecstim_attr_set, \
comm=env.comm, io_size=env.io_size)
for (input_path, input_ns) in spike_input_source_loc ]
for cell_spikes_iter in cell_spikes_iters:
spikes_output_dict.update(dict(list(cell_spikes_iter)))
if rank == 0:
logger.info('*** Writing spike trains for population %s: %s' %
(pop_name, str(spikes_output_dict)))
write_cell_attributes(write_selection_file_path, pop_name, spikes_output_dict, \
namespace=env.spike_input_ns, **write_kwds)
def write_connection_selection(env,
write_selection_file_path,
populations=None,
write_kwds={}):
"""
Loads NeuroH5 connectivity file, and writes the corresponding
synapse and network connection mechanisms for the selected postsynaptic cells.
:param env: an instance of the `dentate.Env` class
"""
if 'comm' not in write_kwds:
write_kwds['comm'] = env.comm
if 'io_size' not in write_kwds:
write_kwds['io_size'] = env.io_size
connectivity_file_path = env.connectivity_file_path
forest_file_path = env.forest_file_path
rank = int(env.comm.Get_rank())
nhosts = int(env.comm.Get_size())
syn_attrs = env.synapse_attributes
if populations is None:
pop_names = sorted(env.cell_selection.keys())
else:
pop_names = populations
input_sources = {pop_name: set([]) for pop_name in env.celltypes}
for (postsyn_name, presyn_names) in sorted(viewitems(env.projection_dict)):
gc.collect()
if rank == 0:
logger.info('*** Writing connection selection of population %s' %
(postsyn_name))
if postsyn_name not in pop_names:
continue
gid_range = [
gid for i, gid in enumerate(env.cell_selection[postsyn_name])
if i % nhosts == rank
]
synapse_config = env.celltypes[postsyn_name]['synapses']
weight_dicts = []
has_weights = False
if 'weights' in synapse_config:
has_weights = True
weight_dicts = synapse_config['weights']
if rank == 0:
logger.info('*** Reading synaptic attributes for population %s' %
(postsyn_name))
syn_attributes_iter = scatter_read_cell_attribute_selection(
forest_file_path,
postsyn_name,
selection=gid_range,
namespace='Synapse Attributes',
comm=env.comm,
io_size=env.io_size)
syn_attributes_output_dict = dict(list(syn_attributes_iter))
write_cell_attributes(write_selection_file_path,
postsyn_name,
syn_attributes_output_dict,
namespace='Synapse Attributes',
**write_kwds)
del syn_attributes_output_dict
del syn_attributes_iter
if has_weights:
for weight_dict in weight_dicts:
weights_namespaces = weight_dict['namespace']
if rank == 0:
logger.info(
'*** Reading synaptic weights of population %s from namespaces %s'
% (postsyn_name, str(weights_namespaces)))
for weights_namespace in weights_namespaces:
syn_weights_iter = scatter_read_cell_attribute_selection(
forest_file_path,
postsyn_name,
namespace=weights_namespace,
selection=gid_range,
comm=env.comm,
io_size=env.io_size)
weight_attributes_output_dict = dict(
list(syn_weights_iter))
write_cell_attributes(write_selection_file_path,
postsyn_name,
weight_attributes_output_dict,
namespace=weights_namespace,
**write_kwds)
del weight_attributes_output_dict
del syn_weights_iter
logger.info(
'*** Rank %i: reading connectivity selection from file %s for postsynaptic population: %s: selection: %s'
% (rank, connectivity_file_path, postsyn_name, str(gid_range)))
(graph, attr_info) = scatter_read_graph_selection(connectivity_file_path, selection=gid_range, \
projections=[ (presyn_name, postsyn_name) for presyn_name in sorted(presyn_names) ], \
comm=env.comm, io_size=env.io_size, namespaces=['Synapses', 'Connections'])
for presyn_name in sorted(presyn_names):
gid_dict = {}
edge_count = 0
node_count = 0
if postsyn_name in graph:
if postsyn_name in attr_info and presyn_name in attr_info[
postsyn_name]:
edge_attr_info = attr_info[postsyn_name][presyn_name]
else:
raise RuntimeError('write_connection_selection: missing edge attributes for projection %s -> %s' % \
(presyn_name, postsyn_name))
if 'Synapses' in edge_attr_info and \
'syn_id' in edge_attr_info['Synapses'] and \
'Connections' in edge_attr_info and \
'distance' in edge_attr_info['Connections']:
syn_id_attr_index = edge_attr_info['Synapses']['syn_id']
distance_attr_index = edge_attr_info['Connections'][
'distance']
else:
raise RuntimeError('write_connection_selection: missing edge attributes for projection %s -> %s' % \
(presyn_name, postsyn_name))
edge_iter = compose_iter(lambda edgeset: input_sources[presyn_name].update(edgeset[1][0]), \
graph[postsyn_name][presyn_name])
for (postsyn_gid, edges) in edge_iter:
presyn_gids, edge_attrs = edges
edge_syn_ids = edge_attrs['Synapses'][syn_id_attr_index]
edge_dists = edge_attrs['Connections'][distance_attr_index]
gid_dict[postsyn_gid] = (presyn_gids, {
'Synapses': {
'syn_id': edge_syn_ids
},
'Connections': {
'distance': edge_dists
}
})
edge_count += len(presyn_gids)
node_count += 1
env.comm.barrier()
logger.info(
'*** Rank %d: Writing projection %s -> %s selection: %d nodes, %d edges'
% (rank, presyn_name, postsyn_name, node_count, edge_count))
write_graph(write_selection_file_path, \
src_pop_name=presyn_name, dst_pop_name=postsyn_name, \
edges=gid_dict, comm=env.comm, io_size=env.io_size)
env.comm.barrier()
return input_sources
def write_cell_selection(env,
write_selection_file_path,
populations=None,
write_kwds={}):
"""
Writes out the data necessary to instantiate the selected cells.
:param env: an instance of the `dentate.Env` class
"""
if 'comm' not in write_kwds:
write_kwds['comm'] = env.comm
if 'io_size' not in write_kwds:
write_kwds['io_size'] = env.io_size
rank = int(env.comm.Get_rank())
nhosts = int(env.comm.Get_size())
dataset_path = env.dataset_path
data_file_path = env.data_file_path
if populations is None:
pop_names = sorted(env.cell_selection.keys())
else:
pop_names = populations
for pop_name in pop_names:
gid_range = [
gid for i, gid in enumerate(env.cell_selection[pop_name])
if i % nhosts == rank
]
trees_output_dict = {}
coords_output_dict = {}
num_cells = 0
if (pop_name in env.cell_attribute_info) and (
'Trees' in env.cell_attribute_info[pop_name]):
if rank == 0:
logger.info("*** Reading trees for population %s" % pop_name)
cell_tree_iter, _ = scatter_read_tree_selection(data_file_path, pop_name, selection=gid_range, \
topology=False, comm=env.comm, io_size=env.io_size)
if rank == 0:
logger.info("*** Done reading trees for population %s" %
pop_name)
for i, (gid, tree) in enumerate(cell_tree_iter):
trees_output_dict[gid] = tree
num_cells += 1
assert (len(trees_output_dict) == len(gid_range))
elif (pop_name in env.cell_attribute_info) and (
'Coordinates' in env.cell_attribute_info[pop_name]):
if rank == 0:
logger.info("*** Reading coordinates for population %s" %
pop_name)
cell_attributes_iter = scatter_read_cell_attribute_selection(data_file_path, pop_name, selection=gid_range, \
namespace='Coordinates', comm=env.comm, io_size=env.io_size)
if rank == 0:
logger.info("*** Done reading coordinates for population %s" %
pop_name)
for i, (gid, coords) in enumerate(cell_attributes_iter):
coords_output_dict[gid] = coords
num_cells += 1
if rank == 0:
logger.info(
"*** Writing cell selection for population %s to file %s" %
(pop_name, write_selection_file_path))
append_cell_trees(write_selection_file_path, pop_name,
trees_output_dict, **write_kwds)
write_cell_attributes(write_selection_file_path,
pop_name,
coords_output_dict,
namespace='Coordinates',
**write_kwds)
env.comm.barrier()