convert2nwb.py

'''
Convert to NWB

Run this script to convert derived data generated by the Brain-wide
infra-slow dynamics study at UoL to the Neurodata Without Borders (NWB)
file format.

The script works by loading general, animal, and recording session
metadata from nwbParams, nwbAnimalParams, nwbSessionParams, respectively.
It then locates the derived data MAT files for each animal and converts
them into derived data NWB files dividing the data per recording session.
The derived data include spiking, waveform, pupil area size fluctuation,
and total facial movement data.

The conversion pipeline depends on the specific structure of derived data
MAT files used in this study. The way the pipeline is organised is
dictated by the fact that the conversion procedure was adopted late in
the study. Ideally NWB file format should be adopted early in the study.

You can use this pipeline to get an idea of how to convert your own
ecephys data to the NWB file format to store spiking and behavioural
data combined with some metadata.
'''

from datetime import datetime
from dateutil.tz import tzlocal

import numpy as np
from pynwb import NWBFile, NWBHDF5IO, TimeSeries
from pynwb.file import Subject
from pynwb.ecephys import ElectricalSeries, LFP
from pynwb.core import DynamicTable
from pynwb.behavior import PupilTracking, BehavioralTimeSeries

from localFunctions import createElectrodeTable, array2list, getSpikes, reshapeWaveforms, concatenateMat, parsePeriod, markQualitySamples

# Load general NWB parameters
exec(open("nwbParams.py").read())

# Load animal-specific parameters
exec(open("nwbAnimalParams.py").read())

# Load session-specific parameters
exec(open("nwbSessionParams.py").read())

# Generate NWB files for every recording session
if not os.path.isdir(animalDerivedDataFolderNWB):
  os.makedirs(animalDerivedDataFolderNWB, exist_ok=True) # Folder to store animal's converted NWB data
derivedData = animalDerivedDataFile
for iSess in range(len(sessionID)):
  # Assign NWB file fields
  nwb = NWBFile(
    session_description = sessionDescription[iSess],
    identifier = animalID + '_' + sessionID[iSess], 
    session_start_time = sessionStartTime[iSess], 
    experimenter = experimenter,  # optional
    session_id = sessionID[iSess],  # optional
    institution = institution,  # optional
    related_publications = publications,  # optional
    notes = sessionNotes[iSess],  # optional
    lab = lab) # optional

  # Create subject object
  nwb.subject = Subject(
    subject_id = animalID,
    age = age,
    description = description,
    species = species,
    sex = sex)

  # Create electrode tables: Info about each recording channel
  input = {
    "iElectrode": 0,
    "electrodeDescription": electrodeDescription[iSess],
    "electrodeManufacturer": electrodeManufacturer[iSess],
    "nShanks": nShanks[iSess],
    "nChannelsPerShank": nChannelsPerShank[iSess],
    "electrodeLocation": electrodeLocation[iSess],
    "electrodeCoordinates": electrodeCoordinates[iSess],
    "sessionID": sessionID[iSess],
    "electrodeLabel": electrodeLabel[iSess]}
  if probeInserted[iSess][input["iElectrode"]] and endCh[iSess][input["iElectrode"]].any():
    (tbl1, columnLabels, columnDescription) = createElectrodeTable(nwb, input)
  else:
    tbl1 = []; columnLabels = []; columnDescription = []
  
  input["iElectrode"] = 1
  if probeInserted[iSess][input["iElectrode"]] and endCh[iSess][input["iElectrode"]].any():
    (tbl2, columnLabels, columnDescription) = createElectrodeTable(nwb, input)
  else:
    tbl2 = []
    if not len(columnLabels):
      columnLabels = []; columnDescription = []
  
  if len(tbl1) and len(tbl2):
    tbl = array2list(np.append(tbl1, tbl2, axis=0), columnLabels, columnDescription)
  elif len(tbl1) and not len(tbl2):
    tbl = array2list(tbl1, columnLabels, columnDescription)
  elif not len(tbl1) and len(tbl2):
    tbl = array2list(tbl2, columnLabels, columnDescription)
  else:
    tbl = []

  if len(tbl):
    for iColumn in range(len(columnLabels)):
      if columnLabels[iColumn] is not 'location' and columnLabels[iColumn] is not 'group':
        nwb.add_electrode_column(name=columnLabels[iColumn], description=columnDescription[iColumn])
    
    for iElec in range(len(tbl[0].data)):
      nwb.add_electrode(
          channel_id=tbl[0].data[iElec],
          channel_local_index=tbl[1].data[iElec],
          x=tbl[2].data[iElec],
          y=tbl[3].data[iElec],
          z=float(tbl[4].data[iElec]),
          imp=tbl[5].data[iElec],
          location=tbl[6].data[iElec],
          filtering=tbl[7].data[iElec],
          group=tbl[8].data[iElec],
          channel_label=tbl[9].data[iElec],
          probe_label=tbl[10].data[iElec]
      )

    nwb.electrodes.to_dataframe()
  
  # Load spike times from the MAT file
  (spikes, metadata, derivedData, columnLabels, columnDescription) = getSpikes(derivedData, animalID, sessionID[iSess], tbl)

  # Load and reshape unit waveforms
  if probeInserted[iSess][0] and len(spikes) and endCh[iSess][0].any():
    waveformsFile1 = os.path.join(electrodeFolder[iSess][0], 'waveforms.mat')
    if os.path.isfile(waveformsFile1):
      waveformsProbe1 = h5py.File(waveformsFile1,'r')
    else:
      waveformsProbe1 = []
    waveformMeans1 = reshapeWaveforms(waveformsProbe1, 1, metadata)
  else:
    waveformMeans1 = []

  if probeInserted[iSess][1] and len(spikes) and endCh[iSess][1].any():
    waveformsFile2 = os.path.join(electrodeFolder[iSess][1], 'waveforms.mat')
    if os.path.isfile(waveformsFile2):
      waveformsProbe2 = h5py.File(waveformsFile2,'r')
    else:
      waveformsProbe2 = []
    waveformMeans2 = reshapeWaveforms(waveformsProbe2, 2, metadata)
  else:
    waveformMeans2 = []

  waveformMeans = np.squeeze(np.array(waveformMeans1 + waveformMeans2))

  # Store spiking and waveform data inside the nwb object
  if len(spikes):
    for iColumn in range(len(columnLabels)-1):
      nwb.add_unit_column(name=columnLabels[iColumn], description=columnDescription[iColumn])
    for iUnit in range(len(spikes)):
      nwb.add_unit(
        cluster_id=metadata[iUnit,0],
        local_cluster_id=metadata[iUnit,1],
        type=metadata[iUnit,2],
        peak_channel_index=metadata[iUnit,3],
        peak_channel_id=metadata[iUnit,4],
        local_peak_channel_id=metadata[iUnit,5],
        rel_horz_pos=metadata[iUnit,6],
        rel_vert_pos=metadata[iUnit,7],
        isi_violations=metadata[iUnit,8],
        isolation_distance=metadata[iUnit,9],
        area=metadata[iUnit,10],
        probe_id=metadata[iUnit,11],
        electrode_group=metadata[iUnit,12],
        spike_times=spikes[iUnit],
        waveform_mean=waveformMeans[iUnit]
      )

    # Add behavioural data: Pupil area size
    acceptablePeriod = np.array(derivedData.get('dataStruct/eyeData/' + animalID + '_s' + sessionID[iSess] + '/period')).transpose() # Acceptable quality range in seconds
    acceptablePeriod = parsePeriod(acceptablePeriod, derivedData)
    videoFrameTimes = np.array(derivedData.get('dataStruct/eyeData/' + animalID + '_s' + sessionID[iSess] + '/frameTimes')).transpose() # seconds
    pupilAreaSize = np.array(derivedData.get('dataStruct/eyeData/' + animalID + '_s' + sessionID[iSess] + '/pupilArea')).transpose() # pixels^2
    if len(pupilAreaSize.shape) and not (len(pupilAreaSize) == 2 and not pupilAreaSize[0] and not pupilAreaSize[1]):
      acceptableSamples = markQualitySamples(acceptablePeriod, videoFrameTimes)
      pupilAreaSize = TimeSeries(
        name='TimeSeries',
        data=pupilAreaSize.squeeze(),
        timestamps=videoFrameTimes.squeeze(),
        unit='pixels^2',
        control=acceptableSamples.squeeze(),
        control_description='low quality samples that should be excluded from analyses; acceptable quality samples',
        description='Pupil area size over the recording session measured in pixels^2. ' +
          'Acceptable quality period starting and ending times are given by data_continuity parameter. ' +
          'The full data range can be divided into multiple acceptable periods.'
      )
      behaviour_module = nwb.create_processing_module(name="behavior", description="contains behavioral data")
      pupilTracking = PupilTracking(pupilAreaSize)
      behaviour_module.add(pupilTracking)
    
    # Add behavioural data: Total movement of the facial area
    acceptablePeriod = np.array(derivedData.get('dataStruct/motionData/' + animalID + '_s' + sessionID[iSess] + '/period')).transpose() # Acceptable quality range in seconds
    acceptablePeriod = parsePeriod(acceptablePeriod, derivedData)
    videoFrameTimes = np.array(derivedData.get('dataStruct/motionData/' + animalID + '_s' + sessionID[iSess] + '/frameTimes')).transpose() # seconds
    totalFaceMovement = np.array(derivedData.get('dataStruct/motionData/' + animalID + '_s' + sessionID[iSess] + '/sa')).transpose() # z-scored change in the frame pixels' content with respect to the previous frame
    if len(totalFaceMovement.shape) and not (len(totalFaceMovement) == 2 and not totalFaceMovement[0] and not totalFaceMovement[1]):
      acceptableSamples = markQualitySamples(acceptablePeriod, videoFrameTimes)
      totalFaceMovement = TimeSeries(
        name='TimeSeries',
        data=totalFaceMovement.squeeze(),
        timestamps=videoFrameTimes.squeeze(),
        unit='a.u.',
        control=acceptableSamples.squeeze(),
        control_description='low quality samples that should be excluded from analyses; acceptable quality samples',
        description='Z-scored change in the frame pixels'' content with respect to the previous frame. ' +
          'It measures the total movement of objects inside the video.'
      )
      if not ('behaviour_module' in locals() or 'behaviour_module' in globals()):
        behaviour_module = nwb.create_processing_module(name="behavior", description="contains behavioral data")
      behavioralTimeSeries = BehavioralTimeSeries(totalFaceMovement)
      behaviour_module.add(behavioralTimeSeries)
  else:
    acceptablePeriod = []; videoFrameTimes = []; pupilAreaSize = []; totalFaceMovement = [];

  # Save the NWB file for the session
  if iSess < 9:
    nwbFilename = os.path.join(animalDerivedDataFolderNWB, 'ecephys_session_0' + str(iSess+1) + '.nwb')
  else:
    nwbFilename = os.path.join(animalDerivedDataFolderNWB, 'ecephys_session_' + str(iSess+1) + '.nwb')
  with NWBHDF5IO(nwbFilename, "w") as io:
    io.write(nwb)
  
  # Delete variables as a precaution
  try:
    del behaviour_module, acceptableSamples
  except:
    print('behaviour_module does not exist')
  del nwb, input, tbl, spikes, metadata, columnLabels, columnDescription, waveformMeans
  del acceptablePeriod, videoFrameTimes, pupilAreaSize, totalFaceMovement