# valentyn1997/causaltransformer architecture

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

- Interactive diagram: https://gitdiagram.com/valentyn1997/causaltransformer
- Repository: https://github.com/valentyn1997/causaltransformer

## Overview

#### Purpose and execution
CausalTransformer is a Python research-training repository for estimating longitudinal counterfactual outcomes under time-varying treatment confounding. Its executable entry points are model-family-specific runners: `runnables/train_multi.py`, `runnables/train_enc_dec.py`, `runnables/train_gnet.py`, `runnables/train_msm.py`, and `runnables/train_rmsn.py`.

#### Configuration-driven experiment assembly
Hydra composes a root experiment configuration with a selected dataset and backbone, including saved hyperparameter variants. Central configuration is `config/config.yaml`; dataset selections are in `config/dataset/cancer_sim.yaml`, `config/dataset/mimic3_synthetic.yaml`, and `config/dataset/mimic3_real.yaml`; backbone defaults are under `config/backbone/`. This is the main control boundary between runners and implementations.

#### Longitudinal dataset layer
A shared collection abstraction coordinates cohort splits and sequence-oriented data consumed by models: `src/data/dataset_collection.py`. Data originates from:
- Synthetic tumor-growth trajectories: `src/data/cancer_sim/cancer_simulation.py`, `src/data/cancer_sim/dataset.py`.
- MIMIC-III semi-synthetic trajectories: `src/data/mimic_iii/semi_synthetic_dataset.py`.
- MIMIC-III real-world outcome data: `src/data/mimic_iii/real_dataset.py`.
MIMIC source extraction is loaded from local processed data through `src/data/mimic_iii/load_data.py`; expected external input is `data/processed/all_hourly_data.h5`.

#### Core temporal-model contract
`src/models/time_varying_model.py` is the likely shared abstraction for temporal causal prediction, allowing runners to train and evaluate different estimators against a common longitudinal-data interface. Shared model operations and neural primitives sit in `src/models/utils.py`, `src/models/utils_lstm.py`, and `src/models/utils_transformer.py`, separating estimator-specific causal logic from recurrent/Transformer mechanics.

#### Transformer causal estimators
The repository’s primary contribution is implemented in `src/models/ct.py` (Causal Transformer) and `src/models/edct.py` (encoder-decoder Causal Transformer). These models predict future outcomes from historical covariates and treatments; configurations support balancing through counterfactual domain confusion or gradient reversal. Architecture and training variants are selected via `config/backbone/ct.yaml` and `config/backbone/edct.yaml`.

#### Baseline estimator families
Comparable causal baselines are independently implemented: Counterfactual Recurrent Network (`src/models/crn.py`), G-Net (`src/models/gnet.py`), Marginal Structural Models (`src/models/msm.py`), and Recurrent Marginal Structural Networks (`src/models/rmsn.py`). Runner choice is meaningful: encoder-decoder training serves CT/EDCT/CRN, while G-Net, MSM, and RMSN have dedicated workflows. Backbone declarations include `config/backbone/crn.yaml`, `config/backbone/gnet.yaml`, and `config/backbone/rmsn.yaml`.

#### Runtime, experiment infrastructure, and outputs
The material runtime is Python with PyTorch Lightning for training lifecycle/device execution, Hydra for CLI configuration and multiruns, and MLflow for experiment tracking. GPU selection is supplied externally through `CUDA_VISIBLE_DEVICES`; MLflow requires an independently running tracking server. Dependencies are declared in `requirements.txt`, packaging metadata in `setup.py`, and processed external datasets are intentionally excluded except for the placeholder `data/processed/.gitkeep`.

Main entry points: [`src/models/time_varying_model.py`](https://github.com/valentyn1997/causaltransformer/blob/main/src/models/time_varying_model.py) (Temporal model contract), [`runnables/train_gnet.py`](https://github.com/valentyn1997/causaltransformer/blob/main/runnables/train_gnet.py) (Baseline runners), [`config/backbone/ct.yaml`](https://github.com/valentyn1997/causaltransformer/blob/main/config/backbone/ct.yaml) (Backbone configs), [`config/config.yaml`](https://github.com/valentyn1997/causaltransformer/blob/main/config/config.yaml) (Root Hydra config)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_experiment["Experiment assembly"]
  node_train_multi["Multi-model runner<br/>[train_multi.py]"]
  node_train_enc_dec["Encoder-decoder runner<br/>[train_enc_dec.py]"]
  node_baseline_runners["Baseline runners<br/>[train_gnet.py]"]
  node_root_config["Root Hydra config<br/>configuration<br/>[config.yaml]"]
  node_dataset_configs["Dataset configs<br/>configuration<br/>[cancer_sim.yaml]"]
  node_backbone_configs["Backbone configs<br/>configuration<br/>[ct.yaml]"]
end

subgraph group_data["Longitudinal data"]
  node_dataset_collection["Dataset collection<br/>data abstraction"]
  node_cancer_data["Cancer simulation<br/>synthetic data"]
  node_mimic_semi["MIMIC semi-synthetic<br/>dataset"]
  node_mimic_real["MIMIC real outcomes<br/>dataset<br/>[real_dataset.py]"]
  node_mimic_loader["MIMIC source loader<br/>data loader<br/>[load_data.py]"]
  node_mimic_hdf[("Processed MIMIC data<br/>external dataset")]
end

subgraph group_models["Causal estimators"]
  node_temporal_contract["Temporal model contract<br/>model abstraction"]
  node_model_primitives["Model primitives<br/>model utilities<br/>[utils.py]"]
  node_transformers["CT and EDCT<br/>causal estimator<br/>[ct.py]"]
  node_crn["CRN baseline<br/>causal estimator<br/>[crn.py]"]
  node_gnet["G-Net baseline<br/>causal estimator<br/>[gnet.py]"]
  node_msm["MSM baseline<br/>causal estimator<br/>[msm.py]"]
  node_rmsn["RMSN baseline<br/>causal estimator<br/>[rmsn.py]"]
end

subgraph group_runtime["Training runtime"]
  node_training_runtime{{"Lightning, MLflow, GPU<br/>runtime infrastructure"}}
end

node_root_config -->|"selects"| node_dataset_configs
node_root_config -->|"selects"| node_backbone_configs
node_train_multi -->|"composes"| node_root_config
node_train_enc_dec -->|"composes"| node_root_config
node_baseline_runners -->|"composes"| node_root_config
node_dataset_configs -->|"builds"| node_dataset_collection
node_cancer_data -->|"supplies trajectories"| node_dataset_collection
node_mimic_semi -->|"supplies trajectories"| node_dataset_collection
node_mimic_real -->|"supplies outcomes"| node_dataset_collection
node_mimic_hdf -->|"loaded from"| node_mimic_loader
node_mimic_loader -->|"provides source data"| node_mimic_semi
node_mimic_loader -->|"provides source data"| node_mimic_real
node_dataset_collection -->|"provides sequences"| node_temporal_contract
node_temporal_contract -->|"implemented by"| node_transformers
node_temporal_contract -->|"implemented by"| node_crn
node_temporal_contract -->|"implemented by"| node_gnet
node_temporal_contract -->|"implemented by"| node_msm
node_temporal_contract -->|"implemented by"| node_rmsn
node_model_primitives -->|"supports"| node_transformers
node_model_primitives -->|"supports"| node_crn
node_backbone_configs -->|"configures"| node_transformers
node_backbone_configs -->|"configures"| node_crn
node_backbone_configs -->|"configures"| node_gnet
node_backbone_configs -->|"configures"| node_rmsn
node_train_enc_dec -->|"trains"| node_transformers
node_train_enc_dec -->|"trains"| node_crn
node_baseline_runners -->|"trains"| node_gnet
node_baseline_runners -->|"trains"| node_msm
node_baseline_runners -->|"trains"| node_rmsn
node_train_multi -->|"executes on"| node_training_runtime
node_train_enc_dec -->|"executes on"| node_training_runtime
node_baseline_runners -->|"executes on"| node_training_runtime

click node_train_multi "https://github.com/valentyn1997/causaltransformer/blob/main/runnables/train_multi.py"
click node_train_enc_dec "https://github.com/valentyn1997/causaltransformer/blob/main/runnables/train_enc_dec.py"
click node_baseline_runners "https://github.com/valentyn1997/causaltransformer/blob/main/runnables/train_gnet.py"
click node_root_config "https://github.com/valentyn1997/causaltransformer/blob/main/config/config.yaml"
click node_dataset_configs "https://github.com/valentyn1997/causaltransformer/blob/main/config/dataset/cancer_sim.yaml"
click node_backbone_configs "https://github.com/valentyn1997/causaltransformer/blob/main/config/backbone/ct.yaml"
click node_dataset_collection "https://github.com/valentyn1997/causaltransformer/blob/main/src/data/dataset_collection.py"
click node_cancer_data "https://github.com/valentyn1997/causaltransformer/blob/main/src/data/cancer_sim/cancer_simulation.py"
click node_mimic_semi "https://github.com/valentyn1997/causaltransformer/blob/main/src/data/mimic_iii/semi_synthetic_dataset.py"
click node_mimic_real "https://github.com/valentyn1997/causaltransformer/blob/main/src/data/mimic_iii/real_dataset.py"
click node_mimic_loader "https://github.com/valentyn1997/causaltransformer/blob/main/src/data/mimic_iii/load_data.py"
click node_temporal_contract "https://github.com/valentyn1997/causaltransformer/blob/main/src/models/time_varying_model.py"
click node_model_primitives "https://github.com/valentyn1997/causaltransformer/blob/main/src/models/utils.py"
click node_transformers "https://github.com/valentyn1997/causaltransformer/blob/main/src/models/ct.py"
click node_crn "https://github.com/valentyn1997/causaltransformer/blob/main/src/models/crn.py"
click node_gnet "https://github.com/valentyn1997/causaltransformer/blob/main/src/models/gnet.py"
click node_msm "https://github.com/valentyn1997/causaltransformer/blob/main/src/models/msm.py"
click node_rmsn "https://github.com/valentyn1997/causaltransformer/blob/main/src/models/rmsn.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_train_multi,node_train_enc_dec,node_baseline_runners,node_root_config,node_dataset_configs,node_backbone_configs toneBlue
class node_dataset_collection,node_cancer_data,node_mimic_semi,node_mimic_real,node_mimic_loader,node_mimic_hdf toneAmber
class node_temporal_contract,node_model_primitives,node_transformers,node_crn,node_gnet,node_msm,node_rmsn toneMint
class node_training_runtime toneRose
```

## Components

### Experiment assembly

Hydra selects dataset and estimator variants.

- [Multi-model runner](https://github.com/valentyn1997/causaltransformer/blob/main/runnables/train_multi.py): runner · Entry point for multi-model training. · `runnables/train_multi.py`
- [Encoder-decoder runner](https://github.com/valentyn1997/causaltransformer/blob/main/runnables/train_enc_dec.py): runner · Trains CT, EDCT, and CRN workflows. · `runnables/train_enc_dec.py`
- [Baseline runners](https://github.com/valentyn1997/causaltransformer/blob/main/runnables/train_gnet.py): runner · Dedicated G-Net, MSM, and RMSN entry points. · `runnables/train_gnet.py`
- [Root Hydra config](https://github.com/valentyn1997/causaltransformer/blob/main/config/config.yaml): configuration · Composes experiment, dataset, and backbone selections. · `config/config.yaml`
- [Dataset configs](https://github.com/valentyn1997/causaltransformer/blob/main/config/dataset/cancer_sim.yaml): configuration · Cancer and MIMIC-III dataset selections. · `config/dataset/cancer_sim.yaml`
- [Backbone configs](https://github.com/valentyn1997/causaltransformer/blob/main/config/backbone/ct.yaml): configuration · CT, EDCT, CRN, G-Net, and RMSN variants. · `config/backbone/ct.yaml`

### Longitudinal data

Shared cohort and trajectory data layer.

- [Dataset collection](https://github.com/valentyn1997/causaltransformer/blob/main/src/data/dataset_collection.py): data abstraction · Coordinates cohort splits and model-ready sequences. · `src/data/dataset_collection.py`
- [Cancer simulation](https://github.com/valentyn1997/causaltransformer/blob/main/src/data/cancer_sim/cancer_simulation.py): synthetic data · Generates tumor-growth trajectories and dataset records. · `src/data/cancer_sim/cancer_simulation.py`
- [MIMIC semi-synthetic](https://github.com/valentyn1997/causaltransformer/blob/main/src/data/mimic_iii/semi_synthetic_dataset.py): dataset · Builds semi-synthetic longitudinal trajectories. · `src/data/mimic_iii/semi_synthetic_dataset.py`
- [MIMIC real outcomes](https://github.com/valentyn1997/causaltransformer/blob/main/src/data/mimic_iii/real_dataset.py): dataset · Provides real-world outcome data. · `src/data/mimic_iii/real_dataset.py`
- [MIMIC source loader](https://github.com/valentyn1997/causaltransformer/blob/main/src/data/mimic_iii/load_data.py): data loader · Loads locally processed hourly MIMIC data. · `src/data/mimic_iii/load_data.py`
- **Processed MIMIC data**: external dataset · Externally supplied processed source data.

### Causal estimators

Common temporal contract and estimator families.

- [Temporal model contract](https://github.com/valentyn1997/causaltransformer/blob/main/src/models/time_varying_model.py): model abstraction · Shared interface for longitudinal causal prediction. · `src/models/time_varying_model.py`
- [Model primitives](https://github.com/valentyn1997/causaltransformer/blob/main/src/models/utils.py): model utilities · Shared operations plus LSTM and Transformer mechanics. · `src/models/utils.py`
- [CT and EDCT](https://github.com/valentyn1997/causaltransformer/blob/main/src/models/ct.py): causal estimator · Transformer estimators for historical treatment and outcome prediction. · `src/models/ct.py`
- [CRN baseline](https://github.com/valentyn1997/causaltransformer/blob/main/src/models/crn.py): causal estimator · `src/models/crn.py`
- [G-Net baseline](https://github.com/valentyn1997/causaltransformer/blob/main/src/models/gnet.py): causal estimator · `src/models/gnet.py`
- [MSM baseline](https://github.com/valentyn1997/causaltransformer/blob/main/src/models/msm.py): causal estimator · `src/models/msm.py`
- [RMSN baseline](https://github.com/valentyn1997/causaltransformer/blob/main/src/models/rmsn.py): causal estimator · `src/models/rmsn.py`

### Training runtime

Execution, tracking, and external inputs.

- **Lightning, MLflow, GPU**: runtime infrastructure · PyTorch Lightning executes training; MLflow records experiments.

## Connections

- Root Hydra config → Dataset configs: selects
- Root Hydra config → Backbone configs: selects
- Multi-model runner → Root Hydra config: composes
- Encoder-decoder runner → Root Hydra config: composes
- Baseline runners → Root Hydra config: composes
- Dataset configs → Dataset collection: builds
- Cancer simulation → Dataset collection: supplies trajectories
- MIMIC semi-synthetic → Dataset collection: supplies trajectories
- MIMIC real outcomes → Dataset collection: supplies outcomes
- Processed MIMIC data → MIMIC source loader: loaded from
- MIMIC source loader → MIMIC semi-synthetic: provides source data
- MIMIC source loader → MIMIC real outcomes: provides source data
- Dataset collection → Temporal model contract: provides sequences
- Temporal model contract → CT and EDCT: implemented by
- Temporal model contract → CRN baseline: implemented by
- Temporal model contract → G-Net baseline: implemented by
- Temporal model contract → MSM baseline: implemented by
- Temporal model contract → RMSN baseline: implemented by
- Model primitives → CT and EDCT: supports
- Model primitives → CRN baseline: supports
- Backbone configs → CT and EDCT: configures
- Backbone configs → CRN baseline: configures
- Backbone configs → G-Net baseline: configures
- Backbone configs → RMSN baseline: configures
- Encoder-decoder runner → CT and EDCT: trains
- Encoder-decoder runner → CRN baseline: trains
- Baseline runners → G-Net baseline: trains
- Baseline runners → MSM baseline: trains
- Baseline runners → RMSN baseline: trains
- Multi-model runner → Lightning, MLflow, GPU: executes on
- Encoder-decoder runner → Lightning, MLflow, GPU: executes on
- Baseline runners → Lightning, MLflow, GPU: executes on