adesomics.io.metadata

  1"""Sample metadata for RNA-seq experiments."""
  2
  3
  4import re
  5from dataclasses import dataclass, field
  6
  7import pandas as pd
  8
  9# columns that describe the study as a whole, not individual samples
 10_STUDY_LEVEL = [
 11    "Organism", "strain", "substrain", "BioProject", "SRA Study",
 12    "Instrument", "Platform", "Assay Type", "LibraryLayout",
 13    "LibrarySelection", "LibrarySource", "Center Name", "culture_type",
 14]
 15
 16# bookkeeping columns with no biological meaning
 17_DROP = [
 18    "Bytes", "Bases", "AvgSpotLen", "Consent", "ReleaseDate", "create_date",
 19    "version", "DATASTORE filetype", "DATASTORE provider", "DATASTORE region",
 20]
 21
 22_ID_CANDIDATES = ["Run", "run_accession", "sample_id"]
 23
 24
 25# columns holding a free-text sample description (pysradb / GEO style)
 26_TITLE_CANDIDATES = ["sample_title", "title", "Sample Name", "sample_name"]
 27
 28# marks a structured SRA Run Selector export
 29_SRA_MARKERS = {"Run", "BioProject", "Assay Type"}
 30
 31
 32def _slug(text) -> str:
 33    """'Acetoacetate/LiCl(10mM)' -> 'Acetoacetate_LiCl_10mM'"""
 34    return re.sub(r"[^0-9A-Za-z]+", "_", str(text)).strip("_")
 35
 36
 37def _strip_replicate(title) -> str:
 38    """'del atoC, M9, LiCl, rep1' -> 'del_atoC_M9_LiCl'"""
 39    without = re.sub(r"[,;]?\s*rep(licate)?\s*\.?\s*\d+\s*$", "", str(title),
 40                     flags=re.I)
 41    return _slug(without)
 42
 43
 44# columns that are never experimental factors, even at low cardinality
 45_NON_FACTOR = {
 46    "Assay Type", "LibrarySelection", "LibrarySource", "Platform",
 47    "LibraryLayout", "Sample Name", "Library Name", "Experiment", "BioSample",
 48    "source_name", "strain",
 49}
 50
 51# preferred factor columns, tried in this order
 52_FACTOR_PRIORITY = ["genotype", "media", "supplement", "treatment",
 53                    "condition_", "time", "temperature", "carbon_source"]
 54
 55
 56def _detect_factors(df):
 57    """Pick columns that vary but are not per-sample identifiers.
 58
 59    A factor has more than one level but fewer than one level per sample,
 60    and is not on the non-factor blocklist.
 61    """
 62    n = len(df)
 63    usable = [
 64        c for c in df.columns
 65        if c not in _NON_FACTOR and 1 < df[c].nunique() < n
 66    ]
 67    # keep known factor names first, then anything else usable
 68    ordered = [c for c in _FACTOR_PRIORITY if c in usable]
 69    ordered += [c for c in usable if c not in ordered]
 70
 71    if not ordered:
 72        raise ValueError(
 73            "Could not detect any experimental factor. Pass factors=(...) "
 74            f"explicitly. Columns available: {list(df.columns)}"
 75        )
 76    return tuple(ordered)
 77
 78
 79def _pick(columns, candidates, what):
 80    for c in candidates:
 81        if c in columns:
 82            return c
 83    raise ValueError(
 84        f"No column for {what}; expected one of {candidates}, have {list(columns)}"
 85    )
 86
 87
 88@dataclass(repr=False, eq=False)
 89class Metadata:
 90    """Per-sample annotation plus study-level constants.
 91
 92    table: index = run IDs, columns = experimental factors
 93    study: values identical across all samples (organism, platform, ...)
 94    """
 95
 96    table: pd.DataFrame
 97    study: dict = field(default_factory=dict)
 98    group_col: str = "condition"
 99
100    def __post_init__(self):
101        self.table = self.table.copy()
102        if self.table.index.duplicated().any():
103            dupes = self.table.index[self.table.index.duplicated()].unique()
104            raise ValueError(f"Duplicate sample IDs: {list(dupes)}")
105        if self.group_col not in self.table.columns:
106            raise ValueError(
107                f"Column '{self.group_col}' missing; have {list(self.table.columns)}"
108            )
109
110    # --- constructors ----------------------------------------------------
111
112    @classmethod
113    def from_sra_run_table(
114        cls,
115        path,
116        factors=None,
117        id_col=None,
118        assay=None,
119    ) -> "Metadata":
120        """Read SraRunTable.csv from the NCBI SRA Run Selector.
121
122        factors
123            Columns that define a condition. If None, they are detected
124            automatically: every non-constant, non-ID column that looks like
125            an experimental factor.
126        assay
127            Keep only rows whose 'Assay Type' matches (e.g. 'RNA-Seq'), for
128            run tables that mix RNA-Seq and ChIP-Seq.
129        """
130        df = pd.read_csv(path, sep=None, engine="python")
131
132        if assay is not None and "Assay Type" in df.columns:
133            before = len(df)
134            df = df[df["Assay Type"] == assay]
135            print(f"[Metadata] assay filter '{assay}': kept {len(df)}/{before} runs")
136            if df.empty:
137                raise ValueError(
138                    f"No rows with Assay Type == '{assay}'; "
139                    f"available: {sorted(pd.read_csv(path)['Assay Type'].unique())}"
140                )
141
142        id_col = id_col or next(
143            (c for c in _ID_CANDIDATES if c in df.columns), None
144        )
145        if id_col is None:
146            raise ValueError(
147                f"No run-ID column found; expected one of {_ID_CANDIDATES}"
148            )
149        df = df.set_index(id_col).drop(columns=_DROP, errors="ignore")
150
151        # split off values that are constant across the whole study
152        study = {}
153        for col in list(df.columns):
154            if col in _STUDY_LEVEL and df[col].nunique(dropna=False) == 1:
155                study[col] = df[col].iloc[0]
156                df = df.drop(columns=col)
157
158        # normalise the biological factors that we recognise
159        if "genotype" not in df.columns and "source_name" in df.columns:
160            df["genotype"] = (
161                df["source_name"]
162                .str.replace("MG1655", "", regex=False)
163                .str.strip()
164                .replace("", "wt")
165            )
166        if "genotype" in df.columns:
167            df["genotype"] = df["genotype"].map(_slug)
168        if "supplement" in df.columns:
169            df["supplement"] = df["supplement"].replace({"--": "none"}).map(_slug)
170        if "media" in df.columns:
171            df["media"] = df["media"].map(_slug)
172
173        if factors is None:
174            factors = _detect_factors(df)
175            print(f"[Metadata] detected factors: {list(factors)}")
176
177        missing = [f for f in factors if f not in df.columns]
178        if missing:
179            raise ValueError(
180                f"Factors {missing} not in table; have {list(df.columns)}"
181            )
182
183        df["condition"] = df[list(factors)].astype(str).agg("_".join, axis=1)
184        return cls(df, study=study)
185
186    @classmethod
187    def from_groups(cls, samples, groups, name="condition") -> "Metadata":
188        """Manual fallback: one group label per sample."""
189        if isinstance(groups, dict):
190            groups = [groups[s] for s in samples]
191        if len(groups) != len(samples):
192            raise ValueError(f"{len(groups)} groups for {len(samples)} samples")
193        return cls(
194            pd.DataFrame({name: list(groups)}, index=list(samples)),
195            group_col=name,
196        )
197
198    # --- access ----------------------------------------------------------
199
200    @property
201    def groups(self) -> pd.Series:
202        return self.table[self.group_col]
203
204    @property
205    def samples(self) -> list:
206        return list(self.table.index)
207
208    @property
209    def organism(self) -> str | None:
210        return self.study.get("Organism")
211
212    def factor_levels(self) -> dict:
213        """{'genotype': ['wt', 'del_atoC', ...], 'media': [...]}
214
215        Skips constants and ID-like columns (one distinct value per sample).
216        """
217        return {
218            c: sorted(self.table[c].unique())
219            for c in self.table.columns
220            if c != self.group_col and 1 < self.table[c].nunique() < len(self.table)
221        }
222
223    def replicates(self) -> pd.Series:
224        return self.groups.value_counts()
225
226    def align(self, counts: pd.DataFrame) -> "Metadata":
227        """Reorder rows to match count-matrix columns; tolerates BAM paths."""
228        ids = [
229            m.group() if (m := re.search(r"[SED]RR\d+", str(c))) else str(c)
230            for c in counts.columns
231        ]
232        missing = set(ids) - set(self.table.index)
233        if missing:
234            raise ValueError(f"Not in metadata: {sorted(missing)}")
235
236        sub = self.table.loc[ids].copy()
237        sub.index = counts.columns
238        return Metadata(sub, study=self.study, group_col=self.group_col)
239
240    def subset_samples(self, samples) -> "Metadata":
241        """Keep the given sample IDs, in the given order."""
242        missing = set(samples) - set(self.table.index)
243        if missing:
244            raise ValueError(f"Not in metadata: {sorted(missing)}")
245        return Metadata(self.table.loc[list(samples)], study=self.study,
246                        group_col=self.group_col)
247
248    def subset(self, keep=None, **criteria) -> "Metadata":
249        """subset(keep=[...]) by group, or subset(media='M9') by factor."""
250        mask = pd.Series(True, index=self.table.index)
251        if keep is not None:
252            mask &= self.groups.isin(keep)
253        for col, want in criteria.items():
254            want = want if isinstance(want, (list, tuple, set)) else [want]
255            mask &= self.table[col].isin(want)
256        return Metadata(self.table[mask], study=self.study,
257                        group_col=self.group_col)
258
259    def __len__(self):
260        return len(self.table)
261
262    def __repr__(self):
263        org = self.organism or "unknown organism"
264        return (f"<Metadata: {len(self)} samples, "
265                f"{self.groups.nunique()} groups, {org}>")