Analyzing ComBatGAM behavior with imbalance across sites

Imports

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from uniharmony.datasets import make_multisite_classification
from uniharmony.datasets import Covariate, CovariateSiteDistribution

from uniharmony import verbosity
verbosity("warning")
from uniharmony.combat import ComBatGAM

sns.set_theme(style="whitegrid")

Data generation

X, y, sites, covars = make_multisite_classification(n_features=2, signal_type="blobs",
                                                    covariates=["age", "sex"],
                                                    site_effect_strength=10)

df = pd.DataFrame({"Class": y, "Site": sites, "Age": covars["age"],
                   "Feature1":X[:,0], "sex": covars["sex"]})


plt.figure(figsize=[10, 6])
plt.title("Features vs age/sex distribution")
sns.scatterplot(df, y="Age", x="Feature1", hue="Site")
plt.grid(axis="y", color="black", alpha=0.5, linestyle="--")

Caution

Note that we are harmonizing the whole dataset, which must be avoided in ML scenarios. This is just to illustrate the effect of harmonization.

Harmonization

combat = ComBatGAM()
combat.fit(X, sites, smooth_covariates=covars["age"] )
X_harmonized = combat.transform(X, sites, smooth_covariates=covars["age"])

Plotting

df_orig = pd.DataFrame(X, columns=["Feature1", "Feature2"])
df_orig["Site"] = sites
df_orig["Age"] = covars["age"]
df_orig["Phase"] = "Original"

df_harm = pd.DataFrame(X_harmonized, columns=["Feature1", "Feature2"])
df_harm["Site"] = sites
df_harm["Age"] = covars["age"]
df_harm["Phase"] = "Harmonized"


fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharex=True, sharey=True)
sns.scatterplot(data=df_orig, x="Feature1", y="Age", hue="Site", alpha=0.6, ax=axes[0])
axes[0].set_title("Original data by site")
axes[0].grid(alpha=0.3, color="black", linestyle="--")

sns.scatterplot(data=df_harm, x="Feature1", y="Age", hue="Site", alpha=0.6, ax=axes[1])
axes[1].set_title("Harmonized data by site")
axes[1].grid(alpha=0.3, color="black", linestyle="--")
plt.tight_layout()

Preserving the target as covariate

Caution

This is also wrong in ML context, where you don’t have access to the full dataset but may be a good option for statistical analysis.

combat = ComBatGAM()
# This is the key line: we need to include the target variable as a covariate
# to preserve its relationship with the features during harmonization.

combat.fit(X, sites, smooth_covariates=covars["age"])
X_harmonized = combat.transform(X, sites, smooth_covariates=covars["age"])

df_orig = pd.DataFrame(X, columns=["Feature1", "Feature2"])
df_orig["Site"] = sites
df_orig["Target"] = y

df_orig["Phase"] = "Original"

df_harm = pd.DataFrame(X_harmonized, columns=["Feature1", "Feature2"])
df_harm["Site"] = sites
df_harm["Target"] = y

df_harm["Phase"] = "Harmonized"

Plotting

# Plot data distribution by site before and after harmonization
fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharex=True, sharey=True)
sns.scatterplot(data=df_orig, x="Feature1", y="Feature2", hue="Site", style="Target", alpha=0.6, ax=axes[0])
axes[0].set_title("Original data by site")
sns.scatterplot(data=df_harm, x="Feature1", y="Feature2", hue="Site", style="Target",alpha=0.6, ax=axes[1])
axes[1].set_title("Harmonized data by site")
plt.tight_layout()