v3.0.0 (#23)

* make singlecam smoother robust to int/float smooth param

* update version

eks/__init__.py

@@ -1,3 +1,3 @@
 from eks import *
 
-__version__ = '2.1.1'
+__version__ = '3.0.0'

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,
 ):

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]

tests/test_singlecam_smoother.py

@@ -11,36 +11,86 @@
 
 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']]
     markers_list = [
         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,
         smooth_param=smooth_param,
         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():