Identifying a spatial code of ligand-receptor pairs with a genetic algorithm¶

import anndata

rng = np.random.default_rng(0)

n_celltypes = 12
cells_per_type = 40
positions = np.linspace(0, 100, n_celltypes)
celltypes = ['CT-{}'.format(i + 1) for i in range(n_celltypes)]

coords, labels = [], []
for name, x in zip(celltypes, positions):
    coords.append(np.column_stack([rng.normal(x, 2.5, cells_per_type),
                                   rng.normal(0, 2.5, cells_per_type)]))
    labels += [name] * cells_per_type

coords = np.vstack(coords)
adata = anndata.AnnData(X=np.zeros((len(coords), 1), dtype=float),
                        obs=pd.DataFrame({'celltype': labels},
                                         index=['cell-{}'.format(i) for i in range(len(coords))]))
adata.obsm['spatial'] = coords
adata

fig, ax = plt.subplots(figsize=(9, 2.4))
for name in celltypes:
    mask = adata.obs['celltype'] == name
    ax.scatter(adata.obsm['spatial'][mask, 0], adata.obsm['spatial'][mask, 1], s=8, label=name)
ax.set_xlabel('x'); ax.set_ylabel('y'); ax.set_title('Single cells coloured by cell type')
ax.legend(ncol=6, fontsize=7, loc='upper center', bbox_to_anchor=(0.5, -0.35))
plt.show()

n_informative, n_noise = 12, 48

genes, profiles, pairs = [], [], []
informative_centres = rng.uniform(0, 100, n_informative)   # kept for section 6

for k in range(n_informative):
    centre = informative_centres[k]
    bump = 200 * np.exp(-((positions - centre) ** 2) / (2 * 22.0 ** 2))
    for tag in ('L', 'R'):
        genes.append('{}info{}'.format(tag, k))
        profiles.append(bump * rng.uniform(0.8, 1.2, n_celltypes))
    pairs.append(('Linfo{}'.format(k), 'Rinfo{}'.format(k)))

for k in range(n_noise):
    for tag in ('L', 'R'):
        genes.append('{}noise{}'.format(tag, k))
        profiles.append(rng.uniform(0, 200, n_celltypes))
    pairs.append(('Lnoise{}'.format(k), 'Rnoise{}'.format(k)))

rnaseq = pd.DataFrame(np.vstack(profiles), index=genes, columns=celltypes)
lr_pairs = pd.DataFrame(pairs, columns=['A', 'B'])

print('expression :', rnaseq.shape)
print('LR pairs   :', lr_pairs.shape[0], '({} informative, {} noise)'.format(n_informative, n_noise))
rnaseq.iloc[:4, :6].round(1)

centroids = c2c.spatial.celltype_centroids(adata, group_col='celltype', spatial_key='spatial')
centroids.round(2).head()

for method in ['centroid', 'min', 'mean', 'median', 'max']:
    d = c2c.spatial.celltype_distances(adata, group_col='celltype', method=method)
    print('{:<9} CT-1 to CT-12 = {:6.2f}'.format(method, d.loc['CT-1', 'CT-12']))

reference_distances = c2c.spatial.celltype_distances(adata, group_col='celltype',
                                                     method='centroid')
reference_distances.iloc[:5, :5].round(1)

import scipy.spatial
from cell2cell.analysis import correlation_fitness
from cell2cell.core import PreparedCCIScorer
from cell2cell.preprocessing import bidirectional_index
from cell2cell.core.interaction_space import InteractionSpace

analysis_setup = {'communication_score': 'expression_thresholding',
                  'cci_score': 'bray_curtis',
                  'cci_type': 'undirected'}
cutoff_setup = {'type': 'constant_value', 'parameter': 50}

bi_ppi = c2c.preprocessing.bidirectional_ppi_for_cci(lr_pairs, verbose=False)
interaction_space = InteractionSpace(rnaseq_data=rnaseq, ppi_data=bi_ppi,
                                     gene_cutoffs=cutoff_setup,
                                     communication_score=analysis_setup['communication_score'],
                                     cci_score=analysis_setup['cci_score'],
                                     cci_type=analysis_setup['cci_type'], verbose=False)

scorer = PreparedCCIScorer(interaction_space)
source = bidirectional_index(lr_pairs, verbose=False)

reference_vector = scipy.spatial.distance.squareform(reference_distances.values, checks=False)
all_pairs = np.ones(len(lr_pairs))
baseline_distance = scorer.distance_batch(all_pairs[source][None, :])[0]
baseline = correlation_fitness(scipy.spatial.distance.squareform(baseline_distance, checks=False),
                        reference_vector)

print('Spearman correlation using all {} pairs: {:.4f}'.format(len(lr_pairs), baseline))

results = c2c.analysis.optimize_lr_pairs(rnaseq_data=rnaseq,
                                         ppi_data=lr_pairs,
                                         reference_distances=reference_distances,
                                         cutoff_setup=cutoff_setup,
                                         analysis_setup=analysis_setup,
                                         population_size=60,
                                         generations=40,
                                         runs=10,
                                         random_state=888)

summary = pd.DataFrame([{'run': k,
                         'objective': results[k]['obj_fn'],
                         'pairs kept': results[k]['n_selected'],
                         'dropped': '{:.0%}'.format(results[k]['drop_fraction'])}
                        for k in natsorted(results) if k.startswith('run')])
summary

for key in sorted(k for k in results if k.startswith('run')):
    kept = lr_pairs.loc[np.asarray(results[key]['ppi_data'], dtype=bool)]
    informative = kept['A'].str.startswith('Linfo').mean()
    print('{}: objective {:.4f} | {:>2} pairs kept | {:.0%} of them informative'
          .format(key, results[key]['obj_fn'], len(kept), informative))

best = results['best_ppi_data']
print('\nBest run:', results['best_run'], '-> objective {:.4f}'.format(results['best_obj_fn']))
best

runs = natsorted(k for k in results if k.startswith('run'))
objectives = [results[k]['obj_fn'] for k in runs]

fig, ax = plt.subplots(figsize=(5, 3))
ax.axhline(baseline, color='grey', ls='--', label='all pairs')
ax.plot(range(1, len(runs) + 1), objectives, 'o-', color='steelblue', label='selected pairs')
ax.set_xlabel('run'); ax.set_ylabel('|Spearman| vs reference')
ax.set_xticks(range(1, len(runs) + 1)); ax.legend()
plt.tight_layout(); plt.show()

best_mask = np.zeros(len(lr_pairs))
best_mask[np.asarray(results[results['best_run']]['ppi_data'], dtype=bool)] = 1.0
best_distance = scorer.distance_batch(best_mask[source][None, :])[0]

fig, axes = plt.subplots(1, 2, figsize=(9, 3.6), sharey=True)
for ax, matrix, title in zip(axes, [baseline_distance, best_distance],
                             ['All pairs', 'Selected pairs']):
    vector = scipy.spatial.distance.squareform(matrix, checks=False)
    rho = correlation_fitness(vector, reference_vector)
    ax.scatter(reference_vector, vector, s=12, alpha=0.6, color='steelblue')
    ax.set_title('{} ($|\\rho|$ = {:.3f})'.format(title, rho))
    ax.set_xlabel('reference distance')
axes[0].set_ylabel('CCI distance')
plt.tight_layout(); plt.show()

def make_donor(seed):
    """The same informative pairs, re-measured with independent noise."""
    rng_d = np.random.default_rng(seed)
    profiles = []
    for k in range(n_informative):
        centre = informative_centres[k]
        bump = 200 * np.exp(-((positions - centre) ** 2) / (2 * 22.0 ** 2))
        for _ in ('L', 'R'):
            profiles.append(bump * rng_d.uniform(0.8, 1.2, n_celltypes))
    for _ in range(n_noise):
        for _ in ('L', 'R'):
            profiles.append(rng_d.uniform(0, 200, n_celltypes))
    donor = pd.DataFrame(np.vstack(profiles), index=genes, columns=celltypes)

    jitter = coords + rng_d.normal(0, 1.5, coords.shape)
    donor_adata = anndata.AnnData(X=np.zeros((len(jitter), 1)), obs=adata.obs.copy())
    donor_adata.obsm['spatial'] = jitter
    donor_ref = c2c.spatial.celltype_distances(donor_adata, group_col='celltype', method='centroid')
    return donor, donor_ref

donors = [make_donor(seed) for seed in (101, 202, 303)]
print('three donors, each {} genes x {} cell types'.format(*donors[0][0].shape))

def informative_fraction(selected):
    return selected['A'].str.startswith('Linfo').mean() if len(selected) else 0.0

# Settings that do not depend on the objective. `cutoff_setup` and `analysis_setup`
# are deliberately kept out: passing them alongside an `objective` is rejected,
# because the objective already carries them.
search = dict(ppi_data=lr_pairs, population_size=60, generations=40,
              runs=2, random_state=888)
setups = dict(cutoff_setup=cutoff_setup, analysis_setup=analysis_setup)

per_donor = []
for i, (donor_rnaseq, donor_ref) in enumerate(donors, start=1):
    r = c2c.analysis.optimize_lr_pairs(rnaseq_data=donor_rnaseq,
                                       reference_distances=donor_ref, **setups, **search)
    kept = r['best_ppi_data']
    per_donor.append(set(map(tuple, kept[['A', 'B']].values)))
    print('donor {}: objective {:.3f} | {:>2} pairs | {:.0%} informative'
          .format(i, r['best_obj_fn'], len(kept), informative_fraction(kept)))

shared = set.intersection(*per_donor)
print('\npairs selected by all three donors: {} of {}'
      .format(len(shared), len(set.union(*per_donor))))

factories = [c2c.analysis.CorrelationObjective(rnaseq_data=donor_rnaseq,
                                               reference_distances=donor_ref,
                                               cutoff_setup=cutoff_setup,
                                               analysis_setup=analysis_setup)
             for donor_rnaseq, donor_ref in donors]

pooled = c2c.analysis.optimize_lr_pairs(
    objective=c2c.analysis.CombinedObjective(factories), **search)

kept = pooled['best_ppi_data']
print('pooled across 3 donors: objective {:.3f} | {:>2} pairs | {:.0%} informative'
      .format(pooled['best_obj_fn'], len(kept), informative_fraction(kept)))

pool = pooled['pool']

def per_donor_fitness(mask, penalty=0.0):
    bound = c2c.analysis.CombinedObjective(factories, sd_penalty=penalty)(pool)
    return bound.evaluate_components(np.atleast_2d(np.asarray(mask, dtype=float)))[:, 0]

penalised = c2c.analysis.optimize_lr_pairs(
    objective=c2c.analysis.CombinedObjective(factories, sd_penalty=1.0), **search)

for label, res in (('no penalty', pooled), ('sd_penalty=1.0', penalised)):
    components = per_donor_fitness(res[res['best_run']]['ppi_data'])
    print('{:<16} per-donor rho {}  mean {:.3f}  sd {:.3f}  | {:.0%} informative'
          .format(label, np.round(components, 3), components.mean(), components.std(),
                  informative_fraction(res['best_ppi_data'])))

rng_lodo = np.random.default_rng(0)

for held_out in range(3):
    train = [f for i, f in enumerate(factories) if i != held_out]
    fitted = c2c.analysis.optimize_lr_pairs(
        objective=c2c.analysis.CombinedObjective(train), **search)

    mask = np.asarray(fitted[fitted['best_run']]['ppi_data'], dtype=float)
    test = factories[held_out](fitted['pool'])
    selected = test(mask[None, :])[0]

    sizes = int(mask.sum())
    random_masks = np.zeros((100, len(mask)))
    for row in random_masks:
        row[rng_lodo.choice(len(mask), sizes, replace=False)] = 1.0
    random_scores = test(random_masks)

    print('held-out donor {}: selected {:.3f} | random sets of {} {:.3f} +/- {:.3f} | all pairs {:.3f}'
          .format(held_out + 1, selected, sizes, random_scores.mean(), random_scores.std(),
                  test(np.ones((1, len(mask))))[0]))