CohortLearn

CohortLearn is a Python library that builds matched exposed and unexposed cohorts from electronic health records (EHR), following the target trial emulation framework.

It solves two common problems:

• Design bias. Follow-up can start at the wrong time, or controls can come from a different period. CohortLearn sets the same time zero for both arms, and picks controls by risk-set sampling.
• Many confounders. Exact matching on all of them leaves empty groups. CohortLearn combines them into one propensity score and matches on it.

You choose the exposure, the outcome and the confounders as ICD codes. Nothing about the clinical question is fixed in the code.

Quick start: reproduce the worked example

The worked example asks whether depression (ICD-10 F32, F33) is associated with a higher hazard of Alzheimer's disease (ICD-10 F00, G30). It runs on synthetic data of 60,000 participants, generated by the repository itself. The numbers have no clinical meaning.

Choose one of the two options below. Docker gives identical numbers on every machine. Conda is easier if you want to edit the code.

Option A: Conda

You will need Git and conda installed:

• Git: https://git-scm.com/downloads
• conda: https://conda-forge.org/download/

1. Download the repository and move into its folder:

   git clone https://github.com/tmi-lab/CohortLearn.git
   cd CohortLearn

2. Create the environment from the environment.yml file:

   conda env create -f environment.yml

3. Check that the environment was created. cohortlearn should appear in the list:

   conda env list

4. Activate the environment:

   conda activate cohortlearn

5. Run the worked example:

   python run_all.py

Option B: Docker

You will need Docker installed: https://docs.docker.com/get-started/get-docker/. You do not need to download the repository. Docker downloads the image for you.

1. Make an empty folder, and open a terminal in it.

2. Run the published image.

   On Windows (PowerShell):

   docker run --rm -v "${PWD}\outputs:/app/outputs" ghcr.io/ghannamzeinab/cohortlearn:v0.4.0

   On macOS or Linux:

   docker run --rm --platform linux/amd64 -v "$(pwd)/outputs:/app/outputs" ghcr.io/ghannamzeinab/cohortlearn:v0.4.0

Expected output

The run takes a few minutes. When it finishes, a new folder named outputs appears in the folder where you ran it. It holds five files:

File	What it shows
results_summary.txt	The main numerical results, in plain text
love_plot.pdf	Standardised mean difference of every covariate, before and after matching
love_plot_short.pdf	The same plot, showing only covariates imbalanced before matching
ps_overlap.pdf	Propensity score distributions of both arms, before and after matching
matched_cohort_12m.csv	The matched cohort, one row per participant

results_summary.txt reports, in order: the question, the synthetic data, the design settings, the size of each arm after matching and the largest standardised mean difference, follow-up and events in each arm, the hazard ratio, and the E-values.

The main result should be:

• Participants in the exposed arm after matching: 8,476
• Participants in the unexposed arm after matching: 24,919
• Largest standardised mean difference after matching: 0.048
• Hazard ratio: 1.91 (95% CI 1.54–2.37)

With Docker, these numbers are identical on every machine. With conda, the last digit can differ between operating systems. The conclusion is the same.

The terminal also prints two validation checks, run on separate synthetic datasets where the true effect is known:

• Validation: true hazard ratio 1.65. The estimate is 1.77 (95% CI 1.47–2.13), so the true value is inside the interval.
• Null: true hazard ratio 1.00. The estimate is 1.14 (95% CI 0.92–1.40), so the true value is inside the interval.

Use CohortLearn in your own analysis

To use CohortLearn in your own analysis:

1. Download the repository, if you have not already:

   git clone https://github.com/tmi-lab/CohortLearn.git
   cd CohortLearn

2. Install it into your Python environment. If you followed Option A, skip this step: the conda environment already includes it.

   pip install -e .

3. Use it in Python:

from cohortlearn import CohortBuilder, PSMCalculator, SurvivalAnalyser
from cohortlearn.survival import full_rank_covariates

CONF = {
    "exposure_codes": ["F32", "F33"],
    "outcome_codes":  ["F00", "G30"],
    "confounder_codes": {
        "hypertension": ["I10", "I11"],
        "anxiety":      ["F40", "F41"],
    },
    "demographic_confounders": {"age": True, "sex": True, "bmi": True,
                                "education": True, "risk_allele": False,
                                "lookback": True},
    "age_range": (18, 110),
}

builder = CohortBuilder(washout_months=12, require_observation=True,
                        min_lookback_years=2, **CONF)
builder.attach_dataframes(icd10_long=icd10_long_df,
                          demographics=demographics_df,
                          bmi=bmi_df)

cases    = builder.build_case_cohort()
controls = builder.build_control_cohort(cases)
master   = builder.build_master_df(cases, controls)

psm = PSMCalculator(master, exact_match_cols=["Sex"], ratio=3,
                    consort=builder.consort)
psm.fit()
psm.match()

psm.balance_table()
psm.plot_smd()
psm.plot_smd(only_imbalanced=True)
psm.plot_overlap()
psm.strata_table()

probe = SurvivalAnalyser(psm.matched_cohort, study_end="2024-12-31")
probe.prepare()
keep, dropped = full_rank_covariates(probe)

sa = SurvivalAnalyser(psm.matched_cohort, study_end="2024-12-31", covariates=keep)
sa.prepare()
sa.fit(adjusted=True, cluster=True)
sa.summary(exposure_only=True)
sa.e_value()

The scripts can also be run one at a time:

python scripts/generate_synthetic_data.py --n 60000 --out-dir data
python scripts/run_pipeline_example.py --data-dir data --out-dir outputs

To sweep several washout windows:

python scripts/run_pipeline_example.py --data-dir data --sweep 12 24 36 60 120

Input format

Flat tables. No common data model is required.

File	Rows	Columns
diagnoses	one per coded diagnosis	id, code, date
demographics	one per participant	id, Sex, YOB, Education, Date of Death
bmi (optional)	one per participant	id, BMI
genotype (optional)	one per participant	id, plus one genotype column

The diagnosis table also defines each participant's observation span, taken as their first and last coded dates. Identifier fields named eid or person_id are renamed to id automatically.

Repository layout

src/cohortlearn/     the library
  cohort.py          CohortBuilder, CONSORTTracker, rename_id
  matching.py        PSMCalculator
  survival.py        SurvivalAnalyser, full_rank_covariates
scripts/
  generate_synthetic_data.py
  run_pipeline_example.py
tests/               smoke tests
run_all.py           worked example and validation checks
environment.yml      conda environment
Dockerfile           Docker environment
requirements.lock    pinned package versions
data/                generated, not tracked
outputs/             generated, not tracked

What the library does

Cohort construction. Time zero for a participant in the exposed arm is their first exposure code. Time zero for a participant in the unexposed arm is drawn from the observed exposed time zeros, and is accepted only if that person is alive, under observation and unexposed on that date. Controls with the outcome on or before that date are then removed as prevalent cases. Participants who are unexposed early and exposed later contribute their unexposed time and are censored at exposure onset. A configurable washout window removes outcomes that were already developing at time zero. An optional minimum look-back requires both arms to have enough record history before time zero, and look-back length can enter the propensity model.

Matching. Propensity scores come from regularised logistic regression, fitted to tight convergence. Matching is greedy nearest neighbour on the logit score, without replacement, 1:k, inside exact strata for sex and a two-year index-period bucket, with a caliper of 0.2 standard deviations. Distances are compared on an exact integer grid with stable tie-breaking, so matches do not depend on floating-point rounding.

Balance checks. Standardised mean differences and variance ratios before and after matching, a Love plot (full and short), propensity overlap plots, a stratum composition table, and a look-back report comparing record history between arms.

Outcome stage. Cox regression with target-trial censoring, optional standard errors clustered on participant, a proportional hazards check, and E-values for unmeasured confounding.

Synthetic dataset

scripts/generate_synthetic_data.py writes a synthetic dataset in the input format above. Exposure is generated from a logistic function of the covariates, so there is real confounding to remove. Diagnosis dates are set from a drawn age at diagnosis rather than from a calendar window, so timing depends on the person. Records stop at death, so no code is dated after a participant's death. The generator is seeded, so the same settings always give the same data.

Two outcome models are available through --outcome-model:

• logistic (default) sets the exposure effect on the odds of ever having the outcome. No true hazard ratio is defined, so the Cox estimate can be checked for direction only. The worked example uses this model.
• hazard draws outcome times from a proportional hazards model with time-varying exposure, so the true hazard ratio is known and the Cox estimate can be checked against it. The validation checks use this model.

The size of the effect is set with --true-effect, on the log scale. Each run also writes generation_info.json, which records the true effect and its scale.

The synthetic dataset and the worked-example matched cohort are archived on Zenodo (version 2.0.0): https://doi.org/10.5281/zenodo.23072614.

Tests

With the conda environment active, run:

pip install pytest
pytest -q

Scope and limitations

• Time-varying confounding is not supported.
• Diagnoses are assumed to be ICD-coded. Other vocabularies need mapping first.
• One matching method only: propensity score nearest-neighbour matching.

Licence

MIT. See LICENSE.