Nb01 Data Extraction
Jupyter notebook from the AlphaFold MSA Depth as a Lens on the Bacterial Annotation Gap project.
NB01 — MSA Depth × Core/Accessory/Singleton Data Extraction¶
Requires: BERDL JupyterHub (Spark)
Purpose: Join gene_cluster × bakta_annotations × alphafold_msa_depths to produce a per-gene-cluster summary with MSA depth and annotation quality flags.
Output: data/gc_msa_summary.csv (written to MinIO, then downloaded locally)
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
gc = spark.table("kbase_ke_pangenome.gene_cluster").select(
"gene_cluster_id", "gtdb_species_clade_id",
"is_core", "is_auxiliary", "is_singleton"
)
print(f"Total gene clusters: {gc.count():,}")
gc.groupBy("is_core", "is_auxiliary", "is_singleton").count().orderBy("count", ascending=False).show(10)
Total gene clusters: 132,531,501
+-------+------------+------------+--------+ |is_core|is_auxiliary|is_singleton| count| +-------+------------+------------+--------+ | true| false| false|62062686| | false| true| true|50203195| | false| true| false|20265620| +-------+------------+------------+--------+
In [3]:
from pyspark.sql.functions import regexp_replace, col
bakta = spark.table("kbase_ke_pangenome.bakta_annotations").select(
"gene_cluster_id", "hypothetical", "ec", "kegg_orthology_id", "uniref100"
).filter(
"uniref100 IS NOT NULL AND uniref100 NOT LIKE 'UniRef100_UPI%'"
)
bakta = bakta.withColumn(
"uniprot_accession",
regexp_replace(col("uniref100"), "UniRef100_", "")
)
print(f"Bakta rows with real UniProt accession: {bakta.count():,}")
Bakta rows with real UniProt accession: 38,804,903
In [4]:
msa = spark.table("kescience_alphafold.alphafold_msa_depths").select(
"uniprot_accession", "msa_depth"
)
# Join: gene_cluster -> bakta -> alphafold_msa_depths
joined = (
gc
.join(bakta, on="gene_cluster_id", how="inner")
.join(msa, on="uniprot_accession", how="inner")
)
print(f"Joined rows (gene clusters with MSA depth): {joined.count():,}")
joined.printSchema()
joined.show(5, truncate=True)
Joined rows (gene clusters with MSA depth): 38,051,842
root |-- uniprot_accession: string (nullable = true) |-- gene_cluster_id: string (nullable = true) |-- gtdb_species_clade_id: string (nullable = true) |-- is_core: boolean (nullable = true) |-- is_auxiliary: boolean (nullable = true) |-- is_singleton: boolean (nullable = true) |-- hypothetical: boolean (nullable = true) |-- ec: string (nullable = true) |-- kegg_orthology_id: string (nullable = true) |-- uniref100: string (nullable = true) |-- msa_depth: integer (nullable = true)
+-----------------+-------------------+---------------------+-------+------------+------------+------------+----+-----------------+--------------------+---------+ |uniprot_accession| gene_cluster_id|gtdb_species_clade_id|is_core|is_auxiliary|is_singleton|hypothetical| ec|kegg_orthology_id| uniref100|msa_depth| +-----------------+-------------------+---------------------+-------+------------+------------+------------+----+-----------------+--------------------+---------+ | A0A009HIR4| NZ_KB851202.1_61| s__Acinetobacter_...| false| true| true| false|NULL| NULL|UniRef100_A0A009HIR4| 9675| | A0A009HIR4|NZ_JFEM01000045.1_7| s__Acinetobacter_...| false| true| false| false|NULL| NULL|UniRef100_A0A009HIR4| 9675| | A0A009MJC6|NZ_NGGG01000116.1_3| s__Acinetobacter_...| false| true| false| false|NULL| NULL|UniRef100_A0A009MJC6| 18012| | A0A009MJC6|NZ_BKYM01000029.1_2| s__Acinetobacter_...| false| true| true| false|NULL| NULL|UniRef100_A0A009MJC6| 18012| | A0A009MJC6|NZ_NGFA01000100.1_1| s__Acinetobacter_...| false| true| false| false|NULL| NULL|UniRef100_A0A009MJC6| 18012| +-----------------+-------------------+---------------------+-------+------------+------------+------------+----+-----------------+--------------------+---------+ only showing top 5 rows
In [5]:
from pyspark.sql.functions import avg, percentile_approx, count, sum as spark_sum
agg = joined.groupBy("is_core", "is_auxiliary", "is_singleton").agg(
count("*").alias("n_clusters"),
avg(col("msa_depth").cast("double")).alias("mean_msa_depth"),
percentile_approx(col("msa_depth").cast("double"), 0.5).alias("median_msa_depth"),
percentile_approx(col("msa_depth").cast("double"), 0.1).alias("p10_msa_depth"),
percentile_approx(col("msa_depth").cast("double"), 0.9).alias("p90_msa_depth"),
spark_sum((col("msa_depth").cast("double") < 10).cast("int")).alias("n_very_low_msa"),
spark_sum(col("hypothetical").cast("int")).alias("n_hypothetical"),
spark_sum((col("ec").isNotNull()).cast("int")).alias("n_has_ec"),
spark_sum((col("kegg_orthology_id").isNotNull()).cast("int")).alias("n_has_kegg")
)
agg.show(truncate=False)
# Also save the full joined table (sampled for local analysis)
# Write aggregate to MinIO
OUTPUT_PATH = "s3a://cdm-lake/tenant-general-warehouse/microbialdiscoveryforge/projects/alphafold_msa_annotation/data/gc_msa_agg.csv"
agg.toPandas().to_csv("/tmp/gc_msa_agg.csv", index=False)
print("Saved aggregate to /tmp/gc_msa_agg.csv — download and commit to data/")
+-------+------------+------------+----------+------------------+----------------+-------------+-------------+--------------+--------------+--------+----------+ |is_core|is_auxiliary|is_singleton|n_clusters|mean_msa_depth |median_msa_depth|p10_msa_depth|p90_msa_depth|n_very_low_msa|n_hypothetical|n_has_ec|n_has_kegg| +-------+------------+------------+----------+------------------+----------------+-------------+-------------+--------------+--------------+--------+----------+ |false |true |true |7095643 |8040.5719132149125|5299.0 |25.0 |19203.0 |392959 |979300 |766880 |1070890 | |true |false |false |25571299 |11617.805913770748|15308.0 |334.0 |19500.0 |415733 |979912 |5557862 |6502189 | |false |true |false |5384900 |8104.200634366469 |5527.0 |32.0 |19192.0 |245002 |622748 |588587 |860594 | +-------+------------+------------+----------+------------------+----------------+-------------+-------------+--------------+--------------+--------+----------+
Saved aggregate to /tmp/gc_msa_agg.csv — download and commit to data/