From 6d69687476d9256ddf51ceb09d10ec935021c0e1 Mon Sep 17 00:00:00 2001 From: Matt Whiteway Date: Fri, 13 Dec 2024 17:47:26 -0500 Subject: [PATCH] v3.0.0 (#23) * make singlecam smoother robust to int/float smooth param * update version --- eks/__init__.py | 2 +- eks/singlecam_smoother.py | 13 ++++-- tests/test_ibl_pupil_smoother.py | 51 ++++++++++++++++------ tests/test_singlecam_smoother.py | 72 +++++++++++++++++++++++++++----- 4 files changed, 111 insertions(+), 27 deletions(-) diff --git a/eks/__init__.py b/eks/__init__.py index 4533b0d..004dab9 100644 --- a/eks/__init__.py +++ b/eks/__init__.py @@ -1,3 +1,3 @@ from eks import * -__version__ = '2.1.1' +__version__ = '3.0.0' diff --git a/eks/singlecam_smoother.py b/eks/singlecam_smoother.py index 074393b..7bdf743 100644 --- a/eks/singlecam_smoother.py +++ b/eks/singlecam_smoother.py @@ -405,14 +405,22 @@ def nll_loss_sequential_scan(s, cov_mats, cropped_ys, m0s, S0s, Cs, As, Rs, ense # Optimize smooth_param if smooth_param is not None: - s_finals = [smooth_param] if isinstance(smooth_param, float) else smooth_param + if isinstance(smooth_param, float): + s_finals = [smooth_param] + elif isinstance(smooth_param, int): + s_finals = [float(smooth_param)] + else: + s_finals = smooth_param else: guesses = [] cropped_ys = [] for k in range(n_keypoints): current_guess = compute_initial_guesses(ensemble_vars[:, k, :]) guesses.append(current_guess) - cropped_ys.append(crop_frames(ys[k], s_frames)) + if s_frames is None or len(s_frames) == 0: + cropped_ys.append(ys[k]) + else: + cropped_ys.append(crop_frames(ys[k], s_frames)) cropped_ys = np.array(cropped_ys) # Concatenation of this list along dimension 0 @@ -524,7 +532,6 @@ def inner_smooth_min_routine_parallel(y, m0, S0, A, Q, C, R): # Note: This should only be run on GPUs # ------------------------------------------------------------------------------------------------ - def singlecam_smooth_min_parallel( smooth_param, cov_mats, observations, initial_means, initial_covariances, Cs, As, Rs, ): diff --git a/tests/test_ibl_pupil_smoother.py b/tests/test_ibl_pupil_smoother.py index 39190d1..ea5d7b8 100644 --- a/tests/test_ibl_pupil_smoother.py +++ b/tests/test_ibl_pupil_smoother.py @@ -172,6 +172,14 @@ def test_add_mean_to_array_single_row(): def test_ensemble_kalman_smoother_ibl_pupil(): + def _check_outputs(df, params, nlls): + assert isinstance(df, pd.DataFrame), "first return arg should be a DataFrame" + assert df.shape[0] == 100, "markers_df should have 100 rows" + assert len(params) == 2, "Expected 2 smooth parameters" + assert params[0] < 1, "Expected diameter smoothing parameter to be less than 1" + assert params[1] < 1, "Expected COM smoothing parameter to be less than 1" + assert isinstance(nlls, list), "Expected nll_values to be a list" + # Create mock data columns = [ 'pupil_top_r_x', 'pupil_top_r_y', 'pupil_bottom_r_x', 'pupil_bottom_r_y', @@ -181,22 +189,41 @@ def test_ensemble_kalman_smoother_ibl_pupil(): pd.DataFrame(np.random.randn(100, 8), columns=columns), pd.DataFrame(np.random.randn(100, 8), columns=columns), ] - smooth_params = [0.5, 0.5] s_frames = [(1, 20)] - # Run the function with mocked data + # Run with fixed smooth params + smooth_params = [0.5, 0.5] smoothed_df, smooth_params_out, nll_values = ensemble_kalman_smoother_ibl_pupil( markers_list, smooth_params, s_frames, avg_mode='mean', var_mode='var', ) + _check_outputs(smoothed_df, smooth_params_out, nll_values) + assert smooth_params == smooth_params_out - # Assertions - assert isinstance(smoothed_df, pd.DataFrame), "first return arg should be a DataFrame" - - # Verify the shape of the output DataFrames - assert smoothed_df.shape[0] == 100, "markers_df should have 100 rows" - - # Check if the smooth parameters and NLL values are correctly returned - assert len(smooth_params_out) == 2, "Expected 2 smooth parameters" - assert isinstance(nll_values, list), "Expected nll_values to be a list" + # Run with [None, None] smooth params + smoothed_df, smooth_params_out, nll_values = ensemble_kalman_smoother_ibl_pupil( + markers_list, [None, None], s_frames, avg_mode='mean', var_mode='var', + ) + _check_outputs(smoothed_df, smooth_params_out, nll_values) - print("All tests passed successfully.") + # Run with None smooth params + smoothed_df, smooth_params_out, nll_values = ensemble_kalman_smoother_ibl_pupil( + markers_list, None, s_frames, avg_mode='mean', var_mode='var', + ) + _check_outputs(smoothed_df, smooth_params_out, nll_values) + + # CURRENTLY NOT SUPPORTED: fix one smooth param + + # Run with diameter smooth param + # smooth_params = [0.9, None] + # smoothed_df, smooth_params_out, nll_values = ensemble_kalman_smoother_ibl_pupil( + # markers_list, [0.9, None], s_frames, avg_mode='mean', var_mode='var', + # ) + # _check_outputs(smoothed_df, smooth_params_out, nll_values) + # assert smooth_params[0] == smooth_params_out[0] + + # Run with COM smooth param + # smoothed_df, smooth_params_out, nll_values = ensemble_kalman_smoother_ibl_pupil( + # markers_list, [None, 0.9], s_frames, avg_mode='mean', var_mode='var', + # ) + # _check_outputs(smoothed_df, smooth_params_out, nll_values) + # assert smooth_params[1] == smooth_params_out[1] diff --git a/tests/test_singlecam_smoother.py b/tests/test_singlecam_smoother.py index 1eddec0..ac9b828 100644 --- a/tests/test_singlecam_smoother.py +++ b/tests/test_singlecam_smoother.py @@ -11,6 +11,18 @@ def test_ensemble_kalman_smoother_singlecam(): + def _check_outputs(df, params): + # Basic checks to ensure the function runs and returns expected types + assert isinstance(df_smoothed, pd.DataFrame), \ + f"Expected first return value to be a pd.DataFrame, got {type(df_smoothed)}" + assert isinstance(s_finals, (list, np.ndarray)), \ + "Expected s_finals to be a list or an ndarray" + + # Check for different return values in the correct level of the columns + for v in ['x', 'y', 'likelihood']: + assert v in df_smoothed.columns.get_level_values('coords'), \ + "Expected 'likelihood' in DataFrame columns at the 'coords' level" + # Create mock data keypoint_names = ['kp1', 'kp2'] columns = [f'{kp}_{coord}' for kp in keypoint_names for coord in ['x', 'y', 'likelihood']] @@ -18,11 +30,11 @@ def test_ensemble_kalman_smoother_singlecam(): pd.DataFrame(np.random.randn(100, len(columns)), columns=columns), pd.DataFrame(np.random.randn(100, len(columns)), columns=columns), ] - smooth_param = 0.1 s_frames = None blocks = [] - # Call the smoother function + # run with fixed smooth param (float) + smooth_param = 0.1 df_smoothed, s_finals = ensemble_kalman_smoother_singlecam( markers_list=markers_list, keypoint_names=keypoint_names, @@ -30,17 +42,55 @@ def test_ensemble_kalman_smoother_singlecam(): s_frames=s_frames, blocks=blocks, ) + _check_outputs(df_smoothed, s_finals) + assert s_finals == [smooth_param] - # Basic checks to ensure the function runs and returns expected types - assert isinstance(df_smoothed, pd.DataFrame), \ - f"Expected first return value to be a pd.DataFrame, got {type(df_smoothed)}" - assert isinstance(s_finals, (list, np.ndarray)), \ - "Expected s_finals to be a list or an ndarray" + # run with fixed smooth param (int) + smooth_param = 5 + df_smoothed, s_finals = ensemble_kalman_smoother_singlecam( + markers_list=markers_list, + keypoint_names=keypoint_names, + smooth_param=smooth_param, + s_frames=s_frames, + blocks=blocks, + ) + _check_outputs(df_smoothed, s_finals) + assert s_finals == [smooth_param] - # Check for different return values in the correct level of the columns - for v in ['x', 'y', 'likelihood']: - assert v in df_smoothed.columns.get_level_values('coords'), \ - "Expected 'likelihood' in DataFrame columns at the 'coords' level" + # run with fixed smooth param (single-entry list) + smooth_param = [0.1] + df_smoothed, s_finals = ensemble_kalman_smoother_singlecam( + markers_list=markers_list, + keypoint_names=keypoint_names, + smooth_param=smooth_param, + s_frames=s_frames, + blocks=blocks, + ) + _check_outputs(df_smoothed, s_finals) + assert s_finals == smooth_param + + # run with fixed smooth param (list) + smooth_param = [0.1, 0.4] + df_smoothed, s_finals = ensemble_kalman_smoother_singlecam( + markers_list=markers_list, + keypoint_names=keypoint_names, + smooth_param=smooth_param, + s_frames=s_frames, + blocks=blocks, + ) + _check_outputs(df_smoothed, s_finals) + assert np.all(s_finals == smooth_param) + + # run with None smooth param + smooth_param = None + df_smoothed, s_finals = ensemble_kalman_smoother_singlecam( + markers_list=markers_list, + keypoint_names=keypoint_names, + smooth_param=smooth_param, + s_frames=s_frames, + blocks=blocks, + ) + _check_outputs(df_smoothed, s_finals) def test_adjust_observations():