Source code for allensdk.brain_observatory.ecephys.write_nwb.nwb_writer

import inspect
import logging
import os
from pathlib import Path
from typing import Union

from pynwb import NWBHDF5IO, NWBFile

from allensdk.brain_observatory.ecephys.behavior_ecephys_session import \
    BehaviorEcephysSession
from allensdk.brain_observatory.nwb.nwb_utils import NWBWriter
from allensdk.core import JsonReadableInterface, NwbReadableInterface, \
    NwbWritableInterface




class BehaviorEcephysNwbWriter(NWBWriter):
    """NWB Writer for behavior ecephys. Same as `NWBWriter` except also
    writes probe NWB files separately """
    def __init__(
            self,
            session_nwb_filepath: str,
            session_data: dict,
            serializer: Union[
                JsonReadableInterface,
                NwbReadableInterface,
                NwbWritableInterface]):
        super().__init__(
            nwb_filepath=session_nwb_filepath,
            session_data=session_data,
            serializer=serializer
        )



    def write_nwb(self, **kwargs):
        """Tries to write nwb to disk. If it fails, the filepath has ".error"
        appended

        Parameters
        ----------
        kwargs: kwargs sent to `from_json`, `from_nwb`, `to_nwb`

        """
        from_json_kwargs = {
            k: v for k, v in kwargs.items()
            if k in inspect.signature(self._serializer.from_json).parameters}
        json_session = self._serializer.from_json(
            session_data=self._session_data, **from_json_kwargs)

        try:
            nwbfile = self._write_nwb(
                session=json_session, **kwargs)
            self._compare_sessions(nwbfile=nwbfile,
                                   loaded_session=json_session,
                                   **kwargs)
            os.rename(self.nwb_filepath_inprogress, self._nwb_filepath)
        except Exception as e:
            if os.path.isfile(self.nwb_filepath_inprogress):
                os.rename(self.nwb_filepath_inprogress,
                          self._nwb_filepath_error)
            raise e


    def _write_nwb(
            self,
            session: BehaviorEcephysSession,
            **kwargs) -> NWBFile:
        to_nwb_kwargs = {
            k: v for k, v in kwargs.items()
            if k in inspect.signature(self._serializer.to_nwb).parameters}
        session_nwbfile, probe_nwbfile_map = session.to_nwb(**to_nwb_kwargs)

        os.makedirs(Path(self.nwb_filepath_inprogress).parent, exist_ok=True)

        logging.info(f'Writing session NWB file to '
                     f'{self.nwb_filepath_inprogress}')
        with NWBHDF5IO(self.nwb_filepath_inprogress, 'w') as nwb_file_writer:
            nwb_file_writer.write(session_nwbfile)
        logging.info(f'Wrote session NWB file to '
                     f'{self.nwb_filepath_inprogress}')

        for probe_name, probe_nwbfile in probe_nwbfile_map.items():
            probe_id = [p.id for p in session.get_probes_obj()
                        if p.name == probe_name][0]
            if probe_nwbfile is not None:
                probe_nwb_path = Path(self._nwb_filepath).parent / \
                    f'lfp_probe_{probe_id}.nwb'
                logging.info(f'Writing probe NWB file to '
                             f'{probe_nwb_path}')
                with NWBHDF5IO(probe_nwb_path, 'w') as nwb_file_writer:
                    nwb_file_writer.write(probe_nwbfile)
                logging.info(f'Wrote probe NWB file to '
                             f'{probe_nwb_path}')
        return session_nwbfile