msnetloader 0.0.0__tar.gz

msnetloader-0.0.0/PKG-INFO

@@ -0,0 +1,26 @@
+Metadata-Version: 2.4
+Name: msnetloader
+Version: 0.0.0
+Summary: A Python API for seamless integration of π-MSNet into AI workflows
+Author-email: chengxin dai <chengxin2024@126.com>
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/daichengxin/MSnet/
+Project-URL: Repository, https://github.com/daichengxin/MSnet/
+Keywords: sdrf,python,multiomics,proteomics
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: torch==2.6.0
+Requires-Dist: pandas
+Requires-Dist: numpy
+Requires-Dist: pyarrow==17.0.0
+Requires-Dist: tensorflow
+Requires-Dist: duckdb
+
+# MSNetLoader
+A Python API for seamless integration of π-MSNet into AI workflows

msnetloader-0.0.0/README.md

@@ -0,0 +1,2 @@
+# MSNetLoader
+A Python API for seamless integration of π-MSNet into AI workflows

msnetloader-0.0.0/msnetloader/__init__.py

@@ -0,0 +1,6 @@
+
+__version__ = "0.0.1"
+
+__all__ = [
+    "__version__",
+]

msnetloader-0.0.0/msnetloader/denovo_loader.py

@@ -0,0 +1,128 @@
+import torch
+import numpy as np
+from torch.utils.data import IterableDataset
+import duckdb
+
+
+class DeNovoIterableDataset(IterableDataset):
+
+    def __init__(self, parquet_path, max_peaks=150, batch_size=32,
+                 min_consensus_support=None,
+                 max_pep=None):
+
+        con = duckdb.connect()
+        self.min_consensus_support = min_consensus_support
+        self.max_pep = max_pep
+
+        conditions = []
+        params = []
+
+        if self.min_consensus_support is not None:
+            conditions.append("consensus_support >= ?")
+            params.append(self.min_consensus_support)
+
+        if self.max_pep is not None:
+            conditions.append("posterior_error_probability <= ?")
+            params.append(self.max_pep)
+
+        where_clause = ""
+        if conditions:
+            where_clause = "WHERE " + " AND ".join(conditions)
+
+        query = f"""
+                SELECT
+                    peptidoform,
+                    exp_mass_to_charge AS precursor_mz,
+                    precursor_charge AS charge,
+                    mz_array,
+                    consensus_support,
+                    posterior_error_probability,
+                    intensity_array
+                FROM parquet_scan(?)
+                {where_clause}
+                """
+
+        self.cursor = con.execute(query, [parquet_path] + params)
+
+        self.batch_size = batch_size
+        self.max_peaks = max_peaks
+
+    def __iter__(self):
+        reader = self.cursor.fetch_record_batch(self.batch_size)
+
+        for batch in reader:
+            if batch.num_rows == 0:
+                continue
+            yield self.process_batch(batch)
+
+    # -----------------------------
+    def process_batch(self, batch):
+
+        peptidoform = batch["peptidoform"].to_pylist()
+        charges = batch["charge"].to_pylist()
+        mz_list = batch["mz_array"].to_pylist()
+        int_list = batch["intensity_array"].to_pylist()
+        precursor_mz = batch["precursor_mz"].to_pylist()
+        consensus_supports = batch["consensus_support"].to_pylist()
+        peps = batch["posterior_error_probability"].to_pylist()
+
+        spectra_out = []
+        seq_out = []
+        charge_out = []
+        precursor_out = []
+
+        # -----------------------------
+        # per spectrum processing
+        # -----------------------------
+        for i in range(len(peptidoform)):
+            mz = np.asarray(mz_list[i], dtype=np.float32)
+            intensity = np.asarray(int_list[i], dtype=np.float32)
+
+            if len(mz) == 0:
+                continue
+
+            # -------------------------
+            # top-k peaks
+            # -------------------------
+            if len(mz) > self.max_peaks:
+                idx = np.argsort(intensity)[-self.max_peaks:]
+                mz = mz[idx]
+                intensity = intensity[idx]
+
+            # -------------------------
+            # normalize
+            # -------------------------
+            max_int = intensity.max() if len(intensity) > 0 else 1.0
+            intensity = intensity / max_int
+
+            spectrum = np.stack([mz, intensity], axis=1)
+
+            spectra_out.append(torch.tensor(spectrum, dtype=torch.float32))
+            seq_out.append(peptidoform[i])
+            charge_out.append(charges[i])
+            precursor_out.append(precursor_mz[i])
+
+        return {
+            "spectrum": spectra_out,
+            "sequence": seq_out,
+            "precursor_mz": torch.tensor(precursor_out, dtype=torch.float32),
+            "charge": torch.tensor(charge_out, dtype=torch.long),
+        }
+
+
+if __name__ == '__main__':
+    """Test dataset + sampler + dataloader pipeline."""
+    file_path = ["D:/gitrepo/MSnet/MSNetLoader/tests/test_data/PXD014877-Akkermansia_muciniphilia-MSNet.parquet",
+                 "D:/gitrepo/MSnet/MSNetLoader/tests/test_data/PXD014877_Clostridium_Bolteae-MSNet.parquet"]
+    dataset = DeNovoIterableDataset(file_path)
+    from torch.utils.data import DataLoader
+    loader = DataLoader(
+        dataset,
+        batch_size=None,
+        num_workers=0,
+        pin_memory=False
+    )
+
+    batch = next(iter(loader))
+    print(batch)
+

msnetloader-0.0.0/msnetloader/denovo_tf.py

@@ -0,0 +1,143 @@
+import tensorflow as tf
+import numpy as np
+import duckdb
+
+
+class DeNovoTFDataset:
+
+    def __init__(self, parquet_path, max_peaks=150, batch_size=32,
+                 min_consensus_support=None,
+                 max_pep=None):
+
+        con = duckdb.connect()
+        self.min_consensus_support = min_consensus_support
+        self.max_pep = max_pep
+
+        conditions = []
+        params = []
+
+        if self.min_consensus_support is not None:
+            conditions.append("consensus_support >= ?")
+            params.append(self.min_consensus_support)
+
+        if self.max_pep is not None:
+            conditions.append("posterior_error_probability <= ?")
+            params.append(self.max_pep)
+
+        where_clause = ""
+        if conditions:
+            where_clause = "WHERE " + " AND ".join(conditions)
+
+        query = f"""
+                 SELECT
+                     peptidoform,
+                     exp_mass_to_charge AS precursor_mz,
+                     precursor_charge AS charge,
+                     mz_array,
+                     consensus_support,
+                     posterior_error_probability,
+                     intensity_array
+                 FROM parquet_scan(?)
+                 {where_clause}
+                 """
+
+        self.cursor = con.execute(query, [parquet_path] + params)
+
+        self.batch_size = batch_size
+        self.max_peaks = max_peaks
+
+    # =========================================================
+    # Generator (核心)
+    # =========================================================
+    def generator(self):
+        reader = self.cursor.fetch_record_batch(self.batch_size)
+        for batch in reader:
+            if batch.num_rows == 0:
+                continue
+            result = self.process_batch(batch)
+            for i in range(len(result["sequence"])):
+                yield (
+                    result["spectrum"][i],
+                    result["sequence"][i],
+                    result["precursor_mz"][i],
+                    result["charge"][i],
+                )
+
+    # =========================================================
+    # TF Dataset接口
+    # =========================================================
+    def get_dataset(self):
+
+        output_signature = (
+            tf.TensorSpec(shape=(None, 2), dtype=tf.float32),  # spectrum
+            tf.TensorSpec(shape=(), dtype=tf.string),          # sequence
+            tf.TensorSpec(shape=(), dtype=tf.float32),         # precursor_mz
+            tf.TensorSpec(shape=(), dtype=tf.int32),           # charge
+        )
+
+        ds = tf.data.Dataset.from_generator(
+            self.generator,
+            output_signature=output_signature
+        )
+
+        return ds
+
+    # =========================================================
+    # batch处理（基本不变）
+    # =========================================================
+    def process_batch(self, batch):
+
+        peptidoform = batch["peptidoform"].to_pylist()
+        charges = batch["charge"].to_pylist()
+        mz_list = batch["mz_array"].to_pylist()
+        int_list = batch["intensity_array"].to_pylist()
+        precursor_mz = batch["precursor_mz"].to_pylist()
+        consensus_supports = batch["consensus_support"].to_pylist()
+        peps = batch["posterior_error_probability"].to_pylist()
+
+        spectra_out = []
+        seq_out = []
+        charge_out = []
+        precursor_out = []
+
+        for i in range(len(peptidoform)):
+
+            if not self.filter_by_consensus_support(consensus_supports[i]):
+                continue
+
+            if not self.filter_by_pep(peps[i]):
+                continue
+
+            mz = np.asarray(mz_list[i], dtype=np.float32)
+            intensity = np.asarray(int_list[i], dtype=np.float32)
+
+            if len(mz) == 0:
+                continue
+
+            # -------------------------
+            # top-k peaks
+            # -------------------------
+            if len(mz) > self.max_peaks:
+                idx = np.argsort(intensity)[-self.max_peaks:]
+                mz = mz[idx]
+                intensity = intensity[idx]
+
+            # -------------------------
+            # normalize
+            # -------------------------
+            max_int = intensity.max() if len(intensity) > 0 else 1.0
+            intensity = intensity / max_int
+
+            spectrum = np.stack([mz, intensity], axis=1)
+
+            spectra_out.append(spectrum.astype(np.float32))
+            seq_out.append(peptidoform[i].encode("utf-8"))  # TF需要bytes
+            charge_out.append(np.int32(charges[i]))
+            precursor_out.append(np.float32(precursor_mz[i]))
+
+        return {
+            "spectrum": spectra_out,
+            "sequence": seq_out,
+            "precursor_mz": precursor_out,
+            "charge": charge_out,
+        }

msnetloader-0.0.0/msnetloader/ms2_loader.py

@@ -0,0 +1,207 @@
+import torch
+import numpy as np
+from torch.utils.data import IterableDataset
+import duckdb
+import re
+
+
+class MS2TorchDataset(IterableDataset):
+
+    def __init__(self, parquet_path, batch_size=8, ion_types=("b", "y"), charges=(1, 2),
+                 min_consensus_support=None,
+                 max_pep=None
+                 ):
+
+        con = duckdb.connect()
+        self.min_consensus_support = min_consensus_support
+        self.max_pep = max_pep
+
+        conditions = []
+        params = []
+
+        if self.min_consensus_support is not None:
+            conditions.append("consensus_support >= ?")
+            params.append(self.min_consensus_support)
+
+        if self.max_pep is not None:
+            conditions.append("posterior_error_probability <= ?")
+            params.append(self.max_pep)
+
+        where_clause = ""
+        if conditions:
+            where_clause = "WHERE " + " AND ".join(conditions)
+
+        query = f"""
+        SELECT
+            sequence,
+            peptidoform,
+            precursor_charge AS charge,
+            cv_params.Instrument AS instrument,
+            CAST(cv_params."Collision Energy" AS DOUBLE) AS nce,
+            ion_type_array,
+            charge_array,
+            intensity_array
+        FROM parquet_scan(?)
+        {where_clause}
+        ORDER BY length(sequence)
+        """
+
+        self.cursor = con.execute(query, [parquet_path] + params)
+        self.batch_size = batch_size
+
+        self.ion_types = set(ion_types)
+        self.charges = set(charges)
+
+        self.channel_map = {
+            ("b", 1): 0,
+            ("b", 2): 1,
+            ("y", 1): 2,
+            ("y", 2): 3,
+        }
+
+        self.active_channels = [
+            self.channel_map[(t, z)]
+            for t in ion_types
+            for z in charges
+            if (t, z) in self.channel_map
+        ]
+
+    # -----------------------------
+    def __iter__(self):
+        reader = self.cursor.fetch_record_batch(self.batch_size)
+
+        for batch in reader:
+            if batch.num_rows == 0:
+                continue
+            yield self.process_batch(batch)
+
+    # -----------------------------
+    def process_batch(self, batch):
+        sequences = batch["sequence"].to_pylist()
+        peptidoform = batch["peptidoform"].to_pylist()
+        charges = batch["charge"].to_pylist()
+        nces = batch["nce"].to_pylist()
+        instruments = batch["instrument"].to_pylist()
+        fragments = batch["ion_type_array"].to_pylist()
+        fragment_charges = batch["charge_array"].to_pylist()
+        intensities = batch["intensity_array"].to_pylist()
+
+        targets = self.build_batch_fragments(
+            sequences,
+            fragments,
+            fragment_charges,
+            intensities
+        )
+
+        charge_tensor = torch.tensor(charges, dtype=torch.long)
+        nce_tensor = torch.tensor(nces, dtype=torch.float32)
+
+        return {
+            "peptide": peptidoform,
+            "charge": charge_tensor,
+            "nce": nce_tensor,
+            "instruments": instruments,
+            "targets": targets
+        }
+
+    # -----------------------------
+    def build_batch_fragments(
+        self,
+        sequences,
+        fragments_list,
+        frag_charges_list,
+        intensity_list
+    ):
+
+        B = len(sequences)
+        Lmax = max(len(s) for s in sequences)
+
+        out = np.zeros((B, Lmax - 1, 4), dtype=np.float32)
+
+        # -----------------------------
+        for b in range(B):
+
+            ions = fragments_list[b]
+            charges = frag_charges_list[b]
+            ints = intensity_list[b]
+
+            if len(ions) == 0:
+                continue
+
+            ions = np.asarray(ions)
+            charges = np.asarray(charges)
+            ints = np.asarray(ints, dtype=np.float32)
+
+            valid = (ions != None)
+            ions = ions[valid]
+            charges = charges[valid]
+            ints = ints[valid]
+
+            if len(ions) == 0:
+                continue
+
+            # remove neutral loss
+            mask = np.char.find(ions.astype(str), "-") == -1
+            ions = ions[mask]
+            charges = charges[mask]
+            ints = ints[mask]
+
+            if len(ions) == 0:
+                continue
+
+            # -----------------------------
+            # parse ion string safely
+            # -----------------------------
+            ion_str = ions.astype(str)
+
+            ion_type = np.array([x[0] for x in ion_str])
+
+            # SAFE regex position parsing
+            pos = np.array([
+                int(re.findall(r"\d+", x)[0]) if re.findall(r"\d+", x) else -1
+                for x in ion_str
+            ])
+
+            seq_len = len(sequences[b])
+
+            valid = (pos >= 1) & (pos < seq_len)
+
+            ion_type = ion_type[valid]
+            pos = pos[valid] - 1
+            charges = charges[valid]
+            ints = ints[valid]
+
+            if len(pos) == 0:
+                continue
+
+            # -----------------------------
+            # channel mapping (fixed)
+            # -----------------------------
+            ch = np.full(len(ion_type), -1, dtype=np.int32)
+
+            for (t, z), c in self.channel_map.items():
+                if t in self.ion_types and z in self.charges:
+                    ch[(ion_type == t) & (charges == z)] = c
+
+            valid_ch = ch >= 0
+
+            pos = pos[valid_ch]
+            ch = ch[valid_ch]
+            ints = ints[valid_ch]
+
+            if len(pos) == 0:
+                continue
+
+            # -----------------------------
+            # intensity normalize
+            # -----------------------------
+            max_int = ints.max() if len(ints) > 0 else 1.0
+            ints = ints / max_int
+
+            # -----------------------------
+            # scatter
+            # -----------------------------
+            out[b, pos, ch] += ints
+
+        # -----------------------------
+        return torch.from_numpy(out[:, :, self.active_channels])