from berdl_notebook_utils.setup_spark_session import get_spark_session
spark = get_spark_session()

from pyspark.sql.functions import count, countDistinct, col

ipr = spark.table("kbase_ke_pangenome.interproscan_domains").select(
    "gene_cluster_id", "analysis", "ipr_acc"
)

domain_counts = ipr.groupBy("gene_cluster_id").agg(
    count("*").alias("n_domain_hits"),
    countDistinct("ipr_acc").alias("n_distinct_ipr"),
    countDistinct("analysis").alias("n_analysis_tools"  )
)

print(f"Gene clusters with any domain annotation: {domain_counts.count():,}")
domain_counts.describe().show()

Gene clusters with any domain annotation: 111,035,431

+-------+----------------+------------------+------------------+------------------+
|summary| gene_cluster_id|     n_domain_hits|    n_distinct_ipr|  n_analysis_tools|
+-------+----------------+------------------+------------------+------------------+
|  count|       111035431|         111035431|         111035431|         111035431|
|   mean|            NULL| 7.504839874039846|3.2789845972678755| 5.042309368799586|
| stddev|            NULL|6.8237080828582615| 2.493564313545704|2.7728723999597196|
|    min|AAPI01000021.1_1|                 1|                 0|                 1|
|    max|WYES01000176.1_1|               657|                48|                15|
+-------+----------------+------------------+------------------+------------------+

domain_counts.groupBy(
    (col("n_domain_hits") == 0).alias("no_hits"),
).count().show()

# Bin distribution
from pyspark.sql.functions import when
domain_counts.withColumn(
    "hit_bin",
    when(col("n_domain_hits") == 0, "0")
    .when(col("n_domain_hits") <= 2, "1-2")
    .when(col("n_domain_hits") <= 5, "3-5")
    .when(col("n_domain_hits") <= 10, "6-10")
    .otherwise("11+")
).groupBy("hit_bin").count().orderBy("hit_bin").show()

+-------+---------+
|no_hits|    count|
+-------+---------+
|  false|111035431|
+-------+---------+

+-------+--------+
|hit_bin|   count|
+-------+--------+
|    1-2|21472948|
|    11+|24848031|
|    3-5|30095497|
|   6-10|34618955|
+-------+--------+

# domain_counts has >1B rows serialized — toPandas() on the full set exceeds
# spark.driver.maxResultSize (1 GiB). Save a sample and aggregate instead.

# 1. 100K-row sample (for local correlation analysis in NB03)
sample_pd = domain_counts.limit(100_000).toPandas()
sample_pd.to_csv("/tmp/gc_domain_sample.csv", index=False)
print(f"Saved sample: {len(sample_pd):,} rows → /tmp/gc_domain_sample.csv")

# 2. Aggregate stats by hit bin (small, safe to toPandas)
from pyspark.sql.functions import avg, percentile_approx, when

hit_bin = when(col("n_domain_hits") <= 2, "1-2") \
    .when(col("n_domain_hits") <= 5, "3-5") \
    .when(col("n_domain_hits") <= 10, "6-10") \
    .otherwise("11+")

agg_pd = (
    domain_counts
    .withColumn("hit_bin", hit_bin)
    .groupBy("hit_bin")
    .agg(
        count("*").alias("n_clusters"),
        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("hit_bin")
    .toPandas()
)
agg_pd.to_csv("/tmp/gc_domain_agg.csv", index=False)
print(f"Saved aggregate: {len(agg_pd)} rows → /tmp/gc_domain_agg.csv")
print(agg_pd)

Saved sample: 100,000 rows → /tmp/gc_domain_sample.csv

Saved aggregate: 4 rows → /tmp/gc_domain_agg.csv
  hit_bin  n_clusters  mean_n_domain_hits  median_n_domain_hits  \
0     1-2    21472948            1.493300                     1   
1     11+    24848031           16.676197                    14   
2     3-5    30095497            4.030053                     4   
3    6-10    34618955            7.671533                     7   

   mean_n_distinct_ipr  
0             0.823410  
1             6.275686  
2             2.148593  
3             3.633877

Nb02 Domain Annotation

NB02 — Domain Annotation Richness × MSA Depth¶