From 198341087fdf2b3bbfaa0162539eb390dabbf6b5 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Mon, 1 Jul 2024 16:33:40 -0700 Subject: [PATCH] fix tests --- mantis/cli/estimate_stabilization.py | 6 ++- .../tests/test_analysis/test_stabilization.py | 40 +++++++------------ 2 files changed, 19 insertions(+), 27 deletions(-) diff --git a/mantis/cli/estimate_stabilization.py b/mantis/cli/estimate_stabilization.py index ac43050a..c152e038 100644 --- a/mantis/cli/estimate_stabilization.py +++ b/mantis/cli/estimate_stabilization.py @@ -67,8 +67,6 @@ def estimate_position_focus( def get_mean_z_positions(dataframe_path: Path, verbose: bool = False) -> None: - import matplotlib.pyplot as plt - df = pd.read_csv(dataframe_path) # Sort the DataFrame based on 'time_idx' @@ -79,13 +77,17 @@ def get_mean_z_positions(dataframe_path: Path, verbose: bool = False) -> None: # Get the mean of positions for each time point average_focus_idx = df.groupby("time_idx")["focus_idx"].mean().reset_index() + if verbose: + import matplotlib.pyplot as plt + # Get the moving average of the focus_idx plt.plot(average_focus_idx["focus_idx"], linestyle="--", label="mean of all positions") plt.xlabel('Time index') plt.ylabel('Focus index') plt.legend() plt.savefig(dataframe_path.parent / "z_drift.png") + return average_focus_idx["focus_idx"].values diff --git a/mantis/tests/test_analysis/test_stabilization.py b/mantis/tests/test_analysis/test_stabilization.py index f487ea01..efbbc16b 100644 --- a/mantis/tests/test_analysis/test_stabilization.py +++ b/mantis/tests/test_analysis/test_stabilization.py @@ -1,40 +1,30 @@ +import io + import numpy as np import pandas as pd -from mantis.cli.estimate_stabilization import estimate_position_focus, get_mean_z_positions - - -def test_estimate_position_focus(): - # Create input data - z_positions = [1, 2, 3, 4, 5] - focus_scores = [0.1, 0.2, 0.3, 0.4, 0.5] - - # Call the function - result = estimate_position_focus(z_positions, focus_scores) - - # Check the result - assert isinstance(result, tuple) - assert len(result) == 2 - assert isinstance(result[0], float) - assert isinstance(result[1], float) +from mantis.cli.estimate_stabilization import get_mean_z_positions def test_mean_z_positions(): # Create input data df = pd.DataFrame( { - "channel": ["GFP"], - "channel_idx": [0, 0, 0, 0, 0, 0, 0, 0], - "time_min": [1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0], - "focal_idx": [20, 25, 19, 16, 24, 29, 21, 18], - "position_idx": [1, 2, 3, 4, 1, 2, 3, 4], + "position": ['0/2/000000'] * 4 + ['0/2/001002'] * 4, + "time_idx": [0, 1, 2, 3] * 2, + "channel": ["GFP"] * 8, + "focus_idx": [20, 25, np.nan, 16, 24, 29, 21, 18], } ) + # Create a pretend file + s_buf = io.StringIO() + df.to_csv(s_buf) + s_buf.seek(0) + # Call the function - result = get_mean_z_positions(df, 0) - pos_1 = np.array([20, 25, 19, 16]).mean() - pos_2 = np.array([24, 29, 21, 18]).mean() + result = get_mean_z_positions(s_buf, verbose=False) + correct_result = np.array([22, 27, 25, 17]) # Check the result - assert result == [pos_1, pos_2] + assert all(result == correct_result)