"""Knowledge graph schema: node and edge data classes.""" from __future__ import annotations from dataclasses import dataclass, field from enum import Enum from typing import Optional class DomainTag(str, Enum): """High-level domain classification for concepts.""" NEUROANATOMY = "neuroanatomy" DISEASE = "disease" GENE = "gene" NEUROTRANSMITTER = "neurotransmitter" DRUG = "drug" COGNITIVE_FUNCTION = "cognitive_function" CELL_TYPE = "cell_type" BIOMARKER = "biomarker" PARADIGM = "paradigm" # experimental paradigm (BrainMap) CONNECTIVITY = "connectivity" # functional/structural connections IMAGING_FEATURE = "imaging_feature" # cortical thickness, volume, FA, FC, SUVR, etc. DATASET_VARIABLE = "dataset_variable" # genetics, environment, medication, etc. # Phase 1.5 Experiment infrastructure (atlas/modality/dataset/ml_model) # + reserved RECIPE tag (former Phase 4.3, removed 2026-05-13 but kept # in UMLS-skip set for forward compat). RECIPE = "recipe" # reserved ATLAS = "atlas" # brain parcellation (ATLAS:*) MODALITY = "modality" # imaging/data modality (MODALITY:*) DATASET = "dataset" # research dataset (DATASET:*) ML_MODEL = "ml_model" # ML architecture (MODEL:*) # Brain decoding stimuli & psychological-state targets VISUAL_STIMULUS = "visual_stimulus" # image/video stimulus (NSD/BOLD5000/SEED-DV) EMOTION = "emotion" # affective state label (SEED family) VIGILANCE = "vigilance" # alertness/drowsiness label (SEED-VIG) # Treatment-effect / clinical-outcome variables (Phase 1, atom-algebra rollout # 2026-05-20): a finer-grained sibling of DATASET_VARIABLE used specifically # to tag nodes that represent quantitative drug-response or therapy outcomes # (e.g. ΔADAS-Cog at 12 mo, ΔMDS-UPDRS, NIHSS at 6 mo, responder/non-responder # labels). Distinguishes treatment-outcome variables from generic individual-data # variables (age, APOE-ε4 status, lifestyle) when both share dataset_variable. TREATMENT_OUTCOME = "treatment_outcome" # Concept-level anchors for INDIVIDUAL_DATA (Aging, APOE, Big-5 traits, …), # seeded by ingestion.individual_data_anchors. Distinct from dataset_variable # (which is the UKB/ADNI/HCP host category hub) so the bridge layer can # connect concept-side IM/disease/gene nodes to dataset variables in two # hops without polluting either side's domain. INDIVIDUAL_DATA_ANCHOR = "individual_data_anchor" class SemanticType(str, Enum): """UMLS semantic types relevant to neuroscience.""" DISEASE_OR_SYNDROME = "T047" MENTAL_DYSFUNCTION = "T048" NEOPLASTIC_PROCESS = "T191" BODY_PART_ORGAN = "T023" BODY_LOCATION = "T029" CELL = "T025" NEUROTRANSMITTER = "T116" AMINO_ACID_PEPTIDE = "T116" # overlaps with neurotransmitter in UMLS PHARMACOLOGIC_SUBSTANCE = "T121" GENE_OR_GENOME = "T028" INTELLECTUAL_PRODUCT = "T170" @dataclass class ConceptNode: """A concept node in the knowledge graph.""" id: str # unique identifier (CUI, or custom like "NN:1234") preferred_name: str # standard display name semantic_types: list[str] = field(default_factory=list) # TUI codes domain_tags: list[str] = field(default_factory=list) # DomainTag values source_vocab: str = "" # originating vocabulary (MeSH, NeuroNames, etc.) definition: str = "" # text definition aliases: list[str] = field(default_factory=list) # synonyms / alternate names external_ids: dict[str, str] = field(default_factory=dict) # cross-references atlas_mapping: Optional[dict] = None # MNI coords, atlas region ID, etc. metadata: dict = field(default_factory=dict) # catch-all for extra info def to_dict(self) -> dict: return { "id": self.id, "preferred_name": self.preferred_name, "semantic_types": self.semantic_types, "domain_tags": self.domain_tags, "source_vocab": self.source_vocab, "definition": self.definition, "aliases": self.aliases, "external_ids": self.external_ids, "atlas_mapping": self.atlas_mapping, "metadata": self.metadata, } @classmethod def from_dict(cls, d: dict) -> ConceptNode: return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__}) RELATION_TYPES = { # taxonomic / structural "is_a", # A is a subtype of B "part_of", # A is anatomical part of B "has_part", # inverse of part_of # causal / functional "causes", # A causes B "associated_with", # A is associated with B (loose) "predisposes", # A increases risk of B # therapeutic "treats", # A treats B "contraindicated_for", # A is contraindicated for B # molecular / genetic "gene_associated_with_disease", "protein_encoded_by", "modulates", # A modulates activity of B "binds_to", # A binds to receptor B # neuroanatomy "projects_to", # A projects neural connections to B "connects_to", # structural connectivity A-B "activates", # A functionally activates B "coactivates", # A and B co-activate (BrainMap) # evidence "supported_by", "contradicts", "about", # claim predicates (from paper extraction) "reduces", "increases", "correlates_with", "is_biomarker_of", "is_risk_factor_for", "is_associated_with", "predicts", "mediates", "inhibits", "distinguishes", # Phase 1.5 Experiment infrastructure edges "supports_modality", # Model → Modality (compat declaration) "provides_modality", # Dataset → Modality (what the dataset contains) # Brain decoding edges (NSD/BOLD5000/SEED-DV/SEED family) "evokes", # visual_stimulus → neuroanatomy (encoding direction) "decoded_from", # visual_stimulus ← neuroanatomy (decoding direction) "elicits", # stimulus → emotion/vigilance (behavioral label) # Outcome / dataset-variable bridges "measures", # rating_scale → disease (the scale measures the disease) "assessed_in", # rating_scale → dataset_variable host "affects_system", # MedDRA SOC → disease umbrella "provides_signal_for", # disease → dataset_variable (anchor → variable) # Therapy reverse edges (closes drug indegree gap) "is_indicated_for", # drug → disease (canonical indication) "is_treated_by", # disease → drug (reverse of treats; for personalised_treatment) # Imaging modality edges (closes IM↔Idv gap) "measured_in_modality", # imaging_marker → dataset_variable (modality host) "modality_provides", # dataset_variable → imaging_marker (reverse) # Outcome closure edges (closes IM/D/Rx → OUTCOME gap) "is_assessed_by", # disease → rating_scale (reverse of measures) "has_adverse_effect", # drug → AE SOC (drug → outcome) } # Edge tier classification — used by `KnowledgeGraph.export_display_subgraph` # to derive a human-facing view of the KG that drops provenance, inverse # mirrors, and admin/bridge edges. The full graph (self.G) is unchanged and # the hypothesis engine still uses every edge. # # Tiers: # "discovery" — direct semantic claim, surface to readers (claim predicates, # canonical pharmacology / connectivity / scale relations) # "skeleton" — taxonomy/hierarchy (is_a, part_of), usually rendered as a # tree, not as graph edges; opt-in via `include_skeleton=True` # "inverse" — inverse mirror of a canonical Tier-1 edge, kept for reverse # path traversal in the hypothesis engine; never display # "bridge" — admin/bridge edge that connects atom families through hub # nodes (dataset-variable, disease-umbrella); never display # "provenance"— claim→subject/object 'about' edges, support links; never display # "deprecated"— removed Phase 4.3 relations kept in RELATION_TYPES for # forward-compat; never display # # Relations not listed here default to "discovery" so the long-tail of # claim-extracted predicates ("improves", "disrupts", …) is still surfaced. EDGE_TIER: dict[str, str] = { # Tier 1 — discovery "reduces": "discovery", "increases": "discovery", "correlates_with": "discovery", "causes": "discovery", "is_biomarker_of": "discovery", "is_risk_factor_for": "discovery", "is_associated_with": "discovery", "associated_with": "discovery", "predicts": "discovery", "mediates": "discovery", "inhibits": "discovery", "distinguishes": "discovery", "contradicts": "discovery", "predisposes": "discovery", "contraindicated_for": "discovery", "gene_associated_with_disease": "discovery", "treats": "discovery", "protein_encoded_by": "discovery", "binds_to": "discovery", "modulates": "discovery", "projects_to": "discovery", "connects_to": "discovery", "activates": "discovery", "coactivates": "discovery", "evokes": "discovery", "elicits": "discovery", "measures": "discovery", # Tier 2 — skeleton "is_a": "skeleton", "part_of": "skeleton", "supports_modality": "skeleton", "provides_modality": "skeleton", # Tier 3 — inverse mirrors of canonical Tier-1 edges "is_treated_by": "inverse", "is_assessed_by": "inverse", "modality_provides": "inverse", "has_part": "inverse", "decoded_from": "inverse", # Tier 3 — admin / bridge "assessed_in": "bridge", "affects_system": "bridge", "provides_signal_for": "bridge", "measured_in_modality": "bridge", "has_adverse_effect": "bridge", # Tier 3 — provenance "about": "provenance", "supported_by": "provenance", } DISPLAY_TIERS_DEFAULT: frozenset[str] = frozenset({"discovery"}) # Claim-specific predicates (extracted from papers) CLAIM_PREDICATES = { "reduces", # A reduces B "increases", # A increases B "correlates_with", # A correlates with B "causes", # A causes B "is_biomarker_of", # A is a biomarker for B "is_risk_factor_for", # A is a risk factor for B "treats", # A treats B "modulates", # A modulates B "activates", # A activates B "inhibits", # A inhibits B "predicts", # A predicts B "mediates", # A mediates the relationship between B and C "is_associated_with", # A is associated with B "distinguishes", # A distinguishes B from C } @dataclass class Edge: """A directed edge in the knowledge graph.""" source_id: str # source ConceptNode.id target_id: str # target ConceptNode.id relation_type: str # one of RELATION_TYPES source: str = "" # provenance: 'NeuroNames', 'MeSH', 'DisGeNET', etc. confidence: float = 1.0 # 0.0-1.0 evidence_ref: str = "" # citation or reference metadata: dict = field(default_factory=dict) def to_dict(self) -> dict: return { "source_id": self.source_id, "target_id": self.target_id, "relation_type": self.relation_type, "source": self.source, "confidence": self.confidence, "evidence_ref": self.evidence_ref, "metadata": self.metadata, } @classmethod def from_dict(cls, d: dict) -> Edge: return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__}) @dataclass class Evidence: """Experimental evidence supporting a scientific claim.""" study_type: str = "" # "fMRI", "lesion", "meta-analysis", "GWAS", "animal_model" methodology: str = "" # "resting-state FC", "voxel-based morphometry", "DTI", ... p_value: Optional[float] = None effect_size: Optional[float] = None # Cohen's d, r, OR, beta effect_metric: str = "" # "Cohen's d", "r", "OR", "beta", "AUC" sample_size: Optional[int] = None replicability: str = "single_study" # "replicated", "single_study", "controversial" direction: str = "" # "positive", "negative" def to_dict(self) -> dict: return { "study_type": self.study_type, "methodology": self.methodology, "p_value": self.p_value, "effect_size": self.effect_size, "effect_metric": self.effect_metric, "sample_size": self.sample_size, "replicability": self.replicability, "direction": self.direction, } @classmethod def from_dict(cls, d: dict) -> Evidence: return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__}) @dataclass class PaperRef: """Reference to a source paper.""" pmid: str = "" # PubMed ID doi: str = "" title: str = "" authors: str = "" year: Optional[int] = None journal: str = "" def to_dict(self) -> dict: return { "pmid": self.pmid, "doi": self.doi, "title": self.title, "authors": self.authors, "year": self.year, "journal": self.journal, } @classmethod def from_dict(cls, d: dict) -> PaperRef: return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__}) @dataclass class Claim: """A structured scientific claim extracted from a paper. A claim is both stored as a node (for detailed querying) and generates simplified edges (for multi-hop traversal). """ id: str # CLM:uuid subject_id: str # ConceptNode.id in the graph subject_name: str # human-readable subject name predicate: str # one of CLAIM_PREDICATES object_id: str # ConceptNode.id in the graph object_name: str # human-readable object name negated: bool = False # "X does NOT affect Y" confidence: float = 0.5 # overall confidence 0-1 evidence: Evidence = field(default_factory=Evidence) source_paper: PaperRef = field(default_factory=PaperRef) raw_text: str = "" # original sentence from paper metadata: dict = field(default_factory=dict) def to_dict(self) -> dict: return { "id": self.id, "subject_id": self.subject_id, "subject_name": self.subject_name, "predicate": self.predicate, "object_id": self.object_id, "object_name": self.object_name, "negated": self.negated, "confidence": self.confidence, "evidence": self.evidence.to_dict(), "source_paper": self.source_paper.to_dict(), "raw_text": self.raw_text, "metadata": self.metadata, } @classmethod def from_dict(cls, d: dict) -> Claim: d = d.copy() if "evidence" in d and isinstance(d["evidence"], dict): d["evidence"] = Evidence.from_dict(d["evidence"]) if "source_paper" in d and isinstance(d["source_paper"], dict): d["source_paper"] = PaperRef.from_dict(d["source_paper"]) return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__}) def to_edge(self) -> Edge: """Convert claim to a simplified graph edge for traversal.""" return Edge( source_id=self.subject_id, target_id=self.object_id, relation_type=self.predicate, source=f"claim:{self.source_paper.pmid or self.id}", confidence=self.confidence, evidence_ref=self.source_paper.title, metadata={"claim_id": self.id, "negated": self.negated}, )