Sample filter

examples/example.py

@@ -26,7 +26,7 @@
 
 # subset
 project_meta_dataframe = recount_meta_dataframe.filter(
-    (pl.col("project").is_in(["SRP009615"]))
+    pl.col("project").is_in(["SRP009615"]) & pl.col("external_id").is_in(["SRR389077"])
 )
 
 print(project_meta_dataframe)
@@ -46,6 +46,7 @@
 )
 
 asyncio.run(project.cache())
+
 project_metadata = project.load(Dtype.METADATA)
 
 print(project_metadata)

src/pyrecount/accessor.py

@@ -86,6 +86,7 @@ def get_project_urls(self, dtype) -> List[str]:
             dtype=dtype,
             annotation=self.annotation,
             project_ids=self.project_ids,
+            # only needed for bigwig
             sample=self.sample,
             jxn_format=self.jxn_format,
         )
@@ -150,6 +151,11 @@ def _metadata_load(self) -> pl.DataFrame:
 
                 fpath = urlparse(url).path.lstrip("/")
                 df = pl.read_csv(fpath, separator="\t", infer_schema=False)
+                df = (
+                    df.filter(pl.col("external_id").is_in(self.sample))
+                    if self.sample
+                    else df
+                )
                 dfs_for_project.append(df)
 
             if not dfs_for_project:
@@ -272,7 +278,7 @@ def _bw_load(self) -> pl.DataFrame:
 
         return pl.DataFrame(project_urls, schema=["url"])
 
-    def _read_gtf(self, rpath: str) -> pl.DataFrame:
+    def _gtf_read(self, rpath: str) -> pl.DataFrame:
         annotation_dataframe = pl.read_csv(
             rpath,
             comment_prefix="#",
@@ -319,32 +325,47 @@ def _read_gtf(self, rpath: str) -> pl.DataFrame:
             ]
         )
 
-    def _read_counts(self, rname: str):
-        counts_dataframe = pl.read_csv(
+    def _counts_read(self, rname: str, colname: Optional[str] = None):
+        df = pl.read_csv(
             rname,
             comment_prefix="#",
             separator="\t",
         )
-        return counts_dataframe
 
-    def _gene_load(self) -> pl.DataFrame:
+        first_col = df.columns[0]
+
+        if not self.sample:
+            return df
+
+        keep = [first_col] + self.sample
+        missing = set(keep) - set(df.columns)
+        if missing:
+            raise KeyError(f"Missing columns in counts file: {missing}")
+        return df.select(keep)
+
+    def _gene_load(self) -> tuple[pl.DataFrame, pl.DataFrame]:
+        annotation = None
+        counts = None
         for url in self.get_project_urls(Dtype.GENE):
             fpath = urlparse(url).path.lstrip("/")
             if self.annotation.value in fpath:
                 if any(fpath.endswith(ext) for ext in Extensions.GENE.value):
-                    annotation = self._read_gtf(fpath)
+                    annotation = self._gtf_read(fpath)
                 if fpath.endswith(f"{self.annotation.value}.gz"):
-                    counts = self._read_counts(fpath)
+                    counts = self._counts_read(fpath, "gene_id")
+
+        if annotation is None or counts is None:
+            raise RuntimeError("Missing gene annotation or counts file")
         return annotation, counts
 
     def _exon_load(self) -> pl.DataFrame:
         for url in self.get_project_urls(Dtype.EXON):
             fpath = urlparse(url).path.lstrip("/")
             if self.annotation.value in fpath:
                 if any(url.endswith(ext) for ext in Extensions.EXON.value):
-                    annotation = self._read_gtf(fpath)
+                    annotation = self._gtf_read(fpath)
                 if url.endswith(f"{self.annotation.value}.gz"):
-                    counts = self._read_counts(fpath)
+                    counts = self._counts_read(fpath)
                     # TODO: extract first column (chromosome|start_1base|end_1ba…)
                     exon_colname = counts.columns[0]
                     exon_fields = ["chrom", "start", "end", "strand"]

tests/test_accessor.py

@@ -6,7 +6,7 @@
 from pyrecount.models import Dtype, Annotation
 
 # TODO: transform raw counts
-# TODO: materialize requested columns: e.g., async polars for study id column
+# TODO: jxn dataframe headers, external_id not rail id
 # TODO: multi-project support for exon, gene dtypes
 # TODO: expand sra attributes
 # TODO: expose Lazyframes