diff --git a/src/pyrecount/accessor.py b/src/pyrecount/accessor.py index 5445840..a8728df 100644 --- a/src/pyrecount/accessor.py +++ b/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: diff --git a/tests/test_accessor.py b/tests/test_accessor.py index 10defcb..bdaf962 100644 --- a/tests/test_accessor.py +++ b/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",