from __future__ import annotations import json import logging from pathlib import Path from typing import Any import numpy as np from rdkit import Chem, RDLogger from rdkit.Chem import AllChem, Descriptors from .registry import AromaRegistry from .thermo.geometry import extract_3d_shape_features from .thermo.naturals import OPENPOM_DIM, SHAPE_DIM, resolve_natural_oil_vectors logger = logging.getLogger("pino.embeddings") class RealOpenPOMUnavailableError(RuntimeError): """Raised when a caller requires real POM features but no model is loaded.""" STRUCTURAL_DIM = OPENPOM_DIM # 138 (Morgan fallback structural block) REAL_POM_DIM = 256 # genuine OpenPOM v1.0.0 penultimate embedding (ablation-proven) SHAPE_DIM = SHAPE_DIM # 7 PHYSICAL_DIM = 2 # molecular_weight, logp # --- morgan_2048_rp arm (round-7 panel response): 2048-bit radius-2 Morgan + # 10 physicochemical descriptors = 2058-D structural block; full token 2071-D; # fixed seeded random projection down to the OpenPOM arm's combined width so # both trained arms carry an IDENTICAL input-projection parameter count # (parameter-parity control requested by the evaluation reviewer). MORGAN2048_STRUCTURAL_DIM = 2058 MORGAN2048_TOKEN_DIM = MORGAN2048_STRUCTURAL_DIM + SHAPE_DIM + PHYSICAL_DIM + 4 # 2071 MORGAN2048_RP_SEED = 20260725 FUNCTIONAL_DIM = 4 # ester, alcohol, ketone, aldehyde COMBINED_DIM = STRUCTURAL_DIM + SHAPE_DIM + PHYSICAL_DIM + FUNCTIONAL_DIM # 151 COMBINED_DIM_REAL_POM = REAL_POM_DIM + SHAPE_DIM + PHYSICAL_DIM + FUNCTIONAL_DIM # 269 _REPO_ROOT = Path(__file__).resolve().parents[2] _POM_DIR = _REPO_ROOT / "artifacts" / "pom" _DISJOINT_DIR = _REPO_ROOT / "artifacts" / "pom_disjoint" # Mute RDKit's internal C++ warning streams globally. rd_logger = RDLogger.logger() rd_logger.setLevel(RDLogger.CRITICAL) # UNIFAC subgroup ID mapping to functional families (Dortmund-UNIFAC convention). # These IDs are stable in the thermo.unifac DOUFSG parameter set. FUNCTIONAL_GROUP_IDS = { "ester": {21, 22, 23}, "alcohol": {14, 15, 16}, "ketone": {18, 19}, "aldehyde": {20}, } def _functional_group_flags(unifac_groups: dict[str, Any] | str | None) -> np.ndarray: """Return a 4-D binary vector for ester, alcohol, ketone, aldehyde.""" flags = np.zeros(FUNCTIONAL_DIM, dtype=np.float32) if not unifac_groups: return flags if isinstance(unifac_groups, str): try: unifac_groups = json.loads(unifac_groups) except Exception: return flags group_ids = set() for key in unifac_groups: try: group_ids.add(int(key)) except ValueError: # If keys are textual names (legacy), look for matching subgroup names. k = str(key).lower() if "coo" in k or "ester" in k: flags[0] = 1.0 if "oh" in k: flags[1] = 1.0 if "ketone" in k or "co" == k: flags[2] = 1.0 if "cho" in k or "aldehyde" in k: flags[3] = 1.0 return flags if group_ids & FUNCTIONAL_GROUP_IDS["ester"]: flags[0] = 1.0 if group_ids & FUNCTIONAL_GROUP_IDS["alcohol"]: flags[1] = 1.0 if group_ids & FUNCTIONAL_GROUP_IDS["ketone"]: flags[2] = 1.0 if group_ids & FUNCTIONAL_GROUP_IDS["aldehyde"]: flags[3] = 1.0 return flags class OlfactoryEmbeddingEngine: """ Write-through cached embedding engine for 151-D olfactory molecule tokens. Combines a 138-D OpenPOM-style structural vector, a 7-D 3D shape/chirality descriptor, and a 6-D physical-functional payload (molecular weight, LogP, ester/alcohol/ketone/aldehyde flags) from the SQLite AromaRegistry. Missing 3D features are computed once and written back to the registry. Empty or invalid inputs short-circuit to a zero fallback without invoking RDKit. """ def __init__( self, use_fallback: bool = True, global_median_logp: float = 2.15, structural_source: str = "morgan", ) -> None: self.use_fallback = use_fallback self.structural_source = structural_source self.use_real_pom = structural_source == "openpom_256" self.use_morgan_2048_rp = structural_source == "morgan_2048_rp" self.use_morgan_2048_direct = structural_source == "morgan_2048_direct" self.use_disjoint_pom = structural_source == "disjoint_256" self.structural_dim = ( REAL_POM_DIM if (self.use_real_pom or self.use_disjoint_pom) else MORGAN2048_STRUCTURAL_DIM if (self.use_morgan_2048_rp or self.use_morgan_2048_direct) else STRUCTURAL_DIM ) self.shape_dim = SHAPE_DIM self.physical_dim = PHYSICAL_DIM self.functional_dim = FUNCTIONAL_DIM self.embedding_dim = ( COMBINED_DIM_REAL_POM if (self.use_real_pom or self.use_morgan_2048_rp or self.use_disjoint_pom) else MORGAN2048_TOKEN_DIM if self.use_morgan_2048_direct else COMBINED_DIM ) self._rp_matrix: np.ndarray | None = None self.registry = AromaRegistry() self._pom_index: dict[str, int] = {} self._pom_matrix: np.ndarray | None = None if self.use_real_pom: self._load_real_pom() if self.use_disjoint_pom: self._load_disjoint_pom() self.global_median_logp = self._compute_median_logp() if global_median_logp is None else global_median_logp if not self.use_fallback: self._initialize_real_openpom() def _load_real_pom(self) -> None: """Load the genuine OpenPOM v1.0.0 ensemble (256-dim, 10-model mean). Indexed by CAS (with some SMILES:-prefixed keys). Fail closed if the provenance-manifest asset is absent so a Morgan feature is never silently substituted for POM.""" emb_path = _POM_DIR / "pom_embeddings.npy" idx_path = _POM_DIR / "pom_index.json" if not (emb_path.exists() and idx_path.exists()): raise RealOpenPOMUnavailableError( f"Genuine POM asset missing under {_POM_DIR}; cannot honour structural_source='openpom_256'" ) self._pom_matrix = np.load(emb_path).astype(np.float32) idx = json.loads(idx_path.read_text()) ids = idx.get("ids", []) self._pom_index = {self._normalized_cas(k): i for i, k in enumerate(ids)} def _load_disjoint_pom(self) -> None: """Load the benchmark-disjoint MPNN encoder embeddings (256-dim). Same CAS-indexed layout as the genuine-POM asset, produced by scripts/build_disjoint_asset_and_gate.py from the round-7 disjoint encoder (mattbitzesty/pino-disjoint-encoder). Fail closed if absent: a Morgan feature must never be silently substituted for the disjoint arm's representation. Reuses _pom_index/_pom_matrix slots (the two arms are never active simultaneously).""" emb_path = _DISJOINT_DIR / "disjoint_embeddings.npy" idx_path = _DISJOINT_DIR / "disjoint_index.json" if not (emb_path.exists() and idx_path.exists()): raise RealOpenPOMUnavailableError( f"Disjoint-encoder asset missing under {_DISJOINT_DIR}; cannot honour structural_source='disjoint_256'" ) self._pom_matrix = np.load(emb_path).astype(np.float32) idx = json.loads(idx_path.read_text()) ids = idx.get("ids", []) self._pom_index = {self._normalized_cas(k): i for i, k in enumerate(ids)} def _compute_median_logp(self) -> float: rows = self.registry._conn.execute( "SELECT logp FROM aroma_chemicals WHERE logp IS NOT NULL AND logp != 0" ).fetchall() if not rows: return 2.15 return float(np.median([r[0] for r in rows])) def _initialize_real_openpom(self) -> None: # Fail closed. Silently substituting Morgan features here makes a # Morgan-vs-POM ablation compare the same representation under two names. raise RealOpenPOMUnavailableError( "Real OpenPOM model is not configured; use_fallback=False cannot be satisfied" ) def _normalized_cas(self, cas: str | None) -> str: return str(cas).replace("NATURAL:", "").strip() if cas else "" def _is_valid_smiles(self, smiles: str) -> bool: if not smiles or not isinstance(smiles, str): return False s = smiles.strip() return bool(s) and s != "None" and not s.startswith("SMILES:") and not s.startswith("NATURAL:") def get_combined_structural_token(self, cas: str, smiles: str) -> np.ndarray: """Main entry point. Returns a 151-dimensional float32 array.""" return self.get_embedding(smiles, cas=cas) def _extract_physical_features(self, record: dict[str, Any]) -> np.ndarray: """Return a 2-D normalized [MW, LogP] vector.""" mw = float(record.get("molecular_weight", 150.0)) logp_raw = record.get("logp") logp = float(logp_raw) if logp_raw is not None else self.global_median_logp return np.array([mw / 300.0, logp / 5.0], dtype=np.float32) def _extract_functional_features(self, record: dict[str, Any]) -> np.ndarray: """Return a 4-D binary flag vector for ester/alcohol/ketone/aldehyde.""" return _functional_group_flags(record.get("unifac_groups")) def get_embedding(self, smiles: str, cas: str | None = None) -> np.ndarray: """Main entry point. Returns a float32 embedding (dim = self.embedding_dim).""" normalized_cas = self._normalized_cas(cas) is_natural = bool(cas and str(cas).startswith("NATURAL:")) if normalized_cas: try: record = self.registry.get(normalized_cas) if record: if self.use_real_pom or self.use_disjoint_pom: # POM/disjoint arms: never reuse the registry's 138-dim Morgan block. structural = self._compute_structural_embedding(smiles, cas) elif self.use_morgan_2048_rp or self.use_morgan_2048_direct: # 2048-bit arms: always fresh-compute the wide fingerprint; # the registry cache only holds the 138-D block. structural = self._compute_structural_embedding(smiles, cas) else: structural = np.asarray(record.get("openpom_embedding") or [0.0] * STRUCTURAL_DIM, dtype=np.float32) if structural.shape[0] != STRUCTURAL_DIM or (not is_natural and np.count_nonzero(structural) == 0): structural = self._compute_structural_embedding(smiles, cas) shape = np.asarray(record.get("shape_3d_features") or [0.0] * SHAPE_DIM, dtype=np.float32) if shape.shape[0] != SHAPE_DIM or (not is_natural and np.count_nonzero(shape) == 0): shape = self._compute_shape_features(smiles, cas) physical = self._extract_physical_features(record) functional = self._extract_functional_features(record) combined = np.concatenate([structural, shape, physical, functional]).astype(np.float32) combined = self._maybe_project(combined) norm = np.linalg.norm(combined) return combined / norm if norm > 0 else combined if is_natural: structural = self._compute_structural_embedding(smiles, cas) shape = self._compute_shape_features(smiles, cas) physical = np.array([0.5, self.global_median_logp / 5.0], dtype=np.float32) functional = _functional_group_flags(None) combined = np.concatenate([structural, shape, physical, functional]).astype(np.float32) combined = self._maybe_project(combined) norm = np.linalg.norm(combined) return combined / norm if norm > 0 else combined except Exception as exc: logger.debug("Registry lookup failed for %s: %s", normalized_cas, exc) # Fallback: compute structural/shape features on the fly. structural = self._compute_structural_embedding(smiles, cas) shape = self._compute_shape_features(smiles, cas) physical = np.array([0.5, self.global_median_logp / 5.0], dtype=np.float32) functional = _functional_group_flags(None) combined = np.concatenate([structural, shape, physical, functional]).astype(np.float32) combined = self._maybe_project(combined) norm = np.linalg.norm(combined) return combined / norm if norm > 0 else combined def _maybe_project(self, combined: np.ndarray) -> np.ndarray: """morgan_2048_rp arm: fixed seeded random projection 2071-D -> 269-D. The projection is generated from a fixed seed on first use (no learned parameters, no checkpoint dependency), so every process computes the identical map. Output width matches the OpenPOM arm's combined token, giving the two trained arms an identical input-projection parameter count (capacity parity for the representation comparison). The morgan_2048_direct arm skips this projection: its 2071-D token feeds a LEARNED input projection directly (round-10 panel response: the RP control constrains the usable subspace; the direct arm tests whether learning from the raw 2048 bits changes the ordering).""" if not self.use_morgan_2048_rp: return combined if self._rp_matrix is None: rng = np.random.default_rng(MORGAN2048_RP_SEED) self._rp_matrix = ( rng.standard_normal((combined.shape[0], COMBINED_DIM_REAL_POM)).astype(np.float32) / np.sqrt(float(combined.shape[0])) ) return combined @ self._rp_matrix def _compute_structural_embedding(self, smiles: str, cas: str | None = None) -> np.ndarray: normalized_cas = self._normalized_cas(cas) # 2048-bit arms: always the wide fingerprint block (never the cached # 138-D registry block, which is a different representation). if self.use_morgan_2048_rp or self.use_morgan_2048_direct: return self._morgan_structural_wide(smiles) # Genuine OpenPOM arm / disjoint-encoder arm: authoritative 256-dim # embedding for known CAS (both assets share the CAS-indexed layout and # are loaded into the same slots; exactly one arm is active at a time). if self.use_real_pom or self.use_disjoint_pom: if normalized_cas and normalized_cas in self._pom_index and self._pom_matrix is not None: return self._pom_matrix[self._pom_index[normalized_cas]] # CAS absent from the asset: real Morgan block, zero-padded # to the 256-dim structural slot so the combined dim stays constant. morgan = self._morgan_structural(smiles) padded = np.zeros(self.structural_dim, dtype=np.float32) padded[: morgan.shape[0]] = morgan return padded # Morgan arm: live SQLite cache, natural-oil blend, then Morgan compute. if normalized_cas: try: record = self.registry.get(normalized_cas) if record: openpom_embedding = record.get("openpom_embedding") if openpom_embedding and len(openpom_embedding) > 0: return np.asarray(openpom_embedding, dtype=np.float32) except Exception as exc: logger.debug("Registry lookup failed for %s: %s", normalized_cas, exc) if normalized_cas and ((cas or "").startswith("NATURAL:") or normalized_cas.startswith(("8", "6", "9"))): try: all_records = self.registry.all_records() openpom, _ = resolve_natural_oil_vectors(normalized_cas, all_records) if any(openpom): return np.asarray(openpom, dtype=np.float32) except Exception as exc: logger.debug("Natural oil resolution failed for %s: %s", normalized_cas, exc) return self._morgan_structural(smiles) def _morgan_structural_wide(self, smiles: str) -> np.ndarray: """Deterministic 2048-bit radius-2 Morgan + 10 physicochemical descriptors (2058-D). The standard-width fingerprint control: identical construction to the preregistered 128-bit baseline, only fpSize differs.""" if not self._is_valid_smiles(smiles): return np.zeros(MORGAN2048_STRUCTURAL_DIM, dtype=np.float32) mol = Chem.MolFromSmiles(smiles) if not mol: return np.zeros(MORGAN2048_STRUCTURAL_DIM, dtype=np.float32) try: fpgen = AllChem.GetMorganGenerator(radius=2, fpSize=2048) fp_array = np.array(fpgen.GetFingerprint(mol), dtype=np.float32) except AttributeError: fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=2048) fp_array = np.array(fp, dtype=np.float32) physchem = np.array([ Descriptors.MolWt(mol), Descriptors.ExactMolWt(mol), Descriptors.MolLogP(mol), Descriptors.HeavyAtomCount(mol), Descriptors.NumHAcceptors(mol), Descriptors.NumHDonors(mol), Descriptors.NumRotatableBonds(mol), Descriptors.FractionCSP3(mol), Descriptors.TPSA(mol), Descriptors.LabuteASA(mol), ], dtype=np.float32) return np.concatenate([fp_array, physchem]).astype(np.float32) def _morgan_structural(self, smiles: str) -> np.ndarray: """Deterministic 128-bit Morgan + 10 physicochemical descriptors (138-D). Used as the Morgan arm's structural block and as the fallback for CAS absent from the genuine-POM asset.""" if not self._is_valid_smiles(smiles): return np.zeros(STRUCTURAL_DIM, dtype=np.float32) mol = Chem.MolFromSmiles(smiles) if not mol: return np.zeros(STRUCTURAL_DIM, dtype=np.float32) try: fpgen = AllChem.GetMorganGenerator(radius=2, fpSize=128) fp_array = np.array(fpgen.GetFingerprint(mol), dtype=np.float32) except AttributeError: fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=128) fp_array = np.array(fp, dtype=np.float32) physchem = np.array([ Descriptors.MolWt(mol), Descriptors.ExactMolWt(mol), Descriptors.MolLogP(mol), Descriptors.HeavyAtomCount(mol), Descriptors.NumHAcceptors(mol), Descriptors.NumHDonors(mol), Descriptors.NumRotatableBonds(mol), Descriptors.FractionCSP3(mol), Descriptors.TPSA(mol), Descriptors.LabuteASA(mol), ], dtype=np.float32) return np.concatenate([fp_array, physchem]).astype(np.float32) def _compute_shape_features(self, smiles: str, cas: str | None = None) -> np.ndarray: """Fetch cached 3D shape features or compute them on the fly and write through.""" normalized_cas = self._normalized_cas(cas) zero_shape = np.zeros(self.shape_dim, dtype=np.float32) # 1. Query the live SQLite cache. if normalized_cas: try: record = self.registry.get(normalized_cas) if record: shape_3d = record.get("shape_3d_features") if shape_3d and len(shape_3d) > 0 and any(v != 0.0 for v in shape_3d): return np.asarray(shape_3d, dtype=np.float32) except Exception as exc: logger.debug("Registry lookup failed for %s: %s", normalized_cas, exc) # Natural oils may be resolvable from a constituent blend. if normalized_cas and (normalized_cas.startswith("8") or normalized_cas.startswith("6") or normalized_cas.startswith("9")): try: all_records = self.registry.all_records() _, shape = resolve_natural_oil_vectors(normalized_cas, all_records) if any(shape): return np.asarray(shape, dtype=np.float32) except Exception as exc: logger.debug("Natural oil shape resolution failed for %s: %s", normalized_cas, exc) # 2. String guardrail: invalid/empty SMILES short-circuit to zero fallback. if not self._is_valid_smiles(smiles): return zero_shape # 3. Just-In-Time compute 3D features exactly once. try: features = extract_3d_shape_features(smiles) except Exception as exc: logger.debug("3D shape feature extraction failed for %s: %s", smiles, exc) return zero_shape shape_3d_vec = list(features) if any(v != 0.0 for v in shape_3d_vec): # 4. Write through to the persistent SQLite cache. if normalized_cas: try: record = self.registry.get(normalized_cas) or {"cas": normalized_cas} record["shape_3d_features"] = shape_3d_vec self.registry.add(record) except Exception as exc: logger.debug("Failed to persist 3D features for %s: %s", normalized_cas, exc) return np.asarray(shape_3d_vec, dtype=np.float32) def _compute_real_openpom(self, smiles: str) -> np.ndarray: raise NotImplementedError("Real OpenPOM model not yet loaded") # Module-level engine is lazily instantiated so that import-time side effects # (e.g., opening the SQLite registry) do not run until the first actual use. _DEFAULT_ENGINE: OlfactoryEmbeddingEngine | None = None def _get_default_engine() -> OlfactoryEmbeddingEngine: global _DEFAULT_ENGINE if _DEFAULT_ENGINE is None: _DEFAULT_ENGINE = OlfactoryEmbeddingEngine(use_fallback=True) return _DEFAULT_ENGINE def get_openpom_embedding(smiles: str, dim: int = COMBINED_DIM, cas: str | None = None) -> np.ndarray: """Convenience accessor used by the PIMT pipeline.""" engine = _get_default_engine() embedding = engine.get_embedding(smiles, cas=cas) if embedding.shape[0] != dim: raise ValueError(f"Embedding dimension mismatch: {embedding.shape[0]} != {dim}") return embedding