doc improvements

Author: ianhi

docs/abs_rel_time_example.ipynb

@@ -13,6 +13,27 @@ Custom xarray indexes for keeping multiple coordinates in sync across shared dim
 This library provides custom [xarray Index](https://docs.xarray.dev/en/stable/internals/how-to-create-custom-index.html) implementations that automatically constrain related dimensions when you select on any one of them.
 
 
+### DimensionInterval
+
+The DimensionInterval provides the ability to performantly store arbitrary intervals over a continuous coordinate. Like a multiindex but more generalized. See the [comparison with MultiIndex](alt-multiindex.ipynb) for an understanding of the comparison.
+
+![diagram of possible sel calls for DimensionInterval](images/generic-intervals.png.excalidraw.png)
+
+See the [Multi-Interval Example](multi_interval_example.ipynb) for a detailed walkthrough.
+
+### AbsoluteRelative Index
+
+Provides the ability to work with both absolute or relative coord (e.g. time) for trialed data.
+![diagram of two possible abs-rel indexes](images/abs-rel.png.excalidraw.png)
+
+See the [Absolute vs Relative Time Example](abs_rel_time_example.ipynb) for a detailed walkthrough.
+
+This then enables more advanced use cases such as building multiple time reference frames without having to shuffle the underlying data. This makes a task such as time-locking a low cost operation:
+
+![diagram of timelocking](images/event-locking.png.excalidraw.png)
+
+See the [Time-Locking Example](time-locking.ipynb) for a demonstration of event-locked analysis.
+
 ### Use Cases
 
 - **Speech/audio data** with hierarchical annotations (words, phonemes, time)

docs/alt-multiindex.ipynb

@@ -4,23 +4,7 @@
    "cell_type": "markdown",
    "id": "cell-0",
    "metadata": {},
-   "source": [
-    "# Multi-Interval Index Example\n",
-    "\n",
-    "This notebook demonstrates how `DimensionInterval` enables automatic cross-slicing between multiple interval types over a shared continuous dimension.\n",
-    "\n",
-    "## Use Case: Speech Data\n",
-    "\n",
-    "Imagine you have speech data with:\n",
-    "- A **continuous time dimension** (e.g., audio samples at regular intervals)\n",
-    "- **Word intervals** - each word spans a range of time\n",
-    "- **Phoneme intervals** - each phoneme spans a smaller range of time within words\n",
-    "\n",
-    "TODO: add an explanatory image\n",
-    "\n",
-    "\n",
-    "When you select a specific word, you want the time and phoneme dimensions to automatically constrain to only the overlapping values. This is exactly what `DimensionInterval` provides."
-   ]
+   "source": "# Multi-Interval Index Example\n\nThis notebook demonstrates how `DimensionInterval` enables automatic cross-slicing between multiple interval types over a shared continuous dimension.\n\n![Diagram of possible sel calls for DimensionInterval](images/generic-intervals.png.excalidraw.png)\n\nWhen you select a specific word, you want the time and phoneme dimensions to automatically constrain to only the overlapping values. This is exactly what `DimensionInterval` provides.\n\n::::{note}\nThere are two ways to encode intervals with `DimensionInterval`:\n1. **Pandas IntervalIndex** - Used in this notebook, intervals are encoded directly as `pd.IntervalIndex` objects\n2. **Onset/Duration format** - Intervals are specified as separate onset and duration coordinates, see the [Onset/Duration Example](onset_duration_example.ipynb)\n::::\n\n::::{seealso}\nFor a comparison of `DimensionInterval` with xarray's built-in `MultiIndex`, see the [MultiIndex Comparison](alt-multiindex.ipynb) notebook.\n::::"
   },
   {
    "cell_type": "code",
@@ -43,6 +27,20 @@
     "from linked_indices import DimensionInterval"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "271be299-c765-488f-804d-5322552e41e0",
+   "metadata": {},
+   "source": [
+    "\n",
+    "## Use Case: Speech Data\n",
+    "\n",
+    "Imagine you have speech data with:\n",
+    "- A **continuous time dimension** (e.g., audio samples at regular intervals)\n",
+    "- **Word intervals** - each word spans a range of time\n",
+    "- **Phoneme intervals** - each phoneme spans a smaller range of time within words\n"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "cell-2",
@@ -5444,7 +5442,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -5463,4 +5461,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

docs/images/abs-rel.png.excalidraw.png

@@ -24,6 +24,7 @@ dev = [
   "isort>=7.0.0",
   "jupyterlab>=4.5.0",
   "jupyterlab-code-formatter>=3.0.2",
+  "jupyterlab-git>=0.51.3",
   "jupyterlab-myst>=2.4.2",
   "matplotlib>=3.10.7",
   "netcdf4>=1.7.3",
@@ -32,6 +33,7 @@ dev = [
   "pytest-cov>=7.0.0",
   "ruff>=0.14.8",
   "scipy>=1.16.3",
+  "xarray-fancy-repr>=0.0.2",
   "zarr>=3.1.5",
 ]
 docs = [

docs/images/event-locking.png.excalidraw.png

@@ -512,39 +512,51 @@ def onset_duration_dataset() -> "xr.Dataset":
 
 
 def trial_based_dataset(
-    n_trials: int = 5,
-    trial_length: float = 10.0,
+    n_trials: int = 3,
+    trial_length: float = 5.0,
     sample_rate: int = 100,
     trial_labels: list[str] | None = None,
     seed: int | None = 42,
+    mode: str = "stacked",
 ) -> "xr.Dataset":
     """
     Create a dataset with trial-based data and both absolute and relative time.
 
     This is useful for testing AbsoluteRelativeIndex.
 
-    The dataset has dimensions (trial, rel_time) with:
-    - rel_time: relative time within each trial (0 to trial_length)
-    - trial: trial labels
-    - abs_time: 2D coordinate (trial, rel_time) mapping to absolute time
+    Supports two modes:
+    - "stacked" (default): 2D array with dimensions (trial, rel_time). Each trial
+      has the same relative time coordinates but different absolute time ranges.
+    - "linear": 1D array with dimension (abs_time). All trials concatenated into
+      a single continuous stream indexed by absolute time, with trial as a
+      1D coordinate indicating which trial each timepoint belongs to.
+
+    By default, creates 3 trials with distinct waveforms:
+    - Trial 1 ("cosine"): cosine wave
+    - Trial 2 ("square"): square wave
+    - Trial 3 ("sawtooth"): sawtooth wave
 
     Parameters
     ----------
     n_trials : int
-        Number of trials. Default: 5
+        Number of trials. Default: 3
     trial_length : float
-        Duration of each trial in seconds. Default: 10.0
+        Duration of each trial in seconds. Default: 5.0
     sample_rate : int
         Samples per second within each trial. Default: 100
     trial_labels : list[str] | None
-        Labels for each trial. If None, uses ["trial_0", "trial_1", ...].
+        Labels for each trial. If None, uses ["cosine", "square", "sawtooth"]
+        for 3 trials, or ["trial_0", "trial_1", ...] for other counts.
     seed : int | None
         Random seed for reproducibility. None for random.
+    mode : str
+        Either "stacked" (2D with trial × rel_time) or "linear" (1D with abs_time).
+        Default: "stacked"
 
     Returns
     -------
     xr.Dataset
-        Dataset with structure:
+        For mode="stacked":
             Dimensions: (trial: n_trials, rel_time: trial_length * sample_rate)
             Coordinates:
               * trial     (trial) str - trial labels
@@ -554,6 +566,16 @@ def trial_based_dataset(
             Data variables:
                 data      (trial, rel_time) float64 - simulated signal
 
+        For mode="linear":
+            Dimensions: (abs_time: n_trials * trial_length * sample_rate)
+            Coordinates:
+              * abs_time   (abs_time) float64 - absolute time
+                rel_time   (abs_time) float64 - relative time within each trial
+                trial      (abs_time) str - trial label for each timepoint
+                trial_onset (abs_time) float64 - onset time of each trial
+            Data variables:
+                data       (abs_time) float64 - simulated signal
+
     Examples
     --------
     >>> from linked_indices.example_data import trial_based_dataset
@@ -566,8 +588,16 @@ def trial_based_dataset(
     0.0
     >>> float(ds.abs_time[1, 0])  # Second trial starts at t=5
     5.0
+
+    >>> ds_linear = trial_based_dataset(mode="linear")
+    >>> dict(ds_linear.dims)
+    {'abs_time': 1500}
     """
     import xarray as xr
+    from scipy import signal
+
+    if mode not in ("stacked", "linear"):
+        raise ValueError(f"mode must be 'stacked' or 'linear', got '{mode}'")
 
     if seed is not None:
         np.random.seed(seed)
@@ -578,7 +608,10 @@ def trial_based_dataset(
 
     # Trial labels
     if trial_labels is None:
-        trial_labels = [f"trial_{i}" for i in range(n_trials)]
+        if n_trials == 3:
+            trial_labels = ["cosine", "square", "sawtooth"]
+        else:
+            trial_labels = [f"trial_{i}" for i in range(n_trials)]
     elif len(trial_labels) != n_trials:
         raise ValueError(
             f"trial_labels length ({len(trial_labels)}) must match n_trials ({n_trials})"
@@ -587,25 +620,63 @@ def trial_based_dataset(
     # Trial onsets (absolute time when each trial starts)
     trial_onsets = np.arange(n_trials) * trial_length
 
-    # Absolute time is a 2D array: abs_time[trial, rel_time_idx] = trial_onset + rel_time
-    abs_time_2d = trial_onsets[:, np.newaxis] + rel_times[np.newaxis, :]
+    # Generate distinct waveforms for each trial
+    freq = 0.5  # Base frequency in Hz
+    data_2d = np.zeros((n_trials, n_samples))
 
-    # Generate different signal for each trial (sine, square, sawtooth, etc.)
-    data = np.zeros((n_trials, n_samples))
     for i in range(n_trials):
-        freq = 1.0 + i * 0.5  # Different frequency per trial
-        phase = i * np.pi / 4  # Different phase
-        data[i] = np.sin(2 * np.pi * freq * rel_times + phase)
-        data[i] += 0.1 * np.random.randn(n_samples)  # Add noise
+        waveform_type = i % 3  # Cycle through cosine, square, sawtooth
+        if waveform_type == 0:
+            # Cosine wave
+            data_2d[i] = np.cos(2 * np.pi * freq * rel_times)
+        elif waveform_type == 1:
+            # Square wave
+            data_2d[i] = signal.square(2 * np.pi * freq * rel_times)
+        else:
+            # Sawtooth wave
+            data_2d[i] = signal.sawtooth(2 * np.pi * freq * rel_times)
+
+    if mode == "stacked":
+        # 2D mode: (trial, rel_time)
+        # Absolute time is a 2D array: abs_time[trial, rel_time_idx] = trial_onset + rel_time
+        abs_time_2d = trial_onsets[:, np.newaxis] + rel_times[np.newaxis, :]
+
+        ds = xr.Dataset(
+            {"data": (("trial", "rel_time"), data_2d)},
+            coords={
+                "trial": trial_labels,
+                "rel_time": rel_times,
+                "abs_time": (("trial", "rel_time"), abs_time_2d),
+                "trial_onset": ("trial", trial_onsets),
+            },
+        )
+    else:
+        # Linear mode: (abs_time,)
+        # Concatenate all trials into a single 1D array
+        data_1d = data_2d.flatten()
+
+        # Absolute time is continuous across all trials
+        abs_time_1d = np.concatenate(
+            [trial_onsets[i] + rel_times for i in range(n_trials)]
+        )
 
-    ds = xr.Dataset(
-        {"data": (("trial", "rel_time"), data)},
-        coords={
-            "trial": trial_labels,
-            "rel_time": rel_times,
-            "abs_time": (("trial", "rel_time"), abs_time_2d),
-            "trial_onset": ("trial", trial_onsets),
-        },
-    )
+        # Relative time repeats for each trial
+        rel_time_1d = np.tile(rel_times, n_trials)
+
+        # Trial label for each timepoint
+        trial_1d = np.repeat(trial_labels, n_samples)
+
+        # Trial onset for each timepoint
+        trial_onset_1d = np.repeat(trial_onsets, n_samples)
+
+        ds = xr.Dataset(
+            {"data": (("abs_time",), data_1d)},
+            coords={
+                "abs_time": abs_time_1d,
+                "rel_time": ("abs_time", rel_time_1d),
+                "trial": ("abs_time", trial_1d),
+                "trial_onset": ("abs_time", trial_onset_1d),
+            },
+        )
 
     return ds