adesomics.io.count_matrix

  1"""
  2Reads and validates gene count matrices from two formats:
  3  - featureCounts output (.txt / .tsv): has a leading '#' comment line and
  4    metadata columns (Chr, Start, End, Strand, Length) that get dropped.
  5  - Plain CSV (.csv): genes x samples, first column is the gene ID
  6    (any column name), remaining columns are samples.
  7
  8In both cases the result is a DataFrame with genes as rows and samples
  9as columns. Use CountMatrix.from_file() and it will detect the format
 10automatically based on the file extension.
 11"""
 12
 13import pandas as pd
 14from pathlib import Path
 15
 16
 17_FEATURE_COUNTS_SUFFIXES = {".txt", ".tsv"}
 18_CSV_SUFFIXES = {".csv"}
 19_FEATURE_COUNTS_META_COLS = {"Chr", "Start", "End", "Strand", "Length"}
 20
 21
 22class CountMatrix:
 23    """
 24    Represents a gene count matrix.
 25
 26    Attributes
 27    ----------
 28    counts : pd.DataFrame
 29        genes × samples, index is the gene ID column (name preserved from file).
 30    source_path : Path or None
 31    annotations : pd.DataFrame or None
 32    """
 33
 34    def __init__(self, counts, source_path=None, annotations=None, lengths=None, sample_names=None):
 35        self.counts = counts
 36        self.source_path = source_path
 37        self.annotations = annotations
 38        self.lengths = lengths
 39        self.sample_names = sample_names if sample_names is not None else list(counts.columns)
 40
 41    # ------------------------------------------------------------------
 42    # Constructors
 43    # ------------------------------------------------------------------
 44
 45    @classmethod
 46    def from_file(cls, path: str | Path) -> "CountMatrix":
 47        """
 48        Auto-detect format from file extension and load.
 49
 50        .txt / .tsv  →  featureCounts output
 51        .csv         →  plain CSV (genes x samples)
 52
 53        Parameters
 54        ----------
 55        path : str or Path
 56
 57        Raises
 58        ------
 59        FileNotFoundError
 60        ValueError  – unrecognised extension or parse failure
 61        """
 62        path = Path(path)
 63
 64        if not path.exists():
 65            raise FileNotFoundError(f"Count matrix file not found: {path}")
 66
 67        suffix = path.suffix.lower()
 68
 69        if suffix in _FEATURE_COUNTS_SUFFIXES:
 70            return cls._from_feature_counts(path)
 71        elif suffix in _CSV_SUFFIXES:
 72            return cls._from_csv(path)
 73        else:
 74            raise ValueError(
 75                f"Unrecognised file extension '{suffix}'. "
 76                f"Expected one of: {_FEATURE_COUNTS_SUFFIXES | _CSV_SUFFIXES}"
 77            )
 78
 79    @classmethod
 80    def from_dataframe(cls, df: pd.DataFrame) -> "CountMatrix":
 81        """
 82        Wrap an existing DataFrame (genes as index, samples as columns).
 83        Useful for testing.
 84        """
 85        instance = cls(counts=df.copy(), source_path=None)
 86        instance._validate()
 87        return instance
 88
 89    # ------------------------------------------------------------------
 90    # Format-specific loaders (private)
 91    # ------------------------------------------------------------------
 92
 93    @classmethod
 94    def _from_feature_counts(cls, path):
 95        df = pd.read_csv(path, sep="\t", comment="#", index_col="Geneid")
 96        lengths = df["Length"] if "Length" in df.columns else None   
 97        meta_cols = [c for c in _FEATURE_COUNTS_META_COLS if c in df.columns]
 98        df = df.drop(columns=meta_cols)
 99        df.columns = [Path(c).stem.split("_")[0] for c in df.columns]
100        sample_names = list(df.columns)
101        instance = cls(counts=df, source_path=path, lengths=lengths, sample_names=sample_names)  
102        instance._validate()
103        return instance
104    
105    @classmethod
106    def _from_csv(cls, path: Path) -> "CountMatrix":
107        """
108        Parse a plain CSV where the first column is the gene ID and
109        remaining columns are samples.
110        """
111        try:
112            df = pd.read_csv(path, index_col=0)
113            annotations = df.columns.values
114            print(f"Annotations: {annotations}")
115        except Exception as e:
116            raise ValueError(f"Could not parse CSV file '{path}': {e}") from e
117
118        instance = cls(counts=df, source_path=path, annotations=annotations, sample_names=list(df.columns))
119        instance._validate()
120        return instance
121    
122    def tpm(self):
123        if self.lengths is None:
124            raise ValueError("No gene lengths available! TPM requires featureCounts input.")
125        rpk = self.counts.div(self.lengths / 1000, axis=0)   # reads/kb
126        return rpk.div(rpk.sum(axis=0), axis=1) * 1e6
127
128    # ------------------------------------------------------------------
129    # Validation
130    # ------------------------------------------------------------------
131
132    def _validate(self):
133        df = self.counts
134
135        if df.shape[0] == 0:
136            raise ValueError("Count matrix is empty.")
137
138        if df.index.duplicated().any():
139            dupes = df.index[df.index.duplicated()].tolist()
140            raise ValueError(f"Duplicate gene IDs found: {dupes[:5]}")
141
142        self._validate_sample_names()
143
144        non_numeric = [c for c in df.columns if not pd.api.types.is_numeric_dtype(df[c])]
145        if non_numeric:
146            raise ValueError(f"Non-numeric sample columns: {non_numeric}")
147
148        if (df < 0).any().any():
149            raise ValueError("Count matrix contains negative values.")
150
151    def _validate_sample_names(self):
152        names = self.sample_names
153
154        if not names:
155            raise ValueError(
156                "Sample names are not properly defined: no sample columns were found "
157                "in the count matrix."
158            )
159
160        if len(names) != self.counts.shape[1]:
161            raise ValueError(
162                "Sample names are not properly defined: expected "
163                f"{self.counts.shape[1]} sample name(s), got {len(names)}."
164            )
165
166        missing = [
167            n for n in names
168            if n is None or (isinstance(n, float) and pd.isna(n))
169            or not str(n).strip() or str(n).startswith("Unnamed:")
170        ]
171        if missing:
172            raise ValueError(
173                f"Sample names are not properly defined: missing or blank sample name(s): {missing}"
174            )
175
176        seen = set()
177        duplicates = sorted({n for n in names if n in seen or seen.add(n)})
178        if duplicates:
179            raise ValueError(
180                f"Sample names are not properly defined: duplicate sample name(s): {duplicates}"
181            )
182
183    # ------------------------------------------------------------------
184    # Properties / convenience
185    # ------------------------------------------------------------------
186
187    @property
188    def n_genes(self) -> int:
189        return self.counts.shape[0]
190
191    @property
192    def n_samples(self) -> int:
193        return self.counts.shape[1]
194
195    @property
196    def gene_ids(self) -> list[str]:
197        return self.counts.index.tolist()
198
199    def summary(self) -> str:
200        ids = self.gene_ids
201        if len(ids) <= 10:
202            ids_str = ", ".join(ids)
203        else:
204            ids_str = ", ".join(ids[:5]) + "  ...  " + ", ".join(ids[-3:])
205
206        return (
207            f"CountMatrix: {self.n_genes} genes × {self.n_samples} samples\n"
208            #f"Gene IDs: {ids_str}\n"
209            f"Samples: {self.sample_names}\n"
210            #f"Total counts per sample:\n{self.counts.sum(axis=0).to_string()}"
211        )
212
213    def __repr__(self) -> str:
214        src = self.source_path.name if self.source_path else "in-memory"
215        return f"<CountMatrix {self.n_genes}g × {self.n_samples}s from '{src}'>"