# akifzahin/timegaitnet architecture

> Architecture diagram of the GitHub repository akifzahin/timegaitnet, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/akifzahin/timegaitnet
- Repository: https://github.com/akifzahin/timegaitnet

## Overview

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.

Main entry points: [`models/base_model.py`](https://github.com/akifzahin/timegaitnet/blob/main/models/base_model.py) (Base model), [`models/bimamba.py`](https://github.com/akifzahin/timegaitnet/blob/main/models/bimamba.py) (BiMamba), [`training/trainer.py`](https://github.com/akifzahin/timegaitnet/blob/main/training/trainer.py) (Trainer), [`experiments/run_ablation.py`](https://github.com/akifzahin/timegaitnet/blob/main/experiments/run_ablation.py) (Ablation run)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_config["Config"]
  node_config_model_configs["Model configs<br/>config registry<br/>[model_configs.py]"]
end

subgraph group_models["Models"]
  node_models_base_model["Base model<br/>multitask wrapper<br/>[base_model.py]"]
  node_models_bimamba["BiMamba<br/>temporal backbone<br/>[bimamba.py]"]
  node_models_frequency_aware["Freq-aware<br/>temporal variant<br/>[frequency_aware.py]"]
  node_models_components["Components<br/>building blocks<br/>[components.py]"]
end

subgraph group_training["Training"]
  node_training_trainer["Trainer<br/>training loop<br/>[trainer.py]"]
  node_training_loss["Adaptive loss<br/>multitask loss<br/>[loss.py]"]
end

subgraph group_evaluation["Evaluation"]
  node_evaluation_evaluator["Evaluator<br/>metrics engine<br/>[evaluator.py]"]
  node_evaluation_visualizations["Plots<br/>reporting<br/>[visualizations.py]"]
end

subgraph group_experiments["Experiments"]
  node_experiments_run_ablation["Ablation run<br/>experiment orchestrator<br/>[run_ablation.py]"]
end

subgraph group_examples["Examples"]
  node_examples_train_example["Train example<br/>usage example<br/>[train_example.py]"]
end

subgraph group_utils["Utilities"]
  node_utils_helpers["Helpers<br/>shared utilities<br/>[helpers.py]"]
end

subgraph group_external["External"]
  node_external_pytorch(("PyTorch<br/>runtime"))
  node_external_mamba(("Mamba SSM<br/>dependency"))
  node_external_causalconv(("causal-conv1d<br/>dependency"))
  node_external_dataset[("FOG data<br/>data boundary")]
end

node_config_model_configs -->|"selects"| node_models_base_model
node_config_model_configs -->|"configures"| node_models_bimamba
node_config_model_configs -->|"enables"| node_models_frequency_aware
node_models_base_model -->|"uses"| node_models_components
node_models_base_model -->|"encodes with"| node_models_bimamba
node_models_base_model -.->|"variant"| node_models_frequency_aware
node_models_bimamba -->|"reuses"| node_models_components
node_models_bimamba -->|"depends on"| node_external_mamba
node_models_bimamba -->|"depends on"| node_external_causalconv
node_models_bimamba -->|"built on"| node_external_pytorch
node_models_frequency_aware -->|"extends"| node_models_components
node_models_frequency_aware -->|"built on"| node_external_pytorch
node_training_trainer -->|"trains"| node_models_base_model
node_training_trainer -->|"optimizes"| node_training_loss
node_training_trainer -->|"uses"| node_utils_helpers
node_training_trainer -->|"built on"| node_external_pytorch
node_training_loss -->|"built on"| node_external_pytorch
node_evaluation_evaluator -->|"evaluates"| node_models_base_model
node_evaluation_evaluator -->|"feeds"| node_evaluation_visualizations
node_evaluation_evaluator -->|"uses"| node_utils_helpers
node_experiments_run_ablation -->|"loads"| node_config_model_configs
node_experiments_run_ablation -->|"runs"| node_training_trainer
node_experiments_run_ablation -->|"reports"| node_evaluation_evaluator
node_experiments_run_ablation -->|"exports"| node_evaluation_visualizations
node_experiments_run_ablation -->|"uses"| node_external_dataset
node_examples_train_example -->|"loads"| node_external_dataset
node_examples_train_example -->|"demonstrates"| node_config_model_configs
node_examples_train_example -->|"instantiates"| node_models_base_model
node_examples_train_example -->|"invokes"| node_training_trainer
node_utils_helpers -->|"supports"| node_external_pytorch

click node_config_model_configs "https://github.com/akifzahin/timegaitnet/blob/main/config/model_configs.py"
click node_models_base_model "https://github.com/akifzahin/timegaitnet/blob/main/models/base_model.py"
click node_models_bimamba "https://github.com/akifzahin/timegaitnet/blob/main/models/bimamba.py"
click node_models_frequency_aware "https://github.com/akifzahin/timegaitnet/blob/main/models/frequency_aware.py"
click node_models_components "https://github.com/akifzahin/timegaitnet/blob/main/models/components.py"
click node_training_trainer "https://github.com/akifzahin/timegaitnet/blob/main/training/trainer.py"
click node_training_loss "https://github.com/akifzahin/timegaitnet/blob/main/training/loss.py"
click node_evaluation_evaluator "https://github.com/akifzahin/timegaitnet/blob/main/evaluation/evaluator.py"
click node_evaluation_visualizations "https://github.com/akifzahin/timegaitnet/blob/main/evaluation/visualizations.py"
click node_experiments_run_ablation "https://github.com/akifzahin/timegaitnet/blob/main/experiments/run_ablation.py"
click node_examples_train_example "https://github.com/akifzahin/timegaitnet/blob/main/examples/train_example.py"
click node_utils_helpers "https://github.com/akifzahin/timegaitnet/blob/main/utils/helpers.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_config_model_configs,node_utils_helpers toneBlue
class node_models_base_model,node_models_bimamba,node_models_frequency_aware,node_models_components,node_external_pytorch,node_external_mamba,node_external_causalconv,node_external_dataset toneAmber
class node_training_trainer,node_training_loss toneMint
class node_evaluation_evaluator,node_evaluation_visualizations toneRose
class node_experiments_run_ablation toneIndigo
class node_examples_train_example toneTeal
```

## Components

### Config

Model and hyperparameter selection

- [Model configs](https://github.com/akifzahin/timegaitnet/blob/main/config/model_configs.py): config registry · Central model variant and hyperparameter switchboard · `config/model_configs.py`

### Models

Multitask architecture and components

- [Base model](https://github.com/akifzahin/timegaitnet/blob/main/models/base_model.py): multitask wrapper · Wraps shared features into binary detection, intensity, and forecasting outputs · `models/base_model.py`
- [BiMamba](https://github.com/akifzahin/timegaitnet/blob/main/models/bimamba.py): temporal backbone · Bidirectional Mamba sequence encoder for temporal context · `models/bimamba.py`
- [Freq-aware](https://github.com/akifzahin/timegaitnet/blob/main/models/frequency_aware.py): temporal variant · Adds frequency-sensitive gating for clinically relevant gait dynamics · `models/frequency_aware.py`
- [Components](https://github.com/akifzahin/timegaitnet/blob/main/models/components.py): building blocks · Shared feature extractors, fusion blocks, and task heads · `models/components.py`

### Training

Optimization and loss handling

- [Trainer](https://github.com/akifzahin/timegaitnet/blob/main/training/trainer.py): training loop · Runs forward/backward passes, validation, and checkpointing · `training/trainer.py`
- [Adaptive loss](https://github.com/akifzahin/timegaitnet/blob/main/training/loss.py): multitask loss · Balances the three tasks with dynamic weighting · `training/loss.py`

### Evaluation

Metrics and plots

- [Evaluator](https://github.com/akifzahin/timegaitnet/blob/main/evaluation/evaluator.py): metrics engine · Computes classification, regression, and aggregate benchmark metrics · `evaluation/evaluator.py`
- [Plots](https://github.com/akifzahin/timegaitnet/blob/main/evaluation/visualizations.py): reporting · Generates ROC, PR, confusion, and residual diagnostics · `evaluation/visualizations.py`

### Experiments

Ablation and seeded runs

- [Ablation run](https://github.com/akifzahin/timegaitnet/blob/main/experiments/run_ablation.py): experiment orchestrator · Runs seeded comparisons across datasets and model variants · `experiments/run_ablation.py`

### Examples

Minimal usage entrypoint

- [Train example](https://github.com/akifzahin/timegaitnet/blob/main/examples/train_example.py): usage example · Minimal single-run training workflow · `examples/train_example.py`

### Utilities

Shared helpers

- [Helpers](https://github.com/akifzahin/timegaitnet/blob/main/utils/helpers.py): shared utilities · Seeding, logging, checkpoint, and tensor helpers · `utils/helpers.py`

### External

Dependencies and data boundary

- **PyTorch**: runtime · Core tensor and training framework
- **Mamba SSM**: dependency · State-space sequence modeling backend
- **causal-conv1d**: dependency · Convolution kernel used by the temporal backbone
- **FOG data**: data boundary · External parquet-based wearable sensor dataset and dataset class

## Connections

- Model configs → Base model: selects · Config chooses multitask model variant and shared settings
- Model configs → BiMamba: configures · Controls temporal backbone options
- Model configs → Freq-aware: enables · Turns frequency-aware gating on or off
- Base model → Components: uses · Builds shared feature extractors and task heads
- Base model → BiMamba: encodes with · Uses BiMamba as temporal encoder
- Base model → Freq-aware: variant · Alternative frequency-aware architecture path
- BiMamba → Components: reuses · Composes common blocks from the components module
- BiMamba → Mamba SSM: depends on · Uses Mamba state-space layers
- BiMamba → causal-conv1d: depends on · Uses causal convolution support
- BiMamba → PyTorch: built on · Implemented with PyTorch modules
- Freq-aware → Components: extends · Adds frequency extractor and gated fusion blocks
- Freq-aware → PyTorch: built on · Implemented with PyTorch modules
- Trainer → Base model: trains · Drives model forward and backward passes
- Trainer → Adaptive loss: optimizes · Uses adaptive multitask objective
- Trainer → Helpers: uses · Relies on shared runtime helpers
- Trainer → PyTorch: built on · Training loop uses PyTorch execution
- Adaptive loss → PyTorch: built on · Loss functions use PyTorch tensors
- Evaluator → Base model: evaluates · Scores multitask predictions against targets
- Evaluator → Plots: feeds · Metrics drive diagnostic plots
- Evaluator → Helpers: uses · Shares metric and logging helpers
- Ablation run → Model configs: loads · Selects variants and hyperparameters
- Ablation run → Trainer: runs · Launches repeated training sessions
- Ablation run → Evaluator: reports · Aggregates benchmark metrics
- Ablation run → Plots: exports · Generates comparison figures
- Ablation run → FOG data: uses · Operates on external dataset paths and loaders
- Train example → FOG data: loads · Illustrates dataset to dataloader setup
- Train example → Model configs: demonstrates · Shows how to pick a model configuration
- Train example → Base model: instantiates · Creates the multitask model for a single run
- Train example → Trainer: invokes · Hands the model to the training loop
- Helpers → PyTorch: supports · Utility functions assist tensor-based workflows