Interfacing NWB with mne-python?

[rly]

@nmarkowitz asked on Slack

Has anyone created an object to interface NWB with mne-python? Similar to the RawEDF object that lets you open a file, or in this case a specific acquisition, and load the data in with mne-python functionality including lazy loading? The idea is that you can lazy load large acquisitions into a session while easily enabling you to use mne-python functions on that data

@rly

Not that I know of, but I think it’s a good idea!

@nmarkowitz

I've made functions to convert an acquisition to mne-python objects but only when all data is read in. Hoping to make a full object that will support lazy loading

@h-mayorquin

What did get on the way to support lazy loading?

@nmarkowitz

This is what I have. The load_nwb_settings() function loads in the json file I also attached. It's a file for in-lab usage with settings that we regularly use

import numpy as np
import mne
from mne.io import RawArray
from mne.io.base import BaseRaw
from mne import create_info, Annotations
from mne.channels import make_dig_montage
from pynwb import TimeSeries, NWBFile
from pynwb.ecephys import ElectricalSeries
from epipe.fileio.nwb.utils import load_nwb_settings


# TODO: Allow nwb2mne to pass in freesurfer directory with elec_recon to create montage and specify coordinates (LEPTO or fsaverage)

def _getRoot(nwbDataInterface):
    """Simple function to get the root of a nwb file when a child is passed in"""
    c = nwbDataInterface
    #parent = nwbDataInterface
    ii = 0
    while not isinstance(c,NWBFile):
        ii += 1
        #parent = c
        #c = parent.get_ancestor()
        c = c.get_ancestor()
        if ii >= 1000:
            print('Error in code')
            break

    return c

def nwb2mne(id,nwbfile=None,preload=False, create_montage=True):
    """Function for converting a TimeSeries object to a MNE RawArray object

    Parameters
    ----------
    id : TimeSeries, str
        A TimeSeries object, name or identifier of a TimeSeries object (latter two require the nwbfile argument)
    nwbfile : NWBFile
        The root object. Only needed if id is a str
    preload : bool
        Whether to load all the data

    Returns : RawArray
        An instance of mne.io.RawArray
    -------

    """

    # Get the TimeSeries instance
    if isinstance(id,TimeSeries):
        ts = id
        nwbfile = _getRoot(ts)
    elif isinstance(id,str) and isinstance(nwbfile,NWBFile):
        nwbObjIds = list(nwbfile.objects.keys())
        if id in nwbObjIds:
            ts = nwbfile.objects.get(id)
        else:
            inspection = inspectNwb(nwbfile)
            ts_df = inspection['timeseries']
            if id in ts_df['name'].values:
                objId = ts_df.loc[ts_df['name'] == id, 'id'].values[0]
                ts = nwbfile.objects.get(objId)
            else:
                print('ERROR!')
    else:
        print('ERROR!')

    # Set chan type for MNE
    if ts.name == 'ieeg':
        ch_types = 'seeg'
    elif isinstance(ts,ElectricalSeries):
        ch_types = 'eeg'
    else:
        ch_types = 'misc'

    # Determine sampling rate
    fs = ts.rate if 'rate' in ts.fields.keys() else 'timestamps'

    # Get number of channels and load data
    if len(ts.data.shape) == 2:
        ntpts, nchans = ts.data.shape
    else:
        ntpts = ts.data.size
        nchans = 1

    # Get timepoints
    if fs == 'timestamps':
        tvec = ts.timestamps[()]
        rate = 1/np.median(np.diff(tvec))
        t1 = tvec[0] + ts.starting_time
        #tvec_idx = np.logical_and(tvec>=t1, tvec<=t2)
    elif isinstance(fs,float) or isinstance(fs,int):
        rate = ts.rate
        tvec = np.arange(0,ntpts)/rate
        t1 = tvec[0] + ts.starting_time
        #tvec_idx = np.logical_and(tvec>=t1, tvec<=t2)

    # Get data
    if len(ts.data.shape) == 1:
        data = np.reshape(ts.data[()],(-1,ntpts))
    elif len(ts.data.shape) == 2:
        data = ts.data[()].T
    else:
        print('ERROR!')

    # Set time offset if first sample wasn't taken at time 0
    starting_time = ts.starting_time
    first_samp_offset = int(starting_time * rate)

    # Handle based on type
    if isinstance(ts,ElectricalSeries):

        # Get electrode info
        elecTable = nwbfile.electrodes.to_dataframe()
        ch_idx = ts.electrodes.data[()]
        elecs = elecTable.iloc[ch_idx]

        if 'label' in elecTable.columns:
            ch_names = elecs['label'].to_list()
        else:
            ch_names = ['ch{}'.format(ii+1) for ii in range(len(ch_idx))]

        # Create Info and Raw objects
        info = create_info(ch_names,sfreq=rate,ch_types=ch_types)
        info['line_freq'] = 60
        info['description'] = nwbfile.experiment_description
        raw = RawArray(data, info, first_samp=first_samp_offset)

        # Get bad channels
        annots, bad_chans = _badChanAnnotation(ts,raw.times[-1],return_chans=True)
        #raw.set_annotations(annots)
        raw.info['bads'] = bad_chans

        # Now add the montage
        if all([c in list(elecTable.columns) for c in ["lepto_x", "lepto_y", "lepto_z"]]):
            xyz_cols = ["lepto_x", "lepto_y", "lepto_z"]
        else:
            xyz_cols = ["x","y","z"]

        ch_coords = {}
        for idx, row in elecTable.iterrows():
            ch_name = row["label"]
            if ch_name in raw.ch_names:
                xyz = row[xyz_cols].to_numpy() / 1000
                ch_coords[ch_name] = xyz

        montage = make_dig_montage(ch_pos=ch_coords, coord_frame='mri')
        raw.set_montage(montage)


    else:
        ch_names = ['ch{}'.format(ii+1) for ii in range(nchans)]
        info = create_info(ch_names,sfreq=rate,ch_types=ch_types)
        info['line_freq'] = 60
        info['description'] = nwbfile.experiment_description
        raw = RawArray(data, info, first_samp=first_samp_offset)

    return raw

def _badChanAnnotation(es,end_time,return_chans=False):

    # Load NWB settings
    nwbSettings = load_nwb_settings()
    bad_chan_params = nwbSettings['ecog_file']['fields']
    elecTable = es.electrodes.to_dataframe()

    all_chans = []

    # Initiate Annotations object
    annots = None
    chan_types = list(bad_chan_params.keys())
    ntypes = len(chan_types)

    # Make annotations object and get channels with tags
    for k in bad_chan_params.keys():
        cols = bad_chan_params[k]
        bad_chans = []
        for c in cols:
            if c in list(elecTable.columns):
                chans = elecTable.loc[elecTable[c] == 1,'label'].tolist()
                bad_chans += chans
                all_chans += chans

        bad_chans = np.unique( np.array(bad_chans) ).tolist()
        if annots is None:
            annots = mne.Annotations(onset=0, duration=end_time, description=k, ch_names=[bad_chans])
        elif isinstance(annots,mne.Annotations):
            annots.append(onset=0, duration=end_time, description=k, ch_names=[bad_chans])

    if return_chans:
        all_chans = np.unique(np.array(all_chans)).tolist()
        return annots, all_chans
    else:
        return annots

nwb_settings.json:

{
    // Parameters dealing with the electrodetable
    "electrode_table": {
        "name": "electrodes",
        "description": "metadata about extracellular electrodes",

        // Columns that go into the electrode table
        "columns": {
            "label": {
                "title": "label",
                "search": "(?:Label|Contact|label*)",
                "description": "name of the channel",
                "type": "string",
                "required": true
            },
            "spec": {
                "title": "spec",
                "search": "(?:elecSpec|spec*)",
                "description": "type of channel",
                "type": "string",
                "required": true
            },
            "x": {
                "title": "x",
                "search": "(?:fsaverage_coords_1|fsaverage_1)",
                "description": "the x coordinate of the channel location in freesurfer average (fsaverage) space",
                "default": 999.0,
                "type": "float",
                "required": true
            },
            "y": {
                "title": "y",
                "search": "(?:fsaverage_coords_2|fsaverage_2)",
                "description": "the y coordinate of the channel location in freesurfer average (fsaverage) space",
                "default": 999.0,
                "type": "float",
                "required": true
            },
            "z": {
                "title": "z",
                "search": "(?:fsaverage_coords_3|fsaverage_3)",
                "description": "the z coordinate of the channel location in freesurfer average (fsaverage) space",
                "default": 999.0,
                "type": "float",
                "required": true
            },
            "location": {
                "title": "location",
                "search": "(?:FS_vol|AparcAseg_Atlas)",
                "description": "the location of channel within the subject e.g. brain region",
                "default": "unknown",
                "type": "string",
                "required": true
            },
            "filtering": {
                "title": "filtering",
                "description": "description of hardware filtering",
                "default": "none",
                "type": "string",
                "required": true
            },
            "imp": {
                "title": "imp",
                "description": "the impedance of the channel",
                "default": -1.0,
                "type": "float",
                "required": true
            },
            "group": {
                "title": "group",
                "description": "a reference to the ElectrodeGroup this electrode is a part of",
                "required": true
            },
            "group_name": {
                "title": "group_name",
                "description": "the name of the ElectrodeGroup this electrode is a part of",
                "type": "string",
                "required": true
            },
            "lepto_x": {
                "title": "lepto_x",
                "search": "(?:LEPTO_COORDS_1|LEPTO_1)",
                "description": "the x coordinate of the channel location in native leptomeningeal space",
                "default": 999.0,
                "type": "float",
                "required": false
            },
            "lepto_y": {
                "title": "lepto_y",
                "search": "(?:LEPTO_COORDS_2|LEPTO_2)",
                "description": "the y coordinate of the channel location native leptomeningeal space",
                "default": 999.0,
                "type": "float",
                "required": false
            },
            "lepto_z": {
                "title": "lepto_z",
                "search": "(?:LEPTO_COORDS_3|LEPTO_3)",
                "description": "the z coordinate of the channel location native leptomeningeal space",
                "default": 999.0,
                "type": "float",
                "required": false
            },
            "soz": {
                "title": "soz",
                "search": "soz*",
                "description": "seizure onset zone",
                "type": "uint8",
                "default": 0,
                "required": false
            },
            "spikey": {
                "title": "spikey",
                "search": "(?:epileptic_spikes|spikey*)",
                "description": "displays interictal activity",
                "type": "uint8",
                "default": 0,
                "required": false
            },
            "out": {
                "title": "out",
                "search": "(?:outofthebrain|out*)",
                "description": "is outside the brain",
                "default": 0,
                "type": "uint8",
                "required": false
            },
            "ptd": {
                "title": "ptd",
                "search": "ptd",
                "description": "Proximal Tissue Density i.e. the concentration of gray matter and white matter voxels surround the contact",
                "default": 999,
                "type": "float",
                "required": false
            },
            "hem": {
                "title": "hem",
                "search": "(?:LvsR|Hem)",
                "description": "hemisphere of the brain channel is located on",
                "default": "unknown",
                "type": "string",
                "required": false
            },
            "dk": {
                "title": "dk",
                "search": "(?:Desikan_Killiany|DK_Atlas)",
                "description": "Location of channel according to Desikan-Killiany atlas",
                "default": "unknown",
                "type": "string",
                "required": false
            },
            "destrieux": {
                "title": "destrieux",
                "search": "(?:Destrieux|D_Atlas)",
                "description": "Location of channel according to Destrieux atlas",
                "default": "unknown",
                "type": "string",
                "required": false
            },
            "yeo7": {
                "title": "yeo7",
                "search": "(?:Yeo7|Y7_Atlas)",
                "description": "Location of channel according to Yeo 7-network atlas",
                "default": "NAN",
                "type": "string",
                "required": false
            },
            "yeo17": {
                "title": "yeo17",
                "search": "(?:Yeo17|Y17_Atlas)",
                "description": "Location of channel according to Yeo 17-network atlas",
                "default": "NAN",
                "type": "string",
                "required": false
            },
            "artifact_block": {
                "title": "artifact_block",
                "search": "(?:artifact_block)",
                "description": "contains poor signal quality only during this recording",
                "default": 0,
                "type": "uint8",
                "required": false
            },
            "artifact_patient": {
                "title": "artifact_patient",
                "search": "(?:artifact_patient|bad*)",
                "description": "contains poor signal quality in general for this patient",
                "default": 0,
                "type": "uint8",
                "required": false
            },
            "dist": {
                "title": "dist",
                "search": "Dist",
                "description": "Distance of channel to nearest vertex on cortical surface",
                "default": 999,
                "type": "float",
                "required": false
            }
        }
    },

    // File meta-data
    // Below fields were previously present but then removed
    // experimenter: Lab members of the Human Brain Mapping Lab
    "meta_data": {
        "lab": "Human Brain Mapping Lab at Feinstein-Northwell",
        "notes": "NA",
        "institution": "North Shore University Hospital",
        "data_collection": "NA",
        "session_id": "None",
        "session_description": "None",
        "experiment_description": "None",
        "date": "1992-11-09 00:00:00"
    },

    // Subject info
    "subject": {
        "species": "Homo sapiens",
        "subject_id": "None",
        "sex": "U",
        "age": "Unknown",
        "description": "intracranial patient"
    },

    // Recording devices used
    "devices": {
        "tdt": {
            "name": "PZ5",
            "description": "TDT amplifier",
            "manufacturer": "Tucker-Davis Technologies",
            "search": "(?:TDT|RS4|PZ5)"
        },
        "natus": {
            "name": "Natus",
            "description": "Quantum Headbox",
            "manufacturer": "Natus",
            "search": "(?:Natus*|xltek*|quantum*)"
        }
    },

    // Table that holds events of recording
    "events_table": {
        "name": "Events",
        "table_description": "Preprocessed events",
        "columns": [{
            "title": "start_time",
            "description": "Start time of epoch, in seconds",
            "required": true
        }, {
            "title": "stop_time",
            "description": "Stop time of epoch, in seconds",
            "required": true
        }]
    },

    // Data formats
    "raw": {
        "edf": {
            "annotations_to_ignore": ["Clip Not*", "Montage*", "Gain/Filter Change", "Breakout box reconnected.", "Started Analyzer -*", "XLSpik*", "XLEven*","Video System*","Video Recording*","Persyst*"]
        },
        "tdt": {
            "neuro_channels": [
                ["RAWx"],
                ["EEG1", "EEG2"],
                ["RSn1", "RSn2"]
            ]
        },
        "date": "1992-11-09 00:00:00"
    },

    // Labels that indicate an electrode recorded intracranial activity
    "lfp_specs": ["depth", "hd_depth", "strip", "hd_strip", "grid", "hd_grid", "seeg", "ecog", "hseeg"],

    // How and where to store key acquisitions
    "acquisitions": {
        "ieeg": "/acquisitions/ieeg",
        "references": {
            "avg": {
                "path": "/general/average_ref",
                "description": "structure of the average reference of this recording",
                "table": "/general/average_ref/channels"
            },
            "bipolar": {
                "path": "/general/bipolar_ref",
                "description": "structure of bipolar montage",
                "table": "/general/bipolar_ref/channels"
            }
        }
    },

    // How to handle HBML ecog.mat files
    "ecog_file": {
        "fields": {
            "bad_chans": ["bad","out","artifact_block","artifact_patient"],
            "spike_chans": ["spikey"],
            "szr_onset_chans": ["soz"]
        },
        "electable": "channelinfo"
    },

    // Naming conventions of the electrodes used
    "electrode_names": {
        "Da": "amygdala",
        "Dh": "hippocampus",
        "Dp": "posterior hippocampus",
        "Ti": "inferior temporal gyrus to entorhinal cortex",
        "Tp": "posterior temporal gyrus",
        "Ts": "superior temporal gyrus",
        "Tx": "temporal pole",
        "To": "inferior temporal gyrus to fusiform gyrus",
        "Fo": "orbitofrontal cortex",
        "Fi": "inferior frontal gyrus",
        "Fa": "anterior frontal gyrus",
        "Fx": "frontal pole",
        "Fm": "medial frontal lobe",
        "Fp": "posterior frontal lobe",
        "Pc": "supramarginal gyrus",
        "Ps": "superior parietal lobule",
        "Pi": "inferior parietal lobule",
        "Pp": "posterior parietal lobule",
        "Pm": "medial parietal lobule",
        "Ia": "anterior insula",
        "Im": "middle insula",
        "Ip": "posterior insula",
        "Hi": "mid and posterior cingulum",
        "Hs": "mid to anterior cingulum",
        "O": "occipital lobule",
        "Oa": "anterior occipital",
        "Oi": "inferior occipital",
        "Op": "posterior occipital"
    }
}

@nmarkowitz

From what I understand, if I want a mne object with lazy loading I need to define a new class and define the method _read_segment_file() which is used by all mne.io classes to load data in a similar manner. I can define it in that. For now it's easier to just load everything in so all mne functions work as they should

[nmarkowitz]

@rly what do you think of the current code I have that requires user to input id (or maybe path) of the acquisition they want plus the NWBFile object? I have a separate function that finds and returns information on all instances of a TimeSeries object in the file that can help in this regard.