diff --git a/docs/API/data_preparation/longitudinal_dataset.md b/docs/API/data_preparation/longitudinal_dataset.md
index 446567a..6c70f64 100644
--- a/docs/API/data_preparation/longitudinal_dataset.md
+++ b/docs/API/data_preparation/longitudinal_dataset.md
@@ -10,6 +10,8 @@
- setup_features_group
- feature_groups
- non_longitudinal_features
+ - to_wide
+ - to_long
- convert
- save_data
- set_data
diff --git a/docs/assets/images/tutorials/long_wide_reshape/LongToWide.gif b/docs/assets/images/tutorials/long_wide_reshape/LongToWide.gif
new file mode 100644
index 0000000..abacd2f
Binary files /dev/null and b/docs/assets/images/tutorials/long_wide_reshape/LongToWide.gif differ
diff --git a/docs/assets/images/tutorials/long_wide_reshape/WideToLong.gif b/docs/assets/images/tutorials/long_wide_reshape/WideToLong.gif
new file mode 100644
index 0000000..126645d
Binary files /dev/null and b/docs/assets/images/tutorials/long_wide_reshape/WideToLong.gif differ
diff --git a/docs/tutorials/long_wide_reshape.md b/docs/tutorials/long_wide_reshape.md
new file mode 100644
index 0000000..eeb460f
--- /dev/null
+++ b/docs/tutorials/long_wide_reshape.md
@@ -0,0 +1,156 @@
+---
+icon: lucide/git-compare-arrows
+---
+
+# Reshaping Longitudinal Data: Long ⇄ Wide with `LongitudinalDataset`
+
+!!! tip "Read this first"
+ The [Temporal Dependency tutorial](temporal_dependency.md) explains `features_group` and `non_longitudinal_features`, which are central to how `Sklong` represents longitudinal data. Reading it first will make this tutorial much easier to follow.
+
+`Sklong` works on **wide-format** matrices: one row per subject, one column per (feature, wave) pair. Real-world tabular sources are often stored in **long format**: one row per (subject, wave). The [`LongitudinalDataset`](../API/data_preparation/longitudinal_dataset.md) class ships two methods, `to_wide` and `to_long`, that move between the two layouts. Both validate the inputs, keep `features_group` and `non_longitudinal_features` in sync, and optionally write a CSV alongside.
+
+## The toy cohort
+
+Three patients, three follow-up waves, two longitudinal measurements (`bp`, `chol`), one static covariate (`sex`). The same cohort is used in both directions of the tutorial.
+
+```python
+import numpy as np
+import pandas as pd
+
+long_df = pd.DataFrame({
+ "patient_id": [1, 1, 1, 2, 2, 2, 3, 3, 3],
+ "wave": [0, 1, 2, 0, 1, 2, 0, 1, 2],
+ "bp": [120., 122., 121., 130., 131., 128., 110., 112., 115.],
+ "chol": [5.0, 5.5, 6.0, 4.5, np.nan, 5.0, 6.0, 6.0, 6.5],
+ "sex": ["M","M","M","F","F","F","F","F","F"],
+})
+```
+
+| patient_id | wave | bp | chol | sex |
+|------------|------|-----|------|-----|
+| 1 | 0 | 120 | 5.0 | M |
+| 1 | 1 | 122 | 5.5 | M |
+| 1 | 2 | 121 | 6.0 | M |
+| 2 | 0 | 130 | 4.5 | F |
+| 2 | 1 | 131 | NaN | F |
+| 2 | 2 | 128 | 5.0 | F |
+| ... | ... | ... | ... | ... |
+
+## Part 1: Long to Wide
+
+The animation below walks through every long-format row and shows where each value lands in the wide matrix.
+
+
+ [{ width="100%" loading="lazy" }](../assets/images/tutorials/long_wide_reshape/LongToWide.gif){ .expandable-media__trigger }
+ Click the image to expand it.
+
+
+The corresponding code is a single call:
+
+```python
+from scikit_longitudinal.data_preparation import LongitudinalDataset
+
+dataset = LongitudinalDataset(file_path=None, data_frame=long_df)
+
+wide = dataset.to_wide(
+ id_col="patient_id",
+ time_col="wave",
+ longitudinal_columns=["bp", "chol"],
+ static_columns=["sex"],
+ wave_format="{feature}_w{wave}", # default; shown for clarity
+ output_path=None, # set to a path to dump a CSV
+)
+
+print(wide)
+print("features_group :", dataset.feature_groups())
+print("static columns :", dataset.non_longitudinal_features())
+```
+
+Output:
+
+```
+ sex bp_w0 bp_w1 bp_w2 chol_w0 chol_w1 chol_w2
+patient_id
+1 M 120.0 122.0 121.0 5.0 5.5 6.0
+2 F 130.0 131.0 128.0 4.5 NaN 5.0
+3 F 110.0 112.0 115.0 6.0 6.0 6.5
+
+features_group : [[1, 2, 3], [4, 5, 6]]
+static columns : [0]
+```
+
+A few things to notice:
+
+- `wide` has one row per patient, with `static_columns` placed first, then each longitudinal feature expanded across all observed waves (oldest to newest).
+- Missing schedule slots are filled with `NaN`. No imputation is performed.
+- The dataset's `feature_groups()` and `non_longitudinal_features()` are updated in place, so the wide frame can be passed straight to a `Sklong` estimator or pipeline.
+
+### Common pitfalls
+
+| Situation | What `to_wide` does |
+|-----------|---------------------|
+| Two rows for the same `(patient_id, wave)` | `ValueError("Duplicate (id, time) rows in long dataframe; ...")` |
+| `sex` differs between waves of the same patient | `ValueError("Static columns vary within a subject: ...")` |
+| Column listed twice (or as both id/time and longitudinal) | `ValueError("listed more than once")` / `"cannot be a value/static column..."` |
+| Empty `longitudinal_columns=[]` | `ValueError("longitudinal_columns must list at least one longitudinal column.")` |
+
+## Part 2: Wide to Long
+
+The other direction, displays as follows:
+
+
+ [{ width="100%" loading="lazy" }](../assets/images/tutorials/long_wide_reshape/WideToLong.gif){ .expandable-media__trigger }
+ Click the image to expand it.
+
+
+`to_long` reads the dataset's own `features_group` to drive the un-pivot, so there is no need to re-state the column names. Starting from the same wide cohort built in Part 1 — the `LongitudinalDataset` already carries `features_group` and `non_longitudinal_features` after the `to_wide` call:
+
+```python
+from scikit_longitudinal.data_preparation import LongitudinalDataset
+
+dataset = LongitudinalDataset(file_path=None, data_frame=long_df)
+
+wide = dataset.to_wide(
+ id_col="patient_id",
+ time_col="wave",
+ longitudinal_columns=["bp", "chol"],
+ static_columns=["sex"],
+)
+
+print("features_group :", dataset.feature_groups())
+print("static columns :", dataset.non_longitudinal_features())
+# features_group : [[1, 2, 3], [4, 5, 6]]
+# static columns : [0]
+
+long_again = dataset.to_long(
+ feature_base_names=["bp", "chol"], # column names in the long output
+ id_col="patient_id",
+ time_col="wave",
+ keep_static=True,
+ output_path=None,
+)
+print(long_again.head(8))
+```
+
+Output:
+
+```
+ patient_id wave bp chol sex
+0 1 1 120.0 5.0 M
+1 1 2 122.0 5.5 M
+2 1 3 121.0 6.0 M
+3 2 1 130.0 4.5 F
+4 2 2 131.0 NaN F
+5 2 3 128.0 5.0 F
+6 3 1 110.0 6.0 F
+7 3 2 112.0 6.0 F
+```
+
+Notes:
+
+- Wave labels in the long output are positional (1, 2, 3, ...), one slot per group entry.
+- `dataset.feature_groups()` and `dataset.non_longitudinal_features()` are reset to `None` because they describe a wide layout that no longer exists.
+- `keep_static=False` would drop `sex` from the long frame.
+
+!!! tip "Where next?"
+ Review the [`LongitudinalDataset` API reference](../API/data_preparation/longitudinal_dataset.md) for all available parameters, or move on to the [Algorithm Adaptation tutorial](sklong_explore_your_first_estimator.md) to plug your fresh wide dataset into a longitudinal estimator.
diff --git a/docs/tutorials/overview.md b/docs/tutorials/overview.md
index 5ee495f..c10302c 100644
--- a/docs/tutorials/overview.md
+++ b/docs/tutorials/overview.md
@@ -44,6 +44,14 @@ In order to visualise what the library delivers, the figure below shows the high
[Read the tutorial](sklong_longitudinal_data_format.md)
+- __Long ⇄ Wide Reshape__
+
+ ---
+
+ Pivot messy long-format cohorts into the wide layout `Sklong` expects (and back) using `LongitudinalDataset.to_wide` / `to_long`.
+
+ [Read the tutorial](long_wide_reshape.md)
+
- __Data Preparation: Flatten Temporal Dependency for Scikit-Learn Estimators__
---
diff --git a/scikit_longitudinal/data_preparation/_longitudinal_reshape.py b/scikit_longitudinal/data_preparation/_longitudinal_reshape.py
new file mode 100644
index 0000000..b9177b3
--- /dev/null
+++ b/scikit_longitudinal/data_preparation/_longitudinal_reshape.py
@@ -0,0 +1,239 @@
+# pylint: disable=R0902,R0913,R0914,R0912,R0915
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import pandas as pd
+
+PAD_VALUE = -1
+
+
+def _format_wave_column(template: str, feature: str, wave: int) -> str:
+ """Render the wide-format column name for a given feature/wave pair."""
+ return template.format(feature=feature, wave=wave)
+
+
+def _validate_wide_inputs(
+ columns: Sequence[str],
+ features_group: Sequence[Sequence[int]],
+ non_longitudinal_features: Optional[Sequence[Union[int, str]]],
+) -> Tuple[List[List[int]], List[int]]:
+ """Normalise and validate `features_group` / `non_longitudinal_features` against `columns`."""
+ if not features_group:
+ raise ValueError("features_group must contain at least one group.")
+ n_columns = len(columns)
+ cleaned: List[List[int]] = []
+ seen: set[int] = set()
+ for grp_idx, group in enumerate(features_group):
+ if not isinstance(group, (list, tuple)):
+ raise ValueError(f"features_group[{grp_idx}] must be a list, got {type(group).__name__}.")
+ without_pad = [idx for idx in group if idx != PAD_VALUE]
+ if len(without_pad) < 2:
+ raise ValueError(
+ f"features_group[{grp_idx}] must list at least two waves; got {group}."
+ )
+ for idx in without_pad:
+ if not isinstance(idx, (int, np.integer)):
+ raise ValueError(
+ f"features_group[{grp_idx}] must contain integer indices; got {idx!r}."
+ )
+ if idx < 0 or idx >= n_columns:
+ raise ValueError(
+ f"features_group[{grp_idx}] index {idx} out of range [0, {n_columns})."
+ )
+ if idx in seen:
+ raise ValueError(
+ f"Index {idx} ('{columns[idx]}') appears in more than one group."
+ )
+ seen.add(idx)
+ cleaned.append(list(group))
+
+ if non_longitudinal_features is None:
+ non_long_indices: List[int] = []
+ else:
+ non_long_indices = []
+ for ref in non_longitudinal_features:
+ if isinstance(ref, (int, np.integer)):
+ idx = int(ref)
+ elif isinstance(ref, str):
+ if ref not in columns:
+ raise ValueError(f"non_longitudinal_features: column '{ref}' not found.")
+ idx = list(columns).index(ref)
+ else:
+ raise ValueError(
+ f"non_longitudinal_features must contain int or str, got {type(ref).__name__}."
+ )
+ if idx < 0 or idx >= n_columns:
+ raise ValueError(
+ f"non_longitudinal_features index {idx} out of range [0, {n_columns})."
+ )
+ if idx in seen:
+ raise ValueError(
+ f"Column '{columns[idx]}' is declared in features_group and "
+ "non_longitudinal_features simultaneously."
+ )
+ non_long_indices.append(idx)
+ return cleaned, non_long_indices
+
+
+def _validate_long_inputs(
+ df: pd.DataFrame,
+ id_col: str,
+ time_col: str,
+ longitudinal_columns: Sequence[str],
+ static_columns: Sequence[str],
+) -> None:
+ """Validate that the long-format inputs reference real, non-overlapping columns."""
+ if id_col not in df.columns:
+ raise ValueError(f"id_col '{id_col}' not in dataframe.")
+ if time_col not in df.columns:
+ raise ValueError(f"time_col '{time_col}' not in dataframe.")
+ if not longitudinal_columns:
+ raise ValueError("longitudinal_columns must list at least one longitudinal column.")
+ roles = {id_col, time_col}
+ seen: set[str] = set()
+ for col in list(longitudinal_columns) + list(static_columns):
+ if col not in df.columns:
+ raise ValueError(f"Column '{col}' not in dataframe.")
+ if col in roles:
+ raise ValueError(
+ f"Column '{col}' cannot be a value/static column and id/time column."
+ )
+ if col in seen:
+ raise ValueError(f"Column '{col}' listed more than once.")
+ seen.add(col)
+
+
+def long_to_wide(
+ df: pd.DataFrame,
+ *,
+ id_col: str,
+ time_col: str,
+ longitudinal_columns: Sequence[str],
+ static_columns: Sequence[str] = (),
+ wave_format: str = "{feature}_w{wave}",
+) -> Tuple[pd.DataFrame, List[List[int]], List[int]]:
+ """Pivot a long-format dataframe to wide format.
+
+ Returns the wide dataframe, a `features_group` describing the wave layout,
+ and the column indices of the static (non-longitudinal) columns.
+ """
+ _validate_long_inputs(df, id_col, time_col, longitudinal_columns, static_columns)
+
+ longitudinal_columns = list(longitudinal_columns)
+ static_columns = list(static_columns)
+
+ duplicates = df.duplicated(subset=[id_col, time_col])
+ if duplicates.any():
+ offending = df.loc[duplicates, [id_col, time_col]].head(3).to_dict("records")
+ raise ValueError(
+ f"Duplicate (id, time) rows in long dataframe; got e.g. {offending}."
+ )
+
+ if static_columns:
+ nunique = df.groupby(id_col)[static_columns].nunique(dropna=False)
+ bad = nunique[(nunique > 1).any(axis=1)]
+ if not bad.empty:
+ raise ValueError(
+ f"Static columns vary within a subject: {bad.index.tolist()[:3]}"
+ )
+
+ subject_ids = df[id_col].drop_duplicates().tolist()
+ wave_labels = sorted(df[time_col].dropna().unique().tolist())
+ n_subjects = len(subject_ids)
+ n_waves = len(wave_labels)
+ n_static = len(static_columns)
+ wave_to_pos: Dict = {w: i for i, w in enumerate(wave_labels)}
+
+ wide_columns: List[str] = []
+ feature_groups: List[List[int]] = []
+ for feat in longitudinal_columns:
+ grp: List[int] = []
+ for wave in wave_labels:
+ grp.append(n_static + len(wide_columns))
+ wide_columns.append(_format_wave_column(wave_format, feat, wave))
+ feature_groups.append(grp)
+
+ full_columns = list(static_columns) + wide_columns
+ wide_matrix = np.full((n_subjects, len(full_columns)), np.nan, dtype=object)
+ subject_to_row = {sid: i for i, sid in enumerate(subject_ids)}
+
+ if static_columns:
+ static_first = df.groupby(id_col, sort=False)[static_columns].first()
+ for sid in subject_ids:
+ wide_matrix[subject_to_row[sid], :n_static] = static_first.loc[sid].to_numpy()
+
+ for sid, frame in df.groupby(id_col, sort=False):
+ row = subject_to_row[sid]
+ for _, observation in frame.iterrows():
+ wave_pos = wave_to_pos[observation[time_col]]
+ for feat_idx, col in enumerate(longitudinal_columns):
+ j = n_static + feat_idx * n_waves + wave_pos
+ wide_matrix[row, j] = observation[col]
+
+ wide_df = pd.DataFrame(wide_matrix, columns=full_columns, index=subject_ids)
+ wide_df.index.name = id_col
+ for col in full_columns:
+ try:
+ wide_df[col] = pd.to_numeric(wide_df[col])
+ except (ValueError, TypeError):
+ pass
+
+ return wide_df, feature_groups, list(range(n_static))
+
+
+def wide_to_long(
+ df: pd.DataFrame,
+ *,
+ features_group: Sequence[Sequence[int]],
+ non_longitudinal_features: Optional[Sequence[Union[int, str]]] = None,
+ feature_base_names: Optional[Sequence[str]] = None,
+ id_col: str = "subject_id",
+ time_col: str = "wave",
+ keep_static: bool = True,
+) -> pd.DataFrame:
+ """Reshape a wide dataframe back to long format using `features_group`."""
+ columns = list(df.columns)
+
+ cleaned_groups, static_indices = _validate_wide_inputs(
+ columns, features_group, non_longitudinal_features
+ )
+
+ if non_longitudinal_features is None:
+ in_groups = {idx for grp in cleaned_groups for idx in grp if idx != PAD_VALUE}
+ static_indices = [i for i in range(len(columns)) if i not in in_groups]
+ static_cols = [columns[i] for i in static_indices]
+
+ if feature_base_names is None:
+ feature_base_names = [f"feature_{i}" for i in range(len(cleaned_groups))]
+ elif len(feature_base_names) != len(cleaned_groups):
+ raise ValueError(
+ f"feature_base_names has {len(feature_base_names)} entries but features_group has "
+ f"{len(cleaned_groups)} groups."
+ )
+
+ max_waves = max(len(g) for g in cleaned_groups)
+ pieces: List[pd.DataFrame] = []
+ indices = df.index.tolist()
+ for wave_pos in range(max_waves):
+ wave_frame = pd.DataFrame({id_col: indices, time_col: wave_pos + 1})
+ wave_frame.index = df.index
+ for base, group in zip(feature_base_names, cleaned_groups):
+ if wave_pos < len(group) and group[wave_pos] != PAD_VALUE:
+ wave_frame[base] = df.iloc[:, group[wave_pos]].values
+ else:
+ wave_frame[base] = np.nan
+ if keep_static and static_cols:
+ for col in static_cols:
+ wave_frame[col] = df[col].values
+ pieces.append(wave_frame)
+
+ long_df = pd.concat(pieces, axis=0, ignore_index=True)
+ long_df = long_df.sort_values([id_col, time_col], kind="mergesort").reset_index(drop=True)
+
+ return long_df
+
+
+__all__ = ["long_to_wide", "wide_to_long"]
diff --git a/scikit_longitudinal/data_preparation/longitudinal_dataset.py b/scikit_longitudinal/data_preparation/longitudinal_dataset.py
index 98b5879..5844886 100644
--- a/scikit_longitudinal/data_preparation/longitudinal_dataset.py
+++ b/scikit_longitudinal/data_preparation/longitudinal_dataset.py
@@ -1,13 +1,18 @@
import re
from functools import wraps
from pathlib import Path
-from typing import Any, Callable, List, Optional, Tuple, Union
+from typing import Any, Callable, List, Optional, Tuple, Union # noqa: F401
import arff
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
+from scikit_longitudinal.data_preparation._longitudinal_reshape import (
+ long_to_wide,
+ wide_to_long,
+)
+
# pylint: disable=W0212, R0902, W1514, E1101, R0904
def clean_padding(features_group: List[List[int]]) -> List[List[int]]:
@@ -860,3 +865,150 @@ def sety_test(self, y_test: pd.Series) -> None:
"""
self._y_test = y_test
+
+ def to_wide(
+ self,
+ *,
+ id_col: str,
+ time_col: str,
+ longitudinal_columns: List[str],
+ static_columns: List[str] = (),
+ wave_format: str = "{feature}_w{wave}",
+ output_path: Optional[Union[str, Path]] = None,
+ ) -> pd.DataFrame:
+ """Pivot the dataset's long-format data into wide format.
+
+ One row per `(subject, time)` becomes one row per subject, with each
+ longitudinal column expanded into one column per observed wave
+ (oldest \u2192 newest). The dataset's `data`, `feature_groups`, and
+ `non_longitudinal_features` are updated in place to match the new
+ layout.
+
+ Args:
+ id_col (str): Subject identifier column.
+ time_col (str): Observation time column.
+ longitudinal_columns (List[str]): Columns expanded across waves.
+ static_columns (List[str], optional): Columns kept as-is per subject.
+ Must be constant within each subject. Defaults to `()`.
+ wave_format (str): Python `str.format` template that names every wide column
+ produced from a longitudinal source column. Two placeholders are substituted
+ per cell: `{feature}` is the original long-format column name (e.g. `"bp"`),
+ and `{wave}` is the value taken straight from `time_col` for that observation
+ (rendered with its native type — typically an int like `0`, `1`, `2`, but
+ strings such as `"2008"` work too). Waves are emitted in sorted order, so the
+ template fully determines the wide schema: the default `"{feature}_w{wave}"`
+ yields `bp_w0, bp_w1, bp_w2, chol_w0, ...`; `"{feature}.t{wave}"` yields
+ `bp.t0, bp.t1, ...`; and `"{wave}_{feature}"` flips the order to
+ `0_bp, 1_bp, ...`. The template must contain both placeholders and must
+ produce unique column names across `(feature, wave)` pairs. Defaults to
+ `"{feature}_w{wave}"`.
+ output_path (Optional[Union[str, Path]]): If set, also write the wide dataframe
+ to a CSV file at this path.
+
+ Returns:
+ pd.DataFrame: The newly stored wide-format dataframe.
+
+ Raises:
+ ValueError: If no data is loaded, if `(id_col, time_col)` rows are duplicated,
+ if a static column varies within a subject, or if any referenced column
+ is missing.
+
+ Examples:
+ !!! example "Basic Usage"
+ ```python
+ from scikit_longitudinal.data_preparation import LongitudinalDataset
+
+ dataset = LongitudinalDataset(file_path=None, data_frame=long_df)
+ dataset.to_wide(
+ id_col="pid",
+ time_col="wave",
+ longitudinal_columns=["bp", "chol"],
+ static_columns=["sex"],
+ output_path="./wide.csv",
+ )
+ ```
+ """
+ if self._data is None:
+ raise ValueError("No data is loaded. Load data first.")
+
+ wide_df, feature_groups, non_long = long_to_wide(
+ self._data,
+ id_col=id_col,
+ time_col=time_col,
+ longitudinal_columns=list(longitudinal_columns),
+ static_columns=list(static_columns),
+ wave_format=wave_format,
+ )
+ self._data = wide_df
+ self._feature_groups = feature_groups
+ self._non_longitudinal_features = non_long
+ if output_path is not None:
+ wide_df.to_csv(output_path, index=True, na_rep="")
+ return wide_df
+
+ def to_long(
+ self,
+ *,
+ feature_base_names: Optional[List[str]] = None,
+ id_col: str = "subject_id",
+ time_col: str = "wave",
+ keep_static: bool = True,
+ output_path: Optional[Union[str, Path]] = None,
+ ) -> pd.DataFrame:
+ """Reshape the dataset's wide-format data into long format.
+
+ Uses the dataset's own `feature_groups` (set via `setup_features_group`)
+ to drive the reshape. The dataset's `data` is replaced with the long
+ dataframe and `feature_groups` / `non_longitudinal_features` are
+ cleared, since they no longer apply to a long layout.
+
+ Args:
+ feature_base_names (Optional[List[str]]): Names for the long-format feature
+ columns, one per group. Defaults to `["feature_0", "feature_1", ...]`.
+ id_col (str): Output id column name. Defaults to `"subject_id"`.
+ time_col (str): Output wave column name. Defaults to `"wave"`.
+ keep_static (bool): Whether to repeat static columns on every long row.
+ Defaults to `True`.
+ output_path (Optional[Union[str, Path]]): If set, also write the long dataframe
+ to a CSV file at this path.
+
+ Returns:
+ pd.DataFrame: The newly stored long-format dataframe.
+
+ Raises:
+ ValueError: If no data is loaded or `setup_features_group(...)` has not been
+ called yet.
+
+ Examples:
+ !!! example "Basic Usage"
+ ```python
+ from scikit_longitudinal.data_preparation import LongitudinalDataset
+
+ dataset = LongitudinalDataset('./stroke.csv')
+ dataset.load_data()
+ dataset.setup_features_group("elsa")
+ dataset.to_long(feature_base_names=["bp", "chol"], output_path="./long.csv")
+ ```
+ """
+ if self._data is None:
+ raise ValueError("No data is loaded. Load data first.")
+ if self._feature_groups is None:
+ raise ValueError(
+ "setup_features_group(...) must be called before to_long()."
+ )
+
+ long_df = wide_to_long(
+ self._data,
+ features_group=self._feature_groups,
+ non_longitudinal_features=self._non_longitudinal_features,
+ feature_base_names=feature_base_names,
+ id_col=id_col,
+ time_col=time_col,
+ keep_static=keep_static,
+ )
+ self._data = long_df
+ self._feature_groups = None
+ self._non_longitudinal_features = None
+ if output_path is not None:
+ long_df.to_csv(output_path, index=False, na_rep="")
+ return long_df
diff --git a/scikit_longitudinal/estimators/trees/TpT/TpT_decision_tree.py b/scikit_longitudinal/estimators/trees/TpT/TpT_decision_tree.py
index d6dba37..149c871 100644
--- a/scikit_longitudinal/estimators/trees/TpT/TpT_decision_tree.py
+++ b/scikit_longitudinal/estimators/trees/TpT/TpT_decision_tree.py
@@ -3,15 +3,12 @@
import math
import warnings
from numbers import Real
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Union
import numpy as np
-import pandas as pd
from sklearn.tree import DecisionTreeClassifier, _tree
from sklearn.utils._param_validation import Interval
-from ._preprocessing import long_to_wide
-
class TpTDecisionTreeClassifier(DecisionTreeClassifier):
"""
@@ -25,31 +22,17 @@ class TpTDecisionTreeClassifier(DecisionTreeClassifier):
splitter therefore tends to prefer earlier waves (while allowing later waves deeper in the tree) unless
later observations bring a substantially stronger signal.
- ??? note "LONG vs wide input — *[Soon To Be Deprecated](https://github.com/simonprovost/scikit-longitudinal/issues/64)*"
- TpT internally operates on a **wide** matrix (features expanded over waves). If `assume_long_format=True`,
- the classifier can accept a LONG-format dataframe and will convert it to the expected wide representation
- before fitting (using `id_col`, `time_col`, `duration_col`, `time_step`, and `max_horizon`).
-
- - LONG-format: one row per (subject, time) observation.
- - wide format: one row per subject, with features duplicated across waves.
-
- The conversion fills feature values up to each subject's duration/horizon and leaves NaNs beyond,
- enabling "duration leaves" in the TpT logic.
-
Args:
gamma (float, optional):
Time-penalty rate $\\gamma$ in the factor $e^{-\\gamma \\Delta t}$.
- If not provided, falls back to `threshold_gain` (for backward compatibility).
+ If not provided, falls back to `threshold_gain`.
threshold_gain (float, optional):
- Backward-compatible alias for `gamma`. If both are provided, `gamma` takes precedence.
+ Alias for `gamma`. If both are provided, `gamma` takes precedence.
features_group (List[List[int]], optional):
A list of lists where each inner list contains indices of features corresponding to a specific longitudinal
attribute across different waves. The order within each inner list reflects the temporal sequence, with the
first element being the oldest wave and the last being the most recent. For example, `[[0,1],[2,3]]` indicates
two longitudinal attributes, each with two waves (e.g., 0: oldest, 1: recent; 2: oldest, 3: recent).
- max_horizon (int, optional):
- Maximum temporal horizon for predictions. Limits the time range explored during tree construction.
- If None, no limit is applied.
criterion (str, default="entropy"):
The function to measure the quality of a split. Fixed to "entropy" for this algorithm; do not change.
splitter (str, default="TpT"):
@@ -180,15 +163,8 @@ class TpTDecisionTreeClassifier(DecisionTreeClassifier):
def __init__(
self,
gamma: Optional[float] = None,
- threshold_gain: Optional[float] = None, # deprecated alias for gamma
+ threshold_gain: Optional[float] = None,
features_group: Optional[List[List[int]]] = None,
- max_horizon: Optional[int] = None,
- id_col: Optional[str] = None,
- time_col: Optional[str] = None,
- duration_col: Optional[str] = None,
- time_step: float = 1.0,
- assume_long_format: bool = False,
- long_feature_columns: Optional[List[str]] = None,
criterion: str = "entropy",
splitter: str = "TpT",
max_depth: Optional[int] = None,
@@ -205,21 +181,10 @@ def __init__(
monotonic_cst: Optional[List[int]] = None,
min_penalized_gain: float = 0.0,
):
- # Resolve gamma with backward-compatible alias
_gamma = gamma if gamma is not None else (threshold_gain if threshold_gain is not None else 0.0015)
self.gamma = float(_gamma)
- # Keep attribute name threshold_gain for downstream Cython param compatibility
self.threshold_gain = self.gamma
self.features_group = features_group
- self.max_horizon = max_horizon
- self.id_col = id_col
- self.time_col = time_col
- self.duration_col = duration_col
- self.time_step = float(time_step)
- self.assume_long_format = assume_long_format
- self.long_feature_columns = long_feature_columns
- self._uses_long_format = False
- self._expected_n_features_: Optional[int] = None
self.min_penalized_gain = float(min_penalized_gain)
if monotonic_cst is not None:
@@ -251,138 +216,37 @@ def __init__(
)
def fit(self, X, y, *args, **kwargs):
- """Fit the classifier, optionally preparing LONG-format data on the fly."""
-
- X_prepared, y_prepared = self._prepare_training_data(X, y)
-
- fitted = super().fit(X_prepared, y_prepared, *args, **kwargs)
-
+ """
+ Fit the Time-penalised Trees (TpT) Decision Tree Classifier to the training data.
+
+ This method trains the classifier using the provided training data and labels. It requires the `features_group`
+ parameter to be set, as the time-penalised splitter relies on it to read the wave index of each candidate split.
+
+ Args:
+ X (array-like of shape (n_samples, n_features)):
+ The training input samples in wide format (features expanded over waves).
+ y (array-like of shape (n_samples,)):
+ The target values (class labels).
+ *args:
+ Additional positional arguments passed to the superclass `fit` method.
+ **kwargs:
+ Additional keyword arguments passed to the superclass `fit` method.
+
+ Returns:
+ TpTDecisionTreeClassifier:
+ The fitted classifier instance.
+
+ Raises:
+ ValueError:
+ If `features_group` is not provided, as it is required for longitudinal functionality.
+ """
+ if self.features_group is None:
+ raise ValueError("The features_group parameter must be provided.")
+ fitted = super().fit(X, y, *args, **kwargs)
if self.min_penalized_gain > 0.0 and hasattr(self, "tree_"):
self._prune_penalized_gain()
-
return fitted
- # ------------------------------------------------------------------
- # Data preparation utilities
- # ------------------------------------------------------------------
- def _ensure_series(self, y, X) -> pd.Series:
- if isinstance(y, pd.DataFrame):
- if y.shape[1] != 1:
- raise ValueError("Target dataframe must have exactly one column for classification.")
- y_series = y.iloc[:, 0]
- elif isinstance(y, pd.Series):
- y_series = y
- else:
- y_series = pd.Series(np.asarray(y).ravel(), index=getattr(X, "index", None))
-
- if getattr(X, "shape", None) is not None and len(y_series) != len(X):
- raise ValueError("Target vector length must match number of observations.")
-
- if y_series.isna().any():
- raise ValueError("NaN targets are not supported in TpTDecisionTreeClassifier.")
-
- return y_series
-
- def _prepare_training_data(self, X, y):
- """Return (X_prepared, y_prepared), handling LONG-format input if required."""
-
- long_format = self.assume_long_format or (
- self.features_group is None and isinstance(X, pd.DataFrame)
- )
-
- if not long_format:
- self._uses_long_format = False
- if self.features_group is None:
- raise ValueError(
- "When providing wide-format data you must also set features_group to describe waves."
- )
- return X, y
-
- if not isinstance(X, pd.DataFrame):
- raise ValueError(
- "LONG-format input requires a pandas DataFrame with id/time/duration columns."
- )
-
- if not all([self.id_col, self.time_col, self.duration_col]):
- raise ValueError(
- "id_col, time_col and duration_col must be specified to consume LONG-format data."
- )
-
- y_series = self._ensure_series(y, X)
- wide_data = long_to_wide(
- X,
- id_col=self.id_col,
- time_col=self.time_col,
- duration_col=self.duration_col,
- time_step=self.time_step,
- max_horizon=self.max_horizon,
- feature_columns=self.long_feature_columns,
- )
-
- subject_targets = (
- y_series.groupby(X[self.id_col]).first().rename(index=lambda idx: str(idx))
- )
- reindexed_targets = subject_targets.reindex(wide_data.subject_ids)
- if reindexed_targets.isna().any():
- missing = reindexed_targets[reindexed_targets.isna()].index.tolist()
- raise ValueError(
- "Missing target values for subjects after aggregation: " + ", ".join(map(str, missing))
- )
-
- wide_data.y = reindexed_targets.to_numpy()
-
- self.features_group = wide_data.feature_groups
- self._uses_long_format = True
- self._wide_feature_names_ = wide_data.feature_names
- self._subject_ids_ = wide_data.subject_ids
- self._feature_columns_long_ = list(wide_data.feature_columns)
- self._n_waves_ = len(wide_data.time_indices)
- self._expected_n_features_ = wide_data.X.shape[1]
-
- return wide_data.X, wide_data.y
-
- def _prepare_long_dataset(self, X: pd.DataFrame) -> Tuple[np.ndarray, List[str]]:
- if not self._uses_long_format:
- raise ValueError("This estimator was not fitted on LONG-format data.")
-
- data = long_to_wide(
- X,
- id_col=self.id_col,
- time_col=self.time_col,
- duration_col=self.duration_col,
- time_step=self.time_step,
- max_horizon=self.max_horizon,
- feature_columns=self._feature_columns_long_,
- )
-
- expected_features = getattr(self, "_expected_n_features_", None)
- if expected_features is None:
- expected_features = data.X.shape[1]
- if data.X.shape[1] < expected_features:
- pad = np.full((data.X.shape[0], expected_features - data.X.shape[1]), np.nan)
- X_wide = np.concatenate([data.X, pad], axis=1)
- elif data.X.shape[1] > expected_features:
- X_wide = data.X[:, :expected_features]
- else:
- X_wide = data.X
-
- return X_wide, data.subject_ids
-
- # ------------------------------------------------------------------
- # Prediction helpers
- # ------------------------------------------------------------------
- def predict(self, X): # type: ignore[override]
- if self._uses_long_format and isinstance(X, pd.DataFrame):
- X_wide, _ = self._prepare_long_dataset(X)
- return super().predict(X_wide)
- return super().predict(X)
-
- def predict_proba(self, X): # type: ignore[override]
- if self._uses_long_format and isinstance(X, pd.DataFrame):
- X_wide, _ = self._prepare_long_dataset(X)
- return super().predict_proba(X_wide)
- return super().predict_proba(X)
-
def _more_tags(self):
tags = super()._more_tags()
tags["allow_nan"] = True
diff --git a/scikit_longitudinal/estimators/trees/TpT/TpT_decision_tree_regressor.py b/scikit_longitudinal/estimators/trees/TpT/TpT_decision_tree_regressor.py
index aa103ed..0163071 100644
--- a/scikit_longitudinal/estimators/trees/TpT/TpT_decision_tree_regressor.py
+++ b/scikit_longitudinal/estimators/trees/TpT/TpT_decision_tree_regressor.py
@@ -2,15 +2,11 @@
import warnings
from numbers import Real
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Union
-import numpy as np
-import pandas as pd
from sklearn.tree import DecisionTreeRegressor
from sklearn.utils._param_validation import Interval
-from ._preprocessing import long_to_wide
-
class TpTDecisionTreeRegressor(DecisionTreeRegressor):
"""
@@ -25,42 +21,15 @@ class TpTDecisionTreeRegressor(DecisionTreeRegressor):
tends to prefer earlier waves (while allowing later waves deeper in the tree) unless later observations
bring a substantially stronger signal.
- ??? note "LONG vs wide input — *[Soon To Be Deprecated](https://github.com/simonprovost/scikit-longitudinal/issues/64)*"
- TpT internally operates on a **wide** matrix (features expanded over waves). If `assume_long_format=True`,
- the regressor can accept a LONG-format dataframe and will convert it to the expected wide representation
- before fitting (using `id_col`, `time_col`, `duration_col`, `time_step`, and `max_horizon`).
-
- - LONG-format: one row per (subject, time) observation.
- - wide format: one row per subject, with features duplicated across waves.
-
- The conversion fills feature values up to each subject's duration/horizon and leaves NaNs beyond,
- enabling "duration leaves" in the TpT logic.
-
Args:
gamma (float, optional):
Time-penalty rate $\\gamma$ in $e^{-\\gamma \\Delta t}$. If not provided, falls back to
- `threshold_gain` for backward compatibility.
+ `threshold_gain`.
threshold_gain (float, optional):
- Backward-compatible alias for `gamma`. If both are provided, `gamma` takes precedence. (Internally
+ Alias for `gamma`. If both are provided, `gamma` takes precedence. (Internally
reused to match existing Cython parameter naming.)
features_group (List[List[int]], optional):
- Temporal grouping of feature indices (waves per covariate). Required when using wide-format input.
- If `assume_long_format=True`, this can be inferred/constructed during preprocessing depending on
- how wide features are generated.
- max_horizon (int, optional):
- Optional cap for the horizon considered during LONG-format preprocessing.
- id_col (str, optional):
- Subject identifier column name in LONG-format.
- time_col (str, optional):
- Observation time column name in LONG-format.
- duration_col (str, optional):
- Subject-specific horizon/duration column name in LONG-format.
- time_step (float, default=1.0):
- Temporal discretisation step used to map times to wave indices.
- assume_long_format (bool, default=False):
- If True, interpret `X` as LONG-format and convert to wide prior to fitting.
- long_feature_columns (List[str], optional):
- Subset of LONG-format columns to treat as features.
+ Temporal grouping of feature indices (waves per covariate).
criterion (str, default="friedman_mse"):
Split criterion for regression. The intended criterion is MSE / variance reduction. (Other criteria
may not be supported depending on the current Cython implementation.)
@@ -97,33 +66,21 @@ class TpTDecisionTreeRegressor(DecisionTreeRegressor):
The underlying fitted tree structure.
feature_importances_ (ndarray of shape (n_features,)):
Impurity-based feature importances (variance reduction based).
- _wide_feature_names_ (List[str]):
- Names of generated wide features (when LONG-format preprocessing is enabled).
- _subject_ids_ (List[str]):
- Subject ids aligned with the wide matrix rows (when LONG-format preprocessing is enabled).
Examples:
!!! example "Basic Usage"
```python
- import pandas as pd
from scikit_longitudinal.estimators.trees import TpTDecisionTreeRegressor
- df_long = pd.read_csv("my_longitudinal_dataset.csv")
- y = df_long["target"]
- X = df_long.drop(columns=["target"])
-
+ features_group = [[0, 1], [2, 3]]
reg = TpTDecisionTreeRegressor(
gamma=0.01,
- assume_long_format=True,
- id_col="id",
- time_col="time_point",
- duration_col="duration",
- time_step=1.0,
+ features_group=features_group,
max_depth=4,
random_state=0,
)
- reg.fit(X, y)
- preds = reg.predict(X)
+ reg.fit(X_wide, y)
+ preds = reg.predict(X_wide)
```
"""
@@ -138,13 +95,6 @@ def __init__(
gamma: Optional[float] = None,
threshold_gain: Optional[float] = None,
features_group: Optional[List[List[int]]] = None,
- max_horizon: Optional[int] = None,
- id_col: Optional[str] = None,
- time_col: Optional[str] = None,
- duration_col: Optional[str] = None,
- time_step: float = 1.0,
- assume_long_format: bool = False,
- long_feature_columns: Optional[List[str]] = None,
criterion: str = "friedman_mse",
splitter: str = "TpT",
max_depth: Optional[int] = None,
@@ -163,15 +113,6 @@ def __init__(
self.gamma = float(_gamma)
self.threshold_gain = self.gamma
self.features_group = features_group
- self.max_horizon = max_horizon
- self.id_col = id_col
- self.time_col = time_col
- self.duration_col = duration_col
- self.time_step = float(time_step)
- self.assume_long_format = assume_long_format
- self.long_feature_columns = long_feature_columns
- self._uses_long_format = False
- self._expected_n_features_: Optional[int] = None
if monotonic_cst is not None:
warnings.warn(
@@ -200,117 +141,34 @@ def __init__(
features_group=self.features_group,
)
-
- def _ensure_series(self, y, X) -> pd.Series:
- if isinstance(y, pd.DataFrame):
- if y.shape[1] != 1:
- raise ValueError("Target dataframe must have exactly one column for regression.")
- y_series = y.iloc[:, 0]
- elif isinstance(y, pd.Series):
- y_series = y
- else:
- y_series = pd.Series(np.asarray(y).ravel(), index=getattr(X, "index", None))
-
- if getattr(X, "shape", None) is not None and len(y_series) != len(X):
- raise ValueError("Target vector length must match number of observations.")
- if y_series.isna().any():
- raise ValueError("NaN targets are not supported in TpTDecisionTreeRegressor.")
- return y_series
-
- def _prepare_training_data(self, X, y):
- long_format = self.assume_long_format or (
- self.features_group is None and isinstance(X, pd.DataFrame)
- )
-
- if not long_format:
- self._uses_long_format = False
- if self.features_group is None:
- raise ValueError(
- "When providing wide-format data you must also set features_group to describe waves."
- )
- return X, y
-
- if not isinstance(X, pd.DataFrame):
- raise ValueError(
- "LONG-format input requires a pandas DataFrame with id/time/duration columns."
- )
-
- if not all([self.id_col, self.time_col, self.duration_col]):
- raise ValueError(
- "id_col, time_col and duration_col must be specified to consume LONG-format data."
- )
-
- y_series = self._ensure_series(y, X)
- wide_data = long_to_wide(
- X,
- id_col=self.id_col,
- time_col=self.time_col,
- duration_col=self.duration_col,
- time_step=self.time_step,
- max_horizon=self.max_horizon,
- feature_columns=self.long_feature_columns,
- )
-
- subject_targets = (
- y_series.groupby(X[self.id_col]).first().rename(index=lambda idx: str(idx))
- )
- reindexed_targets = subject_targets.reindex(wide_data.subject_ids)
- if reindexed_targets.isna().any():
- missing = reindexed_targets[reindexed_targets.isna()].index.tolist()
- raise ValueError(
- "Missing target values for subjects after aggregation: " + ", ".join(map(str, missing))
- )
-
- wide_data.y = reindexed_targets.to_numpy(dtype=np.float64)
-
- self.features_group = wide_data.feature_groups
- self._uses_long_format = True
- self._wide_feature_names_ = wide_data.feature_names
- self._subject_ids_ = wide_data.subject_ids
- self._feature_columns_long_ = list(wide_data.feature_columns)
- self._n_waves_ = len(wide_data.time_indices)
- self._expected_n_features_ = wide_data.X.shape[1]
-
- return wide_data.X, wide_data.y
-
- def _prepare_long_dataset(self, X: pd.DataFrame) -> Tuple[np.ndarray, List[str]]:
- if not self._uses_long_format:
- raise ValueError("This estimator was not fitted on LONG-format data.")
-
- data = long_to_wide(
- X,
- id_col=self.id_col,
- time_col=self.time_col,
- duration_col=self.duration_col,
- time_step=self.time_step,
- max_horizon=self.max_horizon,
- feature_columns=self._feature_columns_long_,
- )
-
- expected = getattr(self, "_expected_n_features_", None)
- if expected is None:
- expected = data.X.shape[1]
-
- if data.X.shape[1] < expected:
- pad = np.full((data.X.shape[0], expected - data.X.shape[1]), np.nan)
- X_wide = np.concatenate([data.X, pad], axis=1)
- elif data.X.shape[1] > expected:
- X_wide = data.X[:, :expected]
- else:
- X_wide = data.X
-
- return X_wide, data.subject_ids
-
-
def fit(self, X, y, *args, **kwargs): # type: ignore[override]
- X_prepared, y_prepared = self._prepare_training_data(X, y)
- return super().fit(X_prepared, y_prepared, *args, **kwargs)
-
- def predict(self, X): # type: ignore[override]
- if self._uses_long_format and isinstance(X, pd.DataFrame):
- X_wide, _ = self._prepare_long_dataset(X)
- return super().predict(X_wide)
- return super().predict(X)
+ """
+ Fit the Time-penalised Trees (TpT) Decision Tree Regressor to the training data.
+
+ This method trains the regressor using the provided training data and targets. It requires the `features_group`
+ parameter to be set, as the time-penalised splitter relies on it to read the wave index of each candidate split.
+
+ Args:
+ X (array-like of shape (n_samples, n_features)):
+ The training input samples in wide format (features expanded over waves).
+ y (array-like of shape (n_samples,)):
+ The target values (continuous).
+ *args:
+ Additional positional arguments passed to the superclass `fit` method.
+ **kwargs:
+ Additional keyword arguments passed to the superclass `fit` method.
+
+ Returns:
+ TpTDecisionTreeRegressor:
+ The fitted regressor instance.
+
+ Raises:
+ ValueError:
+ If `features_group` is not provided, as it is required for longitudinal functionality.
+ """
+ if self.features_group is None:
+ raise ValueError("The features_group parameter must be provided.")
+ return super().fit(X, y, *args, **kwargs)
def _more_tags(self):
tags = super()._more_tags()
diff --git a/scikit_longitudinal/estimators/trees/TpT/_preprocessing.py b/scikit_longitudinal/estimators/trees/TpT/_preprocessing.py
deleted file mode 100644
index 6ec7b0f..0000000
--- a/scikit_longitudinal/estimators/trees/TpT/_preprocessing.py
+++ /dev/null
@@ -1,257 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import Dict, Iterable, List, Optional, Sequence, Tuple
-
-import numpy as np
-import pandas as pd
-
-
-@dataclass
-class WideTpTData:
- """Container holding the wide-format representation expected by TpT splitters."""
-
- X: np.ndarray
- y: np.ndarray
- feature_groups: List[List[int]]
- feature_names: List[str]
- subject_ids: List[str]
- time_indices: np.ndarray
- feature_columns: Sequence[str]
-
-
-def _validate_long_dataframe(
- df: pd.DataFrame,
- id_col: str,
- time_col: str,
- duration_col: str,
- feature_columns: Optional[Sequence[str]] = None,
-) -> Tuple[pd.DataFrame, List[str]]:
- """Validate and normalize a LONG-format longitudinal dataframe for TpT preprocessing.
-
- This helper checks that the structural columns are present (subject id, time,
- duration), determines which columns should be treated as covariates, and
- returns a reordered copy of the dataframe sorted by (id, time).
-
- The returned dataframe has columns ordered as::
-
- feature_columns + [id_col, time_col, duration_col]
-
- and is sorted with a stable sort (mergesort) to preserve determinism.
-
- Parameters
- ----------
- df : pandas.DataFrame
- Input LONG-format dataframe containing one row per subject/time observation.
- id_col : str
- Name of the subject identifier column.
- time_col : str
- Name of the observation time column.
- duration_col : str
- Name of the subject-specific duration / horizon column.
- feature_columns : sequence of str, optional
- Columns to use as covariates. If ``None``, all columns except
- ``id_col``, ``time_col`` and ``duration_col`` are used. Any accidental
- inclusion of those structural columns is removed.
-
- Returns
- -------
- ordered : pandas.DataFrame
- A reordered copy of ``df`` restricted to the selected feature columns
- plus the structural columns, sorted by ``(id_col, time_col)``.
- feature_columns : list of str
- The finalized list of covariate column names (in the order used in
- ``ordered``).
-
- Raises
- ------
- ValueError
- If any required structural column is missing, or if no covariate column
- can be determined after excluding ``id_col``, ``time_col`` and
- ``duration_col``.
- """
-
- missing = {id_col, time_col, duration_col} - set(df.columns)
- if missing:
- raise ValueError(
- "Missing required columns for LONG-format input: " + ", ".join(sorted(missing))
- )
-
- disallowed = {id_col, time_col, duration_col}
- if feature_columns is None:
- feature_columns = [
- col
- for col in df.columns
- if col not in disallowed
- ]
- if not feature_columns:
- raise ValueError("No feature columns found after excluding id/time/duration columns.")
- feature_columns = [col for col in feature_columns if col not in disallowed]
- if not feature_columns:
- raise ValueError("After removing id/time/duration columns no feature columns remain.")
-
- ordered = (
- df[feature_columns + [id_col, time_col, duration_col]]
- .copy()
- .sort_values([id_col, time_col], kind="mergesort")
- .reset_index(drop=True)
- )
- return ordered, list(feature_columns)
-
-
-def long_to_wide(
- df: pd.DataFrame,
- *,
- id_col: str,
- time_col: str,
- duration_col: str,
- time_step: float = 1.0,
- max_horizon: Optional[float] = None,
- feature_columns: Optional[Sequence[str]] = None,
- target: Optional[pd.Series] = None,
-) -> WideTpTData:
- """Convert a LONG-format dataframe into the wide matrix required by TpT splitters.
-
- Parameters
- ----------
- df:
- Longitudinal dataframe containing one row per subject/time observation.
- id_col, time_col, duration_col:
- Column names identifying the subject id, the observation time and the
- subject-specific censoring time (horizon).
- time_step:
- Temporal granularity used to discretise times into waves.
- max_horizon:
- Optional cap for the temporal horizon. If omitted, the maximum duration
- observed in ``df`` is used.
- feature_columns:
- Optional subset of columns to treat as features. When omitted, all
- columns except ``id_col``, ``time_col`` and ``duration_col`` are used.
-
- Returns
- -------
- WideTpTData
- Dataclass with the dense feature matrix, feature groups and metadata.
- """
-
- ordered, feature_columns = _validate_long_dataframe(
- df,
- id_col=id_col,
- time_col=time_col,
- duration_col=duration_col,
- feature_columns=feature_columns,
- )
-
- # Discretise times into integer wave indices.
- if time_step <= 0:
- raise ValueError("time_step must be strictly positive.")
-
- time_index = np.floor_divide(
- np.asarray(ordered[time_col], dtype=np.float64) + 1e-9,
- time_step,
- ).astype(np.int64)
-
- duration_index = np.floor_divide(
- np.asarray(ordered[duration_col], dtype=np.float64) + 1e-9,
- time_step,
- ).astype(np.int64)
-
- max_observed_wave = int(np.floor(time_index.max())) if len(time_index) else 0
- if max_horizon is not None:
- max_horizon_wave = int(np.floor(max_horizon / time_step))
- else:
- max_horizon_wave = max_observed_wave
-
- max_time_index = min(max_observed_wave, max_horizon_wave)
- if max_time_index < 0:
- raise ValueError("No valid time horizon could be determined from the input data.")
-
- duration_index = np.minimum(duration_index, max_time_index)
- if max_time_index < 0:
- raise ValueError("No valid time horizon could be determined from the input data.")
-
- n_waves = max_time_index + 1
-
- # Build subject order to ensure deterministic layout.
- subjects: List[str] = ordered[id_col].astype(str).unique().tolist()
- subject_to_row: Dict[str, int] = {subject: idx for idx, subject in enumerate(subjects)}
- n_subjects = len(subjects)
-
- n_features = len(feature_columns)
- X = np.full((n_subjects, n_features * n_waves), np.nan, dtype=np.float64)
- feature_names: List[str] = []
- feature_groups: List[List[int]] = []
-
- # Pre-compute column offsets per feature.
- for feat_idx, feature in enumerate(feature_columns):
- group: List[int] = []
- for wave in range(n_waves):
- column_index = feat_idx * n_waves + wave
- feature_names.append(f"{feature}_t{wave}")
- group.append(column_index)
- feature_groups.append(group)
-
- # Iterate over subjects and fill the matrix with forward-filled measurements
- # until the subject horizon. Beyond the horizon we keep NaN so that TpT can
- # route those samples to duration leaves.
- grouped = ordered.assign(_time_index=time_index, _duration_index=duration_index).groupby(id_col, sort=False)
-
- if target is not None:
- # Align target with the ordered dataframe to preserve subject order.
- target_aligned = target.loc[ordered.index] if isinstance(target, pd.Series) else pd.Series(target)
- else:
- target_aligned = None
-
- subject_targets = np.full(n_subjects, np.nan)
-
- for subject, frame in grouped:
- row_idx = subject_to_row[str(subject)]
- first_original_idx = frame.index[0]
- frame = frame.reset_index(drop=True)
- if target_aligned is not None:
- subject_targets[row_idx] = target_aligned.loc[first_original_idx]
- times = frame.pop("_time_index").to_numpy(dtype=np.int64)
- durations = frame.pop("_duration_index").to_numpy(dtype=np.int64)
- subject_duration = int(durations.max())
- entry_wave = int(times.min())
- initial_values = {
- feature: frame.at[0, feature]
- for feature in feature_columns
- }
-
- # Track latest available values for each feature.
- last_values = {feature: np.nan for feature in feature_columns}
- pointer = 0
-
- for wave in range(n_waves):
- while pointer < len(frame) and times[pointer] <= wave:
- for feature in feature_columns:
- last_values[feature] = frame.at[pointer, feature]
- pointer += 1
-
- if wave > subject_duration:
- # Subject no longer observed: keep NaNs so that the splitter treats
- # them as terminated subjects.
- continue
-
- for feat_idx, feature in enumerate(feature_columns):
- value = last_values[feature]
- if np.isnan(value) and wave < entry_wave:
- value = initial_values.get(feature, np.nan)
- if np.isnan(value):
- continue
- X[row_idx, feat_idx * n_waves + wave] = value
-
- time_indices = np.arange(n_waves, dtype=np.int64)
-
- # Targets are not produced here; the caller is responsible for selecting one
- # label per subject.
- return WideTpTData(
- X=X,
- y=subject_targets if target_aligned is not None else np.empty(n_subjects, dtype=np.float64),
- feature_groups=feature_groups,
- feature_names=feature_names,
- subject_ids=subjects,
- time_indices=time_indices,
- feature_columns=feature_columns,
- )
diff --git a/scikit_longitudinal/tests/test_longitudinal_reshape.py b/scikit_longitudinal/tests/test_longitudinal_reshape.py
new file mode 100644
index 0000000..2999057
--- /dev/null
+++ b/scikit_longitudinal/tests/test_longitudinal_reshape.py
@@ -0,0 +1,453 @@
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from scikit_longitudinal.data_preparation import LongitudinalDataset
+from scikit_longitudinal.data_preparation._longitudinal_reshape import PAD_VALUE
+
+
+def make_long_frame() -> pd.DataFrame:
+ """Three subjects, three waves, two longitudinal features, one static."""
+ return pd.DataFrame(
+ {
+ "subject_id": [1, 1, 1, 2, 2, 2, 3, 3, 3],
+ "time": [0, 1, 2, 0, 1, 2, 0, 1, 2],
+ "bp": [120.0, 122.0, 121.0, 130.0, 131.0, 128.0, 110.0, 112.0, 115.0],
+ "chol": [5.0, 5.5, 6.0, 4.5, 4.7, 5.0, 6.0, 6.0, 6.5],
+ "sex": ["M", "M", "M", "F", "F", "F", "F", "F", "F"],
+ }
+ )
+
+
+def make_wide_frame() -> pd.DataFrame:
+ """Wide-format mirror with two longitudinal attrs over two waves and one static."""
+ return pd.DataFrame(
+ {
+ "bp_w1": [120.0, 130.0, 110.0],
+ "bp_w2": [122.0, 131.0, 112.0],
+ "chol_w1": [5.0, 4.5, 6.0],
+ "chol_w2": [5.5, 5.0, 6.0],
+ "sex": ["M", "F", "F"],
+ },
+ index=[1, 2, 3],
+ )
+
+
+def _dataset(df: pd.DataFrame) -> LongitudinalDataset:
+ return LongitudinalDataset(file_path=None, data_frame=df)
+
+
+# ---------------------------------------------------------------------------
+# LongitudinalDataset.to_wide
+# ---------------------------------------------------------------------------
+
+
+class TestToWide:
+ def test_pivots_and_updates_dataset_state(self):
+ dataset = _dataset(make_long_frame())
+ wide = dataset.to_wide(
+ id_col="subject_id",
+ time_col="time",
+ longitudinal_columns=["bp", "chol"],
+ static_columns=["sex"],
+ )
+ assert dataset.data is wide
+ assert list(wide.columns)[:4] == ["sex", "bp_w0", "bp_w1", "bp_w2"]
+ assert wide.loc[1, "bp_w0"] == 120.0
+ assert wide.loc[2, "chol_w1"] == 4.7
+ assert dataset.feature_groups() == [[1, 2, 3], [4, 5, 6]]
+ assert dataset.non_longitudinal_features() == [0]
+
+ def test_no_static_columns(self):
+ wide = _dataset(make_long_frame()).to_wide(
+ id_col="subject_id",
+ time_col="time",
+ longitudinal_columns=["bp", "chol"],
+ )
+ assert list(wide.columns) == [
+ "bp_w0", "bp_w1", "bp_w2", "chol_w0", "chol_w1", "chol_w2"
+ ]
+
+ def test_handles_unsorted_long_input(self):
+ df = make_long_frame().sample(frac=1.0, random_state=7).reset_index(drop=True)
+ wide = _dataset(df).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp", "chol"], static_columns=["sex"],
+ )
+ # Values must follow (subject, wave) regardless of input row order.
+ assert wide.loc[1, "bp_w0"] == 120.0
+ assert wide.loc[3, "chol_w2"] == 6.5
+
+ def test_subjects_with_disjoint_wave_sets(self):
+ df = pd.DataFrame({
+ "subject_id": [1, 1, 2, 2, 3],
+ "time": [0, 1, 1, 2, 2],
+ "bp": [120.0, 121.0, 130.0, 131.0, 110.0],
+ "sex": ["M", "M", "F", "F", "F"],
+ })
+ wide = _dataset(df).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp"], static_columns=["sex"],
+ )
+ assert list(wide.columns) == ["sex", "bp_w0", "bp_w1", "bp_w2"]
+ assert wide.loc[1, "bp_w0"] == 120.0
+ assert pd.isna(wide.loc[1, "bp_w2"])
+ assert pd.isna(wide.loc[2, "bp_w0"])
+ assert pd.isna(wide.loc[3, "bp_w0"]) and pd.isna(wide.loc[3, "bp_w1"])
+ assert wide.loc[3, "bp_w2"] == 110.0
+
+ def test_string_wave_labels_sort_lexicographically(self):
+ df = make_long_frame().copy()
+ df["time"] = df["time"].map({0: "wA", 1: "wB", 2: "wC"})
+ wide = _dataset(df).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp"], static_columns=["sex"],
+ )
+ assert list(wide.columns) == ["sex", "bp_wwA", "bp_wwB", "bp_wwC"]
+ assert wide.loc[1, "bp_wwA"] == 120.0
+
+ def test_custom_wave_format(self):
+ wide = _dataset(make_long_frame()).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp"], static_columns=["sex"],
+ wave_format="{feature}@t={wave}",
+ )
+ assert "bp@t=0" in wide.columns and "bp@t=2" in wide.columns
+
+ def test_categorical_longitudinal_values_preserved(self):
+ df = pd.DataFrame({
+ "subject_id": [1, 1, 2, 2],
+ "time": [0, 1, 0, 1],
+ "smoke": ["yes", "no", "no", "no"],
+ "sex": ["M", "M", "F", "F"],
+ })
+ wide = _dataset(df).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["smoke"], static_columns=["sex"],
+ )
+ assert wide.loc[1, "smoke_w0"] == "yes"
+ assert wide.loc[2, "smoke_w1"] == "no"
+ # Non-numeric column should not be coerced to float.
+ assert wide["smoke_w0"].dtype == object
+
+ def test_nan_longitudinal_value_preserved_not_treated_as_missing_wave(self):
+ df = make_long_frame().copy()
+ df.loc[(df["subject_id"] == 1) & (df["time"] == 1), "bp"] = np.nan
+ wide = _dataset(df).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp", "chol"], static_columns=["sex"],
+ )
+ assert pd.isna(wide.loc[1, "bp_w1"])
+ # chol still present at wave 1
+ assert wide.loc[1, "chol_w1"] == 5.5
+
+ def test_single_subject_single_wave_each_still_pivots(self):
+ df = pd.DataFrame({
+ "subject_id": [1, 2, 3],
+ "time": [0, 0, 0],
+ "bp": [120.0, 130.0, 110.0],
+ "chol": [5.0, 4.5, 6.0],
+ })
+ wide = _dataset(df).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp", "chol"],
+ )
+ assert list(wide.columns) == ["bp_w0", "chol_w0"]
+ assert wide.loc[2, "bp_w0"] == 130.0
+
+ def test_subject_with_all_nan_longitudinals(self):
+ df = make_long_frame().copy()
+ df.loc[df["subject_id"] == 2, ["bp", "chol"]] = np.nan
+ wide = _dataset(df).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp", "chol"], static_columns=["sex"],
+ )
+ assert wide.loc[2, ["bp_w0", "bp_w1", "bp_w2"]].isna().all()
+ assert wide.loc[2, "sex"] == "F"
+
+ def test_multiple_static_columns(self):
+ df = make_long_frame().copy()
+ df["country"] = "UK"
+ wide = _dataset(df).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp"], static_columns=["sex", "country"],
+ )
+ assert list(wide.columns)[:2] == ["sex", "country"]
+ assert wide.loc[1, "country"] == "UK"
+
+ def test_writes_csv_when_output_path_given(self, tmp_path):
+ out = tmp_path / "wide.csv"
+ _dataset(make_long_frame()).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp", "chol"], static_columns=["sex"],
+ output_path=out,
+ )
+ assert out.exists()
+ reloaded = pd.read_csv(out)
+ assert "bp_w0" in reloaded.columns and "sex" in reloaded.columns
+
+ # ---- error paths -------------------------------------------------------
+
+ def test_rejects_missing_id_col(self):
+ with pytest.raises(ValueError, match="id_col 'pid' not in dataframe"):
+ _dataset(make_long_frame()).to_wide(
+ id_col="pid", time_col="time", longitudinal_columns=["bp"],
+ )
+
+ def test_rejects_missing_time_col(self):
+ with pytest.raises(ValueError, match="time_col 'wave' not in dataframe"):
+ _dataset(make_long_frame()).to_wide(
+ id_col="subject_id", time_col="wave", longitudinal_columns=["bp"],
+ )
+
+ def test_rejects_duplicate_id_time_rows(self):
+ df = pd.concat([make_long_frame(), make_long_frame().iloc[[0]]], ignore_index=True)
+ with pytest.raises(ValueError, match="Duplicate"):
+ _dataset(df).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp", "chol"], static_columns=["sex"],
+ )
+
+ def test_rejects_static_varying_within_subject(self):
+ df = make_long_frame()
+ df.loc[0, "sex"] = "F"
+ with pytest.raises(ValueError, match="Static columns vary"):
+ _dataset(df).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp", "chol"], static_columns=["sex"],
+ )
+
+ def test_rejects_unknown_value_column(self):
+ with pytest.raises(ValueError, match="Column 'mystery' not in dataframe"):
+ _dataset(make_long_frame()).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["mystery"],
+ )
+
+ def test_rejects_unknown_static_column(self):
+ with pytest.raises(ValueError, match="Column 'ghost' not in dataframe"):
+ _dataset(make_long_frame()).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp"], static_columns=["ghost"],
+ )
+
+ def test_rejects_overlap_between_value_and_static(self):
+ with pytest.raises(ValueError, match="listed more than once"):
+ _dataset(make_long_frame()).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp"], static_columns=["bp"],
+ )
+
+ def test_rejects_id_col_as_longitudinal(self):
+ with pytest.raises(ValueError, match="cannot be a value/static column"):
+ _dataset(make_long_frame()).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["subject_id"],
+ )
+
+ def test_rejects_time_col_as_longitudinal(self):
+ with pytest.raises(ValueError, match="cannot be a value/static column"):
+ _dataset(make_long_frame()).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["time"],
+ )
+
+ def test_rejects_id_col_as_static(self):
+ with pytest.raises(ValueError, match="cannot be a value/static column"):
+ _dataset(make_long_frame()).to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp"], static_columns=["subject_id"],
+ )
+
+ def test_rejects_empty_longitudinal_columns(self):
+ with pytest.raises(ValueError, match="longitudinal_columns must list"):
+ _dataset(make_long_frame()).to_wide(
+ id_col="subject_id", time_col="time", longitudinal_columns=[],
+ )
+
+
+# ---------------------------------------------------------------------------
+# LongitudinalDataset.to_long
+# ---------------------------------------------------------------------------
+
+
+class TestToLong:
+ def _setup(self) -> LongitudinalDataset:
+ dataset = _dataset(make_wide_frame())
+ dataset.setup_features_group([["bp_w1", "bp_w2"], ["chol_w1", "chol_w2"]])
+ return dataset
+
+ def test_pivots_and_updates_dataset_state(self):
+ dataset = self._setup()
+ long_df = dataset.to_long(
+ feature_base_names=["bp", "chol"], id_col="pid", time_col="wave",
+ )
+ assert dataset.data is long_df
+ assert dataset.feature_groups() is None
+ assert dataset.non_longitudinal_features() is None
+ assert set(long_df.columns) == {"pid", "wave", "bp", "chol", "sex"}
+ assert long_df.shape[0] == 6
+ first = long_df[(long_df["pid"] == 1) & (long_df["wave"] == 1)].iloc[0]
+ assert first["bp"] == 120.0 and first["sex"] == "M"
+
+ def test_default_feature_base_names(self):
+ dataset = self._setup()
+ long_df = dataset.to_long(id_col="pid", time_col="wave")
+ assert {"feature_0", "feature_1"}.issubset(long_df.columns)
+
+ def test_keep_static_false_drops_static_columns(self):
+ dataset = self._setup()
+ long_df = dataset.to_long(
+ feature_base_names=["bp", "chol"], id_col="pid",
+ time_col="wave", keep_static=False,
+ )
+ assert "sex" not in long_df.columns
+ assert set(long_df.columns) == {"pid", "wave", "bp", "chol"}
+
+ def test_uneven_wave_counts_pad_with_nan(self):
+ wide = make_wide_frame()
+ wide["chol_w3"] = [np.nan, 5.5, 6.5]
+ dataset = _dataset(wide)
+ dataset.setup_features_group(
+ [["bp_w1", "bp_w2"], ["chol_w1", "chol_w2", "chol_w3"]]
+ )
+ long_df = dataset.to_long(
+ feature_base_names=["bp", "chol"], id_col="pid", time_col="wave",
+ )
+ # 3 subjects * 3 wave slots = 9 rows
+ assert long_df.shape[0] == 9
+ wave3 = long_df[long_df["wave"] == 3]
+ assert wave3["bp"].isna().all()
+ # Subject 1 had NaN at chol_w3 originally; subject 2 has 5.5
+ assert pd.isna(wave3.loc[wave3["pid"] == 1, "chol"].iloc[0])
+ assert wave3.loc[wave3["pid"] == 2, "chol"].iloc[0] == 5.5
+
+ def test_pad_value_inside_group(self):
+ # Build raw indices so we can poke a -1 in.
+ wide = make_wide_frame()
+ wide["chol_w3"] = [9.0, 8.0, 7.0]
+ dataset = _dataset(wide)
+ # bp has only 2 waves; pad slot for wave 3.
+ cols = list(wide.columns)
+ bp_group = [cols.index("bp_w1"), cols.index("bp_w2"), PAD_VALUE]
+ chol_group = [cols.index("chol_w1"), cols.index("chol_w2"), cols.index("chol_w3")]
+ dataset._feature_groups = [bp_group, chol_group]
+ long_df = dataset.to_long(
+ feature_base_names=["bp", "chol"], id_col="pid", time_col="wave",
+ )
+ wave3 = long_df[long_df["wave"] == 3]
+ assert wave3["bp"].isna().all() # padded
+ assert wave3["chol"].tolist() == [9.0, 8.0, 7.0]
+
+ def test_static_columns_inferred_from_unused_indices(self):
+ wide = make_wide_frame().copy()
+ wide["country"] = ["UK", "UK", "FR"]
+ dataset = _dataset(wide)
+ dataset.setup_features_group([["bp_w1", "bp_w2"], ["chol_w1", "chol_w2"]])
+ long_df = dataset.to_long(
+ feature_base_names=["bp", "chol"], id_col="pid", time_col="wave",
+ )
+ assert "country" in long_df.columns and "sex" in long_df.columns
+ assert long_df.loc[long_df["pid"] == 3, "country"].iloc[0] == "FR"
+
+ def test_non_default_index_is_carried_into_id_col(self):
+ wide = make_wide_frame().copy()
+ wide.index = ["alpha", "beta", "gamma"]
+ dataset = _dataset(wide)
+ dataset.setup_features_group([["bp_w1", "bp_w2"], ["chol_w1", "chol_w2"]])
+ long_df = dataset.to_long(
+ feature_base_names=["bp", "chol"], id_col="pid", time_col="wave",
+ )
+ assert set(long_df["pid"]) == {"alpha", "beta", "gamma"}
+
+ def test_writes_csv_when_output_path_given(self, tmp_path):
+ out = tmp_path / "long.csv"
+ self._setup().to_long(
+ feature_base_names=["bp", "chol"], id_col="pid",
+ time_col="wave", output_path=out,
+ )
+ assert out.exists()
+ reloaded = pd.read_csv(out)
+ assert {"pid", "wave", "bp", "chol", "sex"}.issubset(reloaded.columns)
+
+ # ---- error paths -------------------------------------------------------
+
+ def test_rejects_when_features_group_missing(self):
+ dataset = _dataset(make_wide_frame())
+ with pytest.raises(ValueError, match="setup_features_group"):
+ dataset.to_long()
+
+ def test_rejects_wrong_length_feature_base_names(self):
+ dataset = self._setup()
+ with pytest.raises(ValueError, match="feature_base_names has 1 entries"):
+ dataset.to_long(feature_base_names=["only_one"])
+
+ def test_rejects_group_with_single_wave(self):
+ dataset = _dataset(make_wide_frame())
+ dataset._feature_groups = [[0]]
+ with pytest.raises(ValueError, match="must list at least two waves"):
+ dataset.to_long()
+
+ def test_rejects_index_out_of_range(self):
+ dataset = _dataset(make_wide_frame())
+ dataset._feature_groups = [[0, 99]]
+ with pytest.raises(ValueError, match="out of range"):
+ dataset.to_long()
+
+ def test_rejects_index_in_two_groups(self):
+ dataset = _dataset(make_wide_frame())
+ dataset._feature_groups = [[0, 1], [1, 2]]
+ with pytest.raises(ValueError, match="appears in more than one group"):
+ dataset.to_long()
+
+
+# ---------------------------------------------------------------------------
+# Round-trip
+# ---------------------------------------------------------------------------
+
+
+class TestRoundTrip:
+ def test_wide_to_long_then_long_to_wide_recovers_values(self):
+ dataset = _dataset(make_wide_frame())
+ dataset.setup_features_group([["bp_w1", "bp_w2"], ["chol_w1", "chol_w2"]])
+ long_df = dataset.to_long(
+ feature_base_names=["bp", "chol"], id_col="pid", time_col="wave",
+ )
+ wide = _dataset(long_df).to_wide(
+ id_col="pid", time_col="wave",
+ longitudinal_columns=["bp", "chol"], static_columns=["sex"],
+ )
+ assert wide.loc[1, "bp_w1"] == 120.0
+ assert wide.loc[2, "chol_w2"] == 5.0
+ assert list(wide["sex"]) == ["M", "F", "F"]
+
+ def test_long_to_wide_then_wide_to_long_recovers_values(self):
+ original = make_long_frame()
+ dataset = _dataset(original)
+ dataset.to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp", "chol"], static_columns=["sex"],
+ )
+ long_df = dataset.to_long(
+ feature_base_names=["bp", "chol"], id_col="subject_id", time_col="time",
+ )
+ # Same number of rows; same totals.
+ assert long_df.shape[0] == original.shape[0]
+ # Wave labels were 0/1/2; to_long emits 1/2/3 (positional).
+ # Compare per-subject sums (order-independent).
+ for sid in original["subject_id"].unique():
+ assert (
+ long_df.loc[long_df["subject_id"] == sid, "bp"].sum()
+ == pytest.approx(original.loc[original["subject_id"] == sid, "bp"].sum())
+ )
+
+ def test_round_trip_preserves_numeric_dtype(self):
+ dataset = _dataset(make_long_frame())
+ dataset.to_wide(
+ id_col="subject_id", time_col="time",
+ longitudinal_columns=["bp", "chol"], static_columns=["sex"],
+ )
+ for col in ["bp_w0", "bp_w1", "chol_w2"]:
+ assert pd.api.types.is_numeric_dtype(dataset.data[col])
diff --git a/zensical.toml b/zensical.toml
index 0d57e31..2e48c8b 100644
--- a/zensical.toml
+++ b/zensical.toml
@@ -25,6 +25,7 @@ nav = [
{ "Temporal Dependency" = "tutorials/temporal_dependency.md" },
{ "Uneven Temporal Setup" = "tutorials/advanced_temporal_setup.md" },
{ "Data Format" = "tutorials/sklong_longitudinal_data_format.md" },
+ { "Long ⇄ Wide Reshape" = "tutorials/long_wide_reshape.md" },
{ "Data Preparation" = "tutorials/sklong_data_preparation_first_exploration.md" },
{ "Algorithm Adaptation" = "tutorials/sklong_explore_your_first_estimator.md" },
{ "Binary & Multiclass" = "tutorials/binary_vs_multiclass.md" },