Nb02B Msa Domain Join
Jupyter notebook from the AlphaFold MSA Depth as a Lens on the Bacterial Annotation Gap project.
NB02b — MSA Depth × Domain Richness: Per-Cluster Join (H3 Spearman)¶
Requires: BERDL JupyterHub (Spark)
Purpose: Join gene_cluster × bakta_annotations × alphafold_msa_depths × interproscan_domains to produce per-cluster (msa_depth, n_domain_hits) pairs for H3 Spearman correlation. All heavy computation stays in Spark; only a manageable sample is collected to the driver.
Outputs:
data/gc_msa_domain_sample.csv— 200K sample of (gene_cluster_id, is_core, msa_depth, n_domain_hits)data/gc_h3_spearman.csv— Spearman ρ and p-value (approximate, computed via rank correlation in Spark)data/gc_msa_domain_agg.csv— mean/median domain hits by MSA depth bin and pangenome class
In [1]:
from berdl_notebook_utils.setup_spark_session import get_spark_session
spark = get_spark_session()
print("Spark session ready:", spark.version)
Spark session ready: 4.0.1
In [2]:
from pyspark.sql.functions import col, count, countDistinct, regexp_replace, avg, percentile_approx, when
# --- Rebuild NB01 join: gc × bakta × msa_depths ---
gc = spark.table("kbase_ke_pangenome.gene_cluster").select(
"gene_cluster_id", "is_core", "is_auxiliary", "is_singleton"
)
bakta = spark.table("kbase_ke_pangenome.bakta_annotations").select(
"gene_cluster_id", "uniref100"
).filter(
"uniref100 IS NOT NULL AND uniref100 NOT LIKE 'UniRef100_UPI%'"
).withColumn(
"uniprot_accession",
regexp_replace(col("uniref100"), "UniRef100_", "")
)
msa = spark.table("kescience_alphafold.alphafold_msa_depths").select(
"uniprot_accession", col("msa_depth").cast("double")
)
gc_msa = (
gc.join(bakta, on="gene_cluster_id", how="inner")
.join(msa, on="uniprot_accession", how="inner")
.select("gene_cluster_id", "is_core", "is_auxiliary", "is_singleton", "msa_depth")
)
print(f"gc_msa rows: {gc_msa.count():,}")
gc_msa rows: 38,051,842
In [3]:
# --- Domain hit counts per gene cluster (from NB02) ---
ipr = spark.table("kbase_ke_pangenome.interproscan_domains").select(
"gene_cluster_id", "ipr_acc"
)
domain_counts = ipr.groupBy("gene_cluster_id").agg(
count("*").alias("n_domain_hits"),
countDistinct("ipr_acc").alias("n_distinct_ipr")
)
print(f"Clusters with domain annotation: {domain_counts.count():,}")
Clusters with domain annotation: 111,035,431
In [4]:
# --- Join msa + domain counts (left join so unannotated clusters get 0) ---
joined = (
gc_msa
.join(domain_counts, on="gene_cluster_id", how="left")
.fillna({"n_domain_hits": 0, "n_distinct_ipr": 0})
)
print(f"Joined rows: {joined.count():,}")
joined.printSchema()
Joined rows: 38,051,842
root |-- gene_cluster_id: string (nullable = true) |-- is_core: boolean (nullable = true) |-- is_auxiliary: boolean (nullable = true) |-- is_singleton: boolean (nullable = true) |-- msa_depth: double (nullable = true) |-- n_domain_hits: long (nullable = false) |-- n_distinct_ipr: long (nullable = false)
In [5]:
# --- Aggregate: mean domain hits by MSA depth bin and pangenome class ---
msa_bin = (
when(col("msa_depth") < 10, "<10")
.when(col("msa_depth") < 100, "10-99")
.when(col("msa_depth") < 1000, "100-999")
.when(col("msa_depth") < 5000, "1000-4999")
.when(col("msa_depth") < 10000,"5000-9999")
.otherwise("10000+")
)
agg = (
joined
.withColumn("msa_depth_bin", msa_bin)
.groupBy("is_core", "is_auxiliary", "is_singleton", "msa_depth_bin")
.agg(
count("*").alias("n_clusters"),
avg("msa_depth").alias("mean_msa_depth"),
avg("n_domain_hits").alias("mean_n_domain_hits"),
percentile_approx("n_domain_hits", 0.5).alias("median_n_domain_hits"),
avg("n_distinct_ipr").alias("mean_n_distinct_ipr")
)
.orderBy("is_core", "msa_depth_bin")
)
agg.show(30, truncate=False)
agg_pd = agg.toPandas()
agg_pd.to_csv("/tmp/gc_msa_domain_agg.csv", index=False)
print(f"Saved aggregate ({len(agg_pd)} rows) → /tmp/gc_msa_domain_agg.csv")
+-------+------------+------------+-------------+----------+------------------+-------------------+--------------------+--------------------+ |is_core|is_auxiliary|is_singleton|msa_depth_bin|n_clusters|mean_msa_depth |mean_n_domain_hits |median_n_domain_hits|mean_n_distinct_ipr | +-------+------------+------------+-------------+----------+------------------+-------------------+--------------------+--------------------+ |false |true |true |10-99 |825175 |40.50513042687915 |0.7809822158935983 |0 |0.20356288060108463 | |false |true |false |10-99 |629614 |41.13062765440413 |0.7457601006330863 |0 |0.21103723868910157 | |false |true |true |100-999 |1153536 |415.84065603500886|1.6372033469263205 |1 |0.728581509376387 | |false |true |false |100-999 |885826 |413.7956765775672 |1.6929069591545067 |1 |0.7778096375586177 | |false |true |true |1000-4999 |1123356 |2590.6987508857387|3.242658605108265 |3 |1.6265395831775502 | |false |true |false |1000-4999 |860794 |2590.482201316459 |3.4283603277903887 |3 |1.7232101989558477 | |false |true |true |10000+ |2887461 |16766.37599607406 |8.829141242080846 |7 |3.970958915116083 | |false |true |false |10000+ |2211048 |16716.35894064715 |9.685333380369851 |8 |4.2858902203841795 | |false |true |false |5000-9999 |552616 |7339.917083472067 |5.043015403100887 |5 |2.482277024190396 | |false |true |true |5000-9999 |713156 |7313.396892124584 |4.736558340671606 |4 |2.3445641626796943 | |false |true |true |<10 |392959 |4.555599948086187 |0.47031369684878066|0 |0.04602770263564392 | |false |true |false |<10 |245002 |4.7487367450061635|0.41519252904057924|0 |0.045701667741487824| |true |false |false |10-99 |1143785 |43.39528058157783 |0.9603727973351636 |0 |0.24207871234541456 | |true |false |false |100-999 |2301137 |442.08899426674725|1.9528750352543112 |1 |0.8521374433595218 | |true |false |false |1000-4999 |3126558 |2693.1647063639953|3.724952167847198 |3 |1.8427878197046081 | |true |false |false |10000+ |15993075 |16769.926620865594|10.834821446157164 |9 |4.600884757934294 | |true |false |false |5000-9999 |2591011 |7483.764757849349 |5.62260445825973 |5 |2.664387762151531 | |true |false |false |<10 |415733 |4.569935511494156 |0.5893061171473036 |0 |0.05898016274868726 | +-------+------------+------------+-------------+----------+------------------+-------------------+--------------------+--------------------+
Saved aggregate (18 rows) → /tmp/gc_msa_domain_agg.csv
In [6]:
# --- H3 Spearman: msa_depth vs n_domain_hits ---
# Compute via rank correlation in Spark ML
from pyspark.ml.stat import Correlation
from pyspark.ml.feature import VectorAssembler
assembler = VectorAssembler(inputCols=["msa_depth", "n_domain_hits"], outputCol="features")
feat_df = assembler.transform(joined.select("msa_depth", "n_domain_hits")).select("features")
corr_matrix = Correlation.corr(feat_df, "features", method="spearman").head()
rho = float(corr_matrix[0].toArray()[0, 1])
print(f"Spearman ρ (msa_depth vs n_domain_hits): {rho:.4f}")
import pandas as pd
result_pd = pd.DataFrame([{
"metric_x": "msa_depth",
"metric_y": "n_domain_hits",
"spearman_rho": rho,
"n_pairs": joined.count()
}])
result_pd.to_csv("/tmp/gc_h3_spearman.csv", index=False)
print(result_pd)
Spearman ρ (msa_depth vs n_domain_hits): 0.7563
metric_x metric_y spearman_rho n_pairs 0 msa_depth n_domain_hits 0.756334 38051842
In [7]:
# --- Save 200K sample for local scatter plot ---
n_total = joined.count()
frac = min(200_000 / n_total, 1.0)
sample_pd = (
joined
.sample(fraction=float(frac), seed=42)
.limit(200_000)
.select("gene_cluster_id", "is_core", "is_auxiliary", "is_singleton",
"msa_depth", "n_domain_hits", "n_distinct_ipr")
.toPandas()
)
sample_pd.to_csv("/tmp/gc_msa_domain_sample.csv", index=False)
print(f"Saved sample: {len(sample_pd):,} rows → /tmp/gc_msa_domain_sample.csv")
Saved sample: 199,895 rows → /tmp/gc_msa_domain_sample.csv
In [8]:
print("\n=== NB02b complete. Copy outputs: ===")
print(" cp /tmp/gc_msa_domain_agg.csv projects/alphafold_msa_annotation/data/")
print(" cp /tmp/gc_h3_spearman.csv projects/alphafold_msa_annotation/data/")
print(" cp /tmp/gc_msa_domain_sample.csv projects/alphafold_msa_annotation/data/")
=== NB02b complete. Copy outputs: === cp /tmp/gc_msa_domain_agg.csv projects/alphafold_msa_annotation/data/ cp /tmp/gc_h3_spearman.csv projects/alphafold_msa_annotation/data/ cp /tmp/gc_msa_domain_sample.csv projects/alphafold_msa_annotation/data/