eclipse-ms 0.1.2__py3-none-any.whl

eclipse_ms/__init__.py

@@ -0,0 +1,45 @@
+"""ECLIPSE: conditional spectrum autoencoder + clustering for MS/MS.
+
+Importing this package does NOT import TensorFlow; TF is loaded lazily the
+first time you build or load a model (e.g. ``load_encoder``) or call
+``embed_spectra``. This keeps imports fast and lets the clustering / consensus
+utilities be used in TF-free environments.
+
+The Keras model classes live in ``eclipse_ms.models`` (importing that submodule
+does import TensorFlow).
+"""
+
+from .config import COND_DIM, Config
+from .preprocessing import bin_spectrum_numpy, build_cond_vector, preprocess
+from .cluster import cluster_latents, score_clusters
+from .consensus import generate_consensus_spectrum, write_mzml
+from .modelhub import (
+    REGISTRY,
+    cache_dir,
+    get_model_file,
+    load_autoencoder,
+    load_encoder,
+)
+from .embed import embed_raw_spectra, embed_spectra
+
+__version__ = "0.1.2"
+
+__all__ = [
+    "__version__",
+    "Config",
+    "COND_DIM",
+    "bin_spectrum_numpy",
+    "build_cond_vector",
+    "preprocess",
+    "load_encoder",
+    "load_autoencoder",
+    "get_model_file",
+    "cache_dir",
+    "REGISTRY",
+    "embed_spectra",
+    "embed_raw_spectra",
+    "cluster_latents",
+    "score_clusters",
+    "generate_consensus_spectrum",
+    "write_mzml",
+]

eclipse_ms/cli.py

@@ -0,0 +1,107 @@
+"""ECLIPSE command-line interface.
+
+Subcommands:
+  embed      Bin + encode spectra from parquet to a latents .npy
+  cluster    Cluster a latents .npy into cluster labels
+  consensus  Build consensus spectra (mzML) from clusters
+
+Training and HPC data-prep live in the repo's ``training/`` scripts, not in the
+installed package.
+"""
+
+from __future__ import annotations
+
+import argparse
+import glob
+import json
+import os
+
+import numpy as np
+
+
+def _cmd_embed(args):
+    import pandas as pd
+
+    from .config import Config
+    from .embed import embed_raw_spectra
+    from .modelhub import load_encoder
+
+    encoder = load_encoder(weights=args.weights, config=args.config)
+
+    files = sorted(glob.glob(os.path.join(args.input, "*.parquet")))
+    print(f"Found {len(files)} parquet files")
+
+    mz, inten, pmz, charge, im = [], [], [], [], []
+    for fp in files:
+        df = pd.read_parquet(fp)
+        for _, row in df.iterrows():
+            mz.append(row["mz_array"])
+            inten.append(row["intensity_array"])
+            pmz.append(float(row.get("precursor_mz", 0.0)))
+            charge.append(int(row.get("precursor_charge", 2)))
+            im.append(float(row.get("ion_mobility", 0.0)))
+            if args.max_spectra and len(mz) >= args.max_spectra:
+                break
+        if args.max_spectra and len(mz) >= args.max_spectra:
+            break
+
+    latents = embed_raw_spectra(encoder, mz, inten, pmz, charge, im, Config, args.batch_size)
+    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
+    np.save(args.output, latents)
+    print(f"Saved {latents.shape} latents to {args.output}")
+
+
+def _cmd_cluster(args):
+    from .cluster import cluster_latents, score_clusters
+
+    latents = np.load(args.input)
+    labels, info = cluster_latents(
+        latents, method=args.method, min_cluster_size=args.min_cluster_size
+    )
+    os.makedirs(args.output, exist_ok=True)
+    np.save(os.path.join(args.output, "cluster_labels.npy"), labels)
+    with open(os.path.join(args.output, "cluster_info.json"), "w") as f:
+        json.dump(info, f, indent=2)
+    scores = score_clusters(latents, labels)
+    scores.to_csv(os.path.join(args.output, "cluster_scores.csv"), index=False)
+    print(f"Clusters: {info.get('n_clusters')}, noise: {info.get('n_noise', 0)}")
+    print(f"Wrote labels, info, and scores to {args.output}")
+
+
+def _cmd_consensus(args):
+    print(
+        "Consensus generation needs spectra grouped by cluster. See "
+        "eclipse_ms.consensus.generate_consensus_spectrum and the example in "
+        "training/consensus_reference.py for the full pipeline."
+    )
+
+
+def main(argv=None):
+    p = argparse.ArgumentParser(prog="eclipse", description="ECLIPSE")
+    sub = p.add_subparsers(dest="command", required=True)
+
+    pe = sub.add_parser("embed", help="Encode spectra to latents")
+    pe.add_argument("-i", "--input", required=True, help="Parquet directory")
+    pe.add_argument("-o", "--output", required=True, help="Output latents .npy")
+    pe.add_argument("--weights", default=None, help="Local encoder weights (.h5)")
+    pe.add_argument("--config", default=None, help="Local encoder config (.json)")
+    pe.add_argument("--batch-size", type=int, default=256)
+    pe.add_argument("--max-spectra", type=int, default=None)
+    pe.set_defaults(func=_cmd_embed)
+
+    pc = sub.add_parser("cluster", help="Cluster latents")
+    pc.add_argument("-i", "--input", required=True, help="latents .npy")
+    pc.add_argument("-o", "--output", required=True, help="Output directory")
+    pc.add_argument("--method", choices=["hdbscan", "kmeans"], default="hdbscan")
+    pc.add_argument("--min-cluster-size", type=int, default=5)
+    pc.set_defaults(func=_cmd_cluster)
+
+    pk = sub.add_parser("consensus", help="Consensus spectra from clusters")
+    pk.set_defaults(func=_cmd_consensus)
+
+    args = p.parse_args(argv)
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

eclipse_ms/cluster.py

@@ -0,0 +1,105 @@
+"""Cluster latent vectors and score the resulting clusters."""
+
+from __future__ import annotations
+
+import time
+from typing import Tuple
+
+import numpy as np
+
+
+def cluster_latents(
+    latents: np.ndarray,
+    method: str = "hdbscan",
+    min_cluster_size: int = 5,
+    pca_dims: int = 100,
+    random_state: int = 42,
+) -> Tuple[np.ndarray, dict]:
+    """Cluster latent vectors with HDBSCAN (preferred) or MiniBatchKMeans.
+
+    Reduces dimensionality with PCA first when ``n_dims > pca_dims``. Falls back
+    to KMeans if HDBSCAN is not installed.
+
+    Returns ``(labels, info)`` where ``labels`` is ``-1`` for HDBSCAN noise.
+    """
+    from sklearn.decomposition import PCA
+
+    n_samples, n_dims = latents.shape
+    info = {"method": method, "n_samples": int(n_samples)}
+
+    if n_dims > pca_dims:
+        pca = PCA(n_components=pca_dims, random_state=random_state)
+        reduced = pca.fit_transform(latents)
+        info["pca_variance_explained"] = float(pca.explained_variance_ratio_.sum())
+    else:
+        reduced = latents
+
+    if method == "hdbscan":
+        try:
+            import hdbscan
+        except ImportError:
+            print("hdbscan not installed (`pip install eclipse-ms[hdbscan]`); using KMeans.")
+            method = "kmeans"
+
+    if method == "hdbscan":
+        import hdbscan
+
+        start = time.time()
+        clusterer = hdbscan.HDBSCAN(
+            min_cluster_size=min_cluster_size,
+            min_samples=3,
+            metric="euclidean",
+            cluster_selection_method="eom",
+            core_dist_n_jobs=-1,
+        )
+        labels = clusterer.fit_predict(reduced)
+        info["time"] = time.time() - start
+        info["n_clusters"] = len(set(labels)) - (1 if -1 in labels else 0)
+        info["n_noise"] = int((labels == -1).sum())
+
+    elif method == "kmeans":
+        from sklearn.cluster import MiniBatchKMeans
+
+        n_clusters = max(2, min(10000, n_samples // 10))
+        start = time.time()
+        kmeans = MiniBatchKMeans(
+            n_clusters=n_clusters, batch_size=1024, random_state=random_state, n_init=3
+        )
+        labels = kmeans.fit_predict(reduced)
+        info["time"] = time.time() - start
+        info["n_clusters"] = len(set(labels))
+        info["n_noise"] = 0
+    else:
+        raise ValueError(f"Unknown method: {method}")
+
+    return labels, info
+
+
+def score_clusters(latents: np.ndarray, labels: np.ndarray) -> "pd.DataFrame":  # noqa: F821
+    """Lightweight per-cluster quality scores from latent geometry.
+
+    For each non-noise cluster, reports size and intra-cluster cohesion
+    (mean distance to centroid; smaller = tighter).
+    """
+    import pandas as pd
+
+    rows = []
+    for c in sorted(set(labels)):
+        if c == -1:
+            continue
+        idx = np.where(labels == c)[0]
+        pts = latents[idx]
+        centroid = pts.mean(axis=0)
+        dists = np.linalg.norm(pts - centroid, axis=1)
+        rows.append(
+            {
+                "cluster": int(c),
+                "size": int(len(idx)),
+                "cohesion_mean_dist": float(dists.mean()),
+                "cohesion_std_dist": float(dists.std()),
+            }
+        )
+    df = pd.DataFrame(rows)
+    if not df.empty:
+        df = df.sort_values("size", ascending=False).reset_index(drop=True)
+    return df

eclipse_ms/config.py

@@ -0,0 +1,53 @@
+"""Centralised configuration for ECLIPSE.
+
+These values define the spectrum binning and the autoencoder architecture. They
+MUST match the configuration used to train the published weights; the binning
+parameters in particular are baked into the model input.
+"""
+
+
+class Config:
+    """Centralized configuration."""
+
+    SEED = 42
+
+    # Binning parameters (autoencoder input)
+    MZ_MIN = 100.0
+    MZ_MAX = 1700.0
+    BIN_SIZE = 0.5
+    N_BINS = int((MZ_MAX - MZ_MIN) / BIN_SIZE)  # 3200
+
+    # Preprocessing
+    RELATIVE_INTENSITY_THRESHOLD = 0.01
+    TOP_N_PEAKS = 100  # keep only the top N most intense peaks (None to disable)
+
+    # Ion mobility normalisation (1/K0 range for peptides)
+    IM_MIN = 0.6
+    IM_MAX = 1.6
+
+    # Precursor features
+    MAX_CHARGE = 6
+    PRECURSOR_MZ_MAX = 1700.0
+
+    # Autoencoder architecture
+    LATENT_DIM = 256
+    AE_EMBED_DIM = 256
+    AE_NUM_HEADS = 8
+    AE_NUM_LAYERS = 4
+    AE_FF_DIM = 512
+    AE_PATCH_SIZE = 16
+    AE_DROPOUT = 0.1
+
+    # Conditioning embedding
+    COND_EMBED_DIM = 256
+
+    # Training - Autoencoder
+    AE_BATCH_SIZE = 256
+    AE_INITIAL_LR = 1e-4
+    AE_EPOCHS = 50
+    AE_WARMUP_EPOCHS = 5
+    WEIGHT_DECAY = 1e-5
+
+
+# Conditioning dimension: one-hot charge (MAX_CHARGE) + precursor_mz (1) + IM (1).
+COND_DIM = Config.MAX_CHARGE + 2

eclipse_ms/consensus.py

@@ -0,0 +1,148 @@
+"""Consensus (averaged) spectrum generation for clusters.
+
+The consensus algorithm is ported from the project's consensus script; the full
+original (with the richer mzML writer and FragPipe-specific helpers) is kept in
+``training/consensus_reference.py`` for reference.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+import numpy as np
+
+
+def generate_consensus_spectrum(
+    spectra: List[Dict],
+    mz_tolerance: float = 0.01,
+    min_peak_occurrence: float = 0.3,
+) -> Optional[Dict]:
+    """Average multiple spectra into a single consensus spectrum.
+
+    Peaks within ``mz_tolerance`` (Da) are binned; a bin is kept if it occurs in
+    at least ``min_peak_occurrence`` of the spectra. Each input dict must have
+    ``mz_array``, ``intensity_array``, ``precursor_mz``, ``charge``,
+    ``ion_mobility``.
+
+    Returns a consensus dict, or ``None`` if nothing passes the filters.
+    """
+    if not spectra:
+        return None
+    n_spectra = len(spectra)
+
+    all_peaks = []
+    for spec in spectra:
+        mz = np.asarray(spec["mz_array"], dtype=float)
+        intensity = np.asarray(spec["intensity_array"], dtype=float)
+        if len(intensity) > 0 and intensity.max() > 0:
+            intensity = intensity / intensity.max()
+        all_peaks.extend(zip(mz, intensity))
+
+    if not all_peaks:
+        return None
+    all_peaks.sort(key=lambda x: x[0])
+
+    bins: List[List] = []
+    current = [all_peaks[0]]
+    for peak in all_peaks[1:]:
+        if peak[0] - current[0][0] <= mz_tolerance:
+            current.append(peak)
+        else:
+            bins.append(current)
+            current = [peak]
+    bins.append(current)
+
+    consensus_mz, consensus_intensity = [], []
+    for bin_peaks in bins:
+        occurrence = min(len(bin_peaks), n_spectra) / n_spectra
+        if occurrence >= min_peak_occurrence:
+            consensus_mz.append(np.mean([p[0] for p in bin_peaks]))
+            consensus_intensity.append(np.mean([p[1] for p in bin_peaks]))
+
+    if not consensus_mz:
+        return None
+
+    consensus_mz = np.array(consensus_mz)
+    consensus_intensity = np.array(consensus_intensity)
+    if consensus_intensity.max() > 0:
+        consensus_intensity = consensus_intensity / consensus_intensity.max() * 10000
+
+    order = np.argsort(consensus_mz)
+    consensus_mz = consensus_mz[order]
+    consensus_intensity = consensus_intensity[order]
+
+    precursor_mz = float(np.mean([s["precursor_mz"] for s in spectra]))
+    ims = [s["ion_mobility"] for s in spectra if s.get("ion_mobility", 0) > 0]
+    ion_mobility = float(np.mean(ims)) if ims else 0.0
+    charges = [
+        s["charge"] for s in spectra if s.get("charge") and not np.isnan(s["charge"])
+    ]
+    charge = int(np.median(charges)) if charges else 2
+
+    return {
+        "mz_array": consensus_mz,
+        "intensity_array": consensus_intensity,
+        "precursor_mz": precursor_mz,
+        "charge": charge,
+        "ion_mobility": ion_mobility,
+        "n_spectra": n_spectra,
+        "n_peaks": len(consensus_mz),
+    }
+
+
+def write_mzml(consensus_spectra: List[Dict], output_path: str) -> None:
+    """Write consensus spectra to a minimal, valid mzML file.
+
+    A compact writer covering MS2 centroid spectra with a single precursor.
+    For the full-featured writer (instrument metadata, etc.), see
+    ``training/consensus_reference.py``.
+    """
+    import base64
+    import struct
+    from datetime import datetime, timezone
+
+    def _b64(arr, dtype):
+        packed = struct.pack(f"<{len(arr)}{'d' if dtype == 64 else 'f'}", *arr)
+        return base64.b64encode(packed).decode("ascii")
+
+    lines = [
+        '<?xml version="1.0" encoding="utf-8"?>',
+        '<mzML xmlns="http://psi.hupo.org/ms/mzml" version="1.1.0">',
+        "  <run id=\"eclipse_consensus\" "
+        f'startTimeStamp="{datetime.now(timezone.utc).isoformat()}">',
+        f'    <spectrumList count="{len(consensus_spectra)}">',
+    ]
+    for idx, s in enumerate(consensus_spectra):
+        mz = np.asarray(s["mz_array"], dtype=np.float64)
+        inten = np.asarray(s["intensity_array"], dtype=np.float32)
+        mz_b64 = _b64(mz, 64)
+        in_b64 = _b64(inten, 32)
+        lines += [
+            f'      <spectrum index="{idx}" id="scan={idx + 1}" '
+            f'defaultArrayLength="{len(mz)}">',
+            '        <cvParam cvRef="MS" accession="MS:1000580" name="MSn spectrum"/>',
+            '        <cvParam cvRef="MS" accession="MS:1000511" name="ms level" '
+            'value="2"/>',
+            "        <precursorList count=\"1\"><precursor><selectedIonList count=\"1\">"
+            "<selectedIon>",
+            '          <cvParam cvRef="MS" accession="MS:1000744" '
+            f'name="selected ion m/z" value="{s["precursor_mz"]}"/>',
+            '          <cvParam cvRef="MS" accession="MS:1000041" '
+            f'name="charge state" value="{s["charge"]}"/>',
+            "        </selectedIon></selectedIonList></precursor></precursorList>",
+            '        <binaryDataArrayList count="2">',
+            f'          <binaryDataArray encodedLength="{len(mz_b64)}">'
+            '<cvParam cvRef="MS" accession="MS:1000523" name="64-bit float"/>'
+            '<cvParam cvRef="MS" accession="MS:1000514" name="m/z array"/>'
+            f"<binary>{mz_b64}</binary></binaryDataArray>",
+            f'          <binaryDataArray encodedLength="{len(in_b64)}">'
+            '<cvParam cvRef="MS" accession="MS:1000521" name="32-bit float"/>'
+            '<cvParam cvRef="MS" accession="MS:1000515" name="intensity array"/>'
+            f"<binary>{in_b64}</binary></binaryDataArray>",
+            "        </binaryDataArrayList>",
+            "      </spectrum>",
+        ]
+    lines += ["    </spectrumList>", "  </run>", "</mzML>"]
+    with open(output_path, "w") as f:
+        f.write("\n".join(lines))
+    print(f"Wrote {len(consensus_spectra)} consensus spectra to {output_path}")

eclipse_ms/embed.py

@@ -0,0 +1,96 @@
+"""Embed spectra into the autoencoder latent space."""
+
+from __future__ import annotations
+
+import time
+from typing import Optional
+
+import numpy as np
+
+from .config import Config
+from .preprocessing import bin_spectrum_numpy, build_cond_vector
+
+
+def embed_spectra(
+    encoder,
+    spectra: np.ndarray,
+    conditioning: np.ndarray,
+    batch_size: int = 256,
+    latent_dim: Optional[int] = None,
+    verbose: bool = True,
+) -> np.ndarray:
+    """Encode pre-binned spectra to latent vectors in batches.
+
+    Args:
+        encoder: a loaded encoder (see :func:`eclipse_ms.modelhub.load_encoder`)
+            or a full autoencoder exposing ``.encode``.
+        spectra: array ``[n, N_BINS]`` of binned spectra (float32).
+        conditioning: array ``[n, cond_dim]`` of conditioning vectors.
+        batch_size: encode this many spectra at a time.
+        latent_dim: output dimension; inferred from the first batch if None.
+
+    Returns:
+        Array ``[n, latent_dim]`` of latent vectors.
+    """
+    import tensorflow as tf
+
+    n = len(spectra)
+    if n != len(conditioning):
+        raise ValueError(f"spectra and conditioning differ in length: {n} vs {len(conditioning)}")
+
+    spectra = np.asarray(spectra, dtype=np.float32)
+    conditioning = np.asarray(conditioning, dtype=np.float32)
+
+    def _encode(x, cond):
+        # Support both an Encoder (callable) and a full AE (has .encode).
+        if hasattr(encoder, "encode"):
+            z = encoder.encode(x, cond, training=False)
+            if isinstance(z, (tuple, list)):
+                z = z[0]
+            return z
+        return encoder((x, cond), training=False)
+
+    if latent_dim is None:
+        probe = _encode(tf.convert_to_tensor(spectra[:1]), tf.convert_to_tensor(conditioning[:1]))
+        latent_dim = int(probe.shape[-1])
+
+    latents = np.zeros((n, latent_dim), dtype=np.float32)
+    start = time.time()
+    for i in range(0, n, batch_size):
+        end = min(i + batch_size, n)
+        z = _encode(
+            tf.convert_to_tensor(spectra[i:end]),
+            tf.convert_to_tensor(conditioning[i:end]),
+        )
+        latents[i:end] = np.asarray(z)
+        if verbose and (i // batch_size + 1) % 100 == 0:
+            rate = end / (time.time() - start)
+            print(f"  encoded {end:,}/{n:,} ({rate:.0f}/s)")
+    if verbose:
+        print(f"  done in {(time.time() - start) / 60:.1f} min")
+    return latents
+
+
+def embed_raw_spectra(
+    encoder,
+    mz_list,
+    intensity_list,
+    precursor_mz,
+    charge,
+    ion_mobility,
+    config=Config,
+    batch_size: int = 256,
+    verbose: bool = True,
+) -> np.ndarray:
+    """Bin + condition raw peak lists, then embed.
+
+    Each of ``precursor_mz``, ``charge``, ``ion_mobility`` is a sequence aligned
+    with ``mz_list`` / ``intensity_list``.
+    """
+    n = len(mz_list)
+    spectra = np.zeros((n, config.N_BINS), dtype=np.float32)
+    cond = np.zeros((n, config.MAX_CHARGE + 2), dtype=np.float32)
+    for i in range(n):
+        spectra[i] = bin_spectrum_numpy(mz_list[i], intensity_list[i], config)
+        cond[i] = build_cond_vector(precursor_mz[i], charge[i], ion_mobility[i], config)
+    return embed_spectra(encoder, spectra, cond, batch_size=batch_size, verbose=verbose)

eclipse_ms/layers.py

@@ -0,0 +1,84 @@
+"""Keras layer building blocks for the ECLIPSE autoencoder.
+
+Ported verbatim from the training code so the registered serialisable layers
+reconstruct with exactly the same weight structure when loading published
+weights.
+"""
+
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras import layers
+
+
+@keras.utils.register_keras_serializable()
+class PatchEmbedding(layers.Layer):
+    """Convert a 1D spectrum into patch embeddings."""
+
+    def __init__(self, embed_dim: int = 256, patch_size: int = 16, **kwargs):
+        super().__init__(**kwargs)
+        self.embed_dim = embed_dim
+        self.patch_size = patch_size
+        self.projection = layers.Dense(embed_dim)
+
+    def call(self, x):
+        batch_size = tf.shape(x)[0]
+        x = tf.reshape(x, [batch_size, -1, self.patch_size])
+        return self.projection(x)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"embed_dim": self.embed_dim, "patch_size": self.patch_size})
+        return config
+
+
+@keras.utils.register_keras_serializable()
+class TransformerBlock(layers.Layer):
+    """Pre-norm transformer block."""
+
+    def __init__(self, embed_dim=256, num_heads=8, ff_dim=512, dropout=0.1, **kwargs):
+        super().__init__(**kwargs)
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.ff_dim = ff_dim
+        self.dropout_rate = dropout
+
+        self.att = layers.MultiHeadAttention(
+            num_heads=num_heads,
+            key_dim=embed_dim // num_heads,
+            dropout=dropout,
+        )
+
+        self.ffn = keras.Sequential(
+            [
+                layers.Dense(ff_dim, activation="gelu"),
+                layers.Dropout(dropout),
+                layers.Dense(embed_dim),
+                layers.Dropout(dropout),
+            ]
+        )
+
+        self.norm1 = layers.LayerNormalization(epsilon=1e-6)
+        self.norm2 = layers.LayerNormalization(epsilon=1e-6)
+
+    def call(self, x, training=False):
+        x_norm = self.norm1(x)
+        attn_out = self.att(x_norm, x_norm, training=training)
+        x = x + attn_out
+
+        x_norm = self.norm2(x)
+        ffn_out = self.ffn(x_norm, training=training)
+        x = x + ffn_out
+
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "embed_dim": self.embed_dim,
+                "num_heads": self.num_heads,
+                "ff_dim": self.ff_dim,
+                "dropout": self.dropout_rate,
+            }
+        )
+        return config