Add optional pure-JAX network_preprocessing module (hyperedge discovery, pruning, motifs, sparse theta, identifiability)

config.toml

@@ -981,6 +981,57 @@ sensitivity_top_curves = 20
 sensitivity_metric = "l2_norm"
 
 
+# ------------------------------------------------------------
+# HYPEREDGE / NETWORK PREPROCESSING
+# ------------------------------------------------------------
+# Optional pure-JAX preprocessing that discovers, scores, prunes, annotates,
+# and exports kinase-site-substrate triplets before networkmodel ODE construction.
+#
+# Default behavior is disabled so existing PhosKinTime workflows remain unchanged.
+[networkmodel.hyperedge_preprocessing]
+
+# Enable optional preprocessing before Index/model construction.
+# false: keep existing networkmodel behavior exactly unchanged.
+# true : run the independent pure-JAX preprocessing module and pass its pruned
+#        kinase network to downstream model construction.
+enable_hyperedge_preprocessing = false
+
+# Subdirectory name used under standard run result folders, e.g.
+# tables/network_preprocessing and plots/network_preprocessing.
+output_subdir = "network_preprocessing"
+
+# Minimum score used during discovery/reporting. Must be >= 0.0.
+hyperedge_discovery_threshold = 0.0
+
+# Minimum score retained during pruning. Must be >= 0.0.
+hyperedge_pruning_threshold = 0.0
+
+# Detect feed-forward loops, kinase cascades, regulatory triangles, and related
+# 3-node motifs from retained triplets.
+motif_detection = true
+
+# Export sparse COO-like theta tensor indices/values as CSV and NPZ.
+sparse_tensor_export = true
+
+# Run structural identifiability preprocessing diagnostics before fitting.
+identifiability_preprocessing = true
+
+# Optional maximum number of retained triplets for memory safety.
+# Use 0 for no explicit cap. Negative values are invalid.
+max_triplets = 0
+
+# Batch size for JAX kernels and chunked preprocessing. Must be positive.
+batch_size = 65536
+
+# Save diagnostic plots for score distributions, motifs, sparsity, and
+# identifiability.
+plot_generation = true
+
+# Save labelled CSV/JSON tables for discovered/pruned triplets, motifs,
+# identifiability diagnostics, and summary statistics.
+csv_export = true
+
+
 # ------------------------------------------------------------
 # INFERENCE / POST-OPTIMIZATION ANALYSES
 # ------------------------------------------------------------

config_loader.py

@@ -177,6 +177,71 @@ class PhosKinConfig:
     # Models metadata
     available_models: tuple[str, ...] = ()
 
+    # Optional pure-JAX hyperedge/network preprocessing
+    enable_hyperedge_preprocessing: bool = False
+    hyperedge_preprocessing_output_subdir: str = "network_preprocessing"
+    hyperedge_discovery_threshold: float = 0.0
+    hyperedge_pruning_threshold: float = 0.0
+    hyperedge_motif_detection: bool = True
+    hyperedge_sparse_tensor_export: bool = True
+    hyperedge_identifiability_preprocessing: bool = True
+    hyperedge_max_triplets: int | None = None
+    hyperedge_batch_size: int = 65536
+    hyperedge_plot_generation: bool = True
+    hyperedge_csv_export: bool = True
+
+
+def _parse_hyperedge_preprocessing_config(cfg: dict[str, Any]) -> dict[str, Any]:
+    """Parse and validate optional pure-JAX hyperedge preprocessing settings."""
+    prep = cfg.get("hyperedge_preprocessing", {}) or {}
+
+    def _bool(name: str, default: bool) -> bool:
+        raw = prep.get(name, default)
+        if isinstance(raw, bool):
+            return raw
+        if isinstance(raw, str):
+            return raw.strip().lower() in {"1", "true", "yes", "y", "on"}
+        return bool(raw)
+
+    def _float_nonnegative(name: str, default: float) -> float:
+        value = float(prep.get(name, default))
+        if value < 0.0:
+            raise ValueError(f"networkmodel.hyperedge_preprocessing.{name} must be non-negative, got {value}")
+        return value
+
+    def _nonnegative_int_cap(name: str, default: int | None) -> int | None:
+        raw = prep.get(name, default)
+        if raw is None:
+            return None
+        value = int(raw)
+        if value < 0:
+            raise ValueError(f"networkmodel.hyperedge_preprocessing.{name} must be non-negative, got {value}")
+        return None if value == 0 else value
+
+    def _positive_int(name: str, default: int) -> int:
+        value = int(prep.get(name, default))
+        if value <= 0:
+            raise ValueError(f"networkmodel.hyperedge_preprocessing.{name} must be positive, got {value}")
+        return value
+
+    output_subdir = str(prep.get("output_subdir", "network_preprocessing")).strip()
+    if not output_subdir:
+        raise ValueError("networkmodel.hyperedge_preprocessing.output_subdir must be a non-empty string")
+
+    return {
+        "enable_hyperedge_preprocessing": _bool("enable_hyperedge_preprocessing", False),
+        "hyperedge_preprocessing_output_subdir": output_subdir,
+        "hyperedge_discovery_threshold": _float_nonnegative("hyperedge_discovery_threshold", 0.0),
+        "hyperedge_pruning_threshold": _float_nonnegative("hyperedge_pruning_threshold", 0.0),
+        "hyperedge_motif_detection": _bool("motif_detection", True),
+        "hyperedge_sparse_tensor_export": _bool("sparse_tensor_export", True),
+        "hyperedge_identifiability_preprocessing": _bool("identifiability_preprocessing", True),
+        "hyperedge_max_triplets": _nonnegative_int_cap("max_triplets", 0),
+        "hyperedge_batch_size": _positive_int("batch_size", 65536),
+        "hyperedge_plot_generation": _bool("plot_generation", True),
+        "hyperedge_csv_export": _bool("csv_export", True),
+    }
+
 
 def load_config_toml(path: str | Path) -> PhosKinConfig:
     path = Path(path)
@@ -313,6 +378,11 @@ def load_config_toml(path: str | Path) -> PhosKinConfig:
     sensitivity_top_curves = int(cfg.get("sensitivity_top_curves", 50))  # key matches TOML
     sensitivity_metric = str(cfg.get("sensitivity_metric", "total_signal"))
 
+    # -------------------------
+    # 12) Optional hyperedge/network preprocessing
+    # -------------------------
+    hyperedge_cfg = _parse_hyperedge_preprocessing_config(cfg)
+
     return PhosKinConfig(
         kinase_net=kinase_net,
         tf_net=tf_net,
@@ -384,4 +454,6 @@ def load_config_toml(path: str | Path) -> PhosKinConfig:
         sensitivity_metric=sensitivity_metric,
 
         available_models=available_models,
+
+        **hyperedge_cfg,
     )

network_preprocessing/__init__.py

@@ -0,0 +1,2 @@
+from .dataclasses import *
+from .api import preprocess_network, preprocess_networkmodel_frames, discover_hyperedges, prune_triplets, detect_motifs, build_sparse_theta, preprocess_identifiability

network_preprocessing/api.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+from pathlib import Path
+import logging
+import jax
+jax.config.update("jax_enable_x64", True)
+import jax.numpy as jnp
+import pandas as pd
+from .dataclasses import *
+from .io_adapters import encode_kinase_network
+from .jax_kernels import score_triplets, pruning_flags, sparse_indices_values
+
+def discover_hyperedges(encoded: EncodedNetwork, config: NetworkPreprocessingConfig) -> TripletTable:
+    score=score_triplets(encoded.edge_weight, encoded.support_count, encoded.site_observed, encoded.kinase_observed, encoded.kinase_node_ids, encoded.substrate_node_ids)
+    return TripletTable(encoded.kinase_ids, encoded.site_ids, encoded.substrate_ids, score, encoded.support_count, jnp.zeros_like(encoded.support_count, dtype=jnp.uint32))
+
+def prune_triplets(triplets: TripletTable, encoded: EncodedNetwork, config: NetworkPreprocessingConfig) -> TripletTable:
+    flags=pruning_flags(triplets.score, triplets.support_count, encoded.site_observed, encoded.kinase_observed, encoded.kinase_node_ids, encoded.substrate_node_ids, float(max(config.discovery_threshold, config.min_triplet_score)), int(config.min_support_count), bool(config.prune_self_loops), bool(config.prune_missing_observations))
+    keep=flags==0
+    if config.max_triplets is not None and int(keep.sum())>config.max_triplets:
+        kept_idx=jnp.where(keep, size=triplets.score.size, fill_value=-1)[0]
+        order=jnp.argsort(jnp.where(keep, -triplets.score, jnp.inf))[:config.max_triplets]
+        mask=jnp.zeros_like(keep).at[order].set(True); keep=keep & mask
+    return TripletTable(triplets.kinase_ids[keep], triplets.site_ids[keep], triplets.substrate_ids[keep], triplets.score[keep], triplets.support_count[keep], flags[keep])
+
+def build_sparse_theta(triplets: TripletTable, shape: tuple[int,int,int]) -> SparseThetaTensor:
+    idx,val=sparse_indices_values(triplets.kinase_ids, triplets.site_ids, triplets.substrate_ids, triplets.score)
+    return SparseThetaTensor(idx,val,shape)
+
+def _common_node_edges(triplets: TripletTable, encoded: EncodedNetwork) -> tuple[list[int], list[int], list[float], tuple[str, ...]]:
+    """Map motif edges from independent kinase/substrate IDs into one protein namespace."""
+    node_labels = tuple(sorted(set(encoded.kinase_labels) | set(encoded.substrate_labels)))
+    node_id = {label: i for i, label in enumerate(node_labels)}
+    src = [node_id[encoded.kinase_labels[int(k)]] for k in jnp.asarray(triplets.kinase_ids)]
+    dst = [node_id[encoded.substrate_labels[int(s)]] for s in jnp.asarray(triplets.substrate_ids)]
+    score = list(map(float, jnp.asarray(triplets.score)))
+    return src, dst, score, node_labels
+
+def detect_motifs(triplets: TripletTable, config: NetworkPreprocessingConfig, encoded: EncodedNetwork | None = None) -> MotifTable:
+    # bounded host orchestration, core edge testing array based; no dense n^3 enumeration
+    if encoded is None:
+        src=list(map(int, jnp.asarray(triplets.kinase_ids))); dst=list(map(int, jnp.asarray(triplets.substrate_ids))); node_labels=()
+    else:
+        src,dst,_,node_labels=_common_node_edges(triplets, encoded)
+    sc=list(map(float, jnp.asarray(triplets.score)))
+    edge={(a,b):s for a,b,s in zip(src,dst,sc) if a!=b}; rows=[]
+    for (a,b),sab in edge.items():
+        for (bb,c),sbc in edge.items():
+            if bb!=b or c in (a,b): continue
+            if (a,c) in edge:
+                rows.append((1,a,b,c,0b111,(sab*sbc*edge[(a,c)])**(1/3)))
+                if len(rows)>=config.max_motifs: break
+        if len(rows)>=config.max_motifs: break
+    if not rows:
+        return MotifTable(jnp.empty(0,jnp.int16),jnp.empty(0,jnp.int32),jnp.empty(0,jnp.int32),jnp.empty(0,jnp.int32),jnp.empty(0,jnp.uint8),jnp.empty(0,jnp.float64),node_labels)
+    arr=jnp.asarray(rows)
+    return MotifTable(arr[:,0].astype(jnp.int16),arr[:,1].astype(jnp.int32),arr[:,2].astype(jnp.int32),arr[:,3].astype(jnp.int32),arr[:,4].astype(jnp.uint8),arr[:,5].astype(jnp.float64),node_labels)
+
+def preprocess_identifiability(theta: SparseThetaTensor, config: NetworkPreprocessingConfig) -> IdentifiabilityDiagnostics:
+    from .jax_kernels import identifiability_kernel
+    retained,gid,norms,red=identifiability_kernel(theta.indices, theta.values)
+    rank=int(jnp.unique(gid, size=gid.size, fill_value=-1).size) if gid.size else 0
+    return IdentifiabilityDiagnostics(retained,gid,rank,red,norms)
+
+def preprocess_network(kinase_network: pd.DataFrame, *, phospho_observations=None, protein_observations=None, rna_observations=None, tf_network=None, config=None, output_dir=None, logger=None):
+    config=config or NetworkPreprocessingConfig(); logger=logger or logging.getLogger(__name__)
+    encoded, grouped=encode_kinase_network(kinase_network, phospho_observations, protein_observations)
+    discovered=discover_hyperedges(encoded, config); pruned=prune_triplets(discovered, encoded, config)
+    theta=build_sparse_theta(pruned, (len(encoded.kinase_labels), len(encoded.site_labels), len(encoded.substrate_labels)))
+    motifs=detect_motifs(pruned, config, encoded) if config.enable_motifs else None
+    ident=preprocess_identifiability(theta, config) if config.enable_identifiability else None
+    summary={"n_discovered": int(discovered.score.size), "n_retained": int(pruned.score.size), "n_pruned": int(discovered.score.size-pruned.score.size), "n_motifs": int(0 if motifs is None else motifs.score.size), "tensor_nnz": int(theta.values.size)}
+    res=NetworkPreprocessingResult(encoded, discovered, pruned, theta, motifs, ident, summary)
+    if output_dir is not None:
+        if config.export_csv or config.export_sparse_tensor:
+            from .export import export_preprocessing_result
+            export_preprocessing_result(
+                res, output_dir, output_subdir=config.output_subdir,
+                include_csv=config.export_csv, include_sparse_tensor=config.export_sparse_tensor,
+            )
+        if config.generate_plots:
+            from .plotting import plot_preprocessing_result
+            plot_preprocessing_result(res, output_dir, output_subdir=config.output_subdir)
+        logger.info("[NetworkPreprocessing] outputs saved under %s", Path(output_dir)/config.output_subdir)
+    return res
+
+def _pruned_networkmodel_frame(result: NetworkPreprocessingResult) -> pd.DataFrame:
+    enc = result.encoded
+    retained = {
+        (int(k), int(s), int(sub))
+        for k, s, sub in zip(result.pruned.kinase_ids, result.pruned.site_ids, result.pruned.substrate_ids)
+    }
+    rows = []
+    for k, s, sub, alpha, support in zip(
+        enc.kinase_ids, enc.site_ids, enc.substrate_ids, enc.edge_weight, enc.support_count
+    ):
+        key = (int(k), int(s), int(sub))
+        if key not in retained:
+            continue
+        site_label = enc.site_labels[key[1]]
+        protein, psite = site_label.split(":", 1)
+        rows.append({
+            "protein": protein,
+            "psite": psite,
+            "kinase": enc.kinase_labels[key[0]],
+            "alpha": float(alpha),
+            "support_count": int(support),
+        })
+    return pd.DataFrame(rows, columns=["protein", "psite", "kinase", "alpha", "support_count"])
+
+
+def preprocess_networkmodel_frames(df_kin, df_tf, df_prot, df_pho, df_rna, *, config=None, output_dir=None, logger=None):
+    res=preprocess_network(df_kin, phospho_observations=df_pho, protein_observations=df_prot, rna_observations=df_rna, tf_network=df_tf, config=config, output_dir=output_dir, logger=logger)
+    return _pruned_networkmodel_frame(res), res

network_preprocessing/dataclasses.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+from dataclasses import dataclass, field
+from typing import Any, Mapping
+import jax
+import jax.numpy as jnp
+
+@dataclass(frozen=True)
+class NetworkPreprocessingConfig:
+    min_triplet_score: float = 0.0
+    discovery_threshold: float = 0.0
+    min_support_count: int = 1
+    max_triplets: int | None = None
+    batch_size: int = 65536
+    dtype: str = "float64"
+    enable_motifs: bool = True
+    enable_identifiability: bool = True
+    prune_self_loops: bool = True
+    prune_missing_observations: bool = False
+    output_subdir: str = "network_preprocessing"
+    export_sparse_tensor: bool = True
+    generate_plots: bool = True
+    export_csv: bool = True
+    dense_motif_node_limit: int = 256
+    max_motifs: int = 10000
+    score_weights: Mapping[str, float] = field(default_factory=dict)
+
+    @staticmethod
+    def from_args(args: Any) -> "NetworkPreprocessingConfig":
+        return NetworkPreprocessingConfig(
+            min_triplet_score=float(getattr(args, "network_preprocessing_min_score", 0.0)),
+            discovery_threshold=float(getattr(args, "network_preprocessing_discovery_threshold", 0.0)),
+            min_support_count=int(getattr(args, "network_preprocessing_min_support", 1)),
+            max_triplets=(None if getattr(args, "network_preprocessing_max_triplets", None) == 0 else getattr(args, "network_preprocessing_max_triplets", None)),
+            batch_size=int(getattr(args, "network_preprocessing_batch_size", 65536)),
+            enable_motifs=bool(getattr(args, "network_preprocessing_enable_motifs", True)),
+            enable_identifiability=bool(getattr(args, "network_preprocessing_enable_identifiability", True)),
+            prune_self_loops=bool(getattr(args, "network_preprocessing_prune_self_loops", True)),
+            prune_missing_observations=bool(getattr(args, "network_preprocessing_prune_missing_observations", False)),
+            output_subdir=str(getattr(args, "network_preprocessing_output_subdir", "network_preprocessing")),
+            export_sparse_tensor=bool(getattr(args, "network_preprocessing_export_sparse_tensor", True)),
+            generate_plots=bool(getattr(args, "network_preprocessing_generate_plots", True)),
+            export_csv=bool(getattr(args, "network_preprocessing_export_csv", True)),
+        )
+
+@dataclass(frozen=True)
+class EncodedNetwork:
+    kinase_ids: jax.Array
+    substrate_ids: jax.Array
+    kinase_node_ids: jax.Array
+    substrate_node_ids: jax.Array
+    site_ids: jax.Array
+    edge_weight: jax.Array
+    support_count: jax.Array
+    site_observed: jax.Array
+    kinase_observed: jax.Array
+    kinase_labels: tuple[str, ...]
+    substrate_labels: tuple[str, ...]
+    site_labels: tuple[str, ...]
+
+@dataclass(frozen=True)
+class TripletTable:
+    kinase_ids: jax.Array; site_ids: jax.Array; substrate_ids: jax.Array
+    score: jax.Array; support_count: jax.Array; flags: jax.Array
+
+@dataclass(frozen=True)
+class SparseThetaTensor:
+    indices: jax.Array; values: jax.Array; shape: tuple[int, int, int]
+
+@dataclass(frozen=True)
+class MotifTable:
+    motif_type: jax.Array; node_a: jax.Array; node_b: jax.Array; node_c: jax.Array
+    edge_mask: jax.Array; score: jax.Array
+    node_labels: tuple[str, ...] = ()
+
+@dataclass(frozen=True)
+class IdentifiabilityDiagnostics:
+    retained_param_mask: jax.Array; group_id: jax.Array; local_rank_estimate: int
+    redundancy_score: jax.Array; design_column_norm: jax.Array
+
+@dataclass(frozen=True)
+class NetworkPreprocessingResult:
+    encoded: EncodedNetwork; discovered: TripletTable; pruned: TripletTable
+    theta: SparseThetaTensor; motifs: MotifTable | None
+    identifiability: IdentifiabilityDiagnostics | None; summary: Mapping[str, Any]