updating plotting functionality

Author: kaeldai

bmtk/analyzer/compartment.py

@@ -0,0 +1,128 @@
+import os
+from six import string_types
+import numpy as np
+
+from bmtk.utils import sonata
+from bmtk.utils.sonata.config import SonataConfig
+from bmtk.utils.reports.compartment import CompartmentReport
+from bmtk.utils.reports.compartment import plotting
+from bmtk.simulator.utils import simulation_reports
+
+
+def _get_report(report_path=None, config=None, report_name=None):
+    if report_path is not None:
+        return report_path, CompartmentReport.load(report=report_path)
+
+    elif config is not None:
+        selected_reports = []
+        sim_reports = simulation_reports.from_config(config)
+        for report in sim_reports:
+            if report.module in ['membrane_report', 'multimeter_report']:
+                rname = report.report_name
+                rfile = report.params['file_name']
+                # TODO: Full path should be determined by config/simulation_reports module
+                rpath = rfile if os.path.isabs(rfile) else os.path.join(report.params['tmp_dir'], rfile)
+                if report_name is not None and report_name == rname:
+                        selected_reports.append((rname, CompartmentReport.load(rpath)))
+                elif report_name is None:
+                    selected_reports.append((rname, CompartmentReport.load(rpath)))
+
+        if len(selected_reports) == 0:
+            msg = 'Could not find a report '
+            msg += '' if report_name is None else 'with report_name "{}"'.format(report_name)
+            msg += ' from configuration file. . Use "report_path" parameter instead.'
+            raise ValueError(msg)
+
+        elif len(selected_reports) > 1:
+            avail_reports = ', '.join(s[0] for s in selected_reports)
+            raise ValueError('Configuration file contained multiple "membrane_reports", use "report_name" or'
+                             '"report_path" to pick which one to plot. Option values: {}'.format(avail_reports))
+
+        else:
+            return selected_reports[0]
+
+    else:
+        raise AttributeError('Could not find a compartment report SONATA file. Please user "config_file" or '
+                             '"report_path" options.')
+
+
+def _find_nodes(population, config=None, nodes_file=None, node_types_file=None):
+    if nodes_file is not None:
+        network = sonata.File(data_files=nodes_file, data_type_files=node_types_file)
+        if population not in network.nodes.population_names:
+            raise ValueError('node population "{}" not found in {}'.format(population, nodes_file))
+        return network.nodes[population]
+
+    elif config is not None:
+        for nodes_grp in config.nodes:
+            network = sonata.File(data_files=nodes_grp['nodes_file'], data_type_files=nodes_grp['node_types_file'])
+            if population in network.nodes.population_names:
+                return network.nodes[population]
+
+    raise ValueError('Could not find nodes file with node population "{}".'.format(population))
+
+
+def plot_traces(report_path=None, config_file=None, report_name=None, population=None, group_by=None,
+                group_excludes=None, nodes_file=None, node_types_file=None,
+                node_ids=None, sections='origin', average=False, times=None, title=None,
+                show_legend=None, show=True):
+    sonata_config = SonataConfig.from_json(config_file) if config_file else None
+    report_name, cr = _get_report(report_path=report_path, config=sonata_config, report_name=report_name)
+
+    if population is None:
+        pops = cr.populations
+        if len(pops) > 1:
+            raise ValueError('Report {} contains more than population of nodes ({}). Use population parameter'.format(
+                report_name, pops
+            ))
+        population = pops[0]
+
+    if title is None:
+        title = '{} ({})'.format(report_name, population)
+
+    # Create node-groups
+    if group_by is not None:
+        node_groups = []
+        nodes = _find_nodes(population=population, config=sonata_config, nodes_file=nodes_file,
+                            node_types_file=node_types_file)
+
+        grouped_df = None
+        for grp in nodes.groups:
+            if group_by in grp.all_columns:
+                grp_df = grp.to_dataframe()
+                grp_df = grp_df[['node_id', group_by]]
+                grouped_df = grp_df if grouped_df is None else grouped_df.append(grp_df, ignore_index=True)
+
+        if grouped_df is None:
+            raise ValueError('Could not find any nodes with group_by attribute "{}"'.format(group_by))
+
+        # Convert from string to list so we can always use the isin() method for filtering
+        if isinstance(group_excludes, string_types):
+            group_excludes = [group_excludes]
+        elif group_excludes is None:
+            group_excludes = []
+
+        for grp_key, grp in grouped_df.groupby(group_by):
+            if grp_key in group_excludes:
+                continue
+            node_groups.append({'node_ids': np.array(grp['node_id']), 'label': grp_key})
+
+        if len(node_groups) == 0:
+            exclude_str = ' excluding values {}'.format(', '.join(group_excludes)) if len(group_excludes) > 0 else ''
+            raise ValueError('Could not find any node-groups using group_by="{}"{}.'.format(group_by, exclude_str))
+
+    else:
+        node_groups = None
+
+    return plotting.plot_traces(
+        report=cr,
+        population=population,
+        node_ids=node_ids,
+        sections=sections,
+        average=average,
+        node_groups=node_groups,
+        times=times,
+        title=title,
+        show_legend=show_legend,
+        show=show
+    )

bmtk/analyzer/spike_trains.py

@@ -1,39 +1,176 @@
+import os
 import numpy as np
 import pandas as pd
-import h5py
+from functools import partial
+from six import string_types
 
-
-import matplotlib.pyplot as plt
-
-
-from bmtk.utils.sonata.config import SonataConfig as ConfigDict
+from bmtk.utils import sonata
+from bmtk.utils.sonata.config import SonataConfig
 from bmtk.utils.reports import SpikeTrains
 from bmtk.utils.reports.spike_trains import plotting
+from bmtk.simulator.utils import simulation_reports
 
 
-def load_spikes_file(config_file=None, spikes_file=None):
-    if spikes_file is not None:
-        return SpikeTrains.load(spikes_file)
+def _find_spikes(spikes_file=None, config_file=None, population=None):
+    candidate_spikes = []
+
+    # Get spikes file(s)
+    if spikes_file:
+        # User has explicity set the location of the spike files
+        candidate_spikes.append(spikes_file)
 
     elif config_file is not None:
-        config = ConfigDict.from_json(config_file)
-        return SpikeTrains.load(config['output']['spikes_file'])
+        # Otherwise search the config.json for all possible output spikes_files. We can use the simulation_reports
+        # module to find any spikes output file specified in config's "output" or "reports" section.
+        config = SonataConfig.from_json(config_file)
+        sim_reports = simulation_reports.from_config(config)
+        for report in sim_reports:
+            if report.module == 'spikes_report':
+                # BMTK can end up output the same spikes file in SONATA, CSV, and NWB format. Try fetching the SONATA
+                # version first, then CSV, and finally NWB if it exists.
+                spikes_sonata = report.params.get('spikes_file', None)
+                spikes_csv = report.params.get('spikes_file_csv', None)
+                spikes_nwb = report.params.get('spikes_file_nwb', None)
+
+                if spikes_sonata is not None:
+                    candidate_spikes.append(spikes_sonata)
+                elif spikes_csv is not None:
+                    candidate_spikes.append(spikes_csv)
+                elif spikes_csv is not None:
+                    candidate_spikes.append(spikes_nwb)
+
+        # TODO: Should we also look in the "inputs" for displaying input spike statistics?
+
+    if not candidate_spikes:
+        raise ValueError('Could not find an output spikes-file. Use "spikes_file" parameter option.')
+
+    # Find file that contains spikes for the specified "population" of nodes. If "population" parameter is not
+    # specified try to guess that spikes that the user wants to visualize.
+    if population is not None:
+        spikes_obj = None
+        for spikes_f in candidate_spikes:
+            st = SpikeTrains.load(spikes_f)
+            if population in st.populations:
+                if spikes_obj is None:
+                    spikes_obj = st
+                else:
+                    spikes_obj.merge(st)
+
+        if spikes_obj is None:
+            raise ValueError('Could not fine spikes file with node population "{}".'.format(population))
+        else:
+            return population, spikes_obj
 
+    else:
+        if len(candidate_spikes) > 1:
+            raise ValueError('Found more than one spike-trains file')
 
-def to_dataframe(config_file, spikes_file=None):
-    spike_trains = load_spikes_file(config_file=config_file, spikes_file=spikes_file)
-    return spike_trains.to_dataframe()
+        spikes_f = candidate_spikes[0]
+        if not os.path.exists(spikes_f):
+            raise ValueError('Did not find spike-trains file {}. Make sure the simulation has completed.'.format(
+                spikes_f))
+
+        spikes_obj = SpikeTrains.load(spikes_f)
+        if len(spikes_obj.populations) > 1:
+            raise ValueError('Spikes file {} contains more than one node population.'.format(spikes_f))
+        else:
+            return spikes_obj.populations[0], spikes_obj
+
+
+def _find_nodes(population, config=None, nodes_file=None, node_types_file=None):
+    if nodes_file is not None:
+        network = sonata.File(data_files=nodes_file, data_type_files=node_types_file)
+        if population not in network.nodes.population_names:
+            raise ValueError('node population "{}" not found in {}'.format(population, nodes_file))
+        return network.nodes[population]
+
+    elif config is not None:
+        for nodes_grp in config.nodes:
+            network = sonata.File(data_files=nodes_grp['nodes_file'], data_type_files=nodes_grp['node_types_file'])
+            if population in network.nodes.population_names:
+                return network.nodes[population]
+
+    raise ValueError('Could not find nodes file with node population "{}".'.format(population))
+
+
+def _plot_helper(plot_fnc, config_file=None, population=None, times=None, title=None, show=True,
+                 group_by=None, group_excludes=None,
+                 spikes_file=None, nodes_file=None, node_types_file=None):
+    sonata_config = SonataConfig.from_json(config_file) if config_file else None
+    pop, spike_trains = _find_spikes(config_file=config_file, spikes_file=spikes_file, population=population)
+
+    # Create the title
+    title = title if title is not None else '{} Nodes'.format(pop)
+
+    # Get start and stop times from config if needed
+    if sonata_config and times is None:
+        times = (sonata_config.tstart, sonata_config.tstop)
+
+    # Create node-groups
+    if group_by is not None:
+        node_groups = []
+        nodes = _find_nodes(population=pop, config=sonata_config, nodes_file=nodes_file,
+                            node_types_file=node_types_file)
+        grouped_df = None
+        for grp in nodes.groups:
+            if group_by in grp.all_columns:
+                grp_df = grp.to_dataframe()
+                grp_df = grp_df[['node_id', group_by]]
+                grouped_df = grp_df if grouped_df is None else grouped_df.append(grp_df, ignore_index=True)
+
+        if grouped_df is None:
+            raise ValueError('Could not find any nodes with group_by attribute "{}"'.format(group_by))
+
+        # Convert from string to list so we can always use the isin() method for filtering
+        if isinstance(group_excludes, string_types):
+            group_excludes = [group_excludes]
+        elif group_excludes is None:
+            group_excludes = []
+
+        for grp_key, grp in grouped_df.groupby(group_by):
+            if grp_key in group_excludes:
+                continue
+            node_groups.append({'node_ids': np.array(grp['node_id']), 'label': grp_key})
+
+    else:
+        node_groups = None
+
+    plot_fnc(spike_trains=spike_trains, node_groups=node_groups, population=pop, times=times, title=title, show=show)
+
+
+def plot_raster(config_file=None, population=None, with_histogram=True, times=None, title=None, show=True,
+                group_by=None, group_excludes=None,
+                spikes_file=None, nodes_file=None, node_types_file=None):
+
+    plot_fnc = partial(plotting.plot_raster, with_histogram=with_histogram)
+    return _plot_helper(plot_fnc,
+                        config_file=config_file, population=population, times=times, title=title, show=show,
+                        group_by=group_by, group_excludes=group_excludes,
+                        spikes_file=spikes_file, nodes_file=nodes_file, node_types_file=node_types_file
+    )
 
 
-def plot_raster(config_file, spikes_file=None):
-    spike_trains = load_spikes_file(config_file=config_file, spikes_file=spikes_file)
-    plotting.plot_raster(spike_trains)
-    plt.show()
+def plot_rates(config_file=None, population=None, smoothing=False, smoothing_params=None, times=None, title=None,
+               show=True, group_by=None, group_excludes=None, spikes_file=None, nodes_file=None, node_types_file=None):
 
+    plot_fnc = partial(plotting.plot_rates, smoothing=smoothing, smoothing_params=smoothing_params)
+    return _plot_helper(plot_fnc,
+                        config_file=config_file, population=population, times=times, title=title, show=show,
+                        group_by=group_by, group_excludes=group_excludes,
+                        spikes_file=spikes_file, nodes_file=nodes_file, node_types_file=node_types_file
+    )
 
-def plot_rates(config_file):
-    spike_trains = load_spikes_file(config_file)
-    plotting.plot_rates(spike_trains)
+
+def plot_rates_boxplot(config_file=None, population=None, times=None, title=None, show=True,
+                       group_by=None, group_excludes=None,
+                       spikes_file=None, nodes_file=None, node_types_file=None):
+
+    plot_fnc = partial(plotting.plot_rates_boxplot)
+    return _plot_helper(plot_fnc,
+        config_file=config_file, population=population, times=times, title=title, show=show,
+        group_by=group_by, group_excludes=group_excludes,
+        spikes_file=spikes_file, nodes_file=nodes_file, node_types_file=node_types_file
+    )
 
 
 def spike_statistics(spikes_file, simulation=None, simulation_time=None, groupby=None, network=None, **filterparams):
@@ -68,3 +205,8 @@ def calc_stats(r):
         return vals_df
     else:
         return spike_counts_df
+
+
+def to_dataframe(config_file, spikes_file=None, population=None):
+    _, spike_trains = _find_spikes(config_file=config_file, spikes_file=spikes_file, population=population)
+    return spike_trains.to_dataframe()

bmtk/analyzer/visualization/spikes.py

@@ -26,13 +26,14 @@
 from six import string_types
 import pandas as pd
 import numpy as np
+import warnings
 import matplotlib.pyplot as plt
 import matplotlib.cm as cmx
 import matplotlib.colors as colors
 import matplotlib.gridspec as gridspec
 
 import bmtk.simulator.utils.config as config
-from bmtk.utils.reports.spike_trains.plotting import plot_raster, plot_rates, plot_raster_cmp
+from bmtk.utils.reports.spike_trains.plotting import plot_raster, plot_rates # , plot_raster_cmp
 
 
 from mpl_toolkits.axes_grid1 import make_axes_locatable
@@ -119,6 +120,7 @@ def parse_line(line):
 
 
 def plot_spikes_config(configure, group_key=None, exclude=[], save_as=None, show_plot=True):
+    warnings.warn('Deprecated: Please use bmtk.analyzer.spike_trains.plot_raster instead.', DeprecationWarning)
     if isinstance(configure, string_types):
         conf = config.from_json(configure)
     elif isinstance(configure, dict):
@@ -135,6 +137,7 @@ def plot_spikes_config(configure, group_key=None, exclude=[], save_as=None, show
 
 def plot_spikes(cells_file, cell_models_file, spikes_file, population=None, group_key=None, exclude=[], save_as=None,
                 show=True, title=None, legend=True, font_size=None):
+    warnings.warn('Deprecated: Please use bmtk.analyzer.spike_trains.plot_raster instead.', DeprecationWarning)
     # check if can be shown and/or saved
     #if save_as is not None:
     #    if os.path.exists(save_as):

bmtk/simulator/bionet/modules/record_cellvars.py

@@ -81,6 +81,7 @@ def __init__(self, tmp_dir, file_name, variable_name, cells=None, gids=None, sec
 
         self._tmp_dir = tmp_dir
 
+        # TODO: Full path should be determined by config/simulation_reports module
         self._file_name = file_name if os.path.isabs(file_name) else os.path.join(tmp_dir, file_name)
         self._all_gids = cells
         self._local_gids = []