diff --git a/examples/example.py b/examples/example.py index 2ef849a..4a5b9ea 100755 --- a/examples/example.py +++ b/examples/example.py @@ -26,7 +26,7 @@ # subset project_meta_dataframe = recount_meta_dataframe.filter( - pl.col("project").is_in(["SRP009615"]) & pl.col("external_id").is_in(["SRR389077"]) + pl.col("project").is_in(["SRP009615"]) ) print(project_meta_dataframe) @@ -61,6 +61,21 @@ print(gene_annotation) print(gene_counts) +scaled_counts = project.scale_mapped_reads( + gene_counts, + target_size=4e7, + L=100, +) + +print(scaled_counts) + +scaled_counts = project.scale_auc( + gene_counts, + target_size=4e7, +) + +print(scaled_counts) + exon_annotation, exon_counts = project.load(Dtype.EXON) print(exon_annotation) diff --git a/src/pyrecount/accessor.py b/src/pyrecount/accessor.py index 9cea22f..42d7a28 100644 --- a/src/pyrecount/accessor.py +++ b/src/pyrecount/accessor.py @@ -47,7 +47,9 @@ class Project: project_ids: List[str] = field(init=False) sample: List[str] = field(init=False) - endpoints: "EndpointConnector" = field(init=False) + endpoints: EndpointConnector = field(init=False) + + _cached_metadata: Optional[pl.DataFrame] = field(init=False, default=None) def __post_init__(self): if not isinstance(self.metadata, pl.DataFrame): @@ -86,7 +88,6 @@ 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, ) @@ -108,13 +109,71 @@ async def cache(self) -> None: makedirs(path.dirname(fpath), exist_ok=True) tasks.append(download_url_to_path(url=url, fpath=fpath)) + # launches all tasks, without rate limiting if tasks: await asyncio.gather(*tasks) + def scale_mapped_reads(self, counts: pl.DataFrame, target_size: float, L: int): + md = self.load_metadata() + + mapped_reads = pl.col("star.all_mapped_reads").cast(pl.Float64) + avg_mapped_len = pl.col("star.average_mapped_length").cast(pl.Float64) + avg_read_len = pl.col("avg_len").cast(pl.Float64) + + # paired-end detection + ratio = (avg_mapped_len / avg_read_len).round(0) + paired_end = ratio == 2 + paired_factor = pl.when(paired_end).then(2).otherwise(1) + + # scale factors + sf = md.select( + [ + pl.col("external_id"), + ( + target_size * L * paired_factor / (mapped_reads * avg_mapped_len**2) + ).alias("sf"), + ] + ) + + sf_map = dict(zip(sf["external_id"], sf["sf"])) + + return counts.with_columns( + [ + (pl.col(c) * sf_map[c]) + for c in counts.select(pl.selectors.numeric()).columns + ] + ) + + def scale_auc(self, counts: pl.DataFrame, target_size: float) -> pl.DataFrame: + md = self.load_metadata() + + auc = pl.col("bc_auc.all_reads_all_bases").cast(pl.Float64) + + # scale factors + sf = md.select( + "external_id", + (target_size / auc).alias("sf"), + ) + + sf_map = dict(zip(sf["external_id"], sf["sf"])) + + return counts.with_columns( + [ + pl.col(c).mul(sf_map[c]).round(0).cast(pl.Int64).alias(c) + for c in counts.columns + if c != "gene_id" + ] + ) + + def load_metadata(self) -> pl.DataFrame: + if self._cached_metadata is None: + self._cached_metadata = self._metadata_load() + return self._cached_metadata + def load(self, dtype) -> Union[pl.DataFrame, Tuple[pl.DataFrame, pl.DataFrame]]: match dtype: case Dtype.METADATA: - return self._metadata_load() + return self.load_metadata() case Dtype.JXN: return self._jxn_load() case Dtype.GENE: @@ -315,9 +374,11 @@ def _gtf_read(self, rpath: str) -> pl.DataFrame: [ annotation_dataframe["attribute"] .map_elements( - lambda x: re.findall(rf'{field} "([^"]*)"', x)[0] - if rf'{field} "' in x - else "", + lambda x: ( + re.findall(rf'{field} "([^"]*)"', x)[0] + if rf'{field} "' in x + else "" + ), return_dtype=pl.Utf8, ) .alias(field) @@ -366,7 +427,6 @@ def _exon_load(self) -> pl.DataFrame: annotation = self._gtf_read(fpath) if url.endswith(f"{self.annotation.value}.gz"): 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"] counts = ( diff --git a/tests/test_accessor.py b/tests/test_accessor.py index 07359e6..15f32fe 100644 --- a/tests/test_accessor.py +++ b/tests/test_accessor.py @@ -5,7 +5,6 @@ from pyrecount.accessor import Metadata, Project from pyrecount.models import Dtype, Annotation -# TODO: transform raw counts # TODO: jxn dataframe headers, external_id not rail id # TODO: multi-project support for exon, gene dtypes # TODO: expand sra attributes @@ -241,20 +240,20 @@ async def test_project_gene_accessor( (pl.col("project").is_in(project_ids)) ) - dtype = Dtype.GENE + dtype = [Dtype.METADATA, Dtype.GENE] gene = Project( metadata=project_meta_dataframe, dbase=dbase, organism=organism, - dtype=[dtype], + dtype=dtype, annotation=annotation, jxn_format=None, root_url=root_url, ) await gene.cache() - gene_annotation, gene_counts = gene.load(dtype) + gene_annotation, gene_counts = gene.load(Dtype.GENE) assert gene_annotation.shape == expected_annotation_shape assert gene_counts.shape == expected_counts_shape