GEM-FISH — joint Hi-C + FISH 3-D reconstruction

:func:uchrom.recon.fish.reconstruct_gem_fish is an independent PyTorch reimplementation of Abbas et al. 2019 (Nature Communications, doi:10.1038/s41467-019-10005-6). It solves a three-stage optimisation that fuses bulk Hi-C contacts with FISH-measured pairwise distances into a single 3-D chromosome model.

Pipeline (per chromosome):

  1. TAD partition — call the Dixon directionality-index caller on the Hi-C matrix.

  2. Stage 1 — TAD-level gradient descent on ``C_g = C_1 + λ_E·C_2

    • λ_F·C_3`` (Hi-C KL + polymer prior + FISH inter-TAD distance).

  3. Stage 2 — intra-TAD per-TAD descent on C_t = C_1 + λ_E·C_2 + λ_R·C_4 (intra-TAD Hi-C KL + polymer prior + FISH Rg target).

  4. Stage 3 — assembly — Kabsch-style rotation of each intra-TAD cloud around its Stage-1 centre.

The notebook has two parts. The first runs end-to-end on a fully synthetic chain so you can see the pipeline’s mechanics with a known ground truth. The second downloads Bintu et al. 2018 (Science, the follow-up to the paper’s original Wang 2016 FISH) per-trace imaging data for IMR90 chr21:18.6-20.6 Mb, pairs it with the real Rao et al. 2014 IMR90 in situ Hi-C of chr21 (GEO GSE63525, shipped as a 3.8 MB slice under example-data/rao2014_imr90/), and runs GEM-FISH on both.

import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.distance import cdist

from uchrom.recon.fish._tad_level import (
    Stage1Params, reconstruct_tad_level,
)
from uchrom.recon.fish._intra_tad import (
    Stage2Params, reconstruct_intra_tad,
)
from uchrom.recon.fish._assembly import assemble
from uchrom.recon.fish._hic import contact_to_distance

Part 1 — Synthetic end-to-end

Generate a known polymer chain organised into 5 tight TADs, fabricate Hi-C contacts from f = 1/(1+d) and FISH distances from the truth with noise, then run each of the three stages and compare to the ground truth at every step.

rng = np.random.default_rng(0)
n_tads, bins_per_tad = 5, 20
truth = []
centre = np.zeros(3)
for t in range(n_tads):
    walk = np.cumsum(rng.normal(0, 0.3, (bins_per_tad, 3)), axis=0)
    truth.append(walk - walk.mean(0) + centre)
    centre = centre + 5.0 * (np.array([1.,0,0]) + 0.3 * rng.normal(0, 1, 3))
truth = np.vstack(truth)
N = truth.shape[0]
tad_bins = [(t*bins_per_tad, (t+1)*bins_per_tad) for t in range(n_tads)]
tad_centres_true = np.array([truth[s:e].mean(0) for (s, e) in tad_bins])

D_true = cdist(truth, truth)
contacts = 1.0 / (1.0 + D_true); np.fill_diagonal(contacts, 0)

# Inter-TAD contacts
inter_contacts = np.zeros((n_tads, n_tads))
for i, (si, ei) in enumerate(tad_bins):
    for j, (sj, ej) in enumerate(tad_bins):
        inter_contacts[i, j] = contacts[si:ei, sj:ej].sum()

# FISH centre distances + noise
F = cdist(tad_centres_true, tad_centres_true)
F = F + rng.normal(0, 0.1, F.shape); np.fill_diagonal(F, 0)

# Rg² targets per TAD (with noise)
rg_targets = []
for (s, e) in tad_bins:
    sub = truth[s:e]
    rg2 = ((sub - sub.mean(0)) ** 2).sum(1).mean()
    rg_targets.append(rg2 * (1 + 0.05 * rng.standard_normal()))

print(f'Ground truth: {N} bins across {n_tads} TADs')

# Stage 1 — TAD-level
tad_centres_recon, s1_info = reconstruct_tad_level(
    inter_contacts, F,
    params=Stage1Params(
        lambda_E=0.05, lambda_F=0.5,
        n_iter=500, n_ensemble=20, device='cpu',
    ),
)
iu = np.triu_indices(n_tads, 1)
corr_s1 = np.corrcoef(cdist(tad_centres_recon, tad_centres_recon)[iu],
                      cdist(tad_centres_true, tad_centres_true)[iu])[0, 1]
print(f'Stage 1 centre-distance correlation: {corr_s1:.3f}')

# Stage 2 — intra-TAD
intra_coords, s2_info = reconstruct_intra_tad(
    contacts, tad_bins,
    target_rg_sq_per_tad=rg_targets,
    params=Stage2Params(
        lambda_E=0.05, lambda_R=0.01,
        n_iter=300, n_ensemble=6, device='cpu',
    ),
)
per_tad_corr = []
for i, (s, e) in enumerate(tad_bins):
    Di = cdist(intra_coords[i], intra_coords[i])
    Dt = cdist(truth[s:e], truth[s:e])
    n = e - s
    ci = np.corrcoef(Di[np.triu_indices(n, 1)],
                     Dt[np.triu_indices(n, 1)])[0, 1]
    per_tad_corr.append(ci)
print(f'Stage 2 mean intra-TAD distance corr: {np.mean(per_tad_corr):.3f}')

# Stage 3 — assembly
tad_mid = np.array([(s + e - 1) / 2 * bins_per_tad for (s, e) in tad_bins])
genomic = np.abs(tad_mid[:, None] - tad_mid[None, :]).astype(np.float64)
gap = contact_to_distance(
    inter_contacts, genomic_distances=genomic,
    alpha=0.25, fallback_c=1.0, fallback_beta=0.5,
)
final = assemble(tad_centres_recon, intra_coords, inter_tad_distances=gap)

corr_final = np.corrcoef(cdist(final, final)[np.triu_indices(N, 1)],
                         D_true[np.triu_indices(N, 1)])[0, 1]
print(f'Final chain distance correlation: {corr_final:.3f}')

# Visualise truth vs reconstruction using uchrom.pl.plot_structure_3d.
# Use PyVista subplot so both tubes render side-by-side in one figure.
import pyvista as pv
from uchrom.pl import plot_structure_3d

# TAD-membership scalar so each TAD is coloured distinctly
tad_scalar = np.concatenate([
    np.full(e - s, t, dtype=np.float64) for t, (s, e) in enumerate(tad_bins)
])

plotter = pv.Plotter(shape=(1, 2), notebook=True, window_size=(1100, 500))

plotter.subplot(0, 0)
plotter.add_text('Ground truth', font_size=10)
plot_structure_3d(
    truth, colour=tad_scalar, cmap='viridis',
    plotter=plotter, show=False,
)

plotter.subplot(0, 1)
plotter.add_text(
    f'GEM-FISH reconstruction (chain corr = {corr_final:.2f})', font_size=10,
)
plot_structure_3d(
    final, colour=tad_scalar, cmap='viridis',
    plotter=plotter, show=False,
)

plotter.show(jupyter_backend='static')

# Distance-matrix comparison
fig, axes = plt.subplots(1, 3, figsize=(13, 4))
vmax = np.percentile(D_true, 98)
for ax, mat, title in zip(
    axes,
    [D_true, cdist(final, final), np.abs(D_true - cdist(final, final))],
    ['Truth', 'Reconstruction', '|Truth − Recon|'],
):
    im = ax.imshow(mat, cmap='viridis_r', origin='lower',
                   vmax=vmax if title != '|Truth − Recon|' else vmax/3)
    ax.set_title(title); ax.set_xlabel('Bin'); ax.set_ylabel('Bin')
    plt.colorbar(im, ax=ax, fraction=0.046)
plt.tight_layout()
plt.show()

Part 2 — Paper-faithful IMR90 chr21 reconstruction

Load the real Bintu 2018 FISH data (auto-downloaded, ~2 MB) and the real Rao et al. 2014 IMR90 in situ Hi-C (GEO GSE63525, GSE63525_IMR90_combined.hic, the map Bintu 2018 compared against). The chr21 slice at 5 kb (raw observed counts, hg19) is shipped in the repo as example-data/rao2014_imr90/IMR90_chr21_5kb_hg19.cool (3.8 MB; rebuilt by python example-data/build_rao2014_slices.py).

Both inputs are put on the same hg19 coordinates: Bintu’s segment 1 starts at hg19 chr21:20,000,032 (hg38 18,627,714, from the Bintu repo README), and the 5 kb Hi-C is summed into 30 kb bins on the Bintu segment grid (bin edges ≡ 20,000,000 mod 30 kb) so each FISH segment maps onto one Hi-C bin. Then run the one-shot reconstruct_gem_fish entry point.

from pathlib import Path
import urllib.request
from uchrom.io import read_bintu_tracing

DATA_URL = (
    'https://raw.githubusercontent.com/BogdanBintu/ChromatinImaging/'
    'master/Data/IMR90_chr21-18-20Mb.csv'
)

def _repo_root():
    for p in [Path.cwd(), *Path.cwd().parents]:
        if (p / 'pyproject.toml').exists():
            return p
    return Path.cwd()

# Auto-download Bintu 2018 (~2 MB) if absent
root = _repo_root()
bintu_csv = root / 'example-data' / 'IMR90_chr21-18-20Mb.csv'
if not bintu_csv.exists():
    print(f'[data] downloading Bintu 2018 (~2 MB) → {bintu_csv}')
    urllib.request.urlretrieve(DATA_URL, bintu_csv)

# Load into ChromData via the Bintu reader, in hg19 to match Rao 2014
START_HG19 = 20_000_032     # segment 1; hg38 18,627,714 (Bintu repo README)
fish_cd = read_bintu_tracing(
    str(bintu_csv), chrom='chr21',
    start_bp=START_HG19,
    resolution_bp=30_000,
)
print(fish_cd)
print(f'  n_traces: {fish_cd.spots["trace_id"].nunique()}')

# Real Rao 2014 IMR90 Hi-C, chr21 at 5 kb (hg19), shipped in the repo
hic_5kb = root / 'example-data' / 'rao2014_imr90' / 'IMR90_chr21_5kb_hg19.cool'
if not hic_5kb.exists():
    raise FileNotFoundError(
        f'{hic_5kb} missing -- rebuild it with '
        '`python example-data/build_rao2014_slices.py --cell IMR90` '
        '(reads GSE63525_IMR90_combined.hic over HTTP, ~30 s)')
print(f'Hi-C: {hic_5kb} ({hic_5kb.stat().st_size / 1e6:.2f} MB)')

import cooler
import pandas as pd
import tempfile

def rebin_on_grid(src, out, chrom='chr21', res=30_000, anchor=START_HG19):
    """Sum a fine cooler into ``res`` bins whose edges fall on ``anchor``'s grid."""
    c = cooler.Cooler(str(src))
    L = int(c.chromsizes[chrom])
    phase = anchor // c.binsize * c.binsize % res       # 20 kb for Bintu chr21
    edges = np.unique(np.r_[0, np.arange(phase, L, res), L])
    bins = pd.DataFrame({'chrom': chrom, 'start': edges[:-1], 'end': edges[1:]})
    px = c.pixels(join=True).fetch(chrom)
    b1 = np.searchsorted(edges, px['start1'].to_numpy(), side='right') - 1
    b2 = np.searchsorted(edges, px['start2'].to_numpy(), side='right') - 1
    agg = (pd.DataFrame({'bin1_id': b1, 'bin2_id': b2, 'count': px['count'].to_numpy()})
           .groupby(['bin1_id', 'bin2_id'], as_index=False)['count'].sum())
    cooler.create_cooler(str(out), bins, agg, assembly='hg19', ordered=True)
    cooler.balance_cooler(cooler.Cooler(str(out)), store=True)   # ICE weights
    return out

hic_cool = rebin_on_grid(hic_5kb, Path(tempfile.mkdtemp()) / 'IMR90_chr21_30kb_hg19.cool')
clr = cooler.Cooler(str(hic_cool))
print(f'{clr.info["nbins"]} bins x 30 kb, {clr.info["sum"]:,} contacts on chr21')

from uchrom.recon.fish import reconstruct_gem_fish, GEMFISHParams

cd, intermediate = reconstruct_gem_fish(
    hic_path=str(hic_cool),
    chrom='chr21',
    resolution=None,             # single-res .cool
    fish_cd=fish_cd,
    params=GEMFISHParams(
        lambda_E_tad=0.5, lambda_F=0.5,
        lambda_E_intra=0.05, lambda_R=0.01,
        stage1_iter=500, stage1_ensemble=20,
        stage2_iter=300, stage2_ensemble=6,
        tad_di_window=50, tad_di_min_size_frac=0.01,
        tad_max_size_frac=0.05,     # cap huge DI-called TADs (p-arm)
        device='auto', verbose=False,
    ),
    return_intermediate=True,
)
print(f'Reconstructed: {cd}')
print(f'  n_tads called: {cd.uns["gem_fish"]["n_tads"]}')
print(f'  final Stage-1 loss: {cd.uns["gem_fish"]["stage1_final_loss"]:.2f}')

# Evaluate: how well does the reconstruction match the FISH
# population-mean pairwise-distance matrix?
from uchrom.fea.distance import _bin_coord_cube, _pairwise_distance_per_trace

df_fish = fish_cd.to_dataframe()
cube, fish_bin_ids, _ = _bin_coord_cube(df_fish, chrom='chr21')
D_fish = np.nanmedian(_pairwise_distance_per_trace(cube), axis=0)

# Map each FISH segment to the Hi-C bin that contains its midpoint
# (one to one: the Hi-C was binned on the Bintu segment grid)
loci = cd.spots_with_loci()
bin_starts, bin_ends = loci['start'].to_numpy(), loci['end'].to_numpy()
fish_mids = np.array([(s + e) / 2 for s, e in fish_bin_ids])
idx = np.array([int(np.flatnonzero((bin_starts <= m) & (bin_ends > m))[0]) for m in fish_mids])
assert len(set(idx)) == len(idx)
D_recon = cdist(cd.coords[idx], cd.coords[idx])

iu_r = np.triu_indices(D_fish.shape[0], 1)
finite = np.isfinite(D_fish[iu_r]) & np.isfinite(D_recon[iu_r])
m_, f_ = D_recon[iu_r][finite], D_fish[iu_r][finite]
corr_real = np.corrcoef(f_, m_)[0, 1]
# Mean relative error |s*d_model - d_FISH| / d_FISH (Abbas 2019's metric),
# with the model rescaled to nm by the least-squares factor s
s_fit = (m_ / f_).sum() / ((m_ / f_) ** 2).sum()
rel_err = np.mean(np.abs(s_fit * m_ - f_) / f_)
print(f'Reconstruction vs FISH median distances: Pearson {corr_real:.3f}, '
      f'mean relative error {rel_err:.3f}')
print(f'  evaluated over {finite.sum()} segment pairs in the Bintu imaged region')
print('  (the same FISH also constrained the model, so this is a fit, '
      'not a held-out score)')

# Visualise.  Scale the recon distance matrix to nm with the fitted
# factor so the two heat-maps are directly comparable.
scale_est = s_fit
fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))

im0 = axes[0].imshow(D_fish, cmap='viridis_r', origin='lower')
axes[0].set_title('FISH population median (nm)')
plt.colorbar(im0, ax=axes[0])

im1 = axes[1].imshow(D_recon * scale_est, cmap='viridis_r', origin='lower',
                      vmax=np.nanpercentile(D_fish, 98))
axes[1].set_title(f'GEM-FISH (Pearson {corr_real:.2f}, rel. err. {rel_err:.2f})')
plt.colorbar(im1, ax=axes[1])

for ax in axes:
    ax.set_xlabel('Bin'); ax.set_ylabel('Bin')
plt.tight_layout()
plt.show()

3-D rendering of the reconstructed chain

:func:uchrom.pl.plot_structure_3d is a PyVista-based renderer (same stack the uchrom.browser uses, without the Qt window) that draws the reconstructed chain as a smooth tube coloured along the genomic direction.

from uchrom.pl import plot_structure_3d

plot_structure_3d(
    cd, chrom='chr21',
    colour='bin', cmap='plasma',
    window_size=(1000, 800),
    notebook=True, jupyter_backend='static',
)

Hi-C contact map vs reconstructed distance matrix

Sanity check: the reconstructed 3-D structure should reproduce the input Hi-C contact pattern — regions with high contact frequency should be close in 3D. We compare side-by-side:

  1. Input Hi-C contact map log10(1 + f_ij) — raw balanced counts from the Rao 2014 IMR90 .cool.

  2. Reconstructed pairwise distance matrix — Euclidean distance between every pair of reconstructed bins.

If the reconstruction is self-consistent with the input, these should show the inverse pattern: bright diagonal bands and loops in the contact map correspond to dark (short-distance) blocks in the distance matrix.

We also compute the Spearman correlation between the balanced contact and 1 / distance over all off-diagonal pairs. On this real IMR90 input it is low (≈ 0.03 over the whole chromosome, ≈ 0.35 within 3 Mb of the diagonal): the TAD-level model plus the single-anchor assembly reproduces the FISH region far better than the genome-wide Hi-C pattern. Treat it as a known limitation of this port (see the caveats below), not as a property of the data.

from uchrom.recon.fish._hic import load_contact_matrix
from scipy.stats import spearmanr
from scipy.spatial.distance import cdist

hic_mat, hic_bin_df = load_contact_matrix(
    str(hic_cool), chrom='chr21', resolution=None, normalize=True,
)
hic_mat = np.nan_to_num(hic_mat, nan=0.0)
D_recon_full = cdist(cd.coords, cd.coords)
n_bins = hic_mat.shape[0]

# Observed / expected: divide each off-diagonal element by the mean at
# its genomic separation.  Reveals structural features independent of
# the distance-decay trend.
def oe_normalize(mat):
    n = mat.shape[0]
    out = np.zeros_like(mat)
    for k in range(1, n):
        diag_vals = np.array([mat[i, i + k] for i in range(n - k)])
        pos = diag_vals[diag_vals > 0]
        if pos.size == 0: continue
        mu = pos.mean()
        for i in range(n - k):
            out[i, i + k] = out[i + k, i] = (
                mat[i, i + k] / mu if mat[i, i + k] > 0 else 0.0
            )
    return out

hic_oe = oe_normalize(hic_mat)
dist_oe = oe_normalize(D_recon_full)  # high distance / expected → far
# For distance matrix, "observed/expected > 1" means farther than
# expected; to make it analogous to contact-frequency we flip the sign.
dist_inv_oe = np.where(dist_oe > 0, 1.0 / dist_oe, 0.0)

fig, axes = plt.subplots(2, 2, figsize=(10, 9))

im00 = axes[0, 0].imshow(np.log10(hic_mat + 1), cmap='Reds', origin='lower')
axes[0, 0].set_title('Input Hi-C (log10 balanced)')
plt.colorbar(im00, ax=axes[0, 0])

im01 = axes[0, 1].imshow(D_recon_full, cmap='viridis_r', origin='lower',
                          vmax=np.percentile(D_recon_full, 98))
axes[0, 1].set_title('Reconstructed distance (recon units)')
plt.colorbar(im01, ax=axes[0, 1])

im10 = axes[1, 0].imshow(np.log10(hic_oe + 1e-2), cmap='RdBu_r', origin='lower',
                          vmin=-1, vmax=1)
axes[1, 0].set_title('Hi-C O/E (log10)')
plt.colorbar(im10, ax=axes[1, 0])

im11 = axes[1, 1].imshow(np.log10(dist_inv_oe + 1e-2), cmap='RdBu_r',
                          origin='lower', vmin=-1, vmax=1)
axes[1, 1].set_title('Reconstructed 1/(distance O/E) (log10)')
plt.colorbar(im11, ax=axes[1, 1])

for ax in axes.flat:
    ax.set_xlabel('Bin'); ax.set_ylabel('Bin')
plt.tight_layout()
plt.show()

# Correlations
iu = np.triu_indices(n_bins, k=1)
ok = (hic_mat[iu] > 0) & (D_recon_full[iu] > 0)
rho_raw, _ = spearmanr(hic_mat[iu][ok], 1.0 / D_recon_full[iu][ok])
ok_oe = (hic_oe[iu] > 0) & (dist_oe[iu] > 0)
rho_oe, _ = spearmanr(hic_oe[iu][ok_oe], 1.0 / dist_oe[iu][ok_oe])
print(f'Spearman(Hi-C raw contact,  1 / recon distance):       {rho_raw:.3f}')
print(f'Spearman(Hi-C O/E,          1 / recon distance O/E):   {rho_oe:.3f}')
print('  (raw captures mostly the distance-decay trend;')
print('   O/E-corrected shows the structural agreement beyond it)')

Scaling up to other chromosomes

rao2014_imr90/IMR90_chr21_5kb_hg19.cool only contains chr21. Slice other chromosomes (or resolutions) of the same Rao 2014 IMR90 map with the builder, which reads only the needed blocks of the 13 GB GEO .hic over HTTP:

python example-data/build_rao2014_slices.py --cell IMR90 --chroms 20 22 --res 5000

Then (FISH for that chromosome, in hg19 coordinates):

cd = reconstruct_gem_fish(
    hic_path='example-data/rao2014_imr90/IMR90_chr20-22_5kb_hg19.cool',
    chrom='chr20',
    resolution=None,           # single-resolution .cool
    fish_cd=fish_cd,           # swap in FISH for that chromosome
    params=GEMFISHParams(device='auto'),
)
cd.write('imr90_chr20_gemfish.chromdata.zarr')

(Do not use 4DN 4DNFI4QQPDMR.mcool for IMR90: it is K562 Hi-C, Ray et al. 2019; see example-data/README.md.)

Notes and caveats

  • λ weights in the paper (λ_E = 5 × 10¹², λ_F = 1 × 10⁻⁸) are calibrated for raw Hi-C counts without row-normalisation. The PyTorch C_1 here row-normalises inside the loss, so the effective weights differ — our defaults (λ_E = 0.05, λ_F = 0.01) balance the three terms for unit-scale inputs.

  • Single model vs ensemble: reconstruct_tad_level returns the single best model from an n_ensemble batch by default. Pass return_all=True to get all ensemble members for downstream clustering / averaging.

  • Assembly here uses a single-anchor Kabsch per TAD, which is a simplification of the paper’s gradient-descent assembly. Good enough for the tutorial; for production use you may want to augment the initial rotation with a short local refinement step.

  • FISH coverage: the Bintu 2018 CSV only covers chr21:20.0–21.95 Mb in hg19 (18.6–20.6 Mb in hg38; 65 segments out of the 1 605-bin chromosome). The FISH C_3 / C_4 terms influence only the TADs that fall inside that region; the rest of the chromosome is shaped by Hi-C alone, which is the expected behaviour when FISH coverage is partial.

  • How good is this on real data? With these settings the model matches the Bintu median distances with Pearson ≈ 0.71–0.73 and a mean relative error ≈ 0.33–0.35 (CPU vs GPU run; 30 kb segments, the same FISH that constrained the model — a fit, not a held-out score), while its agreement with the Hi-C outside the FISH region is weak (Spearman ≈ 0.03 genome-wide on chr21). The paper’s 0.14 is a different setup (Wang 2016 TAD-centre FISH, TAD-level distances, paper-tuned weights), scored by benchmarks/fig3/c_gemfish_wang2016.py. The numbers here are a sanity check of the pipeline on real inputs, not a benchmark.