def save_morph_data(self, morph_stats_filename): if not os.path.exists(morph_stats_filename): # load a neuron from an SWC file nrn = nm.load_neuron(self.morph_file) morph_stats = {} morph_stats['cell_id'] = self.cell_id morph_stats['soma_suface'] = nm.get('soma_surface_areas', nrn)[0] morph_stats['soma_radius'] = np.mean(nm.get('soma_radii', nrn)) # Morph stats for nrn_type_ in NEURITES: morph_stats['length' + '.' + str(nrn_type_).split('.')[1]] = np.sum( nm.get('segment_lengths', nrn, neurite_type=nrn_type_)) morph_stats['area' + '.' + str(nrn_type_).split('.')[1]] = sum( mm.segment_area(s) for s in nm.iter_segments( nrn, neurite_filter=tree_type_checker(nrn_type_))) morph_stats[ 'volume' + '.' + str(nrn_type_).split('.')[1]] = sum( mm.segment_volume(s) for s in nm.iter_segments( nrn, neurite_filter=tree_type_checker(nrn_type_))) morph_stats['taper_rate' + '.' + str(nrn_type_).split('.')[1]] = \ np.mean([mm.segment_taper_rate(s) for s in nm.iter_segments( nrn, neurite_filter=tree_type_checker(nrn_type_))]) utility.save_json(morph_stats_filename, morph_stats)
def write_opt_config_file(self, param_write_path, mech_write_path, mech_release_write_path, train_features_write_path, test_features_write_path, protocols_write_path, release_params, release_param_write_path, opt_config_filename='config_file.json', **kwargs): if not os.path.exists(opt_config_filename): path_dict = dict() else: path_dict = utility.load_json(opt_config_filename) path_dict['parameters'] = param_write_path path_dict['mechanism'] = mech_write_path path_dict['released_aa_mechanism'] = mech_release_write_path path_dict['train_features'] = train_features_write_path path_dict['test_features'] = test_features_write_path path_dict['train_protocols'] = protocols_write_path path_dict['released_aa_model_dict'] = release_params path_dict['released_aa_model'] = release_param_write_path path_dict['released_peri_model'] = kwargs.get('released_peri_model') path_dict['released_peri_mechanism'] = kwargs.get( 'released_peri_mechanism') # for config_key,path in kwargs.items(): # path_dict[config_key] = path utility.save_json(opt_config_filename, path_dict)
def script_generator(self, chain_job='chain_job.sh', **kwargs): # Force change of certain config properties job_config = utility.load_json(self.job_config_path) stage_jobconfig = job_config['stage_jobconfig'] highlevel_job_props = job_config['highlevel_jobconfig'] analysis_config = stage_jobconfig['analysis_config'] optim_config = stage_jobconfig['optim_config'] stage_jobconfig = update(stage_jobconfig, dryrun_config) stage_jobconfig['optim_config']['ipyparallel'] = False stage_jobconfig['analysis_config']['ipyparallel'] = False stage_jobconfig['seed'] = [1] utility.save_json(self.job_config_path, job_config) testjob_string = '#!/bin/bash\n' testjob_string += 'set -ex\n' testjob_string += 'source activate %s\n' % highlevel_job_props[ 'conda_env'] testjob_string += 'python %s --input_json %s\n' %\ (optim_config['main_script'], self.job_config_path) testjob_string += 'python %s --input_json %s\n'\ % (analysis_config['main_script'], self.job_config_path) if 'next_stage_job_config' in kwargs.keys(): if bool(kwargs['next_stage_job_config']): testjob_string += 'bash %s\n' % chain_job with open(self.script_name, "w") as shell_script: shell_script.write(testjob_string)
def write_params_opt(self, model_params, model_params_release, base_dir='config/', **kwargs): param_write_path = kwargs.get('param_write_path') or \ base_dir + self.cell_id + '/parameters.json' release_param_write_path = kwargs.get('release_param_write_path') or \ base_dir + self.cell_id + '/release_parameters.json' utility.create_filepath(param_write_path) utility.save_json(param_write_path, model_params) release_params = dict( ) # for comparison with optimized values (strictly for values) if model_params_release: for param_dict_release in model_params_release: param_name = param_dict_release['param_name'] if param_name not in ['ena', 'ek', 'v_init', 'celsius']: release_params[param_name + '.' + param_dict_release['sectionlist']] = \ param_dict_release['value'] # Released parameter file in bpopt format for running simulation utility.save_json(release_param_write_path, model_params_release) else: release_param_write_path = None return param_write_path, release_param_write_path, release_params
def create_sa_bound_peri(self,bpopt_param_bounds_path,sens_param_bounds_path, max_bound = .5): # For parameter sensitivity create a new set of bounds because the permutations may # fall outside the original bounds max_bound = max(max_bound,self.param_range+.1) param_bounds = utility.load_json(bpopt_param_bounds_path) optim_param_bpopt_format = {} param_sens_list = list() for i,param_dict in enumerate(param_bounds): if 'sectionlist' in param_dict.keys(): name_loc = param_dict['param_name'] + '.' + \ param_dict['sectionlist'] if param_dict['param_name'] not in ['ena','ek']: optim_param_bpopt_format[name_loc] = param_dict['value'] lb = param_dict['value']-\ max_bound*abs(param_dict['value']) ub = param_dict['value']+\ max_bound*abs(param_dict['value']) param_dict['bounds'] = [lb,ub] del param_dict['value'] param_sens_list.append(param_dict) utility.save_json(sens_param_bounds_path,param_sens_list) return optim_param_bpopt_format
def create_sa_bound(self,bpopt_param_bounds_path,sens_param_bounds_path, max_bound = .5): # For parameter sensitivity create a new set of bounds because the permutations may # fall outside the original bounds max_bound = max(max_bound,self.param_range+.1) bpopt_section_map = utility.bpopt_section_map param_bounds = utility.load_json(bpopt_param_bounds_path) optim_param = self.optim_param optim_param_bpopt_format = {} if 'genome' in optim_param.keys(): print('The parameter file is in AIBS format') for aibs_param_dict in optim_param['genome']: param_name,param_sect = aibs_param_dict['name'],\ aibs_param_dict['section'] if param_name in ['e_pas','g_pas','Ra']: param_sect = 'all' param_sect = bpopt_section_map[param_sect] optim_param_bpopt_format[param_name+'.'+param_sect]=\ float(aibs_param_dict['value']) else: print('The parameter file is in bluepyopt format') for key,val in optim_param.items(): key_param,key_sect = key.split('.') try: key_sect = bpopt_section_map[key_sect] except: print('Already in bluepyopt format') optim_param_bpopt_format[key_param+'.'+key_sect]=val param_sens_list = list() for i,param_dict in enumerate(param_bounds): bound = param_dict.get('bounds') if bound: name_loc = param_dict['param_name'] + '.' + param_dict['sectionlist'] lb = min(param_dict['bounds'][0], optim_param_bpopt_format[name_loc]-\ max_bound*abs(optim_param_bpopt_format[name_loc])) ub = max(param_dict['bounds'][1], optim_param_bpopt_format[name_loc]+\ max_bound*abs(optim_param_bpopt_format[name_loc])) param_dict['bounds'] = [lb,ub] param_sens_list.append(param_dict) utility.save_json(sens_param_bounds_path,param_sens_list) return optim_param_bpopt_format
def save_cell_metadata(**cell_metadata): cell_id = cell_metadata["cell_id"] metadata_filename = 'cell_metadata_%s.json' % cell_id # TODO: maybe move this to launch_optimjob? machine_name = socket.gethostname() machine_name = 'aws' if machine_name == 'master' else machine_name cell_metadata['machine'] = machine_name data_source = cell_metadata["data_source"] if data_source == "web": ctc = CellTypesCache( manifest_file=os.path.join(data_dir, 'manifest.json')) cell_metadata.update(get_data_web(cell_id, ctc)) cell_metadata.update(cell_props(cell_id, ctc)) cell_metadata.update(model_props(cell_id)) elif data_source == "lims": cell_metadata.update(get_data_lims(cell_id)) utility.save_json(metadata_filename, cell_metadata) return cell_metadata, metadata_filename
def write_mechanisms_opt(self, model_mechs, model_mechs_release, base_dir='config/', **kwargs): mechanism_write_path = kwargs.get('mechanism_write_path') or \ base_dir + self.cell_id + '/mechanism.json' mechanism_release_write_path = kwargs.get('mechanism_release_write_path') \ or base_dir + self.cell_id + '/mechanism_release.json' utility.create_filepath(mechanism_write_path) utility.save_json(mechanism_write_path, model_mechs) if model_mechs_release: utility.save_json(mechanism_release_write_path, model_mechs_release) else: model_mechs_release = None mechanism_release_write_path = None return mechanism_write_path, mechanism_release_write_path
def get_efeatures(cell_id): from ateamopt.nwb_extractor import NwbExtractor from ateamopt.optim_config_rules import correct_voltage_feat_std from ateam.data import lims import os from ateamopt.utils import utility import shutil acceptable_stimtypes = ['Long Square'] feature_names_path = 'feature_set_all.json' ephys_dir_origin = 'ephys_ttype' efel_feature_path = 'eFEL_features_ttype' lr = lims.LimsReader() cell_efeatures_dir = os.path.join(efel_feature_path, cell_id) cell_protocols_filename = os.path.join(cell_efeatures_dir, 'protocols.json') cell_features_filename = os.path.join(efel_feature_path, cell_id, 'features.json') efeature_filenames = [cell_protocols_filename, cell_features_filename] if not all(os.path.exists(filename_) for filename_ in efeature_filenames): utility.create_filepath(cell_protocols_filename) nwb_path = lr.get_nwb_path_from_lims(int(cell_id), get_sdk_version=True) cell_ephys_dir = os.path.join(ephys_dir_origin, cell_id) nwb_handler = NwbExtractor(cell_id, nwb_path) ephys_data_path, stimmap_filename = nwb_handler.save_cell_data_web( acceptable_stimtypes, ephys_dir=cell_ephys_dir) protocol_dict, feature_dict = nwb_handler.get_efeatures_all( feature_names_path, ephys_data_path, stimmap_filename) feature_dict = correct_voltage_feat_std(feature_dict) utility.save_json(cell_protocols_filename, protocol_dict) utility.save_json(cell_features_filename, feature_dict) shutil.move(stimmap_filename, cell_efeatures_dir) shutil.rmtree(cell_ephys_dir, ignore_errors=True)
def write_ephys_features(self, train_features, test_features, train_protocols, base_dir='config/', **kwargs): cell_name = self.cell_id train_features_write_path = kwargs.get('train_features_write_path') or \ os.path.join(base_dir, cell_name, 'train_features.json') test_features_write_path = kwargs.get('test_features_write_path') \ or os.path.join(base_dir, cell_name, 'test_features.json') train_protocols_write_path = kwargs.get('protocols_write_path') \ or os.path.join(base_dir, cell_name, 'train_protocols.json') utility.create_filepath(train_protocols_write_path) utility.save_json(train_features_write_path, train_features) utility.save_json(test_features_write_path, test_features) utility.save_json(train_protocols_write_path, train_protocols) return train_features_write_path, test_features_write_path,\ train_protocols_write_path
def main(): # Read sensitivity analysis config file sens_config_file = sys.argv[-1] sens_config_dict = utility.load_json(sens_config_file) cell_id = sens_config_dict['Cell_id'] cpu_count = sens_config_dict['cpu_count'] if 'cpu_count'\ in sens_config_dict.keys() else mp.cpu_count() perisomatic_sa = sens_config_dict.get('run_peri_analysis',False) # Parameters to vary (All-active) select_aa_param_path = sens_config_dict['select_aa_param_path'] # knobs # Parameters to vary (Perisomatic) if perisomatic_sa: select_peri_param_path = sens_config_dict['select_peri_param_path'] # knobs select_feature_path = sens_config_dict['select_feature_path'] # knobs param_mod_range = sens_config_dict.get('param_mod_range',.1) # knobs mechanism_path = sens_config_dict['mechanism'] # config files with all the paths for Bluepyopt sim lr = lims.LimsReader() morph_path = lr.get_swc_path_from_lims(int(cell_id)) model_base_path='/allen/aibs/mat/ateam_shared/' \ 'Mouse_Model_Fit_Metrics/{}'.format(cell_id) opt_config_file = os.path.join(model_base_path,'config_file.json') if not os.path.exists(opt_config_file): opt_config = { "morphology": "", "parameters": "config/{}/parameters.json".format(cell_id), "mechanism": "config/{}/mechanism.json".format(cell_id), "protocols": "config/{}/protocols.json".format(cell_id), "all_protocols": "config/{}/all_protocols.json".format(cell_id), "features": "config/{}/features.json".format(cell_id), "peri_parameters": "config/{}/peri_parameters.json".format(cell_id), "peri_mechanism": "config/{}/peri_mechanism.json".format(cell_id) } opt_config_file = os.path.join(os.getcwd(),'config_file.json') utility.save_json(opt_config_file,opt_config) # optimized parameters around which select parameters are varied optim_param_path_aa = '/allen/aibs/mat/ateam_shared/Mouse_Model_Fit_Metrics/'\ '{cell_id}/fitted_params/optim_param_unformatted_{cell_id}.json'.\ format(cell_id = cell_id) if not os.path.exists(optim_param_path_aa): optim_param_path_aa = '/allen/aibs/mat/ateam_shared/Mouse_Model_Fit_Metrics/'\ '{cell_id}/fitted_params/optim_param_{cell_id}_bpopt.json'.\ format(cell_id = cell_id) SA_obj_aa = SA_helper(optim_param_path_aa,select_aa_param_path,param_mod_range, opt_config_file) _,protocol_path,mech_path,feature_path,\ param_bound_path = SA_obj_aa.load_config(model_base_path) # Make sure to get the parameter bounds big enough for BluePyOpt sim sens_param_bound_write_path_aa = "param_sensitivity_aa.json" optim_param_aa = SA_obj_aa.create_sa_bound(param_bound_path, sens_param_bound_write_path_aa) param_dict_uc_aa = SA_obj_aa.create_sens_param_dict() parameters_aa ={key:optim_param_aa[val] for key,val in param_dict_uc_aa.items()} eval_handler_aa = Bpopt_Evaluator(protocol_path, feature_path, morph_path, sens_param_bound_write_path_aa, mech_path, ephys_dir=None, timed_evaluation = False) evaluator_aa = eval_handler_aa.create_evaluator() opt_aa = bpopt.optimisations.DEAPOptimisation(evaluator=evaluator_aa) stim_protocols = utility.load_json(protocol_path) stim_protocols = {key:val for key,val in stim_protocols.items() \ if 'LongDC' in key} stim_dict = {key:val['stimuli'][0]['amp'] \ for key,val in stim_protocols.items()} sorted_stim_tuple= sorted(stim_dict.items(), key=operator.itemgetter(1)) stim_name= sorted_stim_tuple[-1][0] # knobs (the max amp) # Copy compiled modfiles if not os.path.isdir('x86_64'): raise Exception('Compiled modfiles do not exist') efel_features = utility.load_json(select_feature_path) un_features = un.EfelFeatures(features_to_run=efel_features) un_parameters_aa = un.Parameters(parameters_aa) un_parameters_aa.set_all_distributions(un.uniform(param_mod_range)) un_model_aa = un.Model(run=nrnsim_bpopt, interpolate=True, labels=["Time (ms)", "Membrane potential (mV)"], opt=opt_aa,stim_protocols =stim_protocols, param_dict_uc = param_dict_uc_aa, stim_name=stim_name, optim_param=optim_param_aa) # Perform the uncertainty quantification UQ_aa = un.UncertaintyQuantification(un_model_aa, parameters=un_parameters_aa, features=un_features) data_folder = 'sensitivity_data' sa_filename_aa = 'sa_allactive_%s.h5'%cell_id sa_filename_aa_csv = 'sa_allactive_%s.csv'%cell_id sa_data_path_aa = os.path.join(data_folder,sa_filename_aa) sa_aa_csv_path = os.path.join(data_folder,sa_filename_aa_csv) UQ_aa.quantify(seed=0,CPUs=cpu_count,data_folder=data_folder, filename= sa_filename_aa) _ = SA_obj_aa.save_analysis_data(sa_data_path_aa, filepath=sa_aa_csv_path) cell_data_aa = un.Data(sa_data_path_aa) SA_obj_aa.plot_sobol_analysis(cell_data_aa,analysis_path = \ 'figures/sa_analysis_aa_%s.pdf'%cell_id, palette='Set1') # Perisomatic model if perisomatic_sa: try: optim_param_path_peri = None SA_obj_peri = SA_helper(optim_param_path_peri,select_peri_param_path,param_mod_range, opt_config_file) _,_,mech_path_peri,_,\ param_bound_path_peri = SA_obj_peri.load_config(model_base_path, perisomatic=True) sens_param_bound_write_path_peri = "param_sensitivity_peri.json" optim_param_peri = SA_obj_peri.create_sa_bound_peri(param_bound_path_peri, sens_param_bound_write_path_peri) param_dict_uc_peri = SA_obj_peri.create_sens_param_dict() parameters_peri ={key:optim_param_peri[val] for key,val in param_dict_uc_peri.items()} eval_handler_peri = Bpopt_Evaluator(protocol_path, feature_path, morph_path, sens_param_bound_write_path_peri, mech_path_peri, ephys_dir=None, timed_evaluation = False) evaluator_peri = eval_handler_peri.create_evaluator() opt_peri = bpopt.optimisations.DEAPOptimisation(evaluator=evaluator_peri) un_parameters_peri= un.Parameters(parameters_peri) un_parameters_peri.set_all_distributions(un.uniform(param_mod_range)) un_model_peri = un.Model(run=nrnsim_bpopt, interpolate=True, labels=["Time (ms)", "Membrane potential (mV)"], opt=opt_peri,stim_protocols =stim_protocols, param_dict_uc = param_dict_uc_peri, stim_name=stim_name, optim_param=optim_param_peri) UQ_peri = un.UncertaintyQuantification(un_model_peri, parameters=un_parameters_peri, features=un_features) sa_filename_peri = 'sa_perisomatic_%s.h5'%cell_id sa_filename_peri_csv = 'sa_perisomatic_%s.csv'%cell_id sa_data_path_peri = os.path.join(data_folder,sa_filename_peri) sa_peri_csv_path = os.path.join(data_folder,sa_filename_peri_csv) UQ_peri.quantify(seed=0,CPUs=cpu_count,data_folder=data_folder, filename= sa_filename_peri) _ = SA_obj_peri.save_analysis_data(sa_data_path_peri, filepath=sa_peri_csv_path) cell_data_peri = un.Data(sa_data_path_peri) SA_obj_peri.plot_sobol_analysis(cell_data_peri,analysis_path = \ 'figures/sa_analysis_peri_%s.pdf'%cell_id, palette='Set2') except Exception as e: print(e)
def aibs_peri_to_bpopt(self, peri_param_path, base_dir='config/'): peri_params = utility.load_json(peri_param_path) peri_params_release = list() peri_mechs_release = defaultdict(list) peri_mechs_release['all'].append('pas') rev_potential = utility.rev_potential section_map = utility.bpopt_section_map for key, values in peri_params.items(): if key == 'genome': for j in range(len(values)): iter_dict_release = { 'param_name': peri_params[key][j]['name'] } iter_dict_release['sectionlist'] = section_map[ peri_params[key][j]['section']] iter_dict_release['type'] = 'section' iter_dict_release['value'] = float( peri_params[key][j]['value']) iter_dict_release['dist_type'] = 'uniform' if peri_params[key][j]['mechanism'] != '': iter_dict_release['mech'] = peri_params[key][j][ 'mechanism'] iter_dict_release['type'] = 'range' peri_params_release.append(iter_dict_release) elif key == 'passive': for key_pas, val_pas in values[0].items(): if key_pas == 'cm': for pas_param in val_pas: iter_dict_release = { 'param_name': 'cm', 'sectionlist': section_map[pas_param['section']], 'value': pas_param['cm'], 'dist_type': 'uniform', 'type': 'section' } peri_params_release.append(iter_dict_release) else: iter_dict_release = { 'param_name': 'Ra' if key_pas == 'ra' else key_pas, 'sectionlist': 'all', 'value': val_pas, 'dist_type': 'uniform', 'type': 'section' } peri_params_release.append(iter_dict_release) for rev in rev_potential: iter_dict_release = { 'param_name': rev, 'sectionlist': 'somatic', 'dist_type': 'uniform', 'type': 'section' } if rev == 'ena': iter_dict_release['value'] = rev_potential[rev] elif rev == 'ek': iter_dict_release['value'] = rev_potential[rev] peri_params_release.append(iter_dict_release) peri_params_release.append({ "param_name": "celsius", "type": "global", "value": 34 }) peri_params_release.append({ "param_name": "v_init", "type": "global", "value": peri_params['conditions'][0]["v_init"] }) for param_dict in peri_params_release: if 'mech' in param_dict.keys(): if param_dict['mech'] not in peri_mechs_release[ param_dict['sectionlist']]: peri_mechs_release[param_dict['sectionlist']].append( param_dict['mech']) peri_params_write_path = os.path.join(base_dir, self.cell_id, 'peri_parameters.json') peri_mech_write_path = os.path.join(base_dir, self.cell_id, 'peri_mechanism.json') utility.create_filepath(peri_params_write_path) utility.save_json(peri_params_write_path, peri_params_release) utility.save_json(peri_mech_write_path, peri_mechs_release) return peri_params_write_path, peri_mech_write_path
def script_generator(self, chain_job='chain_job.sh', **kwargs): job_config = utility.load_json(self.job_config_path) stage_jobconfig = job_config['stage_jobconfig'] highlevel_job_props = job_config['highlevel_jobconfig'] analysis_flag = kwargs.get( 'analysis') # this means prepare a batch script for analysis if highlevel_job_props['dryrun']: stage_jobconfig = update(stage_jobconfig, dryrun_config) job_config['stage_jobconfig'] = stage_jobconfig utility.save_json(self.job_config_path, job_config) analysis_config = stage_jobconfig['analysis_config'] with open(self.script_template, 'r') as job_template: batchjob_string = job_template.read() jobname = '%s.%s' % (os.path.basename( highlevel_job_props['job_dir']), stage_jobconfig['stage_name']) if analysis_flag: jobname += '.analysis' seed_string = ''.join( ['%s ' % seed_ for seed_ in stage_jobconfig['seed']]) # High level job config batchjob_string = re.sub('conda_env', highlevel_job_props['conda_env'], batchjob_string) batchjob_string = re.sub('jobname', jobname, batchjob_string) batchjob_string = re.sub('jobscript_name', self.script_name, batchjob_string) if highlevel_job_props.get('email'): batchjob_string = re.sub('email', highlevel_job_props['email'], batchjob_string) # Only related to optimization batchjob_string = re.sub('seed_list', seed_string, batchjob_string) batchjob_string = re.sub('analysis_script', analysis_config['main_script'], batchjob_string) batchjob_string = re.sub('job_config_path', self.job_config_path, batchjob_string) # Job config analysis vs optimization if analysis_flag: hpc_job_config = analysis_config batchjob_string = re.sub('# Run[\S\s]*pids', '', batchjob_string) else: hpc_job_config = stage_jobconfig['optim_config'] # Within the batch job script change the analysis launch to batch analysis if analysis_config.get( 'ipyparallel') and stage_jobconfig['run_hof_analysis']: analysis_jobname = kwargs.get('analysis_jobname') batchjob_string = re.sub('# Analyze[\S\s]*.json', 'qsub %s' % analysis_jobname, batchjob_string) # if there is a next stage chain the job if bool(kwargs.get('next_stage_job_config')): batchjob_string += 'bash %s\n' % chain_job hpc_job_parameters = [ 'jobmem', 'ipyparallel_db', 'qos', 'main_script', 'jobtime', 'error_stream', 'output_stream', 'nnodes', 'nprocs', 'nengines' ] for hpc_param in hpc_job_parameters: batchjob_string = re.sub(hpc_param, str(hpc_job_config[hpc_param]), batchjob_string) with open(self.script_name, "w") as batchjob_script: batchjob_script.write(batchjob_string)
def create_optim_job(args): level = logging.getLevelName(args['log_level']) logger.setLevel(level) cty_props = args['cty_config'] cell_id = cty_props['cell_id'] highlevel_job_props = args['job_config']['highlevel_jobconfig'] stage_job_props = args['job_config']['stage_jobconfig'] # Change any paths to absolute path for ii, stage_job_prop in enumerate(stage_job_props): stage_job_props[ii] = convert_paths(stage_job_prop) highlevel_job_props = convert_paths(highlevel_job_props) try: job_dir = os.path.join(os.getcwd(), highlevel_job_props['job_dir']) except: job_dir = os.path.join(os.getcwd(), str(cell_id)) highlevel_job_props['job_dir'] = job_dir utility.create_dirpath(job_dir) os.chdir(job_dir) # Change Working directory cty_config_path = os.path.join('user_config', 'cell_config.json') job_config_path = os.path.join('user_config', 'job_config.json') highlevel_jobconfig_path = 'high_level_job_config.json' stage_tracker_path = 'stage_tracker_config.json' utility.create_filepath(cty_config_path) utility.create_filepath(job_config_path) # Save a copy of the config files utility.save_json(cty_config_path, cty_props) utility.save_json(job_config_path, args['job_config']) try: ateamopt_dir = os.path.join( os.path.dirname(ateamopt.__file__), os.pardir) ateamopt_commitID = subprocess.check_output( ["git", "describe", "--tags"], cwd=ateamopt_dir).strip() ateamopt_commitID = ateamopt_commitID.decode() if isinstance( ateamopt_commitID, bytes) else ateamopt_commitID cty_props['ateamopt_tag'] = ateamopt_commitID except Exception as e: logger.debug(e) try: bluepyopt_dir = os.path.join( os.path.dirname(bluepyopt.__file__), os.pardir) bpopt_commitID = subprocess.check_output( ["git", "describe", "--tags"], cwd=bluepyopt_dir).strip() bpopt_commitID = bpopt_commitID.decode() if isinstance( bpopt_commitID, bytes) else bpopt_commitID cty_props['bluepyopt_tag'] = bpopt_commitID except: pass # pickling consistency depends on pandas version pd_version = pd.__version__ cty_props['pandas_version'] = pd_version cell_metadata_path = glob.glob('cell_metadata*.json') if len(cell_metadata_path) == 0: cty_props.update(highlevel_job_props) cell_metadata, cell_metadata_path = cell_data.save_cell_metadata( **cty_props) morph_stats_filename = 'morph_stats_%s.json' % cell_id morph_handler = MorphHandler( cell_metadata['swc_path'], cell_id=cell_id) morph_handler.save_morph_data(morph_stats_filename) elif len(cell_metadata_path) == 1: cell_metadata_path = cell_metadata_path[0] cell_metadata = utility.load_json(cell_metadata_path) else: raise Exception('More than one metadata files found') # Extract ephys data ephys_dir = highlevel_job_props['ephys_dir'] non_standard_nwb = highlevel_job_props['non_standard_nwb'] feature_stimtypes = highlevel_job_props['feature_stimtypes'] highlevel_job_props['nwb_path'] = cell_metadata['nwb_path'] highlevel_job_props['swc_path'] = cell_metadata['swc_path'] nwb_handler = NwbExtractor(cell_id, nwb_path=highlevel_job_props['nwb_path']) data_source = highlevel_job_props["data_source"] if data_source == "lims": ephys_data_path, stimmap_filename = nwb_handler.save_cell_data(feature_stimtypes, non_standard_nwb=non_standard_nwb, ephys_dir=ephys_dir) else: ephys_data_path, stimmap_filename = nwb_handler.save_cell_data_web(feature_stimtypes, non_standard_nwb=non_standard_nwb, ephys_dir=ephys_dir) feature_names_path = highlevel_job_props['feature_names_path'] protocol_dict, feature_dict = nwb_handler.get_efeatures_all(feature_names_path, ephys_data_path, stimmap_filename) feature_dict = correct_voltage_feat_std(feature_dict) all_protocols_filename = os.path.join(ephys_data_path, 'all_protocols.json') all_features_filename = os.path.join(ephys_data_path, 'all_features.json') utility.save_json(all_protocols_filename, protocol_dict) utility.save_json(all_features_filename, feature_dict) highlevel_job_props['stimmap_file'] = os.path.abspath(stimmap_filename) highlevel_job_props['machine'] = cell_metadata['machine'] highlevel_job_props['log_level'] = args['log_level'] highlevel_job_props['all_features_path'] = all_features_filename highlevel_job_props['all_protocols_path'] = all_protocols_filename highlevel_job_props = convert_paths(highlevel_job_props) utility.save_json(highlevel_jobconfig_path, highlevel_job_props) stage_level_jobconfig = {} stage_level_jobconfig['stage_jobconfig'] = stage_job_props.pop(0) stage_level_jobconfig['highlevel_jobconfig'] = highlevel_job_props utility.save_json(stage_tracker_path, stage_job_props) stage_jobdir = os.path.join(highlevel_job_props['job_dir'], stage_level_jobconfig['stage_jobconfig']['stage_name']) stage_level_jobconfig_path = os.path.join( stage_jobdir, 'stage_job_config.json') utility.create_filepath(stage_level_jobconfig_path) utility.save_json(stage_level_jobconfig_path, stage_level_jobconfig) prepare_jobscript_default = utility.locate_script_file( 'prepare_stagejob.py') analyze_jobscript_default = utility.locate_script_file( 'analyze_stagejob.py') shutil.copy(prepare_jobscript_default, stage_jobdir) shutil.copy(analyze_jobscript_default, stage_jobdir) jobtemplate_path = 'job_templates/chainjob_template.sh' chain_job = ChainSubJob(jobtemplate_path, stage_level_jobconfig_path) chain_job.script_generator() chain_job.run_job()
def main(args): # Job config job_config_path = sys.argv[-1] stage_jobconfig = args['stage_jobconfig'] highlevel_job_props = args['highlevel_jobconfig'] logging.basicConfig(level=highlevel_job_props['log_level']) job_dir = highlevel_job_props['job_dir'] path_to_cell_metadata = glob.glob( os.path.join(job_dir, 'cell_metadata*.json'))[0] stage_tracker_path = os.path.join(job_dir, 'stage_tracker_config.json') cell_metadata = utility.load_json(path_to_cell_metadata) cell_id = cell_metadata['cell_id'] peri_model_id = cell_metadata.get('peri_model_id') released_aa_model_path = cell_metadata.get('model_path_all_active') released_aa_model_id = cell_metadata.get('released_aa_model_id') nwb_path = highlevel_job_props['nwb_path'] swc_path = highlevel_job_props['swc_path'] all_features_path = highlevel_job_props['all_features_path'] all_protocols_path = highlevel_job_props['all_protocols_path'] stage_stimtypes = stage_jobconfig['stage_stimtypes'] stage_feature_names_path = stage_jobconfig['stage_features'] param_bounds_path = stage_jobconfig['stage_parameters'] ap_init_flag = stage_jobconfig['AP_initiation_zone'] ap_init_feature = 'check_AISInitiation' script_repo_dir = stage_jobconfig.get('script_repo_dir') depol_block_check = stage_jobconfig.get('depol_block_check') add_fi_kink = stage_jobconfig.get('add_fi_kink') analysis_parallel = (stage_jobconfig['analysis_config'].get('ipyparallel') and stage_jobconfig['run_hof_analysis'] ) # analysis batch job only for hof analysis param_bound_tolerance = stage_jobconfig.get('adjust_param_bounds_prev') prev_stage_path = stage_jobconfig.get('prev_stage_path') filter_rule_func = getattr(filter_rules, stage_jobconfig['filter_rule']) all_features = utility.load_json(all_features_path) all_protocols = utility.load_json(all_protocols_path) stage_feature_names = utility.load_json( stage_feature_names_path)['features'] # AP init flag is prioritized over feature set file if ap_init_flag == 'soma': if ap_init_feature in stage_feature_names: stage_feature_names.remove(ap_init_feature) elif ap_init_flag == 'axon': if ap_init_feature not in stage_feature_names: stage_feature_names.append(ap_init_feature) select_stim_names = [] for stim_name in all_features.keys(): stim_type = stim_name.rsplit('_', 1)[0] stim_type_aibs = utility.aibs_stimname_map_inv[stim_type] if stim_type_aibs in stage_stimtypes: select_stim_names.append(stim_name) features_dict = defaultdict(lambda: defaultdict(lambda: defaultdict(dict))) for stim_name, stim_dict in all_features.items(): if stim_name in select_stim_names: for loc, loc_features in stim_dict.items(): for feat, val in loc_features.items(): if feat in stage_feature_names: features_dict[stim_name][loc][feat] = [val[0], val[1]] protocols_dict = { proto_key: proto_val for proto_key, proto_val in all_protocols.items() if proto_key in select_stim_names } nwb_handler = NwbExtractor(cell_id, nwb_path=nwb_path) kwargs = { 'depol_block_check': depol_block_check, 'add_fi_kink': add_fi_kink } if depol_block_check: train_features, test_features, train_protocols, DB_proto_dict = filter_rule_func( features_dict, protocols_dict, **kwargs) # also append DB check info to all_protocols json and save all_protocols['DB_check_DC'] = {'stimuli': DB_proto_dict} utility.save_json(all_protocols_path, all_protocols) else: train_features, test_features, train_protocols = filter_rule_func( features_dict, protocols_dict, **kwargs) train_features_path, test_features_path, train_protocols_path = \ nwb_handler.write_ephys_features(train_features, test_features, train_protocols) # Create the parameter bounds for the optimization if prev_stage_path: prev_stage_model_path = os.path.join( prev_stage_path, 'fitted_params', 'optim_param_%s_compact.json' % cell_id) else: prev_stage_model_path = None model_params_handler = AllActive_Model_Parameters( cell_id, swc_path=swc_path, prev_stage_model_path=prev_stage_model_path, released_aa_model_path=released_aa_model_path) model_params, model_params_release = model_params_handler.get_opt_params( param_bounds_path, prev_stage_tolerance=param_bound_tolerance) param_write_path, released_aa_param_write_path, released_aa_params =\ model_params_handler.write_params_opt(model_params, model_params_release) model_mechs, model_mechs_release = model_params_handler.get_opt_mechanism( model_params, model_params_release, param_bounds_path) mech_write_path, mech_release_write_path = model_params_handler.write_mechanisms_opt( model_mechs, model_mechs_release) props = {} if peri_model_id: peri_model_path = cell_metadata['model_path_perisomatic'] peri_params_write_path, peri_mech_write_path = \ model_params_handler.aibs_peri_to_bpopt(peri_model_path) props['released_peri_model'] = peri_params_write_path props['released_peri_mechanism'] = peri_mech_write_path # Config file with all the necessary paths to feed into the optimization # TODO: clarify how this fits into schema model_params_handler.write_opt_config_file( param_write_path, mech_write_path, mech_release_write_path, train_features_path, test_features_path, train_protocols_path, released_aa_params, released_aa_param_write_path, opt_config_filename=job_config_path, **props) # Copy the optimizer scripts in the current directory optimizer_script = stage_jobconfig['optim_config']['main_script'] analysis_script = stage_jobconfig['analysis_config']['main_script'] if script_repo_dir: optimizer_script_repo = os.path.abspath( os.path.join(script_repo_dir, optimizer_script)) optimizer_script_repo = optimizer_script_repo if os.path.exists(optimizer_script_repo)\ else None else: optimizer_script_repo = None optimizer_script_default = utility.locate_script_file(optimizer_script) optimizer_script_path = optimizer_script_repo or optimizer_script_default stage_cwd = os.getcwd() shutil.copy(optimizer_script_path, stage_cwd) next_stage_job_props = utility.load_json(stage_tracker_path) machine = highlevel_job_props['machine'] machine_match_patterns = ['hpc-login', 'aws', 'cori', 'bbp5'] next_stage_jobconfig = {} try: next_stage_jobconfig['stage_jobconfig'] = next_stage_job_props.pop(0) next_stage_jobconfig['highlevel_jobconfig'] = highlevel_job_props next_stage_jobconfig['stage_jobconfig']['prev_stage_path'] = os.getcwd( ) chainjobtemplate_path = 'job_templates/chainjob_template.sh' except: pass utility.save_json(stage_tracker_path, next_stage_job_props) # Create batch jobscript if not any(substr in machine for substr in machine_match_patterns): testJob = test_JobModule('batch_job.sh', job_config_path=job_config_path) testJob.script_generator(next_stage_job_config=next_stage_jobconfig) elif any(pattern in machine for pattern in ['hpc-login', 'aws']): jobtemplate_path = 'job_templates/pbs_jobtemplate.sh' batch_job = PBS_JobModule(jobtemplate_path, job_config_path) if analysis_parallel: batch_job.script_generator(analysis_jobname='analyze_job.sh') # A separate batch job needs to be created in this case analysis_job = PBS_JobModule(jobtemplate_path, job_config_path, script_name='analyze_job.sh') analysis_job.script_generator( analysis=True, next_stage_job_config=next_stage_jobconfig) else: batch_job.script_generator( next_stage_job_config=next_stage_jobconfig) elif any(pattern in machine for pattern in ['cori', 'bbp5']): if 'cori' in machine: jobtemplate_path = 'job_templates/nersc_slurm_jobtemplate.sh' else: jobtemplate_path = 'job_templates/bbp_slurm_jobtemplate.sh' batch_job = Slurm_JobModule(jobtemplate_path, job_config_path) if analysis_parallel: batch_job.script_generator(analysis_jobname='analyze_job.sh') # A separate batch job needs to be created in this case analysis_job = Slurm_JobModule(jobtemplate_path, job_config_path, script_name='analyze_job.sh') analysis_job.script_generator( analysis=True, next_stage_job_config=next_stage_jobconfig) else: batch_job.script_generator( next_stage_job_config=next_stage_jobconfig) if next_stage_jobconfig: stage_jobdir = os.path.join( highlevel_job_props['job_dir'], next_stage_jobconfig['stage_jobconfig']['stage_name']) next_stage_jobconfig_path = os.path.join(stage_jobdir, 'stage_job_config.json') utility.create_filepath(next_stage_jobconfig_path) utility.save_json(next_stage_jobconfig_path, next_stage_jobconfig) prepare_jobscript_default = utility.locate_script_file( 'prepare_stagejob.py') analyze_jobscript_default = utility.locate_script_file(analysis_script) shutil.copy(prepare_jobscript_default, stage_jobdir) shutil.copy(analyze_jobscript_default, stage_jobdir) chain_job = ChainSubJob(chainjobtemplate_path, next_stage_jobconfig_path) chain_job.script_generator()