Source code for allensdk.internal.brain_observatory.util.multi_session_utils
"""Utilities for accessing data across multiple sessions"""
import os
from multiprocessing import Pool
from typing import List, Optional, Set, Callable
from tqdm import tqdm
from allensdk.brain_observatory.behavior.data_files import BehaviorStimulusFile
from allensdk.brain_observatory.behavior.data_files.stimulus_file import \
MalformedStimulusFileError
from allensdk.brain_observatory.behavior.data_objects import BehaviorSessionId
from allensdk.brain_observatory.behavior.data_objects.metadata\
.behavior_metadata.behavior_metadata import \
BehaviorMetadata
from allensdk.brain_observatory.behavior.stimulus_processing import \
get_image_names
from allensdk.internal.api import PostgresQueryMixin
[docs]
def get_images_shown(
behavior_session_ids: List[int],
lims_engine: PostgresQueryMixin,
n_workers: Optional[int] = None
) -> Set[str]:
"""
Gets images shown to mouse during `behavior_session_ids`
Parameters
----------
behavior_session_ids
lims_engine
n_workers
Returns
-------
Set[str]: set of image names shown to mouse in behavior_session_ids
"""
if n_workers is None or n_workers > 1:
# multiprocessing
image_names = multiprocessing_helper(
target=_get_image_names,
behavior_session_ids=behavior_session_ids,
lims_engine=lims_engine,
progress_bar_title='Reading image_names from pkl file',
n_workers=n_workers
)
else:
# single process
image_names = [_get_image_names([behavior_session_id, lims_engine])
for behavior_session_id in behavior_session_ids]
res = set()
for image_name_set in image_names:
for image_name in image_name_set:
res.add(image_name)
return res
[docs]
def multiprocessing_helper(
target: Callable,
progress_bar_title: str,
behavior_session_ids: List[int],
lims_engine: PostgresQueryMixin,
n_workers: Optional[int] = None
):
if n_workers is None:
n_workers = os.cpu_count()
with Pool(n_workers) as p:
res = list(tqdm(
p.imap(target,
zip(
behavior_session_ids,
[lims_engine] * len(behavior_session_ids))
),
total=len(behavior_session_ids),
desc=progress_bar_title))
return res
def _get_session_metadata(*args) -> Optional[BehaviorMetadata]:
"""
Helper function to get session metadata
"""
behavior_session_id, db_conn = args[0]
try:
meta = BehaviorMetadata.from_lims(
behavior_session_id=BehaviorSessionId(behavior_session_id),
lims_db=db_conn)
except MalformedStimulusFileError:
meta = None
return meta
def _get_image_names(*args) -> Set[str]:
"""
Helper function to get image names from behavior stimulus file
"""
behavior_session_id, db_conn = args[0]
behavior_stimulus_file = BehaviorStimulusFile.from_lims(
behavior_session_id=behavior_session_id, db=db_conn)
image_names = get_image_names(
behavior_stimulus_file=behavior_stimulus_file)
return image_names
[docs]
def remove_invalid_sessions(
behavior_sessions: List[BehaviorMetadata],
remove_pretest_sessions: bool = True,
remove_sessions_after_mouse_death_date: bool = True,
remove_aborted_sessions: bool = True,
expected_training_duration: int = 15 * 60,
expected_duration: int = 60 * 60
) -> List[BehaviorMetadata]:
"""
Removes any invalid sessions from `behavior_sessions`
Parameters
----------
behavior_sessions:
List of behavior session metadata
remove_pretest_sessions
Remove any "pretest" session
remove_sessions_after_mouse_death_date
Remove any sessions mistakenly entered that fall after mouse death date
Sessions were loaded into LIMS with the wrong donor_id,
causing there to be sessions associated with some mice
that occur after those mice's recorded death dates. Our
assumption is that the error is with the donor_id rather
than the death date, so we can correct it by filtering
out any sessions that occur on mice that are supposed
to be dead.
remove_aborted_sessions
Remove aborted sessions
expected_training_duration
Expected duration for TRAINING_0 session in seconds
expected_duration
Expected duration for all sessions except TRAINING_0 in seconds
Returns
-------
List[BehaviorMetadata]:
list of behavior sessions with invalid sessions
removed
"""
if remove_pretest_sessions:
behavior_sessions = [x for x in behavior_sessions if not x.is_pretest]
if remove_sessions_after_mouse_death_date:
behavior_sessions = [
x for x in behavior_sessions
if (x.subject_metadata.get_death_date() is None or
x.date_of_acquisition <= x.subject_metadata.get_death_date())]
if remove_aborted_sessions:
training_sessions = \
[x for x in behavior_sessions if x.is_training and
x.get_session_duration() > expected_training_duration]
nontraining_sessions = \
[x for x in behavior_sessions if not x.is_training and
x.get_session_duration() > expected_duration]
behavior_sessions = training_sessions + nontraining_sessions
return behavior_sessions