import glob
import json
import os
from snudda.neurons import NeuronMorphologyExtended
from snudda.utils.snudda_path import snudda_parse_path
[docs]
class NeuronPrototype:
""" Helper class, returns a neuron prototype based on parameter_id, morph_id and modulation_id """
def __init__(self,
neuron_path,
neuron_name,
morphology_path=None,
parameter_path=None,
mechanism_path=None,
modulation_path=None,
reaction_diffusion_path=None,
snudda_data=None,
meta_path=None,
virtual_neuron=False,
load_morphology=True,
verbose=False):
self.verbose = verbose
self.snudda_data = snudda_data
if neuron_path:
self.neuron_path = snudda_parse_path(neuron_path, self.snudda_data)
elif parameter_path:
self.neuron_path = os.path.dirname(snudda_parse_path(parameter_path, self.snudda_data))
elif morphology_path:
# morphology path usually points to "morphology" directory, but it can also point to a specific morphology
# which should be in neuron_path then (old format), but it might be in morphology folder...
if os.path.isdir(snudda_parse_path(morphology_path, self.snudda_data)):
self.neuron_path = os.path.dirname(snudda_parse_path(morphology_path, self.snudda_data))
elif os.path.isfile(snudda_parse_path(morphology_path, self.snudda_data)):
morph_path = os.path.dirname(snudda_parse_path(morphology_path, self.snudda_data))
if "morphology" in morph_path:
self.neuron_path = os.path.dirname(morph_path)
else:
self.neuron_path = morph_path
else:
self.neuron_path = None
if morphology_path:
self.morphology_path = snudda_parse_path(morphology_path, self.snudda_data)
elif self.neuron_path:
self.morphology_path = snudda_parse_path(os.path.join(self.neuron_path, "morphology"), self.snudda_data)
else:
self.morphology_path = None
if mechanism_path:
self.mechanism_path = snudda_parse_path(mechanism_path, self.snudda_data)
elif self.neuron_path:
self.mechanism_path = snudda_parse_path(os.path.join(self.neuron_path, "mechanisms.json"), self.snudda_data)
else:
self.mechanism_path = None
if parameter_path:
self.parameter_path = snudda_parse_path(parameter_path, self.snudda_data)
elif self.neuron_path:
self.parameter_path = snudda_parse_path(os.path.join(self.neuron_path, "parameters.json"), self.snudda_data)
else:
self.parameter_path = None
if meta_path:
self.meta_path = snudda_parse_path(meta_path, self.snudda_data)
elif self.neuron_path:
self.meta_path = snudda_parse_path(os.path.join(self.neuron_path, "meta.json"), self.snudda_data)
else:
self.meta_path = None
if modulation_path:
self.modulation_path = snudda_parse_path(modulation_path, self.snudda_data)
elif self.neuron_path:
self.modulation_path = snudda_parse_path(os.path.join(self.neuron_path, "modulation.json"), self.snudda_data)
if not os.path.exists(self.modulation_path):
self.modulation_path = None
else:
self.modulation_path = None
if isinstance(reaction_diffusion_path, dict):
self.reaction_diffusion_path = reaction_diffusion_path
elif reaction_diffusion_path:
self.reaction_diffusion_path = snudda_parse_path(reaction_diffusion_path, self.snudda_data)
elif self.neuron_path:
self.reaction_diffusion_path = snudda_parse_path(os.path.join(self.neuron_path, "reaction_diffusion.json"),
self.snudda_data)
if not os.path.exists(self.reaction_diffusion_path):
self.reaction_diffusion_path = None
else:
self.reaction_diffusion_path = None
self.neuron_name = neuron_name
self.parameter_info = None
self.meta_info = None
self.modulation_info = None
self.virtual_neuron = virtual_neuron
self.load_morphology = load_morphology
self.morphology_cache = dict()
self.morphology_lookup = dict()
self.load_info()
[docs]
def load_info(self):
""" Reads information about the neuron prototype. """
self.morphology_cache = dict()
self.morphology_lookup = dict()
par_path = self.parameter_path
if self.meta_path and os.path.exists(self.meta_path):
try:
with open(self.meta_path, "r") as fm:
self.meta_info = json.load(fm)
except Exception as e:
print(f"!! Attempting to load {self.meta_path}")
import traceback
print(traceback.format_exc())
raise e
if par_path is None or not os.path.exists(par_path):
if self.verbose:
print(f"Missing parameters.json : {par_path}")
self.parameter_info = None
return
with open(par_path, "r") as f:
self.parameter_info = json.load(f)
if self.meta_info and isinstance(self.parameter_info, dict):
# Check that all parameter_keys are in meta, otherwise remove them
for par_key in list(self.parameter_info.keys()):
if par_key not in self.meta_info:
del self.parameter_info[par_key]
if self.verbose:
print(f"Parameter key {par_key} missing in {self.meta_path}")
# We now expect a dictionary of parameter sets. If it is a list, we convert it to a dictionary
if type(self.parameter_info) == list:
self.parameter_info = {"default": self.parameter_info}
mod_path = self.modulation_path
if mod_path is not None and os.path.exists(mod_path):
with open(mod_path, "r") as f:
self.modulation_info = json.load(f)
else:
self.modulation_info = None
def get_num_morphologies(self, parameter_id=None, parameter_key=None):
if parameter_key is None:
parameter_key = self.get_parameter_key(parameter_id=parameter_id)
elif parameter_id is not None:
assert parameter_key == self.get_parameter_key(parameter_id=parameter_id), \
f"Internal error, parameter_id {parameter_id} does not match parameter_key {parameter_key} specified."
if self.meta_info and parameter_key:
morph_key_list = self.meta_info[parameter_key].keys()
return len(morph_key_list)
else:
return 1
def get_parameter_key(self, parameter_id):
if parameter_id is None:
return None
# assert parameter_id is not None, (f"get_parameter_key: parameter_id is None, for {self.neuron_name}"
# f"\nNeuron path: {self.neuron_path}")
if self.parameter_info:
par_key_list = sorted(list(self.parameter_info.keys()))
par_key = par_key_list[parameter_id % len(par_key_list)]
elif self.meta_info:
if self.verbose:
print(f"Missing parameter file for {self.parameter_path}, using meta.json for parameter key list.")
par_key_list = sorted(list(self.meta_info.keys()))
par_key = par_key_list[parameter_id % len(par_key_list)]
else:
par_key = None
return par_key
def get_morph_key(self, morphology_id, parameter_id=None, parameter_key=None):
if parameter_id is None:
par_key = parameter_key
else:
par_key = self.get_parameter_key(parameter_id=parameter_id)
assert parameter_key is None or par_key == parameter_key, \
(f"parameter_id = {parameter_id} gives parameter_key {par_key}, "
f"different from parameter_key {parameter_key} provided")
if self.meta_info:
assert par_key in self.meta_info, f"Parameter key {par_key} missing in {self.meta_path}"
morph_key_list = sorted(list(self.meta_info[par_key].keys()))
assert len(morph_key_list) > 0, f"No morphologies available for parameter key {par_key} in {self.meta_path}"
morph_key = morph_key_list[morphology_id % len(morph_key_list)]
else:
morph_key = None
return morph_key
def get_modulation_key(self, modulation_id,
parameter_id=None, morphology_id=None,
parameter_key=None, morphology_key=None):
if modulation_id is None:
return None
if self.modulation_info:
if type(self.modulation_info) == list:
# Old format without keys
return None
if parameter_key is None:
parameter_key = self.get_parameter_key(parameter_id)
if morphology_key is None:
morphology_key = self.get_morph_key(parameter_id=None, parameter_key=parameter_key,
morphology_id=morphology_id)
modulation_key_list = self.meta_info[parameter_key][morphology_key]["neuromodulation"]
if len(modulation_key_list) == 0:
print("Modulation key list empty")
print(f"Path: {self.neuron_path}, param: {parameter_key}, morph: {morphology_key}")
import pdb
pdb.set_trace()
# modulation_key_list = list(self.modulation_info.keys())
modulation_key = modulation_key_list[modulation_id % len(modulation_key_list)]
else:
modulation_key = None
return modulation_key
def get_input_parameters(self, parameter_id=None, morphology_id=None, parameter_key=None, morphology_key=None):
if parameter_key is not None and morphology_key is not None:
if self.meta_info and "input" in self.meta_info[parameter_key][morphology_key]:
input_info = self.meta_info[parameter_key][morphology_key]["input"]
else:
input_info = dict()
else:
# Fallback if the user gave ID instead of Keys (will remove this at some point)
par_key = self.get_parameter_key(parameter_id=parameter_id)
morph_key = self.get_morph_key(parameter_id=parameter_id, morphology_id=morphology_id)
if self.meta_info and "input" in self.meta_info[par_key][morph_key]:
input_info = self.meta_info[par_key][morph_key]["input"]
else:
input_info = dict()
return input_info
[docs]
def get_morphology(self, parameter_key=None, morphology_key=None, parameter_id=None, morphology_id=None):
"""
Returns morphology for a given parameter_id, morphology_id combination.
(Each parameter set has a set of morphologies that it is valid for)
"""
# TODO: In the future remove parameter_id and morphology_id and only use keys
if parameter_id is not None:
par_key = self.get_parameter_key(parameter_id=parameter_id)
if parameter_key:
assert par_key == parameter_key, \
f"Mismatch. Parameter ID {parameter_id} has parameter_key {par_key}, " \
f"not {parameter_key} (Is snudda_data set correctly?)"
elif parameter_key:
par_key = parameter_key
else:
par_key = None
if self.meta_info and par_key is not None:
if morphology_id is not None:
morph_key = self.get_morph_key(parameter_id=parameter_id, morphology_id=morphology_id,
parameter_key=par_key)
if morphology_key:
if morph_key != morphology_key:
raise KeyError(f"Mismatch: Expected morphology_key {morph_key}, got {morphology_key}")
elif morphology_key is not None:
morph_key = morphology_key
if morphology_key not in self.meta_info[par_key]:
raise KeyError(f"Missing morphology_key {morphology_key}")
else:
raise KeyError(f"morphology_id or morphology_key must be specified if meta_info is used")
morph_path = os.path.join(self.morphology_path, self.meta_info[par_key][morph_key]["morphology"])
assert os.path.isfile(morph_path), f"Morphology file {morph_path} is missing (listed in {self.meta_path})"
elif self.morphology_path and os.path.isfile(self.morphology_path):
# Fallback if morphology file is specified in the path
morph_path = self.morphology_path
morph_key = morphology_key
else:
# No morphology
file_list = glob.glob(os.path.join(self.neuron_path, "*swc"))
morph_key = morphology_key
if len(file_list) > 0:
morph_path = file_list[0]
assert len(file_list) == 1, f"More than one morphology available in {self.neuron_path}"
else:
print(f"No morphology in {self.neuron_path}")
morph_path = None
if morph_path is None:
print(f"morph_path is None for {self.neuron_name} path: {self.neuron_path}. "
f"Is SNUDDA_DATA set correctly? ({os.environ['SNUDDA_DATA'] if 'SNUDDA_DATA' in os.environ else None})")
import pdb
pdb.set_trace()
assert morph_path is not None, \
(f"morph_path is None for {self.neuron_name} path: {self.neuron_path} (neuron_path should be a directory, not a SWC file). "
f"Is SNUDDA_DATA set correctly? ({os.environ['SNUDDA_DATA']})")
return morph_path, morph_key
[docs]
def get_parameters(self, parameter_key=None, parameter_id=None):
"""
Returns neuron parameter set
Args:
parameter_key (str) : parameter key for the parameter set
parameter_id (int) : parameter ID, which of the parameter sets to use
Returns:
One parameter set as a list of dictionaries
"""
if self.parameter_info:
if parameter_key is not None:
assert parameter_key in self.parameter_info, \
f"Parameter key {parameter_key} not in parameter_info {self.parameter_info}. ERROR, no key: {parameter_key}"
par_set = self.parameter_info[parameter_key]
else:
assert parameter_id is not None, "If parameter_key is not set, then parameter_id must be set"
par_key_list = list(self.parameter_info.keys())
par_key = par_key_list[parameter_id % len(par_key_list)]
par_set = self.parameter_info[par_key]
else:
par_set = []
return par_set
[docs]
def get_modulation_parameters(self, modulation_key=None, modulation_id=None):
"""
Returns modulation parameter set, give either ID or key
Args:
modulation_key (str) : modulation key
modulation_id (int) : modulation ID, which of the modulation parameter sets to use
Returns:
One parameter set as a list of dictionaries
"""
if self.modulation_info:
if modulation_key:
modulation_set = self.modulation_info[modulation_key]
elif modulation_id is not None:
possible_keys = [x for x in self.modulation_info.keys()]
modulation_set = self.modulation_info[possible_keys[modulation_id % len(self.modulation_info)]]
else:
assert False, f"You need to specify modulation_key or modulation_id to look up a modulation set"
else:
modulation_set = []
return modulation_set
[docs]
def instantiate(self):
"""
Instantiates all morphologies at once, instead of on demand.
"""
n_par = len(self.parameter_info) if self.parameter_info is not None else 1
for par_id in range(0, n_par):
if self.verbose:
print(f"Instantiates par_id = {par_id}")
par_data = self.get_parameters(parameter_id=par_id)
n_morph = self.get_num_morphologies(parameter_id=par_id)
for morph_id in range(0, n_morph):
morph_path, morph_key = self.get_morphology(parameter_id=par_id, morphology_id=morph_id)
morph_tag = os.path.basename(morph_path)
if morph_tag not in self.morphology_cache:
if self.verbose:
print(f"morph_tag = {morph_tag}")
self.morphology_cache[morph_tag] = NeuronMorphologyExtended(swc_filename=morph_path,
param_data=self.parameter_path,
mech_filename=self.mechanism_path,
reaction_diffusion=self.reaction_diffusion_path,
neuron_path=self.neuron_path,
snudda_data=self.snudda_data,
name=self.neuron_name,
morphology_key=morph_key,
load_morphology=self.load_morphology,
virtual_neuron=self.virtual_neuron)
[docs]
def apply(self, function_name, arguments):
"""
Applies function to all morphology prototypes
"""
res = []
for m in self.morphology_cache.values():
function = getattr(m, function_name)
res.append(function(*arguments))
return res
def all_have_axon(self):
all_axon = True
for m in self.morphology_cache.values():
if 2 not in m.morphology_data["neuron"].sections \
or len(m.morphology_data["neuron"].sections[2]) == 0:
all_axon = False
return all_axon
def all_have_dend(self):
all_dend = True
for m in self.morphology_cache.values():
if 3 not in m.morphology_data["neuron"].sections \
or len(m.morphology_data["neuron"].sections[3]) == 0:
all_dend = False
return all_dend
[docs]
def clone(self, parameter_id=None, morphology_id=None, modulation_id=None,
parameter_key=None, morphology_key=None, modulation_key=None,
position=None, rotation=None, get_cache_original=False,
morphology_path=None):
"""
Creates a clone of the neuron prototype, with given position and rotation.
Use either keys or id (keys have precedence if both given)
Args:
parameter_id (int) : parameter set to use
morphology_id (int) : morphology set to use, a parameter set can have multiple morphology variations
modulation_id (int) : neuro-modulation parameter set to use
parameter_key (str) : parameter key for parameter set in parameters.json
morphology_key (str) : morphology key for morphology in meta.json
modulation_key (str) : modulation key for neuromodulation in modulation.json
position (float,float,float) : position of neuron clone
rotation : rotation (3x3 numpy matrix)
get_cache_original (bool) : return the original rather than a clone
morphology_path : If morphology path is given it can override morphology key
"""
if morphology_path is not None:
morph_path = morphology_path
else:
morph_path, morph_key = self.get_morphology(parameter_key=parameter_key, morphology_key=morphology_key,
parameter_id=parameter_id, morphology_id=morphology_id)
if morphology_key is None:
morphology_key = morph_key
else:
assert morphology_key == morph_key, \
f"Internal mismatch requested morphology_key {morphology_key}, got {morph_key}"
morph_tag = os.path.basename(morph_path)
if morph_tag not in self.morphology_cache:
# TODO: hoc file will depend on both morphology_id and parameter_id, we ignore it for now
self.morphology_cache[morph_tag] = NeuronMorphologyExtended(name=self.neuron_name,
position=None, # This is set further down when using clone
rotation=None,
swc_filename=morph_path,
snudda_data=self.snudda_data,
param_data=self.parameter_path,
mech_filename=self.mechanism_path,
reaction_diffusion=self.reaction_diffusion_path,
neuron_path=self.neuron_path,
parameter_key=parameter_key,
morphology_key=morphology_key,
modulation_key=modulation_key,
load_morphology=self.load_morphology,
virtual_neuron=self.virtual_neuron,
verbose=self.verbose)
if get_cache_original:
assert position is None and rotation is None and modulation_id is None, \
"If get_cache_original is passed position, rotation and modulation_id must be None"
morph = self.morphology_cache[morph_tag]
else:
# Add the keys also for parameter, morphology and modulation
if parameter_key is None:
parameter_key = self.get_parameter_key(parameter_id=parameter_id)
if morphology_key is None:
morphology_key = self.get_morph_key(parameter_key=parameter_key, morphology_id=morphology_id)
if modulation_key is None:
modulation_key = self.get_modulation_key(parameter_key=parameter_key,
morphology_key=morphology_key,
modulation_id=modulation_id)
morph = self.morphology_cache[morph_tag].clone(position=position,
rotation=rotation,
parameter_key=parameter_key,
morphology_key=morphology_key,
modulation_key=modulation_key)
if self.meta_info is not None and "axon_density" in self.meta_info.get(parameter_key, {}).get(morphology_key, {}):
axon_density_type, axon_density, axon_density_bounds = self.meta_info[parameter_key][morphology_key]["axon_density"]
if axon_density_type == "r":
morph.set_axon_voxel_radial_density(density=axon_density,
max_axon_radius=axon_density_bounds)
elif axon_density_type == "xyz":
morph.set_axon_voxel_xyz_density(density=axon_density,
axon_density_bounds_xyz=axon_density_bounds)
else:
raise ValueError(f"Unknown axon_density type {axon_density_type}, {self.neuron_path =}")
return morph