diff --git a/src/graphomotor/features/trails/time.py b/src/graphomotor/features/trails/time.py new file mode 100644 index 0000000..4babc3e --- /dev/null +++ b/src/graphomotor/features/trails/time.py @@ -0,0 +1,52 @@ +"""Feature extraction module for time-based metrics in trails drawing data.""" + +from graphomotor.core import models + + +def calculate_total_error_time(drawing: models.Drawing) -> dict[str, float]: + """Calculate the total time spent making errors. + + A contiguous "error chunk" is any sequence of rows where df["error"] != "E0". + The start and end of each chunk is defined as the midpoint between the last + timestamp with a "correct" entry and the first timestamp of an "error". The total + error time is the sum of the durations of all error chunks. + + Args: + drawing: Drawing object containing drawing data. + + Returns: + Dictionary containing the total time (s) spent in error states. + """ + mask = drawing.data["error"] != "E0" + if not mask.any(): + return {"total_error_time": 0.0} + + error_change = mask.astype(int).diff() + chunk_starts = error_change[error_change == 1].index.tolist() + chunk_ends = error_change[error_change == -1].index.tolist() + + if mask.iloc[0]: + chunk_starts = [0] + chunk_starts + + if mask.iloc[-1]: + chunk_ends = chunk_ends + [len(drawing.data)] + + seconds = drawing.data["seconds"].to_numpy() + total_error_time = 0.0 + + for start_idx, end_idx in zip(chunk_starts, chunk_ends): + start_time = ( + (seconds[start_idx - 1] + seconds[start_idx]) / 2 + if start_idx > 0 + else seconds[0] + ) + + end_time = ( + (seconds[end_idx - 1] + seconds[end_idx]) / 2 + if end_idx < len(seconds) + else seconds[-1] + ) + + total_error_time += end_time - start_time + + return {"total_error_time": float(total_error_time)} diff --git a/tests/unit/test_trails_time.py b/tests/unit/test_trails_time.py new file mode 100644 index 0000000..fca1e1b --- /dev/null +++ b/tests/unit/test_trails_time.py @@ -0,0 +1,73 @@ +"""Tests for trails time.py.""" + +import pandas as pd + +from graphomotor.core import models +from graphomotor.features.trails import time + + +def test_total_error_time_no_errors() -> None: + """Test case with no errors.""" + df = pd.DataFrame( + { + "error": ["E0", "E0", "E0", "E0"], + "seconds": [0, 1, 2, 3], + } + ) + drawing = models.Drawing(data=df, task_name="trails", metadata={"id": "5555555"}) + + result = time.calculate_total_error_time(drawing) + assert result == {"total_error_time": 0.0} + + +def test_single_error_chunk() -> None: + """Test case with a single error chunk.""" + df = pd.DataFrame( + { + "error": ["E0", "E1", "E1", "E0", "E0"], + "seconds": [0.0, 1.0, 3.0, 5.0, 6.0], + } + ) + drawing = models.Drawing(data=df, task_name="trails", metadata={"id": "5555555"}) + + result = time.calculate_total_error_time(drawing) + assert result == {"total_error_time": 3.5} + + +def test_multiple_error_chunks() -> None: + """Test case with multiple error chunks.""" + df = pd.DataFrame( + { + "error": ["E0", "E1", "E1", "E0", "E2", "E2", "E0"], + "seconds": [0, 1, 3, 5, 6, 7, 9], + } + ) + drawing = models.Drawing(data=df, task_name="trails", metadata={"id": "5555555"}) + result = time.calculate_total_error_time(drawing) + assert result == {"total_error_time": 6.0} + + +def test_error_at_end() -> None: + """Test case with an error chunk that goes to the end of the drawing.""" + df = pd.DataFrame( + { + "error": ["E0", "E0", "E2", "E2"], + "seconds": [0.0, 1.0, 2.0, 4.0], + } + ) + drawing = models.Drawing(data=df, task_name="trails", metadata={"id": "5555555"}) + result = time.calculate_total_error_time(drawing) + assert result == {"total_error_time": 2.5} + + +def test_error_at_start() -> None: + """Test case with an error chunk that starts at the beginning of the drawing.""" + df = pd.DataFrame( + { + "error": ["E1", "E1", "E0", "E0"], + "seconds": [0.0, 1.0, 3.0, 4.0], + } + ) + drawing = models.Drawing(data=df, task_name="trails", metadata={"id": "5555555"}) + result = time.calculate_total_error_time(drawing) + assert result == {"total_error_time": 2.0}