Comparison with simple baselines on the simulation dataset

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

import sys
sys.path.append(r'/import/home2/yzengbj/my_project/code')
sys.path.append(r'/import/home2/share/yqzeng/MacSGP/codes')

import warnings
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import scanpy as sc
import anndata as ad
import scipy.io

import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.colors import ListedColormap, BoundaryNorm, TwoSlopeNorm
from matplotlib.ticker import MaxNLocator
import seaborn as sns

from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA

from scipy.stats import pearsonr
from itertools import combinations

import libpysal
from libpysal.weights import KNN
import esda

from spatialdata_plot.pl.utils import set_zero_in_cmap_to_transparent

import SVF
import simulation as construct_st
import MacSGP

DATA_PATH = "/import/home2/share/yqzeng/MacSPG/data/simulation/final" # Raw data
SAVE_PATH = "/import/home2/share/yqzeng/MacSPG/results/simulation/final" # Deconvolution results
os.makedirs(os.path.join(DATA_PATH), exist_ok = True)
os.makedirs(os.path.join(SAVE_PATH), exist_ok = True)

Residual pca

factor_pattern = 'gradient_y'
factor_celltype = [0]

n_spots = 3000
n_genes = 2000
n_celltypes = 2
n_cells = 500
n_covariates = 1
n_loading_genes = 200

n_groups = 4
ref_noise_ratio = 10

library_size = 5000
prop_confound = True

prop_min = 0.1
prop_max = 0.9

same_tau = True

spots_effects = False
platform_effects = False

shifted = False
c = 1

adata_ref, adata_st, reference, logmu = construct_st.simulate_data_2ct(
    seed=20250723,
    n_spots=n_spots, n_genes=n_genes,
    library_size=np.ones(n_spots) * library_size,
    n_celltypes=n_celltypes,
    n_cells=n_cells,
    n_covariates=n_covariates,
    n_groups=n_groups,
    ref_noise_ratio=ref_noise_ratio,
    same_tau=same_tau,
    prop_confound=prop_confound,
    prop_min=prop_min, prop_max=prop_max,
    spot_effects=spots_effects, platform_effects=platform_effects,
    factor=True, n_loading_gene=n_loading_genes,
    factor_celltypes=factor_celltype, factor_pattern=factor_pattern,
    shifted=False
)

MacSGP.utils.Cal_Spatial_Net(adata_st, mode='KNN', k_cutoff=12)

adata_st, adata_basis = MacSGP.utils.preprocess(adata_st,adata_ref,
                                                celltype_ref_col = "celltype",
                                                n_hvg_group = 1000)

model = MacSGP.model.Model_deconv(adata_st, adata_basis, n_layers =4, training_steps=10000)
model.train(step_interval=1000)
adata_st = model.eval()

model = MacSGP.model.Model(adata_st, adata_basis, n_layers=4, training_steps=3000)

model.train(step_interval=600)
adata_result = model.eval()

import numpy as np
import pandas as pd
from sklearn.decomposition import PCA

count = adata_st.obsm['count']
count = pd.DataFrame(count, index=adata_st.obs_names, columns=adata_st.var_names)
count_deconv = adata_st.obsm['proportion'].dot(adata_basis.X)
count_deconv = pd.DataFrame(count_deconv.values, index=adata_st.obs_names, columns=adata_st.var_names)
count_normalized = count.div(count.sum(axis=1), axis=0)
count_residual = count_deconv - count_normalized
count_residual = count_normalized - count_deconv

pearson_residuals = (count_normalized - count_deconv) / np.sqrt(count_deconv + 1e-6)
X_pca = pearson_residuals
k = 5
pca = PCA(n_components=k, svd_solver='auto', random_state=0)

# Factor scores:
# shape: (n_spots, k)
factor_scores = pca.fit_transform(X_pca)
factor_scores_df = pd.DataFrame(
    factor_scores,
    index=count_residual.index,
    columns=[f'PC{i+1}' for i in range(k)]
)

# Loading matrix:
loadings = pca.components_.T
loadings_df = pd.DataFrame(
    loadings,
    index=count_residual.columns,  # 基因名
    columns=[f'PC{i+1}' for i in range(k)]
)

s = 8
proportion = adata_st.obsm['proportion']['celltype_0']
factor_true = adata_st.obsm['simulated_factor']
factor_macsgp = adata_result.obsm['factor']['celltype_0']
factor_pca = factor_scores_df['PC1'].values

sign_m = np.sign(pearsonr(factor_macsgp.values, factor_true.values.reshape(-1))[0])
factor_macsgp_aligned = factor_macsgp * sign_m

sign_p = np.sign(pearsonr(factor_pca, factor_true.values.reshape(-1))[0])
factor_pca_aligned = factor_pca * sign_p

data = proportion

fig, ax = plt.subplots(figsize=(3.5, 3.2))
#norm = mpl.colors.TwoSlopeNorm(vcenter=0.0)
sc = ax.scatter(
    x,
    y,
    c=data,
    s=8,
    cmap='YlOrBr',
    linewidths=0,
    #norm=norm
)
cb   = fig.colorbar(sc, ax=ax, shrink=0.4, pad=0.02, orientation='horizontal')
cb.locator = MaxNLocator(nbins=1)
cb.update_ticks()
cb.ax.tick_params(labelsize=12)
#cb.set_label('Proportion', fontsize=12)

ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['left'].set_linewidth(False)
ax.spines['bottom'].set_linewidth(False)

ax.set_title('Proportion', fontsize=15)
ax.set_xticks([])
ax.set_yticks([])
ax.set_aspect('equal', adjustable='box')
plt.tight_layout()
plt.show()

data = factor_macsgp_aligned

fig, ax = plt.subplots(figsize=(3.5, 3.2))
norm = mpl.colors.TwoSlopeNorm(vcenter=0.0)
sc = ax.scatter(
    x,
    y,
    c=data,
    s=8,
    cmap='bwr',
    linewidths=0,
    #norm=norm
)
cb   = fig.colorbar(sc, ax=ax, shrink=0.4, pad=0.02, orientation='horizontal')
#cb.locator = MaxNLocator(nbins=1)
#cb.update_ticks()
#cb.ax.tick_params(labelsize=0)
cb.set_label('Normalized score', fontsize=12)
# no ticks
cb.set_ticks([])

ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['left'].set_linewidth(False)
ax.spines['bottom'].set_linewidth(False)

ax.set_title('MacSGP', fontsize=15)
ax.set_xticks([])
ax.set_yticks([])
ax.set_aspect('equal', adjustable='box')
plt.tight_layout()

data = factor_pca_aligned

fig, ax = plt.subplots(figsize=(3.5, 3.2))
norm = mpl.colors.TwoSlopeNorm(vcenter=0.0)
sc = ax.scatter(
    x,
    y,
    c=data,
    s=8,
    cmap='bwr',
    linewidths=0,
    #norm=norm
)
cb   = fig.colorbar(sc, ax=ax, shrink=0.4, pad=0.02, orientation='horizontal')
#cb.locator = MaxNLocator(nbins=1)
#cb.update_ticks()
#cb.ax.tick_params(labelsize=0)
cb.set_label('Normalized score', fontsize=12)
# no ticks
cb.set_ticks([])

ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['left'].set_linewidth(False)
ax.spines['bottom'].set_linewidth(False)

ax.set_title('Residual PCA', fontsize=15)
ax.set_xticks([])
ax.set_yticks([])
ax.set_aspect('equal', adjustable='box')
plt.tight_layout()

import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from matplotlib.ticker import MaxNLocator

ct = 'Ext_L23'
data = factor_true.values

fig, ax = plt.subplots(figsize=(3.5, 3.2))
norm = mpl.colors.TwoSlopeNorm(vcenter=0.0)
sc = ax.scatter(
    x,
    y,
    c=data,
    s=6,
    cmap='bwr',
    linewidths=0,
    #norm=norm
)
cb   = fig.colorbar(sc, ax=ax, shrink=0.4, pad=0.02, orientation='horizontal')
#cb.locator = MaxNLocator(nbins=1)
#cb.update_ticks()
#cb.ax.tick_params(labelsize=0)
cb.set_label('Normalized score', fontsize=12)
# no ticks
cb.set_ticks([])

ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['left'].set_linewidth(False)
ax.spines['bottom'].set_linewidth(False)

ax.set_title('Ground Truth', fontsize=15)
ax.set_xticks([])
ax.set_yticks([])
ax.set_aspect('equal', adjustable='box')
plt.tight_layout()

fig, ax = plt.subplots(figsize=(5, 3.2))
bar_w = 0.3
x_pos = np.arange(2)
bars1 = ax.bar(x_pos - bar_w/2, macsgp_scores, bar_w,
               label='MacSGP', color='#4C72B0', edgecolor='white')
bars2 = ax.bar(x_pos + bar_w/2, pca_scores, bar_w,
               label='Residual PCA', color='#DD8452', edgecolor='white')

for bar in list(bars1) + list(bars2):
    h = bar.get_height()
    ax.text(bar.get_x() + bar.get_width()/2, h + 0.01,
            f'{h:.3f}', ha='center', va='bottom', fontsize=8)

ax.set_xticks(x_pos)
ax.set_xticklabels(['Factor\nCoherence', 'Loading\nCoherence'], fontsize=15)
ax.set_ylim(0, 1.1)
ax.set_ylabel('Coherence', fontsize=15)
#ax.set_title('Coherence Comparison', fontsize=15)
ax.legend(fontsize=8, loc='upper right')
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
plt.tight_layout()
plt.show()

Within celltype pca

factor_pattern = 'gradient_y'
factor_celltype = [0]

n_spots = 3000
n_genes = 2000
n_celltypes = 2
n_cells = 500
n_covariates = 1
n_loading_genes = 200

n_groups = 4
ref_noise_ratio = 10

library_size = 5000
prop_confound = True

prop_min = 0.1
prop_max = 0.9

same_tau = True

spots_effects = False
platform_effects = False

shifted = False
c = 1

adata_ref, adata_st, reference, logmu = construct_st.simulate_data_2ct(
    seed=20250723,
    n_spots=n_spots, n_genes=n_genes,
    library_size=np.ones(n_spots) * library_size,
    n_celltypes=n_celltypes,
    n_cells=n_cells,
    n_covariates=n_covariates,
    n_groups=n_groups,
    ref_noise_ratio=ref_noise_ratio,
    same_tau=same_tau,
    prop_confound=prop_confound,
    prop_min=prop_min, prop_max=prop_max,
    spot_effects=spots_effects, platform_effects=platform_effects,
    factor=True, n_loading_gene=n_loading_genes,
    factor_celltypes=factor_celltype, factor_pattern=factor_pattern,
    shifted=False, c = 0.001,
)

Reference generated
Tau generated
Single cell data generated
Coordinates generated
Proportion generated
Covariates generated
mu generated
Lambda generated
Reference data generated
Spatial data generated

MacSGP.utils.Cal_Spatial_Net(adata_st, mode='KNN', k_cutoff=12)

Calculating spatial neighbor graph ...
The graph contains 40806 edges, 3000 spots.
13.602 neighbors per spot on average.

adata_st, adata_basis = MacSGP.utils.preprocess(adata_st,adata_ref,
                                                celltype_ref_col = "celltype",
                                                n_hvg_group = 1000)

Finding highly variable genes...
2000 highly variable genes selected.
Calculate basis for deconvolution...
Preprocess ST data...

model = MacSGP.model.Model_deconv(adata_st, adata_basis, n_layers =4, training_steps=10000)
model.train(step_interval=1000)
adata_st = model.eval()

  0%|          | 9/10000 [00:02<37:24,  4.45it/s]

Step: 0, Loss: -736.1868, d_loss: -743.5078, f_loss: 73.2095

 10%|█         | 1026/10000 [00:11<01:12, 123.85it/s]

Step: 1000, Loss: -930.1940, d_loss: -933.8029, f_loss: 36.0887

 20%|██        | 2021/10000 [00:19<01:05, 121.14it/s]

Step: 2000, Loss: -931.1221, d_loss: -934.7235, f_loss: 36.0140

 30%|███       | 3022/10000 [00:27<00:55, 125.01it/s]

Step: 3000, Loss: -931.2433, d_loss: -934.8381, f_loss: 35.9480

 40%|████      | 4026/10000 [00:36<00:47, 124.64it/s]

Step: 4000, Loss: -931.3071, d_loss: -934.8942, f_loss: 35.8717

 50%|█████     | 5021/10000 [00:44<00:40, 124.13it/s]

Step: 5000, Loss: -931.3336, d_loss: -934.9161, f_loss: 35.8255

 60%|██████    | 6018/10000 [00:52<00:31, 124.46it/s]

Step: 6000, Loss: -931.3356, d_loss: -934.9124, f_loss: 35.7686

 70%|███████   | 7014/10000 [01:00<00:25, 118.34it/s]

Step: 7000, Loss: -931.3438, d_loss: -934.9145, f_loss: 35.7074

 80%|████████  | 8013/10000 [01:08<00:16, 123.90it/s]

Step: 8000, Loss: -931.3475, d_loss: -934.9190, f_loss: 35.7148

 90%|█████████ | 9022/10000 [01:16<00:07, 137.57it/s]

Step: 9000, Loss: -931.3655, d_loss: -934.9211, f_loss: 35.5562

100%|██████████| 10000/10000 [01:23<00:00, 119.20it/s]

adata_st.obs['ct'] = adata_st.obsm['proportion'].idxmax(axis=1)

adata_st_ct = adata_st[adata_st.obs['ct']=='celltype_0'].copy()
sc.tl.pca(adata_st_ct, n_comps=min(10, adata_st_ct.n_obs - 1, adata_st_ct.n_vars - 1))

model = MacSGP.model.Model(adata_st, adata_basis, n_layers=4, training_steps=3000)

model.train(step_interval=600)
adata_result = model.eval()

  0%|          | 9/3000 [00:00<01:06, 44.80it/s]

Step: 0, Loss: -927.6000, d_loss: -934.9209, f_loss: 73.2095, reg_loss: 0.0000

 20%|██        | 614/3000 [00:08<00:33, 71.42it/s]

Step: 600, Loss: -944.7094, d_loss: -951.2840, f_loss: 36.1229, reg_loss: 0.2669

 40%|████      | 1214/3000 [00:16<00:24, 72.88it/s]

Step: 1200, Loss: -949.8668, d_loss: -956.4533, f_loss: 36.0360, reg_loss: 0.2687

 60%|██████    | 1810/3000 [00:24<00:16, 70.20it/s]

Step: 1800, Loss: -950.9908, d_loss: -957.4516, f_loss: 36.0039, reg_loss: 0.2577

 80%|████████  | 2410/3000 [00:32<00:08, 71.62it/s]

Step: 2400, Loss: -950.9956, d_loss: -957.4464, f_loss: 35.9598, reg_loss: 0.2572

100%|██████████| 3000/3000 [00:40<00:00, 74.00it/s]

adata_result.obs['ct'] = adata_result.obsm['proportion'].idxmax(axis=1)
adata_result.obs['pc1'] = 0.0
adata_result.obs.loc[adata_st_ct.obs_names, 'pc1'] = adata_st_ct.obsm['X_pca'][:, 0]
adata_result.var['pc1'] = adata_st_ct.varm['PCs'][:, 0]

# ============================================================
# Global plot style (Nature Communications standard)
# ============================================================
mpl.rcParams.update({
    'font.family': 'Helvetica',   # Use Helvetica font
    'font.size': 7,
    'axes.titlesize': 8,
    'axes.labelsize': 7,
    'xtick.labelsize': 6.5,
    'ytick.labelsize': 6.5,
    'axes.linewidth': 0.5,
    'xtick.major.width': 0.5,
    'ytick.major.width': 0.5,
    'xtick.major.size': 2.5,
    'ytick.major.size': 2.5,
    'pdf.fonttype': 42,           # Embed fonts for vector editing
    'ps.fonttype': 42,
    'savefig.dpi': 300,
    'figure.dpi': 150,
})

# Filter spots with low proportions
factor_true = adata_st.obsm['simulated_factor']
factor_macsgp = adata_result.obsm['factor']['celltype_0']
coords = adata_result.obsm['spatial']
w = KNN.from_array(coords, k=20)
w.transform = 'r'  # Row standardization
bi = esda.Moran_BV(factor_true, factor_macsgp, w, permutations=999)
print('I:', bi.I, 'p-value:', bi.p_sim)

loading_true = adata_result.varm['simulated_loading'][:, 0]
loading_macsgp = adata_result.varm['loading']['celltype_0']
cor = np.corrcoef(loading_true, loading_macsgp)[0, 1]
print('Correlation of loadings:', cor)
macsgp_scores = [bi.I, cor]

# Filter spots with low proportions
factor_true = adata_st.obsm['simulated_factor'].loc[adata_st_ct.obs_names].values.reshape(-1)
factor_pca = adata_st_ct.obsm['X_pca'][:, 0].reshape(-1)  # Use the first principal component as the factor value
# factor_pca = results['celltype_0']['scores'][:, 0]      # Weighted by cell-type proportion to amplify spots with high proportion
coords = adata_st_ct.obsm['spatial']
w = KNN.from_array(coords, k=20)
w.transform = 'r'  # Row standardization
bi = esda.Moran_BV(factor_true, factor_pca, w, permutations=999)
print('I:', bi.I, 'p-value:', bi.p_sim)

loading_true = adata_result.varm['simulated_loading'][:, 0]
loading_pca = adata_result.var['pc1'].values
# loading_pca = results['celltype_0']['loadings'][0]
cor = np.corrcoef(loading_true, loading_pca)[0, 1]
print('Correlation of loadings:', cor)
pca_scores = [abs(bi.I), cor]

# Shared variables (run once)
s = 8
proportion = adata_st.obsm['proportion']['celltype_0']
factor_true = adata_st.obsm['simulated_factor']
factor_macsgp = adata_result.obsm['factor']['celltype_0']
factor_pca = adata_result.obs['pc1'].values

sign_m = np.sign(pearsonr(factor_macsgp.values, factor_true.values.reshape(-1))[0])
factor_macsgp_aligned = factor_macsgp * sign_m

sign_p = np.sign(pearsonr(factor_pca, factor_true.values.reshape(-1))[0])
factor_pca_aligned = factor_pca * sign_p

x = adata_st.obsm['spatial'][:, 0]
y = adata_st.obsm['spatial'][:, 1]

data = proportion

fig, ax = plt.subplots(figsize=(3.5, 3.2))
sc = ax.scatter(
    x,
    y,
    c=data,
    s=8,
    cmap='YlOrBr',
    linewidths=0,
)
cb = fig.colorbar(sc, ax=ax, shrink=0.4, pad=0.02, orientation='horizontal')
cb.locator = MaxNLocator(nbins=1)
cb.update_ticks()
cb.ax.tick_params(labelsize=12)

# Remove spines
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(False)

ax.set_title('Cell type A', fontsize=15)
ax.set_xticks([])
ax.set_yticks([])
ax.set_aspect('equal', adjustable='box')
plt.tight_layout()
plt.show()

data = 1 - proportion

fig, ax = plt.subplots(figsize=(3.5, 3.2))
sc = ax.scatter(
    x,
    y,
    c=data,
    s=8,
    cmap='Blues',
    linewidths=0,
)
cb = fig.colorbar(sc, ax=ax, shrink=0.4, pad=0.02, orientation='horizontal')
cb.locator = MaxNLocator(nbins=1)
cb.update_ticks()
cb.ax.tick_params(labelsize=12)

# Remove spines
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(False)

ax.set_title('Cell type B', fontsize=15)
ax.set_xticks([])
ax.set_yticks([])
ax.set_aspect('equal', adjustable='box')
plt.tight_layout()
plt.show()

new_cmap_2 = set_zero_in_cmap_to_transparent(cmap=sns.color_palette("YlOrBr", as_cmap=True))
colors=new_cmap_2(np.arange(new_cmap_2.N))
colors = colors[0:new_cmap_2.N*3//4]
color2 = colors[new_cmap_2.N//2]

new_cmap_1 = set_zero_in_cmap_to_transparent(cmap=sns.color_palette("Blues", as_cmap=True))
colors=new_cmap_1(np.arange(new_cmap_1.N))
colors = colors[0:new_cmap_1.N*3//4]
color1 = colors[new_cmap_1.N//2]

data = adata_st.obs['ct']

# Convert categorical data to codes and labels
if hasattr(data, 'cat'):
    categories = data.cat.categories.tolist()
    codes = data.cat.codes.values
else:
    categories = sorted(data.unique())
    cat_map = {c: i for i, c in enumerate(categories)}
    codes = np.array([cat_map[v] for v in data])

n_cats = len(categories)

# Discrete colormap — specify colors for the two categories explicitly
color_dict = {
    'celltype_0': color2,   # Cell type A
    'celltype_1': color1,   # Cell type B
}

# Replace with your actual category names and colors
colors = [color_dict[c] for c in categories]
cmap = ListedColormap(colors)

fig, ax = plt.subplots(figsize=(3.5, 3.2))
sc = ax.scatter(
    x,
    y,
    c=codes,
    s=8,
    cmap=cmap,
    linewidths=0,
)

# Use legend instead of colorbar for categorical data
handles = [
    plt.Line2D([0], [0], marker='o', color='w',
               markerfacecolor=colors[i], markersize=6, label=categories[i])
    for i in range(n_cats)
]
ax.legend(
    handles=handles,
    fontsize=10,
    loc='upper center',
    bbox_to_anchor=(0.5, -0.05),
    frameon=False,
    ncol=2,
    handletextpad=0.3,
    columnspacing=1.0,
    labels=['Cell type A', 'Cell type B']
)

# Remove spines
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(False)

ax.set_title('Assigned Cell Type', fontsize=15)
ax.set_xticks([])
ax.set_yticks([])
ax.set_aspect('equal', adjustable='box')
plt.tight_layout()
plt.show()

data = factor_macsgp_aligned

fig, ax = plt.subplots(figsize=(3.5, 3.2))
sc = ax.scatter(
    x,
    y,
    c=data,
    s=8,
    cmap='bwr',
    linewidths=0,
)
cb = fig.colorbar(sc, ax=ax, shrink=0.4, pad=0.02, orientation='horizontal')
cb.set_label('Normalized score', fontsize=12)
cb.set_ticks([])

# Remove spines
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(False)

ax.set_title('MacSGP', fontsize=15)
ax.set_xticks([])
ax.set_yticks([])
ax.set_aspect('equal', adjustable='box')
plt.tight_layout()

# Color background spots gray
data = factor_pca_aligned
mask = adata_result.obs['ct'] != 'celltype_0'  # Only show PCA results for celltype_0; other cells are grayed out

fig, ax = plt.subplots(figsize=(3.5, 3.2))
ax.scatter(
    x[mask],
    y[mask],
    c='lightgray',
    s=8,
    linewidths=0,
)
sc = ax.scatter(
    x[~mask],
    y[~mask],
    c=data[~mask],
    s=8,
    cmap='bwr',
    linewidths=0,
)
cb = fig.colorbar(sc, ax=ax, shrink=0.4, pad=0.02, orientation='horizontal')
cb.set_label('Normalized score', fontsize=12)
cb.set_ticks([])

# Remove spines
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(False)

ax.set_title('Within-cell-type PCA', fontsize=15)
ax.set_xticks([])
ax.set_yticks([])
ax.set_aspect('equal', adjustable='box')
plt.tight_layout()

data = factor_true.values

fig, ax = plt.subplots(figsize=(3.5, 3.2))
sc = ax.scatter(
    x,
    y,
    c=data,
    s=6,
    cmap='bwr',
    linewidths=0,
)
cb = fig.colorbar(sc, ax=ax, shrink=0.4, pad=0.02, orientation='horizontal')
cb.set_label('Normalized score', fontsize=12)
cb.set_ticks([])

# Remove spines
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(False)

ax.set_title('Ground Truth', fontsize=15)
ax.set_xticks([])
ax.set_yticks([])
ax.set_aspect('equal', adjustable='box')
plt.tight_layout()

# Figure 5: Coherence Comparison
fig, ax = plt.subplots(figsize=(5, 3.2))
bar_w = 0.3
x_pos = np.arange(2)
bars1 = ax.bar(x_pos - bar_w/2, macsgp_scores, bar_w,
               label='MacSGP', color='#4C72B0', edgecolor='white')
bars2 = ax.bar(x_pos + bar_w/2, pca_scores, bar_w,
               label='Within-cell-type PCA', color='#DD8452', edgecolor='white')

for bar in list(bars1) + list(bars2):
    h = bar.get_height()
    ax.text(bar.get_x() + bar.get_width()/2, h + 0.01,
            f'{h:.3f}', ha='center', va='bottom', fontsize=8)

ax.set_xticks(x_pos)
ax.set_xticklabels(['Factor\nCoherence', 'Loading\nCoherence'], fontsize=15)
ax.set_ylim(0, 1.1)
ax.set_ylabel('Coherence', fontsize=15)
ax.legend(fontsize=8, loc='upper right')
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
plt.tight_layout()
plt.show()