multiproj gene

src/pyrecount/accessor.py

@@ -255,7 +255,9 @@ def _gtf_read(self, rpath: str) -> pl.DataFrame:
             ]
         )
 
-    def _counts_read(self, rname: str):
+    def _counts_read(
+        self, rname: str, samples: list[str] | None = None
+    ) -> pl.DataFrame:
         df = pl.read_csv(
             rname,
             comment_prefix="#",
@@ -264,10 +266,12 @@ def _counts_read(self, rname: str):
 
         first_col = df.columns[0]
 
-        if not self.sample:
+        sample_ids = samples if samples is not None else self.sample
+
+        if not sample_ids:
             return df
 
-        keep = [first_col] + self.sample
+        keep = [first_col] + sample_ids
         missing = set(keep) - set(df.columns)
         if missing:
             raise KeyError(f"Missing columns in counts file: {missing}")
@@ -352,21 +356,36 @@ async def cache(self) -> None:
 
     def load(self) -> tuple[pl.DataFrame, pl.DataFrame]:
         annotation: pl.DataFrame | None = None
-        counts: pl.DataFrame | None = None
+        all_counts: pl.DataFrame = list()
 
+        # load shared annotation first
         for url in self._urls():
             fpath = urlparse(url).path.lstrip("/")
-
             if self.project.annotation.value in fpath:
                 if any(fpath.endswith(ext) for ext in Extensions.GENE.value):
                     annotation = self.project._gtf_read(fpath)
-                elif fpath.endswith(f"{self.project.annotation.value}.gz"):
-                    counts = self.project._counts_read(fpath)
+                    break
 
-        if annotation is None or counts is None:
+        # load per-project counts
+        for project_id in self.project.project_ids:
+            project_urls = [u for u in self._urls() if project_id in u]
+            project_samples = (
+                self.project.metadata.filter(pl.col("project") == project_id)[
+                    "external_id"
+                ]
+                .unique()
+                .to_list()
+            )
+
+            for url in project_urls:
+                fpath = urlparse(url).path.lstrip("/")
+                if fpath.endswith(f"{self.project.annotation.value}.gz"):
+                    all_counts.append(self.project._counts_read(fpath, project_samples))
+
+        if annotation is None or not all_counts:
             raise RuntimeError("Missing gene annotation or counts file")
 
-        return annotation, counts
+        return annotation, pl.concat(all_counts, how="align")
 
 
 class JunctionLoader:

tests/test_accessor.py

@@ -6,7 +6,7 @@
 from pyrecount.models import Dtype, Annotation
 
 # TODO: jxn dataframe headers, external_id not rail id
-# TODO: multi-project support for exon, gene dtypes
+# TODO: multi-project support for exon dtype
 # TODO: expand sra attributes
 # TODO: expose Lazyframes
 
@@ -112,6 +112,57 @@ async def test_multi_project_jxn_accessor(
     assert jxn_dataframe.shape == expected_shape
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "organism, dbase, project_ids, annotation, expected_annotation_shape, expected_counts_shape",
+    [
+        (
+            "human",
+            "sra",
+            ["SRP009615", "SRP075759"],
+            Annotation.GENCODE_V29,
+            (64836, 21),
+            (64837, 44),
+        ),
+    ],
+)
+async def test_multi_project_gene_accessor(
+    organism,
+    dbase,
+    project_ids,
+    annotation,
+    expected_annotation_shape,
+    expected_counts_shape,
+):
+    root_url = "http://duffel.rail.bio/recount3"
+    recount_metadata = Metadata(organism=organism, root_url=root_url)
+
+    recount_metadata.cache()
+
+    recount_meta_dataframe = recount_metadata.load()
+
+    project_meta_dataframe = recount_meta_dataframe.filter(
+        pl.col("project").is_in(project_ids)
+    )
+
+    dtype = Dtype.GENE
+
+    proj = Project(
+        metadata=project_meta_dataframe,
+        dbase=dbase,
+        organism=organism,
+        annotation=annotation,
+        jxn_format=None,
+        root_url=root_url,
+    )
+
+    await proj.cache(dtype)
+    gene_annotation, gene_counts = proj.load(dtype)
+
+    assert gene_annotation.shape == expected_annotation_shape
+    assert gene_counts.shape == expected_counts_shape
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "organism, dbase, project_ids, annotation, expected_annotation_shape, expected_counts_shape",