{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "provenance": [ { "file_id": "1RMWC7BqiNWUz7snIUBnZdNZbb1ybzh2G", "timestamp": 1765455085369 }, { "file_id": "1CwpzZWAOY7dsipRxUGLw6IqpjUm_GMua", "timestamp": 1761142076411 } ], "last_runtime": {} }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" } }, "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": { "id": "19f7MzRoNJzA", "cellView": "form", "executionInfo": { "status": "ok", "timestamp": 1765454980711, "user_tz": 0, "elapsed": 2159, "user": { "displayName": "Natasha Latysheva", "userId": "06959083956285510029" } } }, "outputs": [], "source": [ "# @title Imports\n", "\n", "import functools\n", "import os\n", "from typing import Callable\n", "\n", "import numpy as np\n", "import pandas as pd\n", "from scipy.stats import pearsonr, spearmanr\n", "from sklearn import metrics\n", "from sklearn.metrics import average_precision_score, roc_auc_score\n" ] }, { "cell_type": "code", "source": [ "# @title Helper functions.\n", "\n", "def calculate_tissue_weighted_metric(\n", " predictions_df: pd.DataFrame,\n", " metric_fn: Callable[[np.ndarray, np.ndarray], float],\n", " prediction_col: str = 'prediction',\n", " target_col: str = 'target',\n", " tissue_col: str = 'tissue'\n", ") -> float:\n", " \"\"\"\n", " Calculates a tissue-weighted mean for a given metric function from a DataFrame.\n", "\n", " Args:\n", " predictions_df: DataFrame containing predictions, targets, and tissue info.\n", " metric_fn: A function that accepts two np.ndarrays (y_true, y_pred)\n", " and returns a single float metric.\n", " Note: For scipy functions that return tuples (e.g., spearmanr),\n", " you must wrap them in a lambda function (see example).\n", " prediction_col: Name of the column with prediction scores.\n", " target_col: Name of the column with target labels/values.\n", " tissue_col: Name of the column indicating the tissue/group.\n", "\n", " Returns:\n", " The calculated tissue-weighted mean metric, or np.nan if unable to calculate.\n", " \"\"\"\n", " if predictions_df is None or predictions_df.empty:\n", " print(\"Warning: predictions_df is empty, returning NaN.\")\n", " return float('nan')\n", "\n", " required_cols = [prediction_col, target_col, tissue_col]\n", " if not all(col in predictions_df.columns for col in required_cols):\n", " raise ValueError(f\"DataFrame must contain columns: {required_cols}\")\n", "\n", " tissue_results = []\n", "\n", " # Group by tissue and calculate the metric for each one.\n", " for tissue_name, group_df in predictions_df.groupby(tissue_col):\n", "\n", " # Drop rows where target or prediction is NaN.\n", " clean_group = group_df.dropna(subset=[target_col, prediction_col])\n", "\n", " num_variants = len(clean_group)\n", "\n", " # Skip tissues with < 2 variants (can't calculate metrics).\n", " if num_variants < 2:\n", " # print(f\"Skipping tissue {tissue_name} (variants={num_variants})\")\n", " continue\n", "\n", " try:\n", " # Calculate the metric for this one tissue using the passed function\n", " metric_value = metric_fn(\n", " clean_group[target_col], clean_group[prediction_col]\n", " )\n", "\n", " # Ensure metric is a valid number.\n", " if not np.isfinite(metric_value):\n", " print(f\"Skipping tissue {tissue_name} (metric_fn returned non-finite value)\")\n", " continue\n", "\n", " tissue_results.append({\n", " 'tissue': tissue_name,\n", " 'metric_value': metric_value,\n", " 'count': num_variants # This is the \"weight\".\n", " })\n", " except ValueError as e:\n", " # This will catch single-class errors for AUROC/AUPRC\n", " # and other potential issues from metric_fn.\n", " print(f\"Could not calculate metric for tissue {tissue_name}: {e}\")\n", " continue # Skip tissue if metric calculation fails.\n", "\n", " if not tissue_results:\n", " print(\"Warning: No tissues had scorable metric values, returning NaN.\")\n", " return float('nan')\n", "\n", " # Create a DataFrame of the per-tissue metrics.\n", " metrics_df = pd.DataFrame(tissue_results)\n", "\n", " # Calculate the final weighted mean.\n", " weighted_sum = (metrics_df['metric_value'] * metrics_df['count']).sum()\n", " total_count = metrics_df['count'].sum()\n", "\n", " if total_count == 0:\n", " print(\"Warning: Total count for weighting is zero, returning NaN.\")\n", " return float('nan')\n", "\n", " weighted_mean_metric = weighted_sum / total_count\n", " return weighted_mean_metric\n", "\n", "\n", "def paqtl_auprc(df):\n", " SEED = 0\n", "\n", " # 1. Separate positives and negatives.\n", " pos = df[df['target'] == 1]\n", " neg = df[df['target'] == 0]\n", "\n", " # 2. Merge on 'PI' to find all valid matched pairs.\n", " matched = pd.merge(\n", " pos,\n", " neg,\n", " on='PI',\n", " how='inner',\n", " suffixes=('_pos', '_neg')\n", " )\n", "\n", " # 3. Group by 'PI' and sample exactly one pair per group.\n", " # This ensures 1:1 matching controlled by PI.\n", " sampled_pairs = matched.groupby('PI').sample(n=1, random_state=SEED)\n", "\n", " # 4. Reconstruct a single dataframe for AUPRC calculation\n", " # We stack the positive and negative parts of the sampled pairs back together.\n", " df_sampled = pd.concat([\n", " sampled_pairs[['prediction_pos', 'target_pos']].rename(\n", " columns={'prediction_pos': 'prediction', 'target_pos': 'target'}\n", " ),\n", " sampled_pairs[['prediction_neg', 'target_neg']].rename(\n", " columns={'prediction_neg': 'prediction', 'target_neg': 'target'}\n", " )\n", " ])\n", " return metrics.average_precision_score(\n", " df_sampled['target'], df_sampled['prediction'])\n", "\n", "auroc_fn = roc_auc_score\n", "auprc_fn = average_precision_score\n", "\n", "spearman_fn = lambda y_true, y_pred: spearmanr(y_true, y_pred)[0]\n" ], "metadata": { "cellView": "form", "id": "x-W_KH8_P19G", "executionInfo": { "status": "ok", "timestamp": 1765454981000, "user_tz": 0, "elapsed": 53, "user": { "displayName": "Natasha Latysheva", "userId": "06959083956285510029" } } }, "execution_count": 2, "outputs": [] }, { "cell_type": "code", "source": [ "# @title Eval configs.\n", "PREDS_PATH = 'https://storage.googleapis.com/alphagenome/evals'\n", "\n", "evals = {\n", " 'clinvar_splice_site_region': {\n", " 'output_type': 'SPLICE_SITE_USAGE;RNA_SEQ;SPLICE_SITE_POSITIONS;SPLICE_JUNCTIONS;SPLICE_SITES',\n", " 'metric_name': 'auprc_max_abs_track_aggregation',\n", " 'metric_fn': lambda x: metrics.average_precision_score(x['target'], x['prediction']),\n", " 'reported_metric': 0.5699,\n", " },\n", " 'clinvar_noncoding': {\n", " 'output_type': 'RNA_SEQ;SPLICE_SITE_POSITIONS;SPLICE_SITE_USAGE;SPLICE_JUNCTIONS;SPLICE_SITES',\n", " 'metric_name': 'auprc_max_abs_track_aggregation',\n", " 'metric_fn': lambda x: metrics.average_precision_score(x['target'], x['prediction']),\n", " 'reported_metric': 0.6588,\n", " },\n", " 'clinvar_missense': {\n", " 'output_type': 'SPLICE_SITE_USAGE;SPLICE_JUNCTIONS;SPLICE_SITES;RNA_SEQ;SPLICE_SITE_POSITIONS',\n", " 'metric_name': 'auprc_max_abs_track_aggregation',\n", " 'metric_fn': lambda x: metrics.average_precision_score(x['target'], x['prediction']),\n", " 'reported_metric': 0.1792,\n", " },\n", " 'sqtl_variant_causality_gene_human': {\n", " 'output_type': 'RNA_SEQ;SPLICE_JUNCTIONS;SPLICE_SITES;SPLICE_SITE_USAGE;SPLICE_SITE_POSITIONS',\n", " 'metric_name': 'tissue_weighted_mean_auprc',\n", " 'metric_fn': functools.partial(calculate_tissue_weighted_metric, metric_fn=auprc_fn),\n", " 'reported_metric': 0.7644,\n", " },\n", " 'mfass_splicing': {\n", " 'output_type': 'SPLICE_SITES_LOGITS;SPLICE_SITE_USAGE;SPLICE_JUNCTIONS',\n", " 'metric_name': 'all_tissues_auprc',\n", " 'metric_fn': lambda x: metrics.average_precision_score(x['target'], x['prediction']),\n", " 'reported_metric': 0.5120,\n", " },\n", " 'eqtl_variant_borzoi_sign_human': {\n", " 'output_type': 'RNA_SEQ',\n", " 'metric_name': 'tissue_weighted_mean_auroc',\n", " 'metric_fn': functools.partial(calculate_tissue_weighted_metric, metric_fn=auroc_fn),\n", " 'reported_metric': 0.810077,\n", " },\n", " 'eqtl_variant_catalogue_causality_gene_balanced_human': {\n", " 'output_type': 'RNA_SEQ',\n", " 'metric_name': 'tissue_weighted_mean_auroc',\n", " 'metric_fn': functools.partial(calculate_tissue_weighted_metric, metric_fn=auroc_fn),\n", " 'reported_metric': 0.713255,\n", " },\n", " 'eqtl_variant_borzoi_coefficient_human': {\n", " 'output_type': 'RNA_SEQ',\n", " 'metric_name': 'tissue_weighted_mean_spearmanr',\n", " 'metric_fn': functools.partial(calculate_tissue_weighted_metric, metric_fn=spearman_fn),\n", " 'reported_metric': 0.500588,\n", " },\n", " 'enhancer_gene_linking_e2g': {\n", " 'output_type': 'RNA_SEQ',\n", " 'metric_name': None,\n", " 'metric_fn': lambda x: metrics.average_precision_score(x['target'], x['prediction']),\n", " 'reported_metric': 0.7490,\n", " },\n", " 'paqtl_variant_causality_human': {\n", " 'output_type': 'RNA_SEQ',\n", " 'metric_name': 'PAS_10000_threshold_average_auprc',\n", " 'metric_fn': paqtl_auprc,\n", " 'reported_metric': 0.6294,\n", " 'notes': 'The reported metric is from many permutations, here we just do 1 resample.',\n", " },\n", " 'caqtl_african_variant_causality_human': {\n", " 'output_type': 'DNASE',\n", " 'metric_name': \"auprc_mean_abs_track_aggregation_ontology_curie:['EFO:0002784']\",\n", " 'metric_fn': lambda x: metrics.average_precision_score(x['target'], x['prediction']),\n", " 'reported_metric': 0.5643,\n", " },\n", " 'caqtl_european_variant_causality_human': {\n", " 'output_type': 'DNASE',\n", " 'metric_name': \"auprc_mean_abs_track_aggregation_ontology_curie:['EFO:0002784']\",\n", " 'metric_fn': lambda x: metrics.average_precision_score(x['target'], x['prediction']),\n", " 'reported_metric': 0.3638,\n", " },\n", " 'dsqtl_yoruba_variant_causality_human': {\n", " 'output_type': 'DNASE',\n", " 'metric_name': \"auprc_mean_abs_track_aggregation_ontology_curie:['EFO:0002784']\",\n", " 'metric_fn': lambda x: metrics.average_precision_score(x['target'], x['prediction']),\n", " 'reported_metric': 0.6308,\n", " },\n", " 'caqtl_african_variant_coefficient_human': {\n", " 'output_type': 'DNASE',\n", " 'metric_name': \"pearsonr_mean_track_aggregation_ontology_curie:['EFO:0002784']\",\n", " 'metric_fn': lambda x: pearsonr(x['target'], x['prediction']).statistic,\n", " 'reported_metric': 0.7368,\n", " },\n", " 'caqtl_european_variant_coefficient_human': {\n", " 'output_type': 'DNASE',\n", " 'metric_name': \"pearsonr_mean_track_aggregation_ontology_curie:['EFO:0002784']\",\n", " 'metric_fn': lambda x: pearsonr(x['target'], x['prediction']).statistic,\n", " 'reported_metric': 0.5916,\n", " },\n", " 'dsqtl_yoruba_variant_coefficient_human': {\n", " 'output_type': 'DNASE',\n", " 'metric_name': \"pearsonr_mean_track_aggregation_ontology_curie:['EFO:0002784']\",\n", " 'metric_fn': lambda x: pearsonr(x['target'], x['prediction']).statistic,\n", " 'reported_metric': 0.8323,\n", " 'notes': 'The sign of the correlation flips depending on which allele is labelled REF vs. ALT.',\n", " },\n", " 'caqtl_microglia_variant_coefficient_human': {\n", " 'output_type': 'DNASE',\n", " 'metric_name': \"pearsonr_mean_track_aggregation_ontology_curie:['CL:0000862']\",\n", " 'metric_fn': lambda x: pearsonr(x['target'], x['prediction']).statistic,\n", " 'reported_metric': 0.6357,\n", " 'notes': 'The sign of the correlation flips depending on which allele is labelled REF vs. ALT.',\n", " },\n", " 'caqtl_smc_variant_coefficient_human': {\n", " 'output_type': 'ATAC',\n", " 'metric_name': \"pearsonr_mean_track_aggregation_ontology_curie:['UBERON:0002079']\",\n", " 'metric_fn': lambda x: pearsonr(x['target'], x['prediction']).statistic,\n", " 'reported_metric': 0.687,\n", " },\n", " 'bqtl_spi1_variant_coefficient_human': {\n", " 'output_type': 'CHIP_TF',\n", " 'metric_name': 'pearsonr_mean_track_aggregation_EFO:0002784_SPI1',\n", " 'metric_fn': lambda x: pearsonr(x['target'], x['prediction']).statistic,\n", " 'reported_metric': 0.549967,\n", " },\n", " 'bqtl_spi1_variant_causality_human': {\n", " 'output_type': 'CHIP_TF',\n", " 'metric_name': 'auprc_mean_abs_track_aggregation_EFO:0002784_SPI1',\n", " 'metric_fn': lambda x: metrics.average_precision_score(x['target'], x['prediction']),\n", " 'reported_metric': 0.4952,\n", " },\n", "}" ], "metadata": { "id": "aokAbm4fP_YS", "executionInfo": { "status": "ok", "timestamp": 1765454981311, "user_tz": 0, "elapsed": 54, "user": { "displayName": "Natasha Latysheva", "userId": "06959083956285510029" } }, "cellView": "form" }, "execution_count": 3, "outputs": [] }, { "cell_type": "code", "source": [ "for eval_name, c in evals.items():\n", " print(f\"\\nEval: {eval_name}\")\n", "\n", " filepath = os.path.join(PREDS_PATH, eval_name + '_predictions' + '.feather')\n", " predictions = pd.read_feather(filepath)\n", "\n", " recomputed_metric = c['metric_fn'](predictions)\n", " print(f\" Reported: {c['reported_metric']}\")\n", " print(f\" Recomputed: {recomputed_metric:.4f}\")\n", "\n", " if c.get('notes'):\n", " print(f\" Notes: {c['notes']}\")" ], "metadata": { "id": "RfREAOVlhLEo", "executionInfo": { "status": "ok", "timestamp": 1765455006636, "user_tz": 0, "elapsed": 25072, "user": { "displayName": "Natasha Latysheva", "userId": "06959083956285510029" } }, "outputId": "55cf9270-786a-45a0-bdc0-b55e28df434d" }, "execution_count": 4, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "\n", "Eval: clinvar_splice_site_region\n", " Reported: 0.5699\n", " Recomputed: 0.5699\n", "\n", "Eval: clinvar_noncoding\n", " Reported: 0.6588\n", " Recomputed: 0.6588\n", "\n", "Eval: clinvar_missense\n", " Reported: 0.1792\n", " Recomputed: 0.1792\n", "\n", "Eval: sqtl_variant_causality_gene_human\n", " Reported: 0.7644\n", " Recomputed: 0.7645\n", "\n", "Eval: mfass_splicing\n", " Reported: 0.512\n", " Recomputed: 0.5124\n", "\n", "Eval: eqtl_variant_borzoi_sign_human\n", " Reported: 0.810077\n", " Recomputed: 0.8101\n", "\n", "Eval: eqtl_variant_catalogue_causality_gene_balanced_human\n", " Reported: 0.713255\n", " Recomputed: 0.7133\n", "\n", "Eval: eqtl_variant_borzoi_coefficient_human\n", " Reported: 0.500588\n", " Recomputed: 0.5006\n", "\n", "Eval: enhancer_gene_linking_e2g\n", " Reported: 0.749\n", " Recomputed: 0.7488\n", "\n", "Eval: paqtl_variant_causality_human\n", " Reported: 0.6294\n", " Recomputed: 0.6288\n", " Notes: The reported metric is from many permutations, here we just do 1 resample.\n", "\n", "Eval: caqtl_african_variant_causality_human\n", " Reported: 0.5643\n", " Recomputed: 0.5641\n", "\n", "Eval: caqtl_european_variant_causality_human\n", " Reported: 0.3638\n", " Recomputed: 0.3637\n", "\n", "Eval: dsqtl_yoruba_variant_causality_human\n", " Reported: 0.6308\n", " Recomputed: 0.6311\n", "\n", "Eval: caqtl_african_variant_coefficient_human\n", " Reported: 0.7368\n", " Recomputed: 0.7367\n", "\n", "Eval: caqtl_european_variant_coefficient_human\n", " Reported: 0.5916\n", " Recomputed: 0.5914\n", "\n", "Eval: dsqtl_yoruba_variant_coefficient_human\n", " Reported: 0.8323\n", " Recomputed: -0.8323\n", " Notes: The sign of the correlation flips depending on which allele is labelled REF vs. ALT.\n", "\n", "Eval: caqtl_microglia_variant_coefficient_human\n", " Reported: 0.6357\n", " Recomputed: -0.6354\n", " Notes: The sign of the correlation flips depending on which allele is labelled REF vs. ALT.\n", "\n", "Eval: caqtl_smc_variant_coefficient_human\n", " Reported: 0.687\n", " Recomputed: 0.6870\n", "\n", "Eval: bqtl_spi1_variant_coefficient_human\n", " Reported: 0.549967\n", " Recomputed: 0.5493\n", "\n", "Eval: bqtl_spi1_variant_causality_human\n", " Reported: 0.4952\n", " Recomputed: 0.4948\n" ] } ] } ] }