From 0606add21a1f6e9f4abb72c68a01bd108c1a71cc Mon Sep 17 00:00:00 2001 From: Eric Denovellis Date: Sun, 9 Jun 2024 11:56:26 -0700 Subject: [PATCH] Refactor tests --- tests/test_unit/test_nwb.py | 315 ++++++++++++++++++------------------ 1 file changed, 161 insertions(+), 154 deletions(-) diff --git a/tests/test_unit/test_nwb.py b/tests/test_unit/test_nwb.py index 7319ea48..060c2e99 100644 --- a/tests/test_unit/test_nwb.py +++ b/tests/test_unit/test_nwb.py @@ -1,8 +1,10 @@ +import datetime + import ndx_pose import numpy as np -import pynwb -import pytest -import xarray as xr +from ndx_pose import PoseEstimation, PoseEstimationSeries, Skeleton, Skeletons +from pynwb import NWBHDF5IO, NWBFile +from pynwb.file import Subject from movement import sample_data from movement.io.nwb import ( @@ -19,8 +21,8 @@ def test_create_pose_and_skeleton_objects(): # Call the function pose_estimation, skeletons = _create_pose_and_skeleton_objects( - ds, - subject="subject1", + ds.sel(individuals="individual1"), + subject="individual1", pose_estimation_series_kwargs=None, pose_estimation_kwargs=None, skeleton_kwargs=None, @@ -34,43 +36,59 @@ def test_create_pose_and_skeleton_objects(): assert len(pose_estimation) == 1 # Assert the length of pose_estimation_series list - assert len(pose_estimation[0].pose_estimation_series) == 2 + assert len(pose_estimation[0].pose_estimation_series) == 12 # Assert the name of the first PoseEstimationSeries - assert pose_estimation[0].pose_estimation_series[0].name == "keypoint1" - - # Assert the name of the second PoseEstimationSeries - assert pose_estimation[0].pose_estimation_series[1].name == "keypoint2" + assert "snout" in pose_estimation[0].pose_estimation_series # Assert the name of the Skeleton - assert skeletons.skeletons[0].name == "subject1_skeleton" + assert "individual1_skeleton" in skeletons.skeletons + + +def create_test_pose_estimation_series( + n_time=100, n_dims=2, keypoint="front_left_paw" +): + data = np.random.rand( + n_time, n_dims + ) # num_frames x (x, y) but can be (x, y, z) + timestamps = np.linspace(0, 10, num=n_time) # a timestamp for every frame + confidence = np.ones((n_time,)) # a confidence value for every frame + reference_frame = "(0,0,0) corresponds to ..." + confidence_definition = "Softmax output of the deep neural network." + + return PoseEstimationSeries( + name=keypoint, + description="Marker placed around fingers of front left paw.", + data=data, + unit="pixels", + reference_frame=reference_frame, + timestamps=timestamps, + confidence=confidence, + confidence_definition=confidence_definition, + ) def test__convert_pose_estimation_series(): # Create a sample PoseEstimationSeries object - pose_estimation_series = ndx_pose.PoseEstimationSeries( - name="keypoint1", - data=np.random.rand(10, 3), - confidence=np.random.rand(10), - unit="pixels", - timestamps=np.arange(10), + pose_estimation_series = create_test_pose_estimation_series( + n_time=100, n_dims=2, keypoint="front_left_paw" ) # Call the function movement_dataset = _convert_pose_estimation_series( pose_estimation_series, - keypoint="keypoint1", - subject_name="subject1", + keypoint="leftear", + subject_name="individual1", source_software="software1", source_file="file1", ) # Assert the dimensions of the movement dataset - assert movement_dataset.dims == { - "time": 10, + assert movement_dataset.sizes == { + "time": 100, "individuals": 1, "keypoints": 1, - "space": 3, + "space": 2, } # Assert the values of the position variable @@ -92,155 +110,144 @@ def test__convert_pose_estimation_series(): "source_software": "software1", "source_file": "file1", } + pose_estimation_series = create_test_pose_estimation_series( + n_time=50, n_dims=3, keypoint="front_left_paw" + ) + + # Assert the dimensions of the movement dataset + assert movement_dataset.sizes == { + "time": 50, + "individuals": 1, + "keypoints": 1, + "space": 3, + } def test_add_movement_dataset_to_nwb_single_file(): - # Create a sample NWBFile - nwbfile = pynwb.NWBFile( - "session_description", "identifier", "session_start_time" - ) - # Create a sample movement dataset - movement_dataset = xr.Dataset( - { - "keypoints": (["keypoints"], ["keypoint1", "keypoint2"]), - "position": (["time", "keypoints"], [[1, 2], [3, 4]]), - "confidence": (["time", "keypoints"], [[0.9, 0.8], [0.7, 0.6]]), - "time": [0, 1], - "individuals": ["subject1"], - } + ds = sample_data.fetch_dataset("DLC_two-mice.predictions.csv") + session_start_time = datetime.datetime.now(datetime.timezone.utc) + nwbfile_individual1 = NWBFile( + session_description="session_description", + identifier="individual1", + session_start_time=session_start_time, + ) + add_movement_dataset_to_nwb( + nwbfile_individual1, ds.sel(individuals=["individual1"]) + ) + assert ( + "PoseEstimation" + in nwbfile_individual1.processing["behavior"].data_interfaces + ) + assert ( + "Skeletons" + in nwbfile_individual1.processing["behavior"].data_interfaces ) - # Call the function - add_movement_dataset_to_nwb(nwbfile, movement_dataset) - # Assert the presence of PoseEstimation and Skeletons in the NWBFile - assert "PoseEstimation" in nwbfile.processing["behavior"] - assert "Skeletons" in nwbfile.processing["behavior"] def test_add_movement_dataset_to_nwb_multiple_files(): - # Create sample NWBFiles - nwbfiles = [ - pynwb.NWBFile( - "session_description1", "identifier1", "session_start_time1" - ), - pynwb.NWBFile( - "session_description2", "identifier2", "session_start_time2" - ), - ] - # Create a sample movement dataset - movement_dataset = xr.Dataset( - { - "keypoints": (["keypoints"], ["keypoint1", "keypoint2"]), - "position": (["time", "keypoints"], [[1, 2], [3, 4]]), - "confidence": (["time", "keypoints"], [[0.9, 0.8], [0.7, 0.6]]), - "time": [0, 1], - "individuals": ["subject1", "subject2"], - } + ds = sample_data.fetch_dataset("DLC_two-mice.predictions.csv") + session_start_time = datetime.datetime.now(datetime.timezone.utc) + nwbfile_individual1 = NWBFile( + session_description="session_description", + identifier="individual1", + session_start_time=session_start_time, + ) + nwbfile_individual2 = NWBFile( + session_description="session_description", + identifier="individual2", + session_start_time=session_start_time, ) - # Call the function - add_movement_dataset_to_nwb(nwbfiles, movement_dataset) - # Assert the presence of PoseEstimation and Skeletons in each NWBFile - for nwbfile in nwbfiles: - assert "PoseEstimation" in nwbfile.processing["behavior"] - assert "Skeletons" in nwbfile.processing["behavior"] + nwbfiles = [nwbfile_individual1, nwbfile_individual2] + add_movement_dataset_to_nwb(nwbfiles, ds) -def test_convert_nwb_to_movement(): - # Create sample NWB files - nwb_filepaths = [ - "/path/to/file1.nwb", - "/path/to/file2.nwb", - "/path/to/file3.nwb", - ] - pose_estimation_series = { - "keypoint1": ndx_pose.PoseEstimationSeries( - name="keypoint1", - data=np.random.rand(10, 3), - confidence=np.random.rand(10), - unit="pixels", - timestamps=np.arange(10), - ), - "keypoint2": ndx_pose.PoseEstimationSeries( - name="keypoint2", - data=np.random.rand(10, 3), - confidence=np.random.rand(10), - unit="pixels", - timestamps=np.arange(10), - ), - } - # Mock the NWBHDF5IO read method - def mock_read(filepath): - nwbfile = pynwb.NWBFile( - "session_description", "identifier", "session_start_time" - ) - - pose_estimation = ndx_pose.PoseEstimation( - name="PoseEstimation", - pose_estimation_series=pose_estimation_series, - description="Pose estimation data", - source_software="software1", - skeleton=ndx_pose.Skeleton( - name="skeleton1", nodes=["node1", "node2"] - ), - ) - behavior_pm = pynwb.ProcessingModule( - name="behavior", description="Behavior data" - ) - behavior_pm.add(pose_estimation) - nwbfile.add_processing_module(behavior_pm) - return nwbfile +def create_test_pose_nwb(identifier="subject1", write_to_disk=False): + # initialize an NWBFile object + nwbfile = NWBFile( + session_description="session_description", + identifier=identifier, + session_start_time=datetime.datetime.now(datetime.timezone.utc), + ) - # Patch the NWBHDF5IO read method with the mock - with pytest.patch("pynwb.NWBHDF5IO.read", side_effect=mock_read): - # Call the function - movement_dataset = convert_nwb_to_movement(nwb_filepaths) + # add a subject to the NWB file + subject = Subject(subject_id=identifier, species="Mus musculus") + nwbfile.subject = subject - # Assert the dimensions of the movement dataset - assert movement_dataset.dims == { - "time": 10, - "individuals": 3, - "keypoints": 2, - "space": 3, - } + skeleton = Skeleton( + name="subject1_skeleton", + nodes=["front_left_paw", "body", "front_right_paw"], + edges=np.array([[0, 1], [1, 2]], dtype="uint8"), + subject=subject, + ) - # Assert the values of the position variable - np.testing.assert_array_equal( - movement_dataset["position"].values, - np.concatenate( - [ - pose_estimation_series["keypoint1"].data[ - :, np.newaxis, np.newaxis, : - ], - pose_estimation_series["keypoint2"].data[ - :, np.newaxis, np.newaxis, : - ], - ], - axis=1, - ), + skeletons = Skeletons(skeletons=[skeleton]) + + # create a device for the camera + camera1 = nwbfile.create_device( + name="camera1", + description="camera for recording behavior", + manufacturer="my manufacturer", ) - # Assert the values of the confidence variable - np.testing.assert_array_equal( - movement_dataset["confidence"].values, - np.concatenate( - [ - pose_estimation_series["keypoint1"].confidence[ - :, np.newaxis, np.newaxis - ], - pose_estimation_series["keypoint2"].confidence[ - :, np.newaxis, np.newaxis - ], - ], - axis=1, - ), + n_time = 100 + n_dims = 2 # 2D data + front_left_paw = create_test_pose_estimation_series( + n_time=n_time, n_dims=n_dims, keypoint="front_left_paw" ) - # Assert the attributes of the movement dataset - assert movement_dataset.attrs == { - "fps": np.nanmedian( - 1 / np.diff(pose_estimation_series["keypoint1"].timestamps) + body = create_test_pose_estimation_series( + n_time=n_time, n_dims=n_dims, keypoint="body" + ) + front_right_paw = create_test_pose_estimation_series( + n_time=n_time, n_dims=n_dims, keypoint="front_right_paw" + ) + + # store all PoseEstimationSeries in a list + pose_estimation_series = [front_left_paw, body, front_right_paw] + + pose_estimation = PoseEstimation( + name="PoseEstimation", + pose_estimation_series=pose_estimation_series, + description=( + "Estimated positions of front paws" "of subject1 using DeepLabCut." ), - "time_units": pose_estimation_series["keypoint1"].timestamps_unit, - "source_software": "software1", - "source_file": None, + original_videos=["path/to/camera1.mp4"], + labeled_videos=["path/to/camera1_labeled.mp4"], + dimensions=np.array( + [[640, 480]], dtype="uint16" + ), # pixel dimensions of the video + devices=[camera1], + scorer="DLC_resnet50_openfieldOct30shuffle1_1600", + source_software="DeepLabCut", + source_software_version="2.3.8", + skeleton=skeleton, # link to the skeleton object + ) + + behavior_pm = nwbfile.create_processing_module( + name="behavior", + description="processed behavioral data", + ) + behavior_pm.add(skeletons) + behavior_pm.add(pose_estimation) + + # write the NWBFile to disk + if write_to_disk: + path = "test_pose.nwb" + with NWBHDF5IO(path, mode="w") as io: + io.write(nwbfile) + else: + return nwbfile + + +def test_convert_nwb_to_movement(): + create_test_pose_nwb(write_to_disk=True) + nwb_filepaths = ["test_pose.nwb"] + movement_dataset = convert_nwb_to_movement(nwb_filepaths) + + assert movement_dataset.sizes == { + "time": 100, + "individuals": 1, + "keypoints": 3, + "space": 2, }