01_entity_analysis.py

"""
GeneRIF Entity Analysis
-----------------------
Analyzes a sample of GeneRIF entries to catalog all entity types
that could become nodes in a knowledge graph.

Approach:
  1. spaCy general NER (en_core_web_sm) — catches ORG, PERSON, DATE, etc.
  2. Rule-based patterns for biomedical entities not covered by general NER:
     - Gene/protein symbols (all-caps, alphanumeric, hyphens)
     - Variants / mutations (p.Val600Glu, rs numbers, SNPs)
     - GO terms, MeSH-style diseases
     - Relationships / predicates (verbs like "inhibits", "activates")
  3. Frequency analysis + manual taxonomy of entity types

Output:
  /results/entity_type_summary.tsv   — entity type counts
  /results/entity_examples.tsv       — top examples per type
"""

import gzip
import re
import csv
import random
from collections import Counter, defaultdict

import spacy

# ── configuration ──────────────────────────────────────────────────────────────
DATA_FILE   = "/data/generifs_basic.gz"
OUT_SUMMARY = "/results/entity_type_summary.tsv"
OUT_EXAMPLES = "/results/entity_examples.tsv"

SAMPLE_SIZE   = 20_000   # human entries to analyse
HUMAN_TAX_ID  = "9606"
RANDOM_SEED   = 42

# ── load spaCy ─────────────────────────────────────────────────────────────────
print("Loading spaCy model...")
nlp = spacy.load("en_core_web_sm", disable=["parser"])
nlp.max_length = 2_000_000

# ── regex patterns for biomedical entities ────────────────────────────────────
PATTERNS = {
    # Gene / protein symbol: 2-10 uppercase letters+digits, optional hyphens
    "GENE_SYMBOL": re.compile(
        r'\b([A-Z][A-Z0-9]{1,9}(?:-[A-Z0-9]+)?)\b'
    ),
    # Mutation – HGVS-like: p.Val600Glu, p.R175H, c.1234A>G
    "MUTATION_HGVS": re.compile(
        r'\b(p\.[A-Z][a-z]{2}\d+[A-Z][a-z]{2}|p\.[A-Z]\d+[A-Z*]|'
        r'c\.\d+[ACGT]>[ACGT])\b'
    ),
    # dbSNP rs number
    "SNP_RS": re.compile(r'\brs\d{4,}\b', re.IGNORECASE),
    # Numeric variant: e.g. "V600E", "R175H", "K65R"
    "MUTATION_SHORT": re.compile(r'\b[A-Z]\d{1,4}[A-Z*]\b'),
    # Gene/protein fusion: GENE1-GENE2
    "FUSION": re.compile(r'\b([A-Z][A-Z0-9]{1,8})-([A-Z][A-Z0-9]{1,8})\b'),
    # miRNA / ncRNA
    "MIRNA": re.compile(r'\bmi(?:R|RNA)-?\d+[a-z]?\b', re.IGNORECASE),
    # lncRNA pattern
    "LNCRNA": re.compile(r'\b(?:lnc|LNC)[A-Z0-9\-]{2,20}\b'),
    # Chemical/drug: looks for CamelCase or all-lower terms ending in -ib, -mab, -nib
    "DRUG": re.compile(
        r'\b\w+(?:inib|tinib|rafenib|mab|zumab|ximab|kinase inhibitor)\b',
        re.IGNORECASE
    ),
    # Cell type: ends in "cell", "cells", "cyte", "cytes", "blast"
    "CELL_TYPE": re.compile(
        r'\b(?:[A-Z]?[a-z]+ )?(?:T|B|NK|stem|mast|dendritic|epithelial|'
        r'endothelial|fibroblast|macrophage|monocyte|neutrophil|'
        r'lymphocyte|neuron|hepatocyte|cardiomyocyte|osteoblast|osteoclast)'
        r'(?:s| cells?| cytes?)?\b'
    ),
    # Pathway: ends in "pathway", "signaling", "cascade", "axis"
    "PATHWAY": re.compile(
        r'\b[\w\-/]+ (?:pathway|signaling|signalling|cascade|axis|network)\b',
        re.IGNORECASE
    ),
    # GO biological process keywords
    "BIO_PROCESS": re.compile(
        r'\b(?:apoptosis|autophagy|proliferation|differentiation|migration|'
        r'invasion|metastasis|angiogenesis|inflammation|transcription|'
        r'translation|phosphorylation|ubiquitination|methylation|acetylation|'
        r'glycosylation|SUMOylation|splicing|expression|regulation|'
        r'activation|inhibition|binding|cleavage|degradation|'
        r'localization|trafficking|secretion|signaling|senescence)\b',
        re.IGNORECASE
    ),
    # Disease: ends in disease, syndrome, cancer, carcinoma, disorder, etc.
    "DISEASE": re.compile(
        r'\b[\w\s\-]+ (?:disease|syndrome|cancer|carcinoma|adenocarcinoma|'
        r'lymphoma|leukemia|tumor|tumour|disorder|deficiency|'
        r'malignancy|neoplasm|fibrosis|dystrophy)\b',
        re.IGNORECASE
    ),
    # Tissue / organ
    "TISSUE_ORGAN": re.compile(
        r'\b(?:liver|kidney|lung|heart|brain|colon|breast|prostate|'
        r'pancreas|thyroid|ovary|testis|muscle|bone|skin|blood|'
        r'spleen|thymus|lymph node|adipose|placenta|retina|'
        r'hippocampus|cortex|cerebellum)\b',
        re.IGNORECASE
    ),
    # Species / organism
    "SPECIES": re.compile(
        r'\b(?:human|mouse|rat|zebrafish|yeast|Drosophila|C\. elegans|'
        r'Arabidopsis|E\. coli|S\. cerevisiae|Xenopus|chicken|pig|rabbit)\b',
        re.IGNORECASE
    ),
    # Relationship predicates (verbs)
    "RELATION_VERB": re.compile(
        r'\b(?:activates?|inhibits?|promotes?|suppresses?|induces?|represses?|'
        r'regulates?|modulates?|upregulates?|downregulates?|phosphorylates?|'
        r'ubiquitinates?|methylates?|acetylates?|cleaves?|binds?|interacts?|'
        r'recruits?|stabilizes?|destabilizes?|localizes?|translocates?|'
        r'mediates?|drives?|blocks?|enhances?|reduces?|increases?|decreases?)\b',
        re.IGNORECASE
    ),
    # Subcellular localization
    "SUBCELLULAR": re.compile(
        r'\b(?:nucleus|cytoplasm|mitochondri[ao]|endoplasmic reticulum|'
        r'Golgi|lysosome|peroxisome|plasma membrane|cell membrane|'
        r'cytosol|nucleolus|centrosome|ribosome|proteasome|chromatin)\b',
        re.IGNORECASE
    ),
    # Experimental method
    "METHOD": re.compile(
        r'\b(?:Western blot|immunoprecipitation|ChIP(?:-seq)?|RNA-seq|'
        r'CRISPR|siRNA|shRNA|knockout|knockin|overexpression|'
        r'qPCR|RT-PCR|ELISA|flow cytometry|mass spectrometry|'
        r'proteomics|genomics|transcriptomics|co-immunoprecipitation|'
        r'yeast two-hybrid|pull-down|luciferase|reporter assay)\b',
        re.IGNORECASE
    ),
    # Genomic region / element
    "GENOMIC_ELEMENT": re.compile(
        r'\b(?:promoter|enhancer|exon|intron|UTR|5\'UTR|3\'UTR|'
        r'CpG island|TATA box|splice site|open reading frame|ORF|'
        r'coding sequence|CDS|regulatory element|transcription factor binding site)\b',
        re.IGNORECASE
    ),
}

# ── relationship verb patterns ─────────────────────────────────────────────────

def extract_entities(text):
    """Return dict of entity_type -> list of matched strings."""
    found = defaultdict(list)

    # spaCy NER
    doc = nlp(text)
    for ent in doc.ents:
        label = f"SPACY_{ent.label_}"
        found[label].append(ent.text)

    # regex patterns
    for etype, pattern in PATTERNS.items():
        matches = pattern.findall(text)
        if matches:
            # flatten tuples from groups
            flat = []
            for m in matches:
                if isinstance(m, tuple):
                    flat.append("-".join(m))
                else:
                    flat.append(m)
            found[etype].extend(flat)

    return found


def load_sample(filepath, tax_id=HUMAN_TAX_ID, n=SAMPLE_SIZE, seed=RANDOM_SEED):
    """Load a random sample of GeneRIF entries for a given tax ID."""
    print(f"Loading sample of {n} entries for tax_id={tax_id}...")
    all_texts = []
    with gzip.open(filepath, "rt", encoding="utf-8") as f:
        for line in f:
            if line.startswith("#"):
                continue
            cols = line.strip().split("\t")
            if len(cols) < 5:
                continue
            if cols[0] == tax_id:
                all_texts.append(cols[4])

    random.seed(seed)
    sample = random.sample(all_texts, min(n, len(all_texts)))
    print(f"  Loaded {len(sample)} entries.")
    return sample


def main():
    texts = load_sample(DATA_FILE)

    type_counts   = Counter()     # total entity mentions per type
    type_examples = defaultdict(Counter)  # top surface forms per type

    print(f"Extracting entities from {len(texts)} entries...")
    for i, text in enumerate(texts):
        if i % 2000 == 0:
            print(f"  {i}/{len(texts)}...")
        entities = extract_entities(text)
        for etype, values in entities.items():
            type_counts[etype] += len(values)
            for v in values:
                v_clean = v.strip()
                if len(v_clean) > 2:
                    type_examples[etype][v_clean] += 1

    # ── write summary ──────────────────────────────────────────────────────────
    print(f"\nWriting {OUT_SUMMARY}...")
    with open(OUT_SUMMARY, "w", newline="") as f:
        w = csv.writer(f, delimiter="\t")
        w.writerow(["entity_type", "total_mentions", "unique_values", "notes"])
        for etype, count in type_counts.most_common():
            unique = len(type_examples[etype])
            w.writerow([etype, count, unique, ""])

    # ── write examples ─────────────────────────────────────────────────────────
    print(f"Writing {OUT_EXAMPLES}...")
    with open(OUT_EXAMPLES, "w", newline="") as f:
        w = csv.writer(f, delimiter="\t")
        w.writerow(["entity_type", "surface_form", "frequency"])
        for etype, counter in sorted(type_examples.items()):
            for form, freq in counter.most_common(20):
                w.writerow([etype, form, freq])

    # ── print summary to console ───────────────────────────────────────────────
    print("\n" + "="*60)
    print(f"{'ENTITY TYPE':<30} {'MENTIONS':>10} {'UNIQUE':>10}")
    print("="*60)
    for etype, count in type_counts.most_common():
        unique = len(type_examples[etype])
        print(f"{etype:<30} {count:>10,} {unique:>10,}")
    print("="*60)
    print(f"\nDone. Results written to /results/")


if __name__ == "__main__":
    main()