Prepare repo for public release and Zenodo archival (KDD 2026 AI4Sci)

.gitignore

@@ -98,7 +98,7 @@ ipython_config.py
 #   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
 #   This is especially recommended for binary packages to ensure reproducibility, and is more
 #   commonly ignored for libraries.
-#uv.lock
+uv.lock
 
 # poetry
 #   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
@@ -205,3 +205,12 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+# Agent scratch
+.claude/
+
+# Data: ship ONLY the curated demo artifact; never raw downloads or transient outputs.
+data/_raw_labels/
+data/*
+!data/README.md
+!data/sleepaccel_demo.npz

.zenodo.json

@@ -0,0 +1,28 @@
+{
+  "title": "StageGuard: Physiologically Constrained Sleep Staging",
+  "description": "Official implementation of the KDD 2026 AI4Sciences paper 'StageGuard: Physiologically Constrained Sleep Staging'. StageGuard is a plug-and-play, backbone-agnostic structured-inference framework that wraps any neural sleep-staging backbone with physiology-informed priors: a differentiable soft transition penalty during training and a semi-Markov constrained decoder at inference. The repository bundles a small subset of Sleep-Accel PSG hypnograms (labels only) under the ODC-By license for the demo notebook.",
+  "license": "MIT",
+  "upload_type": "software",
+  "access_right": "open",
+  "creators": [
+    {"name": "Wang, Juntang", "affiliation": "Zu Chongzhi Center, Duke Kunshan University"},
+    {"name": "Wang, Yihan", "affiliation": "Zu Chongzhi Center, Duke Kunshan University"},
+    {"name": "Wu, Hao", "affiliation": "Zu Chongzhi Center, Duke Kunshan University"},
+    {"name": "Gao, Jiayu", "affiliation": "Zu Chongzhi Center, Duke Kunshan University"},
+    {"name": "Xu, Shixin", "affiliation": "Zu Chongzhi Center, Duke Kunshan University"},
+    {"name": "Zou, Dongmian", "affiliation": "Zu Chongzhi Center, Duke Kunshan University"}
+  ],
+  "keywords": [
+    "sleep staging",
+    "physiological constraints",
+    "semi-Markov model",
+    "temporal consistency",
+    "constrained decoding",
+    "deep learning",
+    "KDD 2026"
+  ],
+  "related_identifiers": [
+    {"identifier": "10.1145/3770855.3818916", "relation": "isSupplementTo", "scheme": "doi"},
+    {"identifier": "10.13026/hmhs-py35", "relation": "references", "scheme": "doi"}
+  ]
+}

CITATION.cff

@@ -0,0 +1,67 @@
+cff-version: 1.2.0
+message: "If you use StageGuard in your research, please cite the KDD 2026 paper below."
+title: "StageGuard: Physiologically Constrained Sleep Staging"
+type: software
+license: MIT
+repository-code: "https://github.com/Q9gJYx/StageGuard"
+url: "https://doi.org/10.1145/3770855.3818916"
+# doi: 10.5281/zenodo.XXXXXXXX   # add the archival Zenodo DOI after the release is deposited
+abstract: >-
+  StageGuard is a plug-and-play, backbone-agnostic structured-inference framework that wraps any neural
+  sleep-staging model with physiology-informed priors: a differentiable soft transition penalty during
+  training and a semi-Markov constrained decoder at inference.
+keywords:
+  - sleep staging
+  - physiological constraints
+  - semi-Markov model
+  - temporal consistency
+  - constrained decoding
+  - deep learning
+authors:
+  - family-names: Wang
+    given-names: Juntang
+    affiliation: "Zu Chongzhi Center, Duke Kunshan University"
+    email: "jw853@duke.edu"
+  - family-names: Wang
+    given-names: Yihan
+    affiliation: "Zu Chongzhi Center, Duke Kunshan University"
+    email: "yw572@duke.edu"
+  - family-names: Wu
+    given-names: Hao
+    affiliation: "Zu Chongzhi Center, Duke Kunshan University"
+    email: "hw412@duke.edu"
+  - family-names: Gao
+    given-names: Jiayu
+    affiliation: "Zu Chongzhi Center, Duke Kunshan University"
+    email: "jg563@duke.edu"
+  - family-names: Xu
+    given-names: Shixin
+    affiliation: "Zu Chongzhi Center, Duke Kunshan University"
+    email: "shixin.xu@duke.edu"
+  - family-names: Zou
+    given-names: Dongmian
+    affiliation: "Zu Chongzhi Center, Duke Kunshan University"
+    email: "dongmian.zou@duke.edu"
+preferred-citation:
+  type: conference-paper
+  title: "StageGuard: Physiologically Constrained Sleep Staging"
+  year: 2026
+  collection-title: "Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining V.2 (KDD '26)"
+  doi: "10.1145/3770855.3818916"
+  conference:
+    name: "ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD)"
+    city: "Jeju Island"
+    country: "KR"
+  authors:
+    - family-names: Wang
+      given-names: Juntang
+    - family-names: Wang
+      given-names: Yihan
+    - family-names: Wu
+      given-names: Hao
+    - family-names: Gao
+      given-names: Jiayu
+    - family-names: Xu
+      given-names: Shixin
+    - family-names: Zou
+      given-names: Dongmian

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2026 Juntang
+Copyright (c) 2026 Juntang Wang
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

README.md

@@ -1,14 +1,45 @@
-# StageGuard
+# [KDD 2026 AI4Sci] StageGuard: Physiologically Constrained Sleep Staging
 
 **Backbone-agnostic physiological constraints for neural sleep staging.**
 
-StageGuard improves any neural sleep staging model by enforcing physiological constraints during both training (soft transition penalty) and inference (Semi-Markov constrained decoding). It works across modalities—EEG, actigraphy, cardiorespiratory signals, and bioradar—with modality-specific configurations.
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Q9gJYx/StageGuard/blob/main/notebooks/demo_stageguard.ipynb)
+[![KDD 2026](https://img.shields.io/badge/KDD-2026%20AI4Sci-1d4ed8.svg)](https://doi.org/10.1145/3770855.3818916)
+[![DOI](https://img.shields.io/badge/DOI-10.1145%2F3770855.3818916-blue.svg)](https://doi.org/10.1145/3770855.3818916)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+[![Python](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org)
+[![PyTorch](https://img.shields.io/badge/PyTorch-2.0%2B-ee4c2c.svg)](https://pytorch.org)
+<!-- After depositing the release on Zenodo, add the archival DOI badge here:
+[![Zenodo](https://zenodo.org/badge/DOI/10.5281/zenodo.XXXXXXXX.svg)](https://doi.org/10.5281/zenodo.XXXXXXXX) -->
+
+Official implementation of the **KDD 2026 AI4Sciences** paper **"StageGuard: Physiologically Constrained
+Sleep Staging"** by Juntang Wang, Yihan Wang, Hao Wu, Jiayu Gao, Shixin Xu, and Dongmian Zou (Zu Chongzhi
+Center, Duke Kunshan University). Equal contribution: Juntang Wang, Yihan Wang, Hao Wu, Jiayu Gao.
+Corresponding author: Dongmian Zou (`dongmian.zou@duke.edu`).
+
+> Automated sleep staging is increasingly used in large-scale studies and downstream scientific analyses to
+> derive sleep-architecture endpoints, including total sleep time, REM latency, sleep efficiency, and
+> bout-duration statistics. Deep learning models achieve epoch-level accuracy approaching inter-rater
+> agreement, yet often produce hypnograms that violate physiological invariants, such as rare state
+> transitions (e.g., direct Wake to REM) or excessively fragmented sequences. Such violations can bias
+> downstream sleep metrics, regardless of overall accuracy. We propose StageGuard, a plug-and-play,
+> backbone-agnostic structured-inference framework that wraps any neural sleep-staging backbone with
+> physiology-informed priors. StageGuard combines (1) a differentiable soft transition penalty that
+> discourages physiologically rare transitions during training, and (2) a semi-Markov constrained decoder
+> with a duration-augmented state space that jointly enforces transition penalties and minimum bout durations
+> at inference. Across six backbones and four datasets, StageGuard reduces the transition-violation rate to
+> physiologically plausible levels and lowers the fragmentation index by 56 to 62%, while maintaining or
+> slightly improving classification accuracy, and improves the accuracy of derived sleep-architecture
+> statistics that are not directly optimized by the method.
+
+## News
+
+- **2026** Accepted to KDD 2026 (AI4Sciences track), Jeju Island, Republic of Korea (Aug 9 to 13, 2026).
 
 ## Key Contributions
 
-1. **Soft Transition Penalty** — A differentiable loss term that discourages physiologically rare stage transitions (e.g., Wake↔REM) during training.
-2. **Semi-Markov Decoder** — An augmented Viterbi decoder that enforces minimum bout durations, penalizes rare transitions, and suppresses flip-flop artifacts at inference time.
-3. **Signal Quality Integration** — Per-epoch signal quality scores modulate decoder confidence, gracefully degrading predictions for noisy segments.
+1. **Soft Transition Penalty** - A differentiable loss term that discourages physiologically rare stage transitions (e.g., Wake to/from REM) during training.
+2. **Semi-Markov Decoder** - An augmented Viterbi decoder that enforces minimum bout durations, penalizes rare transitions, and suppresses flip-flop artifacts at inference time.
+3. **Signal Quality Integration** - Per-epoch signal quality scores modulate decoder confidence, gracefully degrading predictions for noisy segments.
 
 ## Installation
 
@@ -19,8 +50,22 @@ conda activate stageguard
 
 # Or with pip
 pip install -e .
+
+# Optional: notebook/plotting extras
+pip install -e ".[demo]"
 ```
 
+## Demo
+
+Run the interactive demo on Colab (no install, a few seconds on CPU):
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Q9gJYx/StageGuard/blob/main/notebooks/demo_stageguard.ipynb)
+
+It loads real Sleep-Accel PSG hypnograms (shipped under ODC-By), simulates a noisy backbone, and shows the
+semi-Markov decoder driving the transition-violation rate (TVR) and fragmentation index (FI) down while
+holding or improving accuracy, with a before/after hypnogram plot. See
+[`notebooks/demo_stageguard.ipynb`](notebooks/demo_stageguard.ipynb).
+
 ## Quick Start
 
 ```python
@@ -60,55 +105,64 @@ class MyBackbone(nn.Module):
 model = StageGuardWrapper(MyBackbone(), config)
 ```
 
+See [`docs/backbones.md`](docs/backbones.md) for the full interface contract and the `ModalityConfig` schema.
+
 ## Repository Structure
 
 ```
 StageGuard/
-├── configs/               # Modality-specific YAML configurations
-│   ├── mouse_eeg.yaml     # AccuSleep: 3-state, 4s epochs
-│   ├── actigraphy.yaml    # Sleep-Accel: 2-state, 30s epochs
-│   ├── cardiorespiratory.yaml  # SHHS: 3-state, 30s epochs
-│   └── bioradar.yaml      # SLEEPBRL: 3-state, 30s epochs
+├── configs/                      # Modality-specific YAML configurations
+│   ├── mouse_eeg.yaml            # AccuSleep: 3-state, 4s epochs
+│   ├── actigraphy.yaml           # Sleep-Accel: 2-state, 30s epochs
+│   ├── cardiorespiratory.yaml    # SHHS: 3-state, 30s epochs
+│   ├── bioradar.yaml             # SLEEPBRL: 3-state, 30s epochs
+│   └── sleepaccel_demo.yaml      # 3-state Wake/NREM/REM config used by the demo
 ├── stageguard/
-│   ├── config.py           # ModalityConfig dataclass + YAML loader
-│   ├── losses.py           # SoftTransitionPenalty + stageguard_loss
-│   ├── decoder.py          # SemiMarkovDecoder (augmented Viterbi)
-│   ├── wrapper.py          # StageGuardWrapper (backbone + loss + decoder)
-│   ├── metrics.py          # TVR, FI, accuracy, kappa, F1, sleep architecture
-│   ├── sqi.py              # Signal quality index per modality
-│   ├── backbones/          # Backbone implementations + registry
-│   │   ├── base.py         # BackboneBase ABC
-│   │   ├── accusleep.py    # 2-layer CNN
-│   │   └── usleep.py       # U-Net encoder-decoder
-│   └── data/               # Dataset loaders (expect pre-downloaded data)
-│       ├── base.py         # BaseSleepDataset ABC
-│       ├── mouse_eeg.py    # AccuSleep mouse EEG/EMG
-│       ├── actigraphy.py   # Sleep-Accel wrist actigraphy
-│       ├── cardiorespiratory.py  # SHHS cardiorespiratory
-│       └── bioradar.py     # SLEEPBRL bioradar
-├── tests/                  # Comprehensive test suite
-└── examples/
-    └── demo.py             # End-to-end synthetic data demo
+│   ├── config.py                 # ModalityConfig dataclass + YAML loader
+│   ├── losses.py                 # SoftTransitionPenalty + stageguard_loss
+│   ├── decoder.py                # SemiMarkovDecoder (augmented Viterbi)
+│   ├── wrapper.py                # StageGuardWrapper (backbone + loss + decoder)
+│   ├── metrics.py                # TVR, FI, accuracy, kappa, F1, sleep architecture
+│   ├── sqi.py                    # Signal quality index per modality
+│   ├── backbones/                # Backbone implementations + registry
+│   └── data/                     # Dataset loaders (expect pre-downloaded data)
+├── notebooks/
+│   └── demo_stageguard.ipynb     # Colab-runnable demo
+├── data/
+│   ├── sleepaccel_demo.npz       # small Sleep-Accel PSG hypnogram subset (ODC-By)
+│   └── README.md                 # provenance, license, stage-collapse mapping
+├── docs/
+│   ├── datasets.md               # dataset access + loader formats
+│   └── backbones.md              # interface contract + config schema
+├── examples/demo.py              # end-to-end synthetic data demo
+├── tests/                        # test suite
+├── ROADMAP.md
+├── CITATION.cff
+└── .zenodo.json
 ```
 
 ## Datasets
 
-| Dataset | Modality | Classes | Epoch | URL |
-|---------|----------|---------|-------|-----|
-| AccuSleep | Mouse EEG/EMG | 3 (Wake, NREM, REM) | 4s | [Zenodo](https://zenodo.org/records/4079563) |
-| Sleep-Accel | Wrist actigraphy | 2 (Wake, Sleep) | 30s | [PhysioNet](https://physionet.org/content/sleep-accel/1.0.0/) |
-| SHHS | Cardiorespiratory | 3 (Wake, Light, Deep) | 30s | [NSRR](https://sleepdata.org/datasets/shhs) |
+| Dataset | Modality | Classes | Epoch | Access |
+|---------|----------|---------|-------|--------|
+| AccuSleep | Mouse EEG/EMG | 3 (Wake, NREM, REM) | 4s | Open ([Zenodo](https://zenodo.org/records/4079563)) |
+| Sleep-Accel | Wrist actigraphy | 2 (Wake, Sleep) | 30s | Open, ODC-By ([PhysioNet](https://physionet.org/content/sleep-accel/1.0.0/)) |
+| SHHS | Cardiorespiratory | 3 (Wake, Light, Deep) | 30s | Data-use agreement ([NSRR](https://sleepdata.org/datasets/shhs)) |
 | SLEEPBRL | Bioradar | 3 (Wake, Light, Deep) | 30s | Contact authors |
 
-> **Note:** Datasets must be downloaded separately. Some require data use agreements. Each loader provides download instructions via `Dataset.download_instructions()`.
+Datasets must be downloaded separately; some require data use agreements. Each loader provides download
+instructions via `Dataset.download_instructions()`. See [`docs/datasets.md`](docs/datasets.md) for the
+expected on-disk formats. The demo notebook uses a small, redistributable subset of Sleep-Accel PSG
+hypnograms shipped at [`data/sleepaccel_demo.npz`](data/sleepaccel_demo.npz) (see
+[`data/README.md`](data/README.md)).
 
 ## Hyperparameters
 
 | Parameter | Symbol | Default | Description |
 |-----------|--------|---------|-------------|
-| `lambda_trans` | λ | 1.0 | Transition penalty weight |
-| `epsilon` | ε | 5.0 | Rare transition penalty in decoder |
-| `gamma` | γ | 2.0 | Anti-flip-flop penalty |
+| `lambda_trans` | lambda | 1.0 | Transition penalty weight |
+| `epsilon` | epsilon | 5.0 | Rare transition penalty in decoder |
+| `gamma` | gamma | 2.0 | Anti-flip-flop penalty |
 | `k` | k | 5 | Flip-flop lookback window (epochs) |
 | `d_min` | d_min | per-stage | Minimum bout duration (epochs) |
 | `d_max` | d_max | 30 | Maximum tracked duration (epochs) |
@@ -127,6 +181,30 @@ pytest tests/test_decoder.py  # Single module
 python examples/demo.py
 ```
 
+## Acknowledgments
+
+We thank the authors of the datasets used in this work: AccuSleep (Barger et al., 2019), Sleep-Accel (Walch
+et al., 2019, distributed via PhysioNet under ODC-By), SHHS (Quan et al., 1997, via the NSRR), and SLEEPBRL
+(Tataraidze et al., 2015). The reference backbones reimplement architectures from AccuSleep (Barger et al.,
+2019) and U-Sleep (Perslev et al., 2021).
+
+## Citation
+
+If you use StageGuard in your research, please cite:
+
+```bibtex
+@inproceedings{wang2026stageguard,
+  title     = {StageGuard: Physiologically Constrained Sleep Staging},
+  author    = {Wang, Juntang and Wang, Yihan and Wu, Hao and Gao, Jiayu and Xu, Shixin and Zou, Dongmian},
+  booktitle = {Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining V.2 (KDD '26)},
+  year      = {2026},
+  doi       = {10.1145/3770855.3818916}
+}
+```
+
+A machine-readable [CITATION.cff](CITATION.cff) is also provided, which GitHub renders into a "Cite this
+repository" button on the project page.
+
 ## License
 
-MIT
+Released under the MIT License. See [LICENSE](LICENSE) for details.

ROADMAP.md

@@ -0,0 +1,33 @@
+# Roadmap
+
+This file tracks planned work beyond the KDD 2026 camera-ready snapshot (`v1.0.0-kdd2026`). The released
+library is the backbone-agnostic StageGuard core (soft transition penalty + semi-Markov decoder + metrics +
+SQI) with two reference backbones and a runnable demo.
+
+## Additional backbone adapters
+
+The paper evaluates six backbones; two ship here (`accusleep`, `usleep`). Reference adapters for the
+remaining four are planned:
+
+- DeepSleepNet (Supratak et al., 2017)
+- SeqSleepNet (Phan et al., 2019)
+- AttnSleep (Eldele et al., 2021)
+- SleepTransformer (Phan et al., 2022)
+
+Any of these already works today by wrapping the author's own implementation in `StageGuardWrapper`; the only
+requirement is a `(B, T, C)` logits output (see [docs/backbones.md](docs/backbones.md)).
+
+## Reproduction harness
+
+An `experiments/` package with per-dataset entry points (`--mode base / +stageguard`) to reproduce the
+paper's tables, plus exact-version environment locks and per-run seeds.
+
+## Packaging and distribution
+
+- Publish to PyPI (`pip install stageguard`).
+- Archive each tagged release on Zenodo for a citable DOI.
+
+## More modalities and configs
+
+- Additional `configs/*.yaml` for further datasets and stage schemes.
+- Optional helper utilities for converting common PSG label formats into the loaders' expected layouts.

configs/sleepaccel_demo.yaml

@@ -0,0 +1,33 @@
+# Sleep-Accel (demo) -- wrist-PSG hypnograms, 3-state collapse
+# Dataset: https://physionet.org/content/sleep-accel/1.0.0/  (Walch et al. 2019, ODC-By)
+# Demo config for notebooks/demo_stageguard.ipynb. Uses the PSG gold-standard
+# hypnograms collapsed to 3 states (Wake / NREM / REM) at 30-second epochs.
+# NOTE: distinct from configs/actigraphy.yaml (2-state Wake/Sleep) -- this 3-state
+#       scheme showcases the rare-transition and minimum-duration constraints.
+
+dataset_name: Sleep-Accel (demo, 3-state)
+dataset_url: https://physionet.org/content/sleep-accel/1.0.0/
+
+stage_names: [Wake, NREM, REM]
+num_classes: 3
+epoch_sec: 30.0
+
+# Direct Wake<->REM is physiologically rare in humans
+rare_transitions:
+  - [0, 2]  # Wake -> REM
+  - [2, 0]  # REM -> Wake
+
+lambda_trans: 1.0
+
+# Semi-Markov decoder
+d_max: 30
+epsilon: 5.0
+gamma: 2.0
+k: 5
+d_min: [1, 2, 2]  # Wake >= 1 epoch (30s), NREM >= 2 (60s), REM >= 2 (60s);
+                  # chosen so < 2% of true epochs fall in sub-d_min bouts across
+                  # the shipped demo subjects (see data/README.md)
+
+# Signal quality
+sqi_method: acceleration_variance
+sqi_threshold: 0.3