feat: first training attempt

It still sucks. I'm sure the the feature embedding is terribly wrong.

.gitignore

@@ -8,3 +8,7 @@ wheels/
 
 # Virtual environments
 .venv
+
+# Models
+**/*.pt
+*.pt

pyproject.toml

@@ -8,7 +8,10 @@ authors = [
 ]
 requires-python = ">=3.12"
 dependencies = [
-    "torch"
+    "torch",
+    "numpy",
+    "pandas",
+    "loguru"
 ]
 
 [project.scripts]

scripts/train_from_data.sh

@@ -0,0 +1,88 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+DATA_DIR="${ROOT_DIR}/data"
+OUTPUT_FILE="${ROOT_DIR}/sleep-model.pt"
+MODEL_FILE="${OUTPUT_FILE}"
+
+if [[ ${1:-} == "-h" || ${1:-} == "--help" ]]; then
+  cat <<'EOF'
+Usage: scripts/train_from_data.sh [-o OUTPUT_FILE]
+
+Trains the model on complete session triplets found under data/{motion,heart_rate,labels}/.
+EOF
+  exit 0
+fi
+
+if [[ ${1:-} == "-o" || ${1:-} == "--output" ]]; then
+  OUTPUT_FILE="${2:?Missing output file after -o/--output}"
+  MODEL_FILE="${OUTPUT_FILE}"
+fi
+
+shopt -s nullglob
+
+checkpoint_class_count() {
+  python - "$1" <<'PY'
+from pathlib import Path
+import sys
+import torch
+
+payload = torch.load(Path(sys.argv[1]), weights_only=True)
+state_dict = payload.get("model_state_dict", payload)
+print(state_dict["net.2.bias"].shape[0])
+PY
+}
+
+motions=("${DATA_DIR}/motion"/*_acceleration.txt)
+
+subject_ids=()
+
+for motion_file in "${motions[@]}"; do
+  sample_id="$(basename "${motion_file}" _acceleration.txt)"
+  heartrate_file="${DATA_DIR}/heart_rate/${sample_id}_heartrate.txt"
+  label_file="${DATA_DIR}/labels/${sample_id}_labeled_sleep.txt"
+
+  if [[ -f "${heartrate_file}" && -f "${label_file}" ]]; then
+    subject_ids+=("${sample_id}")
+  fi
+done
+
+if [[ ${#subject_ids[@]} -eq 0 ]]; then
+  printf 'No complete data triplets found under %s\n' "${DATA_DIR}" >&2
+  exit 1
+fi
+
+mapfile -t subject_ids < <(printf '%s\n' "${subject_ids[@]}" | sort -u)
+
+printf 'Training on %d subject(s)\n' "${#subject_ids[@]}"
+
+export PYTHONPATH="${ROOT_DIR}/src${PYTHONPATH:+:${PYTHONPATH}}"
+
+for subject_id in "${subject_ids[@]}"; do
+  acceleration_file="${DATA_DIR}/motion/${subject_id}_acceleration.txt"
+  heartrate_file="${DATA_DIR}/heart_rate/${subject_id}_heartrate.txt"
+  label_file="${DATA_DIR}/labels/${subject_id}_labeled_sleep.txt"
+  resume_args=()
+
+  printf 'Training subject %s\n' "${subject_id}"
+
+  if [[ -f "${MODEL_FILE}" ]]; then
+    checkpoint_classes="$(checkpoint_class_count "${MODEL_FILE}")"
+    if [[ "${checkpoint_classes}" == "3" ]]; then
+      resume_args=(-i "${MODEL_FILE}")
+    else
+      printf 'Skipping checkpoint for subject %s: model classes=%s, expected=3\n' "${subject_id}" "${checkpoint_classes}"
+    fi
+  fi
+
+  args=(python -m sleep_detection train -a "${acceleration_file}" -H "${heartrate_file}" -l "${label_file}" -o "${MODEL_FILE}")
+  if [[ ${#resume_args[@]} -gt 0 ]]; then
+    args+=("${resume_args[@]}")
+  fi
+
+  "${args[@]}"
+done
+
+printf 'Saved trained model to %s\n' "${MODEL_FILE}"

src/sleep_detection/main.py

@@ -1,6 +1,344 @@
 """Main Entrypoint codebase."""
 
+import sys
+from argparse import ArgumentParser
+from argparse import Namespace
+from collections.abc import Sequence
+from pathlib import Path
 
-def main():
+import numpy
+import pandas
+import torch
+from loguru import logger
+from torch import Tensor
+from torch.nn import CrossEntropyLoss
+from torch.nn import Module
+from torch.optim import SGD
+
+from sleep_detection.sleep_epoch_classifer import SleepEpochClassifier
+
+SLEEP_STAGE_LABELS = [0, 1, 2]
+
+
+def normalize_sleep_stage(stage) -> float:
+    """Map raw sleep stages to the 3-class target space.
+
+    Raw labels:
+    0 -> 0
+    1, 2, 3, 4 -> 1
+    5, 6 -> 2
+    Other values are treated as unknown.
+    """
+    if pandas.isna(stage):
+        return numpy.nan
+
+    value = int(stage)
+    if value == 0:
+        return 0.0
+    if 1 <= value <= 4:
+        return 1.0
+    if 5 <= value <= 6:
+        return 2.0
+
+    return numpy.nan
+
+
+@logger.catch
+def main() -> None:
     """The Main entrypoint function."""
-    print("Hello from sleep-detection!")
+    cli = parse_args()
+    logger.info("welcome")
+    if cli.command == "train":
+        sessions = load_sessions(cli.a, cli.H, cli.l)
+        epochs, features, target, label_mapping = featureize_sessions(sessions)
+
+        c = SleepEpochClassifier(
+            n_features=features.shape[1],
+            n_classes=len(SLEEP_STAGE_LABELS),
+            hidden_size=8,
+        )
+        if cli.input is not None:
+            c.load(cli.input)
+
+        logger.info(f"Using a {c.__class__.__name__}")
+
+        train(features, target, c)
+        save_network(cli.output, c, label_mapping)
+        return
+
+    if cli.command == "execute":
+        adf = load_acceleration_data(cli.a)
+        hdf = load_heartrate_data(cli.H)
+        ldf = load_labels_data(cli.l)
+
+        df = adf.join(hdf, how="outer").join(ldf, how="outer").sort_index()
+        epochs, features, target, label_mapping = featureize(df)
+
+        c = SleepEpochClassifier(
+            n_features=features.shape[1],
+            n_classes=len(SLEEP_STAGE_LABELS),
+            hidden_size=8,
+        )
+        loaded_label_mapping = c.load(cli.model)
+        logger.info(f"Using a {c.__class__.__name__}")
+        execute(epochs, features, target, c, loaded_label_mapping or label_mapping)
+        return
+
+    raise ValueError(f"Unknown command: {cli.command}")
+
+
+def parse_args() -> Namespace:
+    parser = ArgumentParser(description="Training program.")
+    parser.add_argument("-v", "--verbose", help="Increase verbosity.", action="store_true")
+    parser.add_argument("-V", "--version", action="store_true", help="Show version")
+    command = parser.add_subparsers(dest="command")
+
+    train = command.add_parser("train", help="Train the model.")
+    train.add_argument("-a", type=Path, nargs="+", required=True, help="Acceleration data file(s).")
+    train.add_argument("-H", type=Path, nargs="+", required=True, help="Heartrate data file(s).")
+    train.add_argument("-l", type=Path, nargs="+", required=True, help="Labels data file(s).")
+    train.add_argument("-o", "--output", type=Path, default=None, required=True, help="Output .pt file.")
+    train.add_argument("-i", "--input", type=Path, default=None, required=False, help="Input .pt file.")
+
+    execute = command.add_parser("execute", help="Execute the model.")
+    execute.add_argument("-m", "--model", type=Path, default=None, required=True, help="Input .pr file to use.")
+    execute.add_argument("-a", type=Path, default=None, required=True, help="Acceleration data file.")
+    execute.add_argument("-H", type=Path, default=None, required=True, help="Heartrate data file.")
+    execute.add_argument("-l", type=Path, default=None, required=True, help="Labels data file.")
+
+    return parser.parse_args()
+
+
+def featureize(data: pandas.DataFrame) -> tuple[pandas.DataFrame, Tensor, Tensor, list]:
+    """Convert loaded data into input and target tensors.
+
+    Args:
+        data: Loaded data, should contain the keys: [timestamp, heartrate, motion, x, y, z, stage]
+
+    """
+    logger.info(f"Featurizing the data...")
+    epochs = build_epochs(data)
+
+    epochs["hr_std"] = epochs["hr_std"].fillna(0.0)
+    epochs = epochs.dropna(subset=["hr_mean", "motion_mean", "motion_max", "stage"])
+
+    epochs["seconds_since_start"] = epochs.index.total_seconds()
+
+    CONTEXT = 2  # epochs each side -> 5 epochs total = 25 min window
+    FEATURE_COLS = ["hr_mean", "hr_std", "motion_mean", "motion_max", "seconds_since_start"]
+
+    features = epochs[FEATURE_COLS].to_numpy(dtype=numpy.float32)
+    labels = epochs["stage"].astype(numpy.int64).to_numpy()
+
+    X, y = [], []
+    for i in range(CONTEXT, len(epochs) - CONTEXT):
+        window = features[i - CONTEXT : i + CONTEXT + 1]  # shape: (2*CONTEXT+1, n_feats)
+        X.append(window.flatten())  # -> 1D vector
+        y.append(labels[i])
+
+    if not X:
+        raise ValueError("Not enough complete epochs to build a training set.")
+
+    X = torch.tensor(numpy.stack(X), dtype=torch.float32)  # shape: (n_epochs, (2*CONTEXT+1)*n_feats)
+    y = torch.tensor(y, dtype=torch.long)  # shape: (n_epochs,)
+    return epochs, X, y, SLEEP_STAGE_LABELS
+
+
+def featureize_sessions(dataframes: Sequence[pandas.DataFrame]) -> tuple[pandas.DataFrame, Tensor, Tensor, list]:
+    """Convert multiple loaded sessions into input and target tensors."""
+    epochs_per_session = []
+    for data in dataframes:
+        epochs = build_epochs(data)
+
+        epochs["hr_std"] = epochs["hr_std"].fillna(0.0)
+        epochs = epochs.dropna(subset=["hr_mean", "motion_mean", "motion_max", "stage"])
+        epochs["seconds_since_start"] = epochs.index.total_seconds()
+        epochs_per_session.append(epochs)
+
+    if not epochs_per_session:
+        raise ValueError("No sessions provided.")
+
+    combined_epochs = pandas.concat(epochs_per_session)
+
+    CONTEXT = 2  # epochs each side -> 5 epochs total = 25 min window
+    FEATURE_COLS = ["hr_mean", "hr_std", "motion_mean", "motion_max", "seconds_since_start"]
+
+    X, y = [], []
+    for epochs in epochs_per_session:
+        features = epochs[FEATURE_COLS].to_numpy(dtype=numpy.float32)
+        labels = epochs["stage"].astype(numpy.int64).to_numpy()
+
+        for i in range(CONTEXT, len(epochs) - CONTEXT):
+            window = features[i - CONTEXT : i + CONTEXT + 1]  # shape: (2*CONTEXT+1, n_feats)
+            X.append(window.flatten())  # -> 1D vector
+            y.append(labels[i])
+
+    if not X:
+        raise ValueError("Not enough complete epochs to build a training set.")
+
+    X = torch.tensor(numpy.stack(X), dtype=torch.float32)  # shape: (n_epochs, (2*CONTEXT+1)*n_feats)
+    y = torch.tensor(y, dtype=torch.long)  # shape: (n_epochs,)
+    return combined_epochs, X, y, SLEEP_STAGE_LABELS
+
+
+def load_sessions(
+    acceleration_files: Sequence[Path], heartrate_files: Sequence[Path], label_files: Sequence[Path]
+) -> list[pandas.DataFrame]:
+    """Load multiple session triplets."""
+    if not (len(acceleration_files) == len(heartrate_files) == len(label_files)):
+        raise ValueError("Acceleration, heartrate, and label file counts must match.")
+
+    sessions = []
+    for acceleration_file, heartrate_file, label_file in zip(acceleration_files, heartrate_files, label_files):
+        adf = load_acceleration_data(acceleration_file)
+        hdf = load_heartrate_data(heartrate_file)
+        ldf = load_labels_data(label_file)
+        sessions.append(adf.join(hdf, how="outer").join(ldf, how="outer").sort_index())
+
+    return sessions
+
+
+def build_epochs(data: pandas.DataFrame) -> pandas.DataFrame:
+    """Aggregate raw signals into 5-minute epochs."""
+    EPOCH_LEN = pandas.Timedelta(minutes=5)
+
+    return data.resample(EPOCH_LEN).agg(
+        hr_mean=("hr", "mean"),
+        hr_std=("hr", "std"),
+        motion_mean=("motion", "mean"),
+        motion_max=("motion", "max"),
+        stage=("stage", lambda s: s.dropna().mode().iloc[0] if len(s.dropna()) else numpy.nan),
+    )
+
+
+def execute(epochs: pandas.DataFrame, input: Tensor, target: Tensor, classifier: Module, label_mapping: list) -> None:
+    """Run the model and emit per-epoch predictions."""
+    classifier.eval()
+    with torch.no_grad():
+        output = classifier(input)
+        predicted = output.argmax(dim=1)
+
+    decoded_predicted = [label_mapping[index] for index in predicted.tolist()]
+    decoded_target = [label_mapping[index] for index in target.tolist()]
+
+    report = pandas.DataFrame(
+        {
+            "timestamp": epochs.index[2 : len(epochs) - 2].astype("timedelta64[s]") / numpy.timedelta64(1, "s"),
+            "predicted": decoded_predicted,
+            "actual": decoded_target,
+        }
+    )
+    report.to_csv(sys.stdout, index=False)
+
+    accuracy = (predicted == target).float().mean().item()
+    logger.info(f"accuracy={accuracy:.3f}")
+
+
+def load_acceleration_data(data_file: Path) -> pandas.DataFrame:
+    """Load acceleration data file.
+
+    Expected input format: <timestamp> <x> <y> <z>
+
+    Args:
+        data_file: The filepath for the file containing acceleration data.
+
+    Returns:
+        A pandas data frame with the loaded acceleration data.
+
+    """
+    logger.info(f"Loading acceleration data file {data_file}")
+    df = pandas.read_csv(data_file, sep=r"\s+", names=["timestamp", "x", "y", "z"], header=None)
+    df["timestamp"] = pandas.to_numeric(df["timestamp"], errors="coerce")
+    df = df.dropna(subset=["timestamp"])
+    df = df.sort_values("timestamp").set_index(pandas.to_timedelta(df["timestamp"], unit="s"))
+    df = df.drop(columns=["timestamp"])
+    df.index.name = "timestamp"
+    # df["motion"] = df[["x", "y", "z"]].pow(2).sum(axis=1).pow(0.5)
+    df["motion"] = df[["x", "y", "z"]].diff().pow(2).sum(axis=1).pow(0.5)
+    return df
+
+
+def load_heartrate_data(data_file: Path) -> pandas.DataFrame:
+    """Load heart rate data file.
+
+    Expected input format: <timestamp>,<heartrate>
+
+    Args:
+        data_file: The filepath for the file containing heartrate data.
+
+    Returns:
+        A pandas data frame with the loaded heartrate data.
+
+    """
+    logger.info(f"Loading heartrate date file {data_file}")
+    df = pandas.read_csv(data_file, sep=",", names=["timestamp", "hr"], header=None)
+    df["timestamp"] = pandas.to_numeric(df["timestamp"], errors="coerce")
+    df = df.dropna(subset=["timestamp"])
+    df = df.sort_values("timestamp").set_index(pandas.to_timedelta(df["timestamp"], unit="s"))
+    df = df.drop(columns=["timestamp"])
+    df.index.name = "timestamp"
+    return df
+
+
+def load_labels_data(data_file: Path) -> pandas.DataFrame:
+    """Load label data file.
+
+    Expected input format: <timestamp>,<label>
+
+    Args:
+        data_file: The filepath for the file containing label data.
+
+    Returns:
+        A pandas data frame with the loaded label data.
+
+    """
+    logger.info(f"Loading label date file {data_file}")
+    df = pandas.read_csv(data_file, sep=r"\s+", names=["timestamp", "stage"], header=None)
+    df["timestamp"] = pandas.to_numeric(df["timestamp"], errors="coerce")
+    df["stage"] = pandas.to_numeric(df["stage"], errors="coerce").map(normalize_sleep_stage)
+    df = df.dropna(subset=["timestamp"])
+    df = df.sort_values("timestamp").set_index(pandas.to_timedelta(df["timestamp"], unit="s"))
+    df = df.drop(columns=["timestamp"])
+    df.index.name = "timestamp"
+    return df
+
+
+def train(input: Tensor, target: Tensor, classifier: Module):
+    classifier.train()
+    optimizer = SGD(classifier.parameters(), lr=0.01)
+    criterion = CrossEntropyLoss()
+
+    n_samples = input.shape[0]
+    n_train = max(1, int(n_samples * 0.8))
+    permutation = torch.randperm(n_samples)
+    train_idx = permutation[:n_train]
+    val_idx = permutation[n_train:]
+
+    train_input = input[train_idx]
+    train_target = target[train_idx]
+    val_input = input[val_idx] if len(val_idx) else None
+    val_target = target[val_idx] if len(val_idx) else None
+
+    for epoch in range(100):
+        optimizer.zero_grad()
+        output = classifier(train_input)
+        loss = criterion(output, train_target)
+        loss.backward()
+        optimizer.step()
+
+        if epoch % 10 == 0:
+            message = f"epoch={epoch} train_loss={loss.item():.4f}"
+            if val_input is not None and val_target is not None:
+                classifier.eval()
+                with torch.no_grad():
+                    val_output = classifier(val_input)
+                    val_loss = criterion(val_output, val_target)
+                    val_accuracy = (val_output.argmax(dim=1) == val_target).float().mean().item()
+                classifier.train()
+                message += f" val_loss={val_loss.item():.4f} val_acc={val_accuracy:.3f}"
+            logger.info(message)
+
+
+def save_network(output_file: Path, network: Module, label_mapping: list) -> None:
+    logger.info(f"Saving network {network} to {output_file}")
+    torch.save({"model_state_dict": network.state_dict(), "label_mapping": label_mapping}, output_file)

src/sleep_detection/sleep_epoch_classifer.py

@@ -0,0 +1,40 @@
+from pathlib import Path
+
+import torch
+from loguru import logger
+from torch import Tensor
+from torch.nn import Linear
+from torch.nn import Module
+from torch.nn import ReLU
+from torch.nn import Sequential
+
+
+class SleepEpochClassifier(Module):
+    """Per-epoch sleep stage classifier.
+
+    Input: a fixed-size feature vector summarizing one epoch's context window
+           (e.g. motion mean/max, HR mean/variance, time-of-night, previous
+           predicted label — whatever features you settle on).
+    Output: logits over sleep stage classes for that single epoch.
+
+    No recurrence/attention — sequence smoothing (majority vote / HMM)
+    happens as a separate post-processing step, not inside this network.
+    """
+
+    def __init__(self, n_features: int, n_classes: int, hidden_size: int = 16):
+        super().__init__()
+        self.net = Sequential(
+            Linear(n_features, hidden_size),
+            ReLU(),
+            Linear(hidden_size, n_classes),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        # x shape: (batch, n_features) -> returns (batch, n_classes) logits
+        return self.net(x)
+
+    def load(self, input_file: Path) -> list | None:
+        payload = torch.load(input_file, weights_only=True)
+        state_dict = payload.get("model_state_dict", payload)
+        self.load_state_dict(state_dict)
+        return payload.get("label_mapping")