aind-dynamic-foraging-data-utils 0.1.0__py3-none-any.whl

aind_dynamic_foraging_data_utils/__init__.py

@@ -0,0 +1,3 @@
+"""Init package"""
+
+__version__ = "0.1.0"

aind_dynamic_foraging_data_utils/alignment.py

@@ -0,0 +1,573 @@
+"""
+    Tools for aligning continuous and discrete events
+    functions:
+    get_time_array
+    slice_inds_and_offsets
+    index_of_nearest_value
+    event_triggered_response
+"""
+
+import numpy as np
+import pandas as pd
+
+
+def get_time_array(
+    t_start, t_end, sampling_rate=None, step_size=None, include_endpoint=True
+):  # NOQA E501
+    """
+    A function to get a time array between two specified timepoints at a defined sampling rate  # NOQA E501
+    Deals with possibility of time range not being evenly divisible by desired sampling rate  # NOQA E501
+    Uses np.linspace instead of np.arange given decimal precision issues with np.arange (see np.arange documentation for details)  # NOQA E501
+
+    Parameters:
+    -----------
+    t_start : float
+        start time for array
+    t_end : float
+        end time for array
+    sampling_rate : float
+        desired sampling of array
+        Note: user must specify either sampling_rate or step_size, not both
+    step_size : float
+        desired step size of array
+        Note: user must specify either sampling_rate or step_size, not both
+    include_endpoint : Boolean
+        Passed to np.linspace to calculate relative time
+        If True, stop is the last sample. Otherwise, it is not included.
+            Default is True
+
+    Returns:
+    --------
+    numpy.array
+        an array of timepoints at the desired sampling rate
+
+    Examples:
+    ---------
+    get a time array exclusive of the endpoint
+    >>> t_array = get_time_array(
+        t_start=-1,
+        t_end=1,
+        step_size=0.5,
+        include_endpoint=False
+    )
+
+    np.array([-1., -0.5,  0.,  0.5])
+
+
+    get a time array inclusive of the endpoint
+    >>> t_array = get_time_array(
+        t_start=-1,
+        t_end=1,
+        step_size=0.5,
+        include_endpoint=False
+    )
+
+    np.array([-1., -0.5,  0.,  0.5, 1.0])
+
+
+    get a time array where the range can't be evenly divided by the desired step_size
+    in this case, the time array includes the last timepoint before the desired endpoint
+    >>> t_array = get_time_array(
+        t_start=-1,
+        t_end=0.75,
+        step_size=0.5,
+        include_endpoint=False
+    )
+
+    np.array([-1., -0.5,  0.,  0.5])
+
+
+    Instead of passing the step_size, we can pass the sampling rate
+    >>> t_array = get_time_array(
+        t_start=-1,
+        t_end=1,
+        sampling_rate=2,
+        include_endpoint=False
+    )
+
+    np.array([-1., -0.5,  0.,  0.5])
+    """
+    assert (
+        sampling_rate is not None or step_size is not None
+    ), "must specify either sampling_rate or step_size"  # NOQA E501
+    assert (
+        sampling_rate is None or step_size is None
+    ), "cannot specify both sampling_rate and step_size"  # NOQA E501
+
+    # value as a linearly spaced time array
+    if not step_size:
+        step_size = 1 / sampling_rate
+    # define a time array
+    n_steps = (t_end - t_start) / step_size
+    if n_steps != int(n_steps):
+        # if the number of steps isn't an int, that means it isn't possible
+        # to end on the desired t_after using the defined sampling rate
+        # we need to round down and include the endpoint
+        n_steps = int(n_steps)
+        t_end_adjusted = t_start + n_steps * step_size
+        include_endpoint = True
+    else:
+        t_end_adjusted = t_end
+
+    if include_endpoint:
+        # add an extra step if including endpoint
+        n_steps += 1
+
+    t_array = np.linspace(t_start, t_end_adjusted, int(n_steps), endpoint=include_endpoint)
+
+    return t_array
+
+
+def slice_inds_and_offsets(
+    data_timestamps,
+    event_timestamps,
+    time_window,
+    sampling_rate=None,
+    include_endpoint=False,
+):  # NOQA E501
+    """
+    Get nearest indices to event timestamps, plus ind offsets (start:stop)
+    for slicing out a window around the event from the trace.
+    Parameters:
+    -----------
+    data_timestamps : np.array
+        Timestamps of the datatrace.
+    event_timestamps : np.array
+        Timestamps of events around which to slice windows.
+    time_window : list
+        [start_offset, end_offset] in seconds
+    sampling_rate : float, optional, default=None
+        Sampling rate of the datatrace.
+        If left as None, samplng rate is inferred from data_timestamps.
+
+    Returns:
+    --------
+    event_indices : np.array
+        Indices of events from the timestamps provided.
+    start_ind_offset : int
+    end_ind_offset : int
+    trace_timebase :  np.array
+    """
+    if sampling_rate is None:
+        sampling_rate = 1 / np.diff(data_timestamps).mean()
+
+    event_indices = index_of_nearest_value(data_timestamps, event_timestamps)
+    trace_len = (time_window[1] - time_window[0]) * sampling_rate
+    start_ind_offset = int(time_window[0] * sampling_rate)
+    end_ind_offset = int(start_ind_offset + trace_len) + int(include_endpoint)
+    trace_timebase = np.arange(start_ind_offset, end_ind_offset) / sampling_rate
+
+    return event_indices, start_ind_offset, end_ind_offset, trace_timebase
+
+
+def index_of_nearest_value(data_timestamps, event_timestamps):
+    """
+    The index of the nearest sample time for each event time.
+
+    Parameters:
+    -----------
+    sample_timestamps : np.ndarray of floats
+        sorted 1-d vector of data sample timestamps.
+    event_timestamps : np.ndarray of floats
+        1-d vector of event timestamps.
+
+    Returns:
+    --------
+    event_aligned_ind : np.ndarray of int
+        An array of nearest sample time index for each event times.
+    """
+    insertion_ind = np.searchsorted(data_timestamps, event_timestamps)
+    # is the value closer to data at insertion_ind or insertion_ind-1?
+    ind_diff = data_timestamps[insertion_ind] - event_timestamps
+    ind_minus_one_diff = np.abs(
+        data_timestamps[np.clip(insertion_ind - 1, 0, np.inf).astype(int)] - event_timestamps
+    )
+
+    event_indices = insertion_ind - (ind_diff > ind_minus_one_diff).astype(int)
+    return event_indices
+
+
+def event_triggered_response(  # noqa C901
+    data,
+    t,
+    y,
+    event_times,
+    t_start=None,
+    t_end=None,
+    t_before=None,
+    t_after=None,
+    output_sampling_rate=None,
+    include_endpoint=True,
+    output_format="tidy",
+    interpolate=True,
+    censor=True,
+    censor_times=None,
+    nan_policy="error",
+):  # NOQA E501
+    """
+    Slices a timeseries relative to a given set of event times
+    to build an event-triggered response. From mindscope_utilities (commit af2b70a)
+
+    For example, If we have data such as a measurement of neural activity
+    over time and specific events in time that we want to align
+    the neural activity to, this function will extract segments of the neural
+    timeseries in a specified time window around each event.
+
+    The times of the events need not align with the measured
+    times of the neural data.
+    Relative times will be calculated by linear interpolation.
+
+    Parameters:
+    -----------
+    data: Pandas.DataFrame
+        Input dataframe in tidy format
+        Each row should be one observation
+        Must contains columns representing `t` and `y` (see below)
+    t : string
+        Name of column in data to use as time data
+    y : string
+        Name of column to use as y data
+    event_times: list or array of floats
+        Times of events of interest.
+        Values in column specified by `y` will be sliced and interpolated
+            relative to these times
+    t_start : float
+        start time relative to each event for desired time window
+        e.g.:   t_start = -1 would start the window 1 second before each
+                t_start = 1 would start the window 1 second after each event
+        Note: cannot pass both t_start and t_before
+    t_before : float
+        time before each of event of interest to include in each slice
+        e.g.:   t_before = 1 would start the window 1 second before each event
+                t_before = -1 would start the window 1 second after each event
+        Note: cannot pass both t_start and t_before
+    t_end : float
+        end time relative to each event for desired time window
+        e.g.:   t_end = 1 would end the window 1 second after each event
+                t_end = -1 would end the window 1 second before each event
+        Note: cannot pass both t_end and t_after
+    t_after : float
+        time after each event of interest to include in each slice
+        e.g.:   t_after = 1 would start the window 1 second after each event
+                t_after = -1 would start the window 1 second before each event
+        Note: cannot pass both t_end and t_after
+    output_sampling_rate : float
+        Desired sampling of output.
+        Input data will be interpolated to this sampling rate if interpolate = True (default). # NOQA E501
+        If passing interpolate = False, the sampling rate of the input timeseries will # NOQA E501
+        be used and output_sampling_rate should not be specified.
+    include_endpoint : Boolean
+        Passed to np.linspace to calculate relative time
+        If True, stop is the last sample. Otherwise, it is not included.
+            Default is True
+    output_format : string
+        'wide' or 'tidy' (default = 'tidy')
+        if 'tidy'
+            One column representing time
+            One column representing event_number
+            One column representing event_time
+            One row per observation (# rows = len(time) x len(event_times))
+        if 'wide', output format will be:
+            time as indices
+            One row per interpolated timepoint
+            One column per event,
+                with column names titled event_{EVENT NUMBER}_t={EVENT TIME}
+    interpolate : Boolean
+        if True (default), interpolates each response onto a common timebase
+        if False, shifts each response to align indices to a common timebase
+    censor: Boolean
+        if True (default), censor observations that take place after the next event time
+        if False, do not censor
+    censor_times: list or array or None
+        if None, and censor is True, then use event_times as the censor times
+        if times are provided, then these are the times at which ETR is censored
+    nan_policy: How to handle NaNs in the input data
+        "error": raise an exception if NaNs are present in the time window of an ETR
+        "interpolate": interpolate over NaN values
+        "exclude": exclude any response with a NaN in the response window
+
+    Returns:
+    --------
+    Pandas.DataFrame
+        See description in `output_format` section above
+
+    Examples:
+    ---------
+    An example use case, recover a sinousoid from noise:
+
+    First, define a time vector
+    >>> t = np.arange(-10,110,0.001)
+
+    Now build a dataframe with one column for time,
+    and another column that is a noise-corrupted sinuosoid with period of 1
+    >>> data = pd.DataFrame({
+            'time': t,
+            'noisy_sinusoid': np.sin(2*np.pi*t) + np.random.randn(len(t))*3
+        })
+
+    Now use the event_triggered_response function to get a tidy
+    dataframe of the signal around every event
+
+    Events will simply be generated as every 1 second interval
+    starting at 0, since our period here is 1
+    >>> etr = event_triggered_response(
+            data,
+            x = 'time',
+            y = 'noisy_sinusoid',
+            event_times = np.arange(100),
+            t_start = -1,
+            t_end = 1,
+            output_sampling_rate = 100
+        )
+    Then use seaborn to view the result
+    We're able to recover the sinusoid through averaging
+    >>> import matplotlib.pyplot as plt
+    >>> import seaborn as sns
+    >>> fig, ax = plt.subplots()
+    >>> sns.lineplot(
+            data = etr,
+            x='time',
+            y='noisy_sinusoid',
+            ax=ax
+        )
+    """
+    # ensure that non-conflicting time values are passed
+    assert (
+        t_before is not None or t_start is not None
+    ), "must pass either t_start or t_before"  # noqa: E501
+    assert (
+        t_after is not None or t_end is not None
+    ), "must pass either t_start or t_before"  # noqa: E501
+
+    assert (
+        t_before is None or t_start is None
+    ), "cannot pass both t_start and t_before"  # noqa: E501
+    assert t_after is None or t_end is None, "cannot pass both t_after and t_end"  # noqa: E501
+
+    if interpolate is False:
+        assert (
+            output_sampling_rate is None
+        ), "if interpolation = False, the sampling rate of the input timeseries will be used. Do not specify output_sampling_rate"  # NOQA E501
+
+    # assign time values to t_start and t_end
+    if t_start is None:
+        t_start = -1 * t_before
+    if t_end is None:
+        t_end = t_after
+
+    # ensure that t_end is greater than t_start
+    assert t_end > t_start, "must define t_end to be greater than t_start"
+
+    assert (not censor) or (output_format == "tidy"), "cannot censor data in wide output"
+
+    assert nan_policy in ["error", "interpolate", "exclude"], "unrecognized nan_policy"
+
+    if censor:
+        event_times = np.sort(event_times)
+
+    if output_sampling_rate is None:
+        # if sampling rate is None,
+        # set it to be the mean sampling rate of the input data
+        output_sampling_rate = 1 / np.diff(data[t]).mean()
+
+    # if interpolate is set to True,
+    # we will calculate a common timebase and
+    # interpolate every response onto that timebase
+    if interpolate:
+        # set up a dictionary with key 'time' and
+        t_array = get_time_array(
+            t_start=t_start,
+            t_end=t_end,
+            sampling_rate=output_sampling_rate,
+            include_endpoint=include_endpoint,
+        )
+        data_dict = {"time": t_array}
+
+        # iterate over all event times
+        data_reindexed = data.set_index(t, inplace=False)
+
+        for event_number, event_time in enumerate(np.array(event_times)):
+            # get a slice of the input data surrounding each event time
+            data_slice = data_reindexed[y].loc[
+                event_time + t_start : event_time + t_end
+            ]  # noqa: E501
+
+            # if the slice is empty, we will fill it with NaNs
+            if len(data_slice) == 0:
+                data_dict.update(
+                    {
+                        "event_{}_t={}".format(event_number, event_time): np.full(
+                            len(t_array), np.nan
+                        )
+                    }
+                )
+
+            elif np.any(np.isnan(data_slice)):
+                if nan_policy == "error":
+                    # raise exception
+                    raise Exception("NaN value in data slice, at event time {}".format(event_time))
+                elif nan_policy == "exclude":
+                    # exclude this event
+                    data_dict.update(
+                        {
+                            "event_{}_t={}".format(event_number, event_time): np.full(
+                                len(t_array), np.nan
+                            )
+                        }
+                    )
+                else:
+                    # Interpolate over NaNs
+                    x_data = data_slice[~data_slice.isnull()]
+                    data_slice[:] = np.interp(
+                        data_slice.index.values, x_data.index.values, x_data.values
+                    )
+
+                    # Add to data dict as normal
+                    data_dict.update(
+                        {
+                            "event_{}_t={}".format(event_number, event_time): np.interp(
+                                data_dict["time"],
+                                data_slice.index - event_time,
+                                data_slice.values,
+                            )
+                        }
+                    )
+            else:
+                # update our dictionary to have a new key defined as
+                # 'event_{EVENT NUMBER}_t={EVENT TIME}' and
+                # a value that includes an array that represents the
+                # sliced data around the current event, interpolated
+                # on the linearly spaced time array
+                data_dict.update(
+                    {
+                        "event_{}_t={}".format(event_number, event_time): np.interp(
+                            data_dict["time"],
+                            data_slice.index - event_time,
+                            data_slice.values,
+                        )
+                    }
+                )
+
+        # define a wide dataframe as a dataframe of the above compiled dictionary  # NOQA E501
+        wide_etr = pd.DataFrame(data_dict)
+
+    # if interpolate is False,
+    # we will calculate a common timebase and
+    # shift every response onto that timebase
+    else:
+        (
+            event_indices,
+            start_ind_offset,
+            end_ind_offset,
+            trace_timebase,
+        ) = slice_inds_and_offsets(  # NOQA E501
+            np.array(data[t]),
+            np.array(event_times),
+            time_window=[t_start, t_end],
+            sampling_rate=None,
+            include_endpoint=True,
+        )
+        all_inds = event_indices + np.arange(start_ind_offset, end_ind_offset)[:, None]
+        wide_etr = (
+            pd.DataFrame(
+                data[y].values.T[all_inds],
+                index=trace_timebase,
+                columns=[
+                    "event_{}_t={}".format(event_index, event_time)
+                    for event_index, event_time in enumerate(event_times)
+                ],  # NOQA E501
+            )
+            .rename_axis(index="time")
+            .reset_index()
+        )
+
+    if output_format == "wide":
+        # return the wide dataframe if output_format is 'wide'
+        return wide_etr.set_index("time")
+    elif output_format == "tidy":
+        # if output format is 'tidy',
+        # transform the wide dataframe to tidy format
+        # first, melt the dataframe with the 'id_vars' column as "time"
+        tidy_etr = wide_etr.melt(id_vars="time")
+
+        # add an "event_number" column that contains the event number
+        tidy_etr["event_number"] = (
+            tidy_etr["variable"].map(lambda s: s.split("event_")[1].split("_")[0]).astype(int)
+        )
+
+        # add an "event_time" column that contains the event time ()
+        tidy_etr["event_time"] = tidy_etr["variable"].map(lambda s: s.split("t=")[1]).astype(float)
+
+        # drop the "variable" column, rename the "value" column
+        tidy_etr = tidy_etr.drop(columns=["variable"]).rename(columns={"value": y})
+        # return the tidy event triggered responses
+    if censor:
+        tidy_etr = censor_event_triggered_response(
+            tidy_etr, y, t_start, t_end, event_times, censor_times
+        )
+    return tidy_etr
+
+
+def censor_event_triggered_response(etr, y, t_start, t_end, event_times, censor_times=None):
+    """
+    censors the event triggered response by the immediately preceeding or
+    subsequent event times if that event time is within the (t_start, t_end)
+    time window
+
+    censored timepoints are replaced with NaN, so all data points are still present
+
+    etr: dataframe, event triggered response
+    y: column of the response variable to censor
+    t_start: start of event triggered response window
+    t_end: end of event triggered response window
+    censor: Boolean
+        if True, censor observations that take place after the next event time
+        if False, do not censor
+    censor_times: list or array or None
+        if None, and censor is True, then use event_times as the censor times
+        if times are provided, then these are the times at which ETR is censored
+    """
+
+    if censor_times is None:
+        # Compute when we should censor
+        diff = np.diff(event_times)
+        diff_backward = np.concatenate([[np.inf], diff])
+        diff_forward = np.concatenate([diff, [np.inf]])
+        backward_time = [-np.min([np.abs(t_start), x]) for x in diff_backward]
+        forward_time = [np.min([t_end, x]) for x in diff_forward]
+
+        # double check we have all events
+        assert len(event_times) == len(etr["event_number"].unique()), "event times missing"
+
+        # Censor trials
+        for index, time in enumerate(event_times):
+            vec = (etr["event_number"] == index) & (etr["time"] < backward_time[index])
+            etr.loc[vec, y] = np.nan
+            vec = (etr["event_number"] == index) & (etr["time"] > forward_time[index])
+            etr.loc[vec, y] = np.nan
+
+        return etr
+    else:
+        censor_times = np.sort(censor_times)
+        backward_time = []
+        forward_time = []
+        for e in event_times:
+            before = censor_times[censor_times < e]
+            if len(before) == 0:
+                backward_time.append(t_start)
+            else:
+                backward_time.append(np.max([t_start, before[-1] - e]))
+            after = censor_times[censor_times > e]
+            if len(after) == 0:
+                forward_time.append(t_end)
+            else:
+                forward_time.append(np.min([t_end, after[0] - e]))
+
+        # Censor trials
+        for index, time in enumerate(event_times):
+            vec = (etr["event_number"] == index) & (etr["time"] < backward_time[index])
+            etr.loc[vec, y] = np.nan
+            vec = (etr["event_number"] == index) & (etr["time"] > forward_time[index])
+            etr.loc[vec, y] = np.nan
+
+        return etr

aind_dynamic_foraging_data_utils/nwb_utils.py

@@ -0,0 +1,599 @@
+"""
+Utility functions for processing dynamic foraging data.
+    load_nwb_from_filename
+    unpack_metadata
+    create_single_df_session_inner
+    create_df_session
+    create_single_df_session
+    create_df_trials
+    create_events_df
+    create_fib_df
+"""
+
+import os
+import re
+
+import numpy as np
+import pandas as pd
+from pynwb import NWBHDF5IO
+from hdmf_zarr import NWBZarrIO
+
+
+def load_nwb_from_filename(filename):
+    """
+    Load NWB from file, checking for HDF5 or Zarr
+    if filename is not a string, then return the input, assuming its the NWB file already
+    """
+
+    if type(filename) is str:
+        if os.path.isdir(filename):
+            io = NWBZarrIO(filename, mode="r")
+            nwb = io.read()
+            return nwb
+        elif os.path.isfile(filename):
+            io = NWBHDF5IO(filename, mode="r")
+            nwb = io.read()
+            return nwb
+        else:
+            raise FileNotFoundError(filename)
+    else:
+        # Assuming its already an NWB
+        return filename
+
+
+def unpack_metadata(nwb):
+    """
+    Unpacks metadata as a dictionary attribute, instead of a Dynamic
+    table nested inside a dictionary
+    """
+    nwb.metadata = nwb.scratch["metadata"].to_dataframe().iloc[0].to_dict()
+
+
+def create_single_df_session_inner(nwb):
+    """
+    given a nwb file, output a tidy dataframe
+    """
+    df_trials = nwb.trials.to_dataframe()
+
+    # Reformat data
+    choice_history = df_trials.animal_response.map({0: 0, 1: 1, 2: np.nan}).values
+    reward_history = np.vstack([df_trials.rewarded_historyL, df_trials.rewarded_historyR])
+
+    # -- Session-based table --
+    # - Meta data -
+    session_start_time_from_meta = nwb.session_start_time
+    session_date_from_meta = session_start_time_from_meta.strftime("%Y-%m-%d")
+    subject_id_from_meta = nwb.subject.subject_id
+
+    # TODO, should reprocess old files, and remove this logic
+    if "behavior" in nwb.session_id:
+        splits = nwb.session_id.split("_")
+        subject_id = splits[1]
+        session_date = splits[2]
+        nwb_suffix = splits[3].replace("-", "")
+    else:
+        old_re = re.match(
+            r"(?P<subject_id>\d+)_(?P<date>\d{4}-\d{2}-\d{2})(?:_(?P<n>\d+))?\.json",
+            nwb.session_id,
+        )
+
+        if old_re is not None:
+            # If there are more than one "bonsai sessions" (the trainer clicked "Save" button in the
+            # GUI more than once) in a certain day,
+            # parse nwb_suffix from the file name (0, 1, 2, ...)
+            subject_id, session_date, nwb_suffix = old_re.groups()
+            nwb_suffix = int(nwb_suffix) if nwb_suffix is not None else 0
+        else:
+            # After https://github.com/AllenNeuralDynamics/dynamic-foraging-task/commit/
+            # 62d0e9e2bb9b47a8efe8ecb91da9653381a5f551,
+            # the suffix becomes the session start time. Therefore, I use HHMMSS as the nwb suffix,
+            # which still keeps the order as before.
+
+            # Typical situation for multiple bonsai sessions per day is that the
+            # RAs pressed more than once "Save" button but only started the session once.
+            # Therefore, I should generate nwb_suffix from the bonsai file name
+            # instead of session_start_time.
+            subject_id, session_date, session_json_time = re.match(
+                r"(?P<subject_id>\d+)_(?P<date>\d{4}-\d{2}-\d{2})(?:_(?P<time>.*))\.json",
+                nwb.session_id,
+            ).groups()
+            nwb_suffix = int(session_json_time.replace("-", ""))
+
+    # Ad-hoc bug fixes for some mistyped mouse ID
+    if subject_id in ("689727"):
+        subject_id_from_meta = subject_id
+
+    assert subject_id == subject_id_from_meta, (
+        f"Subject name from the metadata ({subject_id_from_meta}) does not match "
+        f"that from json name ({subject_id})!!"
+    )
+    assert session_date == session_date_from_meta, (
+        f"Session date from the metadata ({session_date_from_meta}) does not match "
+        f"that from json name ({session_date})!!"
+    )
+
+    session_index = pd.MultiIndex.from_tuples(
+        [(subject_id, session_date, nwb_suffix)],
+        names=["subject_id", "session_date", "nwb_suffix"],
+    )
+
+    # Parse meta info
+    # TODO: when generating nwb, put meta info in nwb.scratch and get rid of the regular expression
+    # This could be significantly cleaned up based on new metadata format
+    # But im making it consistent for now
+    # TODO, should reprocess old files, and remove this logic
+    if "behavior" not in nwb.session_id:
+        extra_water, rig = re.search(
+            r"Give extra water.*:(\d*(?:\.\d+)?)? .*?(?:tower|box):(.*)?",
+            nwb.session_description,
+        ).groups()
+        weight_after_session = re.search(
+            r"Weight after.*:(\d*(?:\.\d+)?)?", nwb.subject.description
+        ).groups()[0]
+
+        extra_water = float(extra_water) if extra_water != "" else 0
+        weight_after_session = float(weight_after_session) if weight_after_session != "" else np.nan
+        weight_before_session = float(nwb.subject.weight) if nwb.subject.weight != "" else np.nan
+        user_name = nwb.experimenter[0]
+    else:
+        rig = nwb.scratch["metadata"].box[0]
+        user_name = nwb.experimenter
+        weight_after_session = nwb.scratch["metadata"].weight_after[0]
+        water_during_session = nwb.scratch["metadata"].water_in_session_total[0]
+        weight_before_session = weight_after_session - water_during_session
+        extra_water = nwb.scratch["metadata"].water_in_session_manual[0]
+
+    dict_meta = {
+        "rig": rig,
+        "user_name": user_name,
+        "experiment_description": nwb.experiment_description,
+        "task": nwb.protocol,
+        "session_start_time": session_start_time_from_meta,
+        "weight_before_session": weight_before_session,
+        "weight_after_session": weight_after_session,
+        "water_during_session": weight_after_session - weight_before_session,
+        "water_extra": extra_water,
+    }
+
+    df_session = pd.DataFrame(
+        dict_meta,
+        index=session_index,
+    )
+    # Use hierarchical index (type = {'metadata', 'session_stats'}, variable = {...}, etc.)
+    df_session.columns = pd.MultiIndex.from_product(
+        [["metadata"], dict_meta.keys()], names=["type", "variable"]
+    )
+
+    # - Compute session-level stats -
+    # TODO: Ideally, all these simple stats could be computed in the GUI, and
+    # the GUI sends a copy to the meta session.json file and to the nwb file as well.
+
+    total_trials = len(df_trials)
+    finished_trials = np.sum(~np.isnan(choice_history))
+    reward_trials = np.sum(reward_history)
+
+    reward_rate = reward_trials / finished_trials
+
+    # TODO: add more stats
+    # See code here: https://github.com/AllenNeuralDynamics/map-ephys/blob/
+    # 7a06a5178cc621638d849457abb003151f7234ea/pipeline/foraging_analysis.py#L70C8-L70C8
+    # early_lick_ratio =
+    # double_dipping_ratio =
+    # block_num
+    # mean_block_length
+    # mean_reward_sum
+    # mean_reward_contrast
+    # autowater_num
+    # autowater_ratio
+    #
+    # mean_iti
+    # mean_reward_sum
+    # mean_reward_contrast
+    # ...
+
+    # foraging_eff_func = (
+    #    foraging_eff_baiting if "bait" in nwb.protocol.lower() else foraging_eff_no_baiting
+    # )
+    # foraging_eff, foraging_eff_random_seed = foraging_eff_func(
+    #    reward_rate, p_reward[LEFT, :], p_reward[RIGHT, :]
+    # )
+
+    # -- Add session stats here --
+    dict_session_stat = {
+        "total_trials": total_trials,
+        "finished_trials": finished_trials,
+        "finished_rate": finished_trials / total_trials,
+        "ignore_rate": np.sum(np.isnan(choice_history)) / total_trials,
+        "reward_trials": reward_trials,
+        "reward_rate": reward_rate,
+        # TODO: add more stats here
+    }
+
+    # Generate df_session_stat
+    df_session_stat = pd.DataFrame(dict_session_stat, index=session_index)
+    df_session_stat.columns = pd.MultiIndex.from_product(
+        [["session_stats"], dict_session_stat.keys()],
+        names=["type", "variable"],
+    )
+
+    # -- Add automatic training --
+    if "auto_train_engaged" in df_trials.columns:
+        df_session["auto_train", "curriculum_name"] = (
+            np.nan
+            if df_trials.auto_train_curriculum_name.mode()[0] == "none"
+            else df_trials.auto_train_curriculum_name.mode()[0]
+        )
+        df_session["auto_train", "curriculum_version"] = (
+            np.nan
+            if df_trials.auto_train_curriculum_version.mode()[0] == "none"
+            else df_trials.auto_train_curriculum_version.mode()[0]
+        )
+        df_session["auto_train", "curriculum_schema_version"] = (
+            np.nan
+            if df_trials.auto_train_curriculum_schema_version.mode()[0] == "none"
+            else df_trials.auto_train_curriculum_schema_version.mode()[0]
+        )
+        df_session["auto_train", "current_stage_actual"] = (
+            np.nan
+            if df_trials.auto_train_stage.mode()[0] == "none"
+            else df_trials.auto_train_stage.mode()[0]
+        )
+        df_session["auto_train", "if_overriden_by_trainer"] = (
+            np.nan
+            if all(df_trials.auto_train_stage_overridden.isna())
+            else df_trials.auto_train_stage_overridden.mode()[0]
+        )
+
+        # Add a flag to indicate whether any of the auto train settings were changed
+        # during the training
+        df_session["auto_train", "if_consistent_within_session"] = (
+            len(df_trials.groupby([col for col in df_trials.columns if "auto_train" in col])) == 1
+        )
+    else:
+        for field in [
+            "curriculum_name",
+            "curriculum_version",
+            "curriculum_schema_version",
+            "current_stage_actual",
+            "if_overriden_by_trainer",
+        ]:
+            df_session["auto_train", field] = None
+
+    # -- Merge to df_session --
+    df_session = pd.concat([df_session, df_session_stat], axis=1)
+    return df_session
+
+
+def create_df_session(nwb_filename):
+    """
+    Creates a dataframe where each row is a session
+    nwb_filename can be either a single nwb file, a single filepath
+    or a list of nwb files, or a list of nwb filepaths
+    """
+    if (type(nwb_filename) is not str) and (hasattr(nwb_filename, "__iter__")):
+        dfs = []
+        for nwb_file in nwb_filename:
+            dfs.append(create_single_df_session(nwb_file))
+        return pd.concat(dfs)
+    else:
+        return create_single_df_session(nwb_filename)
+
+
+# % Process nwb and create df_session for every single session
+def create_single_df_session(nwb_filename):
+    """
+    create a dataframe for a single session
+    """
+    nwb = load_nwb_from_filename(nwb_filename)
+
+    df_session = create_single_df_session_inner(nwb)
+
+    df_session.columns = df_session.columns.droplevel("type")
+    df_session = df_session.reset_index()
+    df_session["ses_idx"] = (
+        df_session["subject_id"].values + "_" + df_session["session_date"].values
+    )
+    df_session = df_session.rename(columns={"variable": "session_num"})
+    return df_session
+
+
+def create_df_trials(nwb_filename):
+    """
+    Process nwb and create df_trials for every single session
+    """
+    nwb = load_nwb_from_filename(nwb_filename)
+
+    key_from_acq = [
+        "left_lick_time",
+        "right_lick_time",
+        "left_reward_delivery_time",
+        "right_reward_delivery_time",
+        "FIP_falling_time",
+        "FIP_rising_time",
+    ]
+
+    # Parse subject and session_date
+    if nwb.session_id.startswith("behavior") or nwb.session_id.startswith("FIP"):
+        splits = nwb.session_id.split("_")
+        subject_id = splits[1]
+        session_date = splits[2]
+    else:
+        splits = nwb.session_id.split("_")
+        subject_id = splits[0]
+        session_date = splits[1]
+
+    ses_idx = subject_id + "_" + session_date
+
+    df_ses_trials = nwb.trials.to_dataframe().reset_index()
+    df_ses_trials = df_ses_trials.rename(columns={"id": "trial"})
+    df_ses_trials["ses_idx"] = ses_idx
+
+    # Adjust all times relative to start of the first trial
+    t0 = df_ses_trials.start_time[0]
+    skip_cols = ["right_valve_open_time", "left_valve_open_time"]
+    for col in df_ses_trials.columns:
+        if ("time" in col) and (col not in skip_cols):
+            df_ses_trials[col + "_absolute"] = df_ses_trials[col] - t0
+
+    # Adjust for gaps in trial start/stop, and use the last stop time
+    last_stop = df_ses_trials.iloc[-1]["stop_time_absolute"]
+    df_ses_trials["stop_time_absolute"] = df_ses_trials["start_time_absolute"].shift(
+        -1, fill_value=last_stop
+    )
+
+    # Adjust times relative to go cue
+    for col in df_ses_trials.columns:
+        if (
+            ("time" in col)
+            and ("time_absolute" not in col)
+            and (col != "goCue_start_time")
+            and (col not in skip_cols)
+        ):
+            df_ses_trials.loc[:, col] = (
+                df_ses_trials[col].values - df_ses_trials["goCue_start_time"].values
+            )
+    df_ses_trials["goCue_start_time"] = 0.0
+
+    # Adjust event times relative to trial
+    events_ses = {key: nwb.acquisition[key].timestamps[:] - t0 for key in key_from_acq}
+    for event in [
+        "left_lick_time",
+        "right_lick_time",
+        "left_reward_delivery_time",
+        "right_reward_delivery_time",
+    ]:
+        event_times = events_ses[event]
+        df_ses_trials[event] = df_ses_trials.apply(
+            lambda x: np.round(
+                event_times[
+                    (event_times > (x["goCue_start_time"] + x["goCue_start_time_absolute"]))
+                    & (event_times < (x["stop_time"] + x["goCue_start_time_absolute"]))
+                ]
+                - x["goCue_start_time_absolute"],
+                4,
+            ),
+            axis=1,
+        )
+
+    # Compute time of reward for each trial
+    df_ses_trials["reward_time"] = df_ses_trials.apply(
+        lambda x: np.nanmin(
+            np.concatenate(
+                [
+                    [np.nan],
+                    x["right_reward_delivery_time"],
+                    x["left_reward_delivery_time"],
+                ]
+            )
+        ),
+        axis=1,
+    )
+    df_ses_trials["reward_time_absolute"] = (
+        df_ses_trials["reward_time"] + df_ses_trials["goCue_start_time_absolute"]
+    )
+
+    # Compute time of choice for each trials
+    df_ses_trials["choice_time"] = df_ses_trials.apply(
+        lambda x: np.nanmin(np.concatenate([[np.nan], x["right_lick_time"], x["left_lick_time"]])),
+        axis=1,
+    )
+    df_ses_trials["choice_time_absolute"] = (
+        df_ses_trials["choice_time"] + df_ses_trials["goCue_start_time_absolute"]
+    )
+
+    # Compute boolean of whether animal was rewarded
+    df_ses_trials["reward"] = df_ses_trials.rewarded_historyR.astype(
+        int
+    ) | df_ses_trials.rewarded_historyL.astype(int)
+
+    # Drop columns
+    df_ses_trials = df_ses_trials.drop(
+        columns=[
+            "left_lick_time",
+            "right_lick_time",
+            "left_reward_delivery_time",
+            "right_reward_delivery_time",
+        ]
+    )
+    return df_ses_trials
+
+
+def create_events_df(nwb_filename, adjust_time=True):
+    """
+    returns a tidy dataframe of the events in the nwb file
+
+    adjust_time (bool), set time of first goCue to t=0
+    """
+
+    nwb = load_nwb_from_filename(nwb_filename)
+
+    # Build list of all event types in acqusition, ignore FIP events
+    event_types = set(nwb.acquisition.keys())
+    ignore_types = set(
+        [
+            "FIP_falling_time",
+            "FIP_rising_time",
+            "G_1",
+            "G_2",
+            "Iso_1",
+            "R_1",
+            "R_2",
+            "Iso_2",
+            "G_1_dff-bright",
+            "G_2_dff-bright",
+            "Iso_1_dff-bright",
+            "R_1_dff-bright",
+            "R_2_dff-bright",
+            "Iso_2_dff-bright",
+            "G_1_dff-exp",
+            "G_2_dff-exp",
+            "Iso_1_dff-exp",
+            "R_1_dff-exp",
+            "R_2_dff-exp",
+            "Iso_2_dff-exp",
+            "G_1_dff-poly",
+            "G_2_dff-poly",
+            "Iso_1_dff-poly",
+            "R_1_dff-poly",
+            "R_2_dff-poly",
+            "Iso_2_dff-poly",
+        ]
+    )
+    event_types -= ignore_types
+
+    # Determine time 0
+    t0 = nwb.trials.start_time[0]
+
+    # Iterate over event types and build a dataframe of each
+    events = []
+    for e in event_types:
+        # For each event, get timestamps, data, and label
+        stamps = nwb.acquisition[e].timestamps[:]
+        data = nwb.acquisition[e].data[:]
+        labels = [e] * len(data)
+        if adjust_time:
+            stamps = stamps - t0
+        df = pd.DataFrame({"timestamps": stamps, "data": data, "event": labels})
+        events.append(df)
+
+    # Add keys from trials table
+    # I don't like hardcoding dynamic foraging specific things here.
+    # I think these keys should be added to the stimulus field of the nwb
+    trial_events = ["goCue_start_time"]
+    for e in trial_events:
+        stamps = nwb.trials[:][e].values
+        labels = [e] * len(stamps)
+        if adjust_time:
+            stamps = stamps - t0
+        df = pd.DataFrame({"timestamps": stamps, "event": labels})
+        events.append(df)
+
+    # Build dataframe by concatenating each event
+    df = pd.concat(events)
+    df = df.sort_values(by="timestamps")
+    df = df.dropna(subset="timestamps").reset_index(drop=True)
+
+    # Add trial index for each event
+    trial_starts = nwb.trials.start_time[:] - nwb.trials.start_time[0]
+    last_stop = nwb.trials.stop_time[-1] - nwb.trials.start_time[0]
+    trial_index = []
+    for index, e in df.iterrows():
+        starts = np.where(e.timestamps > trial_starts)[0]
+        if len(starts) == 0:
+            trial_index.append(-1)
+        elif e.timestamps > last_stop:
+            trial_index.append(len(trial_starts))
+        else:
+            trial_index.append(starts[-1])
+    df["trial"] = trial_index
+
+    return df
+
+
+def create_fib_df(nwb_filename, tidy=True, adjust_time=True):
+    """
+    returns a dataframe of the FIB data in the nwb file
+    if tidy, return a tidy dataframe
+    if not tidy, return pivoted by timestamp
+
+    adjust_time (bool), set time of first goCue to t=0
+    """
+
+    nwb = load_nwb_from_filename(nwb_filename)
+
+    # Build list of all FIB events in NWB file
+    nwb_types = set(nwb.acquisition.keys())
+    event_types = set(
+        [
+            "FIP_falling_time",
+            "FIP_rising_time",
+            "G_1",
+            "G_2",
+            "Iso_1",
+            "R_1",
+            "R_2",
+            "Iso_2",
+            "G_1_dff-bright",
+            "G_2_dff-bright",
+            "Iso_1_dff-bright",
+            "R_1_dff-bright",
+            "R_2_dff-bright",
+            "Iso_2_dff-bright",
+            "G_1_dff-exp",
+            "G_2_dff-exp",
+            "Iso_1_dff-exp",
+            "R_1_dff-exp",
+            "R_2_dff-exp",
+            "Iso_2_dff-exp",
+            "G_1_dff-poly",
+            "G_2_dff-poly",
+            "Iso_1_dff-poly",
+            "R_1_dff-poly",
+            "R_2_dff-poly",
+            "Iso_2_dff-poly",
+        ]
+    )
+    event_types = event_types.intersection(nwb_types)
+
+    # If no FIB data available
+    if len(event_types) == 0:
+        return None
+
+    # Determine time 0
+    t0 = nwb.trials.start_time[0]
+
+    # Iterate over event types and build a dataframe of each
+    events = []
+    for e in event_types:
+        # For each event, get timestamps, data, and label
+        stamps = nwb.acquisition[e].timestamps[:]
+        data = nwb.acquisition[e].data[:]
+        labels = [e] * len(data)
+        if adjust_time:
+            stamps = stamps - t0
+        df = pd.DataFrame({"timestamps": stamps, "data": data, "event": labels})
+        events.append(df)
+
+    # Build dataframe by concatenating each event
+    df = pd.concat(events)
+    df = df.sort_values(by="timestamps")
+    df = df.dropna(subset="timestamps").reset_index(drop=True)
+
+    # Add session_idx with subject ID and session date info - JL
+    if nwb.session_id.startswith("behavior") or nwb.session_id.startswith("FIP"):
+        splits = nwb.session_id.split("_")
+        subject_id = splits[1]
+        session_date = splits[2]
+    else:
+        splits = nwb.session_id.split("_")
+        subject_id = splits[0]
+        session_date = splits[1]
+    ses_idx = subject_id + "_" + session_date
+    df["ses_idx"] = ses_idx
+
+    # pivot table based on timestamps
+    if not tidy:
+        df_pivoted = pd.pivot(df, index="timestamps", columns=["event"], values="data")
+        return df_pivoted
+    else:
+        return df

aind_dynamic_foraging_data_utils-0.1.0.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Allen Institute for Neural Dynamics
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

aind_dynamic_foraging_data_utils-0.1.0.dist-info/METADATA

@@ -0,0 +1,125 @@
+Metadata-Version: 2.1
+Name: aind-dynamic-foraging-data-utils
+Version: 0.1.0
+Summary: Functions to help with postprocessing dynamic foraging data.
+Author: Allen Institute for Neural Dynamics
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.7
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: pandas
+Requires-Dist: numpy
+Requires-Dist: pynwb
+Requires-Dist: hdmf-zarr
+Provides-Extra: dev
+Requires-Dist: black ; extra == 'dev'
+Requires-Dist: coverage ; extra == 'dev'
+Requires-Dist: flake8 ; extra == 'dev'
+Requires-Dist: interrogate ; extra == 'dev'
+Requires-Dist: isort ; extra == 'dev'
+Requires-Dist: Sphinx ; extra == 'dev'
+Requires-Dist: furo ; extra == 'dev'
+
+# aind-dynamic-foraging-data-utils
+
+
+[![License](https://img.shields.io/badge/license-MIT-brightgreen)](LICENSE)
+![Code Style](https://img.shields.io/badge/code%20style-black-black)
+[![semantic-release: angular](https://img.shields.io/badge/semantic--release-angular-e10079?logo=semantic-release)](https://github.com/semantic-release/semantic-release)
+![Interrogate](https://img.shields.io/badge/interrogate-100.0%25-brightgreen)
+![Coverage](https://img.shields.io/badge/coverage-100%25-brightgreen?logo=codecov)
+![Python](https://img.shields.io/badge/python->=3.7-blue?logo=python)
+
+
+
+## Scope
+Purpose: Ingests NWB and spits out dataframes with the relevant information. Focused on dynamic foraging. Other tasks can branch and build task-specific utils.
+Inputs are nwbs, outputs are dataframes (tidy and not)
+Dependencies: xarray (includes numpy and pandas), scikit-learn (includes scipy), matplotlib
+
+
+
+
+## Installation
+To use the software, in the root directory, run
+```bash
+pip install -e .
+```
+
+To develop the code, run
+```bash
+pip install -e .[dev]
+```
+
+## Contributing
+
+### Linters and testing
+
+There are several libraries used to run linters, check documentation, and run tests.
+
+- Please test your changes using the **coverage** library, which will run the tests and log a coverage report:
+
+```bash
+coverage run -m unittest discover && coverage report
+```
+
+- Use **interrogate** to check that modules, methods, etc. have been documented thoroughly:
+
+```bash
+interrogate .
+```
+
+- Use **flake8** to check that code is up to standards (no unused imports, etc.):
+```bash
+flake8 .
+```
+
+- Use **black** to automatically format the code into PEP standards:
+```bash
+black .
+```
+
+- Use **isort** to automatically sort import statements:
+```bash
+isort .
+```
+
+### Pull requests
+
+For internal members, please create a branch. For external members, please fork the repository and open a pull request from the fork. We'll primarily use [Angular](https://github.com/angular/angular/blob/main/CONTRIBUTING.md#commit) style for commit messages. Roughly, they should follow the pattern:
+```text
+<type>(<scope>): <short summary>
+```
+
+where scope (optional) describes the packages affected by the code changes and type (mandatory) is one of:
+
+- **build**: Changes that affect build tools or external dependencies (example scopes: pyproject.toml, setup.py)
+- **ci**: Changes to our CI configuration files and scripts (examples: .github/workflows/ci.yml)
+- **docs**: Documentation only changes
+- **feat**: A new feature
+- **fix**: A bugfix
+- **perf**: A code change that improves performance
+- **refactor**: A code change that neither fixes a bug nor adds a feature
+- **test**: Adding missing tests or correcting existing tests
+
+### Semantic Release
+
+The table below, from [semantic release](https://github.com/semantic-release/semantic-release), shows which commit message gets you which release type when `semantic-release` runs (using the default configuration):
+
+| Commit message                                                                                                                                                                                   | Release type                                                                                                    |
+| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------- |
+| `fix(pencil): stop graphite breaking when too much pressure applied`                                                                                                                             | ~~Patch~~ Fix Release, Default release                                                                          |
+| `feat(pencil): add 'graphiteWidth' option`                                                                                                                                                       | ~~Minor~~ Feature Release                                                                                       |
+| `perf(pencil): remove graphiteWidth option`<br><br>`BREAKING CHANGE: The graphiteWidth option has been removed.`<br>`The default graphite width of 10mm is always used for performance reasons.` | ~~Major~~ Breaking Release <br /> (Note that the `BREAKING CHANGE: ` token must be in the footer of the commit) |
+
+### Documentation
+To generate the rst files source files for documentation, run
+```bash
+sphinx-apidoc -o doc_template/source/ src 
+```
+Then to create the documentation HTML files, run
+```bash
+sphinx-build -b html doc_template/source/ doc_template/build/html
+```
+More info on sphinx installation can be found [here](https://www.sphinx-doc.org/en/master/usage/installation.html).

aind_dynamic_foraging_data_utils-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+aind_dynamic_foraging_data_utils/__init__.py,sha256=GumePyX4MIeArBCTdlS-9iWt3heecE1BNf4BjIH2OyI,42
+aind_dynamic_foraging_data_utils/alignment.py,sha256=ccTZwDtQQFSgWTS73P0X-4tDLMOR3o4nkh_rKX_BQwY,21280
+aind_dynamic_foraging_data_utils/nwb_utils.py,sha256=9P711GaJH7IFksC_6CSQj3HiZpM0nLafStT1yNCEv_Q,20575
+aind_dynamic_foraging_data_utils-0.1.0.dist-info/LICENSE,sha256=U0Y7B3gZJHXpjJVLgTQjM8e_c8w4JJpLgGhIdsoFR1Y,1092
+aind_dynamic_foraging_data_utils-0.1.0.dist-info/METADATA,sha256=bnjgQlSsFHqsHt9dDrOWNOuXc5BKdC0xndyWkrsmcnE,5610
+aind_dynamic_foraging_data_utils-0.1.0.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+aind_dynamic_foraging_data_utils-0.1.0.dist-info/top_level.txt,sha256=cxZtQWcebZ8eA0nZpYUavHj14nMd6BtwyrO0DYfpoWg,33
+aind_dynamic_foraging_data_utils-0.1.0.dist-info/RECORD,,

aind_dynamic_foraging_data_utils-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.3.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

aind_dynamic_foraging_data_utils-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+aind_dynamic_foraging_data_utils