akifzahin/timegaitnet

1. What this repo is

TimeGaitNet is a PyTorch-based research codebase for multitask freezing-of-gait (FOG) detection from wearable accelerometer windows. The core system predicts three outputs from the same input window: binary FOG detection, FOG intensity estimation, and future FOG forecasting.

2. High-level architecture

The repository is organized as a pipeline with four main layers:

configuration → model construction → training/loss optimization → evaluation/experiments.

Data loading is expected to happen outside the shown modules via a dataset class referenced in the README, while the repo focuses on model and experiment logic.

3. Model subsystem

The models/ package is the architectural center:

- base_model.py defines the multitask wrapper that produces the three task outputs.

- bimamba.py implements the BiMamba temporal backbone.

- frequency_aware.py adds the frequency-sensitive variant.

- components.py contains reusable building blocks such as feature extractors, fusion blocks, and task heads.

This suggests a modular design where shared encoders feed task-specific prediction heads.

4. Core modeling pattern

The README indicates the model combines:

- multiscale CNN feature extraction for local/windowed sensor patterns,

- bidirectional Mamba state-space encoding for temporal context,

- gated fusion between spatial and temporal features,

- separate heads for each prediction task.

The frequency-aware variant adds a 3–8 Hz band-oriented gate/extractor to emphasize clinically relevant FOG dynamics.

5. Configuration boundary

config/model_configs.py centralizes model variants and hyperparameters. It is the main switchboard for architectural choices such as:

- temporal backbone type,

- whether frequency gating is enabled,

- task-sharing mode,

- dropout,

- task-loss weights.

This makes the repo configuration-driven rather than hardcoding model variants in training scripts.

6. Training subsystem

training/trainer.py likely owns the end-to-end training loop: batching, forward/backward passes, validation, early stopping, and checkpointing.

training/loss.py holds the adaptive multitask loss, which is important because the three tasks are optimized jointly and need dynamic weighting rather than a fixed single-objective loss.

7. Evaluation subsystem

evaluation/evaluator.py is the metrics/benchmarking layer. Based on the README, it computes:

- classification metrics for binary detection and forecasting,

- regression metrics for intensity estimation,

- aggregate statistics across seeds and datasets.

evaluation/visualizations.py generates the diagnostic plots used in the ablation study, such as ROC/PR curves, confusion matrices, and residual plots.

8. Experiment orchestration

experiments/run_ablation.py is the top-level research entry point. It appears to:

- load dataset paths,

- run multiple seeds,

- compare BiMamba vs BiMamba-FreqAware,

- evaluate on two datasets: DAPHNET and TDCSFOG,

- export metrics and plots.

This script is the main integration point connecting config, training, and evaluation.

9. Example entry point

examples/train_example.py is a lightweight usage example for single-run training. It demonstrates the intended flow:

dataset → dataloader → model config → multitask model → batch prediction/training.

This file is helpful for understanding the intended runtime sequence, even if the real experimentation happens in experiments/.

10. Utilities and supporting code

utils/helpers.py likely contains shared helpers such as seeding, logging, checkpoint utilities, metric helpers, or tensor utilities. These are cross-cutting concerns used by training and evaluation without becoming part of the model core.

11. External dependencies and runtime assumptions

The architecture depends materially on:

- PyTorch for all training/inference,

- mamba-ssm and causal-conv1d for the BiMamba backbone,

- CUDA for practical inference/training speed.

Input data is expected to be in parquet-based wearable sensor datasets, and the README references a FOGDataset class that is not present in the file tree, so data ingestion is an external boundary in this snapshot.

12. Data flow summary

The key flow is:

sensor window parquet data → dataset/dataloader → multitask model → three prediction heads → multitask loss during training → evaluator metrics/plots during validation and experiments.

The important architectural boundary is that the model is purely window-based and multitask, while experiment scripts coordinate dataset-specific runs and reporting.

13. Diagramming takeaway

If drawing an architecture diagram, the most important boxes are:

Config, BiMamba/Frequency-Aware Model, Training/Loss, Evaluation/Visualization, and Experiment Runner.

The main arrows should show config feeding model construction, model feeding training and evaluation, and experiments orchestrating repeated seeded runs across datasets.