add make_dataset function for constructing RecallDataset structures from sparse inputs

Author: githubpsyche

jaxcmr/helpers.py

@@ -33,6 +33,7 @@
     "find_max_list_length",
     "apply_by_subject",
     "has_repeats_per_row",
+    "make_dataset",
 ]
 
 
@@ -260,3 +261,119 @@ def check_row(row):
         return jnp.any(valid_repeats)
 
     return vmap(check_row)(arr)
+
+
+def make_dataset(
+    recalls: Integer[Array, "n_trials num_recalled"],
+    pres_itemnos: Integer[Array, "n_trials num_presented"] | None = None,
+    listLength: Integer[Array, "n_trials 1"] | int | None = None,
+    subject: Integer[Array, "n_trials 1"] | int = 1,
+    *,
+    reps: int = 1,
+) -> RecallDataset:
+    """Construct a ``RecallDataset`` from flexible inputs.
+
+    Parameters
+    ----------
+    recalls
+        Within-list recall events (1-indexed, 0 = padding). A 1-D vector is
+        treated as a single trial.
+    pres_itemnos
+        Within-list presentation items. When omitted, generated as
+        ``arange(1, list_length + 1)`` per trial.
+    listLength
+        List length per trial. Inferred from *pres_itemnos* when omitted.
+        When both are provided, an assertion checks compatibility.
+        When neither is provided, inferred as
+        ``max(recalls.shape[1], recalls.max())``.
+    subject
+        Subject identifier per trial. Defaults to ``1`` (single subject).
+    reps
+        Number of times to tile the resulting trials along axis 0.
+
+    Returns
+    -------
+    RecallDataset
+        Dictionary with keys ``recalls``, ``pres_itemnos``, ``listLength``,
+        and ``subject``, all shaped ``(n_trials * reps, ...)``.
+
+    """
+    # -- normalise recalls (required) -----------------------------------------
+    recalls_arr = jnp.atleast_2d(jnp.asarray(recalls, dtype=jnp.int32))
+
+    # -- normalise pres_itemnos (optional) ------------------------------------
+    pres_arr: jnp.ndarray | None = None
+    if pres_itemnos is not None:
+        pres_arr = jnp.atleast_2d(jnp.asarray(pres_itemnos, dtype=jnp.int32))
+
+    # -- normalise listLength (optional) --------------------------------------
+    ll_arr: jnp.ndarray | None = None
+    if listLength is not None:
+        if isinstance(listLength, int):
+            ll_arr = jnp.array([[listLength]], dtype=jnp.int32)
+        else:
+            ll_arr = jnp.asarray(listLength, dtype=jnp.int32).reshape(-1, 1)
+
+    # -- normalise subject ----------------------------------------------------
+    if isinstance(subject, int):
+        subj_arr = jnp.array([[subject]], dtype=jnp.int32)
+    else:
+        subj_arr = jnp.asarray(subject, dtype=jnp.int32).reshape(-1, 1)
+
+    # -- resolve n_trials -----------------------------------------------------
+    sizes: list[int] = []
+    for arr in (recalls_arr, pres_arr, ll_arr, subj_arr):
+        if arr is not None and arr.shape[0] > 1:
+            sizes.append(arr.shape[0])
+    if sizes:
+        n_trials = sizes[0]
+        assert all(s == n_trials for s in sizes), (
+            f"Multi-trial args disagree on n_trials: {sizes}"
+        )
+    else:
+        n_trials = 1
+
+    # -- tile single-trial args to n_trials -----------------------------------
+    if recalls_arr.shape[0] == 1 and n_trials > 1:
+        recalls_arr = jnp.tile(recalls_arr, (n_trials, 1))
+    if pres_arr is not None and pres_arr.shape[0] == 1 and n_trials > 1:
+        pres_arr = jnp.tile(pres_arr, (n_trials, 1))
+    if ll_arr is not None and ll_arr.shape[0] == 1 and n_trials > 1:
+        ll_arr = jnp.tile(ll_arr, (n_trials, 1))
+    if subj_arr.shape[0] == 1 and n_trials > 1:
+        subj_arr = jnp.tile(subj_arr, (n_trials, 1))
+
+    # -- infer / validate list_length -----------------------------------------
+    if pres_arr is not None and ll_arr is not None:
+        assert jnp.all(ll_arr == pres_arr.shape[1]), (
+            f"listLength ({ll_arr.ravel()}) incompatible with "
+            f"pres_itemnos width ({pres_arr.shape[1]})"
+        )
+        list_length = int(pres_arr.shape[1])
+    elif pres_arr is not None:
+        list_length = int(pres_arr.shape[1])
+        ll_arr = jnp.full((n_trials, 1), list_length, dtype=jnp.int32)
+    elif ll_arr is not None:
+        list_length = int(ll_arr[0, 0])
+    else:
+        list_length = int(max(recalls_arr.shape[1], jnp.max(recalls_arr)))
+        ll_arr = jnp.full((n_trials, 1), list_length, dtype=jnp.int32)
+
+    # -- generate default pres_itemnos ----------------------------------------
+    if pres_arr is None:
+        row = jnp.arange(1, list_length + 1, dtype=jnp.int32)
+        pres_arr = jnp.tile(row[None, :], (n_trials, 1))
+
+    # -- apply reps -----------------------------------------------------------
+    if reps > 1:
+        recalls_arr = jnp.tile(recalls_arr, (reps, 1))
+        pres_arr = jnp.tile(pres_arr, (reps, 1))
+        ll_arr = jnp.tile(ll_arr, (reps, 1))
+        subj_arr = jnp.tile(subj_arr, (reps, 1))
+
+    return {
+        "recalls": recalls_arr,
+        "pres_itemnos": pres_arr,
+        "listLength": ll_arr,
+        "subject": subj_arr,
+    }  # type: ignore

tests/test_backreplagrank.py

@@ -13,23 +13,7 @@
     plot_back_rep_lagrank,
 )
 from jaxcmr.analyses.replagrank import replagrank
-from jaxcmr.typing import RecallDataset
-
-
-def _make_dataset(recalls, presentations, subjects=None) -> RecallDataset:
-    recalls = jnp.array(recalls)
-    presentations = jnp.array(presentations)
-    n = recalls.shape[0]
-    if subjects is None:
-        subjects = jnp.zeros(n, dtype=int)
-    else:
-        subjects = jnp.array(subjects)
-    return {
-        "recalls": recalls,
-        "pres_itemnos": presentations,
-        "subject": subjects,
-        "listLength": jnp.full(n, presentations.shape[1]),
-    }
+from jaxcmr.helpers import make_dataset
 
 
 def _rep_dataset():
@@ -45,7 +29,7 @@ def _rep_dataset():
         [3, 1, 2, 4, 5, 0, 0, 0],
         [1, 3, 5, 2, 0, 0, 0, 0],
     ])
-    return _make_dataset(recalls, pres, subjects=[0, 0, 1, 1])
+    return make_dataset(recalls, pres, subject=jnp.array([0, 0, 1, 1]))
 
 
 class TestReversalChangesResults:

tests/test_compound_cueing_crp.py

@@ -22,29 +22,7 @@
     tabulate_trial,
     compound_cueing_crp,
 )
-from jaxcmr.typing import RecallDataset
-
-
-# -----------------------------------------------------------------------------
-# Helpers
-# -----------------------------------------------------------------------------
-
-
-def _make_dataset(
-    recalls: jnp.ndarray,
-    presentations: jnp.ndarray,
-) -> RecallDataset:
-    """Return minimal dataset keyed for the compound cueing analysis."""
-    recalls = jnp.asarray(recalls, dtype=jnp.int32)
-    presentations = jnp.asarray(presentations, dtype=jnp.int32)
-    n_trials, _ = recalls.shape
-    list_length = presentations.shape[1]
-    return {
-        "subject": jnp.ones((n_trials, 1), dtype=jnp.int32),
-        "listLength": jnp.full((n_trials, 1), list_length, dtype=jnp.int32),
-        "pres_itemnos": presentations,
-        "recalls": recalls,
-    }
+from jaxcmr.helpers import make_dataset
 
 
 # -----------------------------------------------------------------------------
@@ -421,7 +399,7 @@ def test_crp_computation_single_trial():
     # Arrange
     presentation = jnp.array([[1, 2, 3, 4, 5, 6, 7, 8, 9, 3, 10, 11]], dtype=jnp.int32)
     trial = jnp.array([[1, 2, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=jnp.int32)
-    dataset = _make_dataset(trial, presentation)
+    dataset = make_dataset(trial, presentation)
 
     # Act
     result = compound_cueing_crp(dataset, min_spacing=6, size=2)
@@ -447,7 +425,7 @@ def test_crp_nan_when_no_opportunities():
     # Arrange
     presentation = jnp.array([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=jnp.int32)
     trial = jnp.array([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=jnp.int32)
-    dataset = _make_dataset(trial, presentation)
+    dataset = make_dataset(trial, presentation)
 
     # Act
     result = compound_cueing_crp(dataset, min_spacing=6, size=2)
@@ -483,7 +461,7 @@ def test_crp_aggregates_across_trials():
         ],
         dtype=jnp.int32,
     )
-    dataset = _make_dataset(trials, presentation)
+    dataset = make_dataset(trials, presentation)
 
     # Act
     result = compound_cueing_crp(dataset, min_spacing=6, size=2)

tests/test_crp.py

@@ -7,6 +7,7 @@
 from matplotlib.figure import Figure
 import pytest
 from jaxcmr.analyses import crp
+from jaxcmr.helpers import make_dataset
 from jaxcmr.typing import RecallDataset
 
 
@@ -28,23 +29,6 @@ def _lag_index(lags_len: int, lag: int) -> int:
     center = lags_len // 2
     return center + lag
 
-def _make_dataset(
-    recalls: jnp.ndarray,
-    presentations: jnp.ndarray,
-) -> RecallDataset:
-    """Return minimal dataset keyed for the CRP analysis."""
-
-    recalls = jnp.asarray(recalls, dtype=jnp.int32)
-    presentations = jnp.asarray(presentations, dtype=jnp.int32)
-    n_trials, _ = recalls.shape
-    list_length = presentations.shape[1]
-    return {
-        "subject": jnp.ones((n_trials, 1), dtype=jnp.int32),
-        "listLength": jnp.full((n_trials, 1), list_length, dtype=jnp.int32),
-        "pres_itemnos": presentations,
-        "recalls": recalls,
-    } # type: ignore
-
 # -----------------------------------------------------------------------------
 # set_false_at_index tests
 # -----------------------------------------------------------------------------
@@ -444,8 +428,8 @@ def test_produces_same_crp_when_repeats_removed():
     presentations = jnp.array([[1, 2, 3, 4]], dtype=jnp.int32)
 
     # Act / When
-    dataset_with_repeat = _make_dataset(trials_with_repeat, presentations)
-    dataset_without_repeat = _make_dataset(trials_without_repeat, presentations)
+    dataset_with_repeat = make_dataset(trials_with_repeat, presentations)
+    dataset_without_repeat = make_dataset(trials_without_repeat, presentations)
 
     with_repeat = crp.crp(dataset_with_repeat, size=1)
     without_repeat = crp.crp(dataset_without_repeat, size=1)
@@ -471,7 +455,7 @@ def test_returns_array_when_single_trial_item():
     presentations = jnp.array([[1]], dtype=jnp.int32)
 
     # Act / When
-    dataset = _make_dataset(trials, presentations)
+    dataset = make_dataset(trials, presentations)
     out = crp.crp(dataset, size=1)
 
     # Assert / Then
@@ -495,7 +479,7 @@ def test_matches_uncompiled_result_when_jitted_with_static_argnames():
     trials = jnp.array([[1, 2, 3], [1, 3, 2]], dtype=jnp.int32)
     presentations = jnp.array([[1, 2, 3], [1, 2, 3]], dtype=jnp.int32)
     size = 1
-    dataset = _make_dataset(trials, presentations)
+    dataset = make_dataset(trials, presentations)
     expected = crp.crp(dataset, size)
 
     # Act / When
@@ -523,7 +507,7 @@ def test_runs_with_larger_size_when_jitted():
     trials = jnp.array([[1, 2, 3]], dtype=jnp.int32)
     presentations = jnp.array([[1, 2, 3]], dtype=jnp.int32)
     size = 2
-    dataset = _make_dataset(trials, presentations)
+    dataset = make_dataset(trials, presentations)
     expected = crp.crp(dataset, size)
 
     # Act / When

tests/test_cue_centered_lagrank.py

@@ -16,28 +16,7 @@
     test_cue_centered_lagrank as run_test,
 )
 from jaxcmr.analyses.lagrank import LagRankTestResult
-from jaxcmr.typing import RecallDataset
-
-
-def _make_dataset(recalls, presentations, cues, should_tabulate, subjects=None) -> RecallDataset:
-    """Wrap arrays into a RecallDataset dict with cue fields."""
-    recalls = jnp.array(recalls)
-    presentations = jnp.array(presentations)
-    cues = jnp.array(cues)
-    should_tab = jnp.array(should_tabulate, dtype=bool)
-    n = recalls.shape[0]
-    if subjects is None:
-        subjects = jnp.zeros(n, dtype=int)
-    else:
-        subjects = jnp.array(subjects)
-    return {
-        "recalls": recalls,
-        "pres_itemnos": presentations,
-        "cue_clips": cues,
-        "_should_tabulate": should_tab,
-        "subject": subjects,
-        "listLength": jnp.full(n, presentations.shape[1]),
-    }
+from jaxcmr.helpers import make_dataset
 
 
 # ---- Tabulation: no cue → skipped ----
@@ -98,7 +77,7 @@ def dataset(self):
             [3, 2, 1, 0],
         ])
         should_tab = recalls > 0
-        return _make_dataset(recalls, pres, cues, should_tab)
+        return {**make_dataset(recalls, pres), "cue_clips": cues, "_should_tabulate": should_tab}
 
     def test_returns_1d(self, dataset):
         """cue_centered_lagrank returns (n_trials,)."""
@@ -140,7 +119,7 @@ def test_subject_shape(self):
         ])
         should_tab = recalls > 0
         subjects = jnp.array([0, 0, 1, 1])
-        dataset = _make_dataset(recalls, pres, cues, should_tab, subjects)
+        dataset = {**make_dataset(recalls, pres, subject=subjects), "cue_clips": cues, "_should_tabulate": should_tab}
         mask = jnp.ones(4, dtype=bool)
         result = subject_cue_centered_lagrank(dataset, mask, size=1)
         assert result.shape == (2,)
@@ -178,7 +157,7 @@ def test_plot_returns_axes(self):
             [1, 2, 0, 0],
         ])
         should_tab = recalls > 0
-        dataset = _make_dataset(recalls, pres, cues, should_tab)
+        dataset = {**make_dataset(recalls, pres), "cue_clips": cues, "_should_tabulate": should_tab}
         mask = jnp.ones(2, dtype=bool)
         ax = plot_cue_centered_lagrank(
             dataset, mask, should_tabulate=should_tab, size=1, labels=["Test"]
@@ -206,7 +185,7 @@ def test_jit_compatible(self):
             [1, 2, 0, 0],
         ])
         should_tab = recalls > 0
-        dataset = _make_dataset(recalls, pres, cues, should_tab)
+        dataset = {**make_dataset(recalls, pres), "cue_clips": cues, "_should_tabulate": should_tab}
         result_nojit = cue_centered_lagrank(dataset, size=1)
         result_jit = jit(cue_centered_lagrank, static_argnames=("size",))(
             dataset, size=1