Fix tidy_anndata to return empty dataframe when there are no scores.

This should fix #39, where we may have an empty anndata object with gene_id and strand obs, but no rows.

PiperOrigin-RevId: 874524624
Change-Id: I053d90491346cba679cd45069fda31104f575964

Author: tomwardio

src/alphagenome/models/variant_scorers.py

@@ -18,6 +18,7 @@
 import dataclasses
 import enum
 import itertools
+import math
 from typing import TypeVar
 
 from alphagenome.models import dna_output
@@ -700,16 +701,18 @@ def tidy_anndata(
       'junction_Start',
       'junction_End',
   ]
-  if 'gene_id' in adata.obs and 'strand' in adata.obs:
+  if math.prod(adata.X.shape) == 0:
+    # Scores are empty, so we return an empty dataframe.
+    return pd.DataFrame()
+  elif 'gene_id' in adata.obs and 'strand' in adata.obs:
     # Scores are for a gene-based scorer.
     obs = adata.obs.rename({'strand': 'gene_strand'}, axis=1)
-    obs['gene_id'] = obs['gene_id'].str.split('.', expand=True)[0]  # Depatch.
+
+    # Remove patch number from gene_id.
+    obs['gene_id'] = obs['gene_id'].str.split('.', expand=True).get(0)
     for col in gene_columns:
       if col not in obs:
         obs[col] = None
-  elif adata.X.shape[0] == 0:
-    # Scores are empty, so we return an empty dataframe.
-    return pd.DataFrame()
   elif adata.X.shape[0] == 1 and adata.obs.empty:
     # Scores are for a non-gene-based scorer.
     obs = pd.DataFrame([[None] * len(gene_columns)], columns=gene_columns)

src/alphagenome/models/variant_scorers_test.py

@@ -877,6 +877,13 @@ def test_tidy_anndata_splicing_sites(self):
         ],
     )
 
+  def test_tidy_anndata_empty(self):
+    adata = anndata.AnnData(
+        X=np.ones((0, 3)), obs=pd.DataFrame(columns=['gene_id', 'strand'])
+    )
+    result = variant_scorers.tidy_anndata(adata)
+    pd.testing.assert_frame_equal(result, pd.DataFrame())
+
   def test_tidy_lists_of_anndata(self):
     # List of AnnDatas.
     anndata_list = [self.adata_gene_centric, self.adata_variant_centric]