From b29253c2da94662779ad9f58b7d464a89e0f0676 Mon Sep 17 00:00:00 2001 From: Tibor Rozsa Date: Thu, 11 Jan 2024 16:47:37 +0100 Subject: [PATCH] Export sim results to json based on Arkheia documentation --- mozaik/controller.py | 36 ++++-- mozaik/tools/json_export.py | 216 +++++++++++++++++++++++++++++++ mozaik/visualization/plotting.py | 29 +++-- 3 files changed, 262 insertions(+), 19 deletions(-) create mode 100644 mozaik/tools/json_export.py diff --git a/mozaik/controller.py b/mozaik/controller.py index 7ee3ccbfb..522337d50 100644 --- a/mozaik/controller.py +++ b/mozaik/controller.py @@ -12,6 +12,7 @@ import time from datetime import datetime import logging +from mozaik.tools.json_export import save_json, get_experimental_protocols, get_recorders, get_stimuli logger = mozaik.getMozaikLogger() @@ -139,20 +140,28 @@ def prepare_workflow(simulation_name, model_class): if mozaik.mpi_comm.rank == mozaik.MPI_ROOT: - # Let's store the full and modified parameters, if we are the 0 rank process - parameters.save(Global.root_directory + "parameters", expand_urls=True) - import pickle - f = open(Global.root_directory+"modified_parameters","wb") - pickle.dump(str(modified_parameters),f) - f.close() + # Store simulation run info, if we are the 0 rank process, + # with several components to be stored/filled in later during the simulation run + sim_info = { + 'submission_date' : None, + 'run_date': datetime.now().strftime('%d/%m/%Y-%H:%M:%S'), + 'simulation_run_name': simulation_run_name, + 'model_name': simulation_name, + "model_description": model_class.__doc__, + 'results': {"$ref": "results.json"}, + 'stimuli': {"$ref": "stimuli.json"}, + 'recorders': {"$ref": "recorders.json"}, + 'experimental_protocols': {"$ref": "experimental_protocols.json"}, + 'parameters': {"$ref": "parameters.json"}, + } + save_json(sim_info, Global.root_directory + 'sim_info.json') + save_json(parameters.to_dict(), Global.root_directory + 'parameters.json') + #save_json(modified_parameters, Global.root_directory + 'modified_parameters.json') + recorders = get_recorders(parameters.to_dict()) + save_json(recorders, Global.root_directory + 'recorders.json') setup_logging() - if mozaik.mpi_comm.rank == mozaik.MPI_ROOT: - # Let's store some basic info about the simulation run - f = open(Global.root_directory+"info","w") - f.write(str({'model_class' : str(model_class), 'model_docstring' : model_class.__doc__,'simulation_run_name' : simulation_run_name, 'model_name' : simulation_name, 'creation_data' : datetime.now().strftime('%d/%m/%Y-%H:%M:%S')})) - f.close() return sim, num_threads, parameters def run_workflow(simulation_name, model_class, create_experiments): @@ -264,4 +273,9 @@ def run_experiments(model,experiment_list,parameters,load_from=None): logger.info('Simulator run time: %.0fs (%d%%)' % (simulation_run_time, int(simulation_run_time /total_run_time * 100))) logger.info('Mozaik run time: %.0fs (%d%%)' % (mozaik_run_time, int(mozaik_run_time /total_run_time * 100))) + experimental_protocols = get_experimental_protocols(data_store) + stimuli = get_stimuli(data_store,parameters.store_stimuli, parameters.input_space) + save_json(experimental_protocols, Global.root_directory + 'experimental_protocols.json') + save_json(stimuli, Global.root_directory + 'stimuli.json') + return data_store diff --git a/mozaik/tools/json_export.py b/mozaik/tools/json_export.py new file mode 100644 index 000000000..4989841a8 --- /dev/null +++ b/mozaik/tools/json_export.py @@ -0,0 +1,216 @@ +import os +import re +import json +import numpy as np +from numpyencoder import NumpyEncoder +from sphinx.util import docstrings + + +PARAMETERS_REGEX = re.compile(".*Parameters.*") +OTHER_PARAMETER_REGEX = re.compile(".*Other\ [pP]arameters\ *\n-{15}-+") +PARAMETER_REGEX = re.compile( + "\s*(?P[^:\s]+)\s*\:\s* (?P[^\n]*)\n\s*(?P[^\n]*)" +) + +def save_json(d, filename): + with open(filename, 'w', encoding='utf-8') as f: + json.dump(d, f, ensure_ascii=False, indent=4, cls=NumpyEncoder) + +def get_params_from_docstring(cls): + params = {} + for cls1 in cls.__mro__: + params.update(parse_docstring(cls1.__doc__)["params"]) + return params + +def parse_docstring(docstring): + """Parse the docstring into its components. + :returns: a dictionary of form + { + "short_description": ..., + "long_description": ..., + "params": [{"name": ..., "doc": ...}, ...], + "returns": ... + } + """ + + short_description = long_description = returns = "" + params = [] + + if docstring: + docstring = "\n".join(docstrings.prepare_docstring(docstring)) + + lines = docstring.split("\n", 1) + short_description = lines[0] + + if len(lines) > 1: + reminder = lines[1].strip() + # params_returns_desc = None + match_parameters = PARAMETERS_REGEX.search(reminder) + if match_parameters: + long_desc_end = match_parameters.start() + long_description = reminder[:long_desc_end].rstrip() + reminder = reminder[long_desc_end:].strip() + + match = OTHER_PARAMETER_REGEX.search(reminder) + + if match: + end = match.start() + if not match_parameters: + long_description = reminder[:end].rstrip() + reminder = reminder[end:].strip() + + if reminder: + params = {} + + for name, tpe, doc in PARAMETER_REGEX.findall(reminder): + params[name] = (tpe, doc) + + if (not match_parameters) and (not match): + long_description = reminder + + return { + "short_description": short_description, + "long_description": long_description, + "params": params, + } + +def get_recorders(parameters): + recorders_docs = [] + for sh in parameters["sheets"].keys(): + for rec in parameters["sheets"][sh]["params"]["recorders"].keys(): + recorder = parameters["sheets"][sh]["params"]["recorders"][rec] + name = recorder["component"].split(".")[-1] + module_path = ".".join(recorder["component"].split(".")[:-1]) + doc_par = get_params_from_docstring( + getattr(__import__(module_path, globals(), locals(), name), name) + ) + p = { + k: (recorder["params"][k], doc_par[k][0], doc_par[k][1]) + for k in recorder["params"].keys() + } + + recorders_docs.append( + { + "code": module_path + "." + name, + "short_description": parse_docstring( + getattr( + __import__(module_path, globals(), locals(), name), name + ).__doc__ + )["short_description"], + "long_description": parse_docstring( + getattr( + __import__(module_path, globals(), locals(), name), name + ).__doc__ + )["long_description"], + "parameters": p, + "variables": recorder["variables"], + "source": sh, + } + ) + return recorders_docs + +def get_experimental_protocols(data_store): + experimental_protocols_docs = [] + for ep in data_store.get_experiment_parametrization_list(): + name = ep[0][8:-2].split(".")[-1] + module_path = ".".join(ep[0][8:-2].split(".")[:-1]) + doc_par = get_params_from_docstring( + getattr(__import__(module_path, globals(), locals(), name), name) + ) + params = eval(ep[1]) + + p = { + k: (params[k], doc_par[k][0], doc_par[k][1]) + if k in doc_par + else params[k] + for k in params.keys() + } + + experimental_protocols_docs.append( + { + "class": module_path + "." + name, + "short_description": parse_docstring( + getattr( + __import__(module_path, globals(), locals(), name), name + ).__doc__ + )["short_description"], + "long_description": parse_docstring( + getattr( + __import__(module_path, globals(), locals(), name), name + ).__doc__ + )["long_description"], + "parameters": p, + } + ) + return experimental_protocols_docs + +from mozaik.tools.mozaik_parametrized import MozaikParametrized + +def reduce_dicts(dicts): + constant = {k : True for k in dicts.keys()} + for d in dicts(): + continue + +#changing_parameters = {} +#unique_stimulus_types = set([MozaikExtendedParameterSet(s).to_dict()["name"] s for s in unique_stimuli]) +# +#for s_type in unique_stimulus_types: +# stim_dicts = [MozaikExtendedParameterSet(s).to_dict() for s in unique_stimuli()] +# stim_dicts = [d for d in stim_dicts if d["name"] == s_type] +# changing_parameters["name"] = reduce_dicts(stim_dicts) +# break + +import imageio + +def get_stimuli(data_store, store_stimuli, input_space): + stim_docs = [] + if not store_stimuli: + return stim_docs + unique_stimuli = [s for s in set(data_store.get_stimuli())] + stim_dir = "stimuli/" + os.makedirs(data_store.parameters.root_directory + stim_dir, exist_ok=True) + for s in unique_stimuli: + sidd = MozaikParametrized.idd(s) + params = sidd.get_param_values() + params = {k: (v, sidd.params()[k].doc) for k, v in params} + + # Only save one trial of each different stimulus + if params["trial"][0] != 0: + continue + + raws = data_store.get_sensory_stimulus([s]) + + if raws == [] or raws[0] == None: + img = np.zeros((50,50)).astype(np.uint8) + raws = [img,img] + else: + raws = raws[0] + + mov_duration = input_space["update_interval"] / 1000.0 if input_space != None else 0.1 + gif_name = params["name"][0] + str(hash(s)) + ".gif" + imageio.mimwrite(data_store.parameters.root_directory + stim_dir + gif_name, raws, duration=mov_duration) + + stim_docs.append( + { + "code": sidd.name, + "short_description": parse_docstring( + getattr( + __import__( + sidd.module_path, globals(), locals(), sidd.name + ), + sidd.name, + ).__doc__ + )["short_description"], + "long_description": parse_docstring( + getattr( + __import__( + sidd.module_path, globals(), locals(), sidd.name + ), + sidd.name, + ).__doc__ + )["long_description"], + "parameters": params, + "movie": stim_dir + gif_name, + } + ) + return stim_docs diff --git a/mozaik/visualization/plotting.py b/mozaik/visualization/plotting.py index 1135bc138..3891742e9 100644 --- a/mozaik/visualization/plotting.py +++ b/mozaik/visualization/plotting.py @@ -66,7 +66,8 @@ logger = mozaik.getMozaikLogger() from builtins import zip - +import json +from mozaik.tools.json_export import save_json class Plotting(ParametrizedObject): """ @@ -186,13 +187,25 @@ def plot(self, params=None): else: # save the analysis plot pylab.savefig(Global.root_directory+self.plot_file_name,transparent=True) - - # and store the record - with open(Global.root_directory+'results','a+') as f: - entry = {'parameters' : self.parameters, 'file_name' : self.plot_file_name, 'class_name' : str(self.__class__)} - f.write(str(entry)+'\n') - f.close() - + # store the entry at the end of a list in a json file + entry = { + 'code': str(self.__class__.__module__) + "." + str(self.__class__.__name__), + 'name': ".".join(self.plot_file_name.split('.')[:-1]), + 'caption': self.caption, + 'parameters': self.parameters, + 'figure': self.plot_file_name + } + + results = [] + results_path = Global.root_directory + 'results.json' + + if os.path.exists(results_path): + with open(results_path, 'r', encoding='utf-8') as f: + results = json.load(f) + + results.append(entry) + save_json(results, results_path) + t2 = time.time() logger.warning(self.__class__.__name__ + ' plotting took: ' + str(t2 - t1) + 'seconds')