diff --git a/.gitignore b/.gitignore index beead43..4061195 100644 --- a/.gitignore +++ b/.gitignore @@ -53,6 +53,9 @@ coverage.xml # Django stuff: *.log + +# Exported graph artifacts +*.graphml local_settings.py # Flask stuff: diff --git a/HISTORY.rst b/HISTORY.rst index e525cc9..ecb4505 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -17,4 +17,124 @@ This release expands GraphFaker’s scope with a new data sources to support gra - Retrieve raw page data (title, summary, content, sections, links, references) via the wikipedia package - Export JSON dumps of article fields -Upgrade now to effortlessly pull in unstructured Wikipedia data \ No newline at end of file +Upgrade now to effortlessly pull in unstructured Wikipedia data + +0.5.0 (2026-08-01) +------------------ + +Realistic graph topology. + +Until now both endpoints of every edge were drawn uniformly at random, so the +synthetic generator produced an Erdos-Renyi graph: Poisson degree distribution, +no hubs, effectively no clustering, no community structure, and attributes +statistically independent of the topology. It was labelled "realistic" but was +not, and anything measured against it was measured against noise. + +Edges are now formed by preferential attachment, triadic closure, and homophily +over latent communities. Measured on 600 nodes / 2400 edges, seed 1: + +=========================== =========== ================= +metric realistic uniform (pre-0.5) +=========================== =========== ================= +degree Fano factor 7.9 1.7 +max degree 88 19 +degree Gini 0.45 0.26 +average clustering 0.180 0.013 +community modularity 0.72 -0.004 +age homophily 0.81 0.01 +isolated nodes 0 4 +=========================== =========== ================= + +* ``graphfaker.metrics`` — ``graph_stats()`` and ``compare_topology()``, so the + claim is checkable rather than asserted. The tests are differential: each + property is compared against ``topology="uniform"``, which reproduces the old + behaviour and is retained solely as a baseline. +* Every node carries a latent ``community``; attributes are drawn from + community-specific distributions, which is what makes homophily possible. +* Places, organizations, events, and products carry a heavy-tailed ``prominence`` + driving popularity. +* ``population`` tracks connectivity and ``employee_count`` now correlates 0.95 + with the ``WORKS_AT`` edges actually present, instead of contradicting them. +* ``industry`` holds an industry rather than ``fake.job()``'s job titles. +* ``LIVES_IN``, ``BORN_IN``, and ``HEADQUARTERED_IN`` are singular. A person + could previously live in four cities. +* ``total_edges`` now means what ``number_of_edges()`` reports. Counting loop + iterations overshot by ~14% when bidirectional friendships dominated and + undershot by ~12% when skipped functional relationships did. +* No isolated nodes. Preferential attachment alone left 24% of a 600-node graph + unreachable, so coverage and top-up passes guarantee a giant component. + +Fixed: + +* GraphML export raised on ``None`` attribute values, which broke export for any + graph small enough to leave a community without a place. + +Note for anyone comparing results across versions: entity resolution is +measurably harder on a realistic graph. The same injected duplicates give +``resolve()`` precision 1.000 on a uniform graph and 0.88 on a realistic one, +because homophily means distinct people share attributes and neighbours. Numbers +produced before 0.5 were flattered by the generator. + +0.4.0 (2026-07-31) +------------------ + +Graph-native entity resolution, reproducible generation, and several fixes. + +New: + +* ``graphfaker.export`` — export connectors that write files rather than + requiring a database driver: ``export_csv`` (key-union headers, so + heterogeneous node types keep their values in the right columns), + ``export_neo4j_csv`` (``neo4j-admin`` typed headers), and ``export_cypher`` + for Neo4j, openCypher, and ISO GQL. Also exposed as ``--format`` on the CLI. +* ``graphfaker.corpus`` — paired text/entity corpora for measuring how many + nodes a graph builder creates per real entity. Documents are clean prose, not + corrupted text; ``Corpus.audit()`` verifies that no surface form could be + claimed by two entities, and ``duplication_report()`` produces the counts. +* ``examples/duplication_experiment.py`` — runs that measurement across several + graph-building frameworks and prints a comparison table. The Cognee adapter is + written against the API verified in cognee 1.4.1 (``add``/``cognify``/``export`` + are coroutines; ``export`` accepts ``format="graphml"``; there is no + ``get_graph_data``). It writes to a run-specific dataset instead of calling + ``cognee.prune``, so an existing local store is not destroyed. +* ``docs/notebooks/duplication_experiment.ipynb`` — the same experiment end to + end, including repairing the graph with ``resolve()``, a threshold sweep + showing the precision/recall tradeoff, and export of the cleaned graph. Runs + without an LLM key using a labelled simulated extraction. +* ``resolve_entities(token_subset_floor=...)`` — floors the attribute score when + one name's tokens are contained in the other's ("Hill" inside "Allison Hill"). + Character ratios score that pair near 0.5, so it was previously discarded + before neighbourhood overlap was consulted. The floor sits below the default + threshold deliberately, so containment alone never merges anything. +* ``graphfaker.resolve`` — find duplicate entities using attribute similarity + **and** neighbourhood overlap, then merge each cluster onto one canonical node + with its edges rewired. Available as ``GraphFaker.resolve()`` or the + standalone ``resolve_entities()`` / ``merge_clusters()``. +* ``evaluate_clusters()`` — pairwise and B-cubed precision/recall/F1 for scoring + a predicted clustering against gold labels you supply. +* Seeding: ``GraphFaker(seed=...)``, ``generate_graph(..., seed=...)``, + ``reseed()``, and ``--seed`` on the CLI. Identical seed and arguments now + produce an identical graph, and seeding no longer touches global + ``random`` state. + +Fixed: + +* **Security:** TLS certificate verification was disabled for all flight-data + downloads. It is now enabled by default; opt out with + ``GRAPHFAKER_INSECURE_TLS=1``, which warns loudly. +* **Packaging:** ``tqdm`` and ``urllib3`` were imported but never declared as + dependencies, so a clean install could fail on ``import graphfaker``. +* The CLI logged through the standard library ``venv`` module's logger by + accident (``from venv import logger``) instead of the package logger. +* ``generate_graph`` reported the wrong valid sources on error ("Use 'random' or + 'osm'"). +* GraphML export now flattens list, tuple, set, and dict attributes instead of + failing on them. +* Three CLI tests were failing on ``main`` because their mocks returned strings + where a graph was required; tests also no longer write artifacts into the + repository root. + +Docs: + +* README no longer advertises unimplemented features (RDF/JSON-LD export, + Neo4j/Kuzu/TigerGraph integration, million-node scale, LLM-driven fetching). \ No newline at end of file diff --git a/PROPOSAL_V2.md b/PROPOSAL_V2.md new file mode 100644 index 0000000..5b4f3a0 --- /dev/null +++ b/PROPOSAL_V2.md @@ -0,0 +1,514 @@ +# GraphFaker v2 — Strategy Proposal + +**Status:** Draft for discussion · **Date:** 2026-07-31 · **Branch:** `new_feature_branch` + +Supersedes the framing in [`Proposal.md`](Proposal.md). All figures were measured on +2026-07-31 and are cited so they can be re-checked. This document went through an +adversarial review that **killed its own original recommendation**; §4 records that, and §0 +states the revised position. The discarded thesis is kept in §3 because the reasoning for +discarding it is the most useful content here. + +--- + +## 0. TL;DR + +**Three findings, in order of confidence.** + +1. **All four ideas currently on the table should be dropped.** Not narrowed — dropped. §2 + gives the evidence for each. This is the highest-confidence conclusion in the document. +2. **The attractive-looking alternative — "GraphFaker becomes the ground-truth substrate for + evaluating graph systems" — also fails.** It survived first-principles reasoning and then + died under evidence. Three independent kills, §4. +3. **What survives is small, and it is not a repositioning.** Ship a **deduplication + capability upstream into an existing framework**, use corruption/ground-truth only as + *internal test fixtures inside that PR*, and never ship it as the product. §6. + +**The single sentence that reorganized this document:** + +> Splink — 2,308★, actively developed, the exact use case — **built a synthetic +> linkage-data generator and then deleted it.** `splink_synthetic_data` is a 404, and +> [#110](https://github.com/moj-analytical-services/splink/issues/110) plus #292, #316, #632 +> are all closed. They ship static CSVs with a `cluster` column instead. + +That is not an unserved market. That is a market that was tried and rejected by the +best-positioned team in the world. + +**Second sentence, nearly as damaging:** the people in acute pain about duplicate entities +never asked to *measure* them. In [microsoft/graphrag#1718](https://github.com/microsoft/graphrag/issues/1718) +— closed as "not planned," the flagship evidence for the whole thesis — **no commenter +mentions benchmarks, ground truth, or evaluation.** They asked for the bug to stop. Revealed +preference on solution-vs-measurement: Splink 2,308★ and Zingg 1.2k★ are *solutions*; the +best-maintained synthetic ER data generator is **2★**. A ~600× ratio. + +**On "lucrative":** no, not directly, and §7 does the arithmetic rather than gesturing. +Direct revenue >$50k/yr: **5–10%**. Reputational return converting to consulting or +employment: **~70%**. If direct revenue is the real goal, it lives in the application layer, +and that is a different project than this one. + +--- + +## 1. Honest audit: what GraphFaker is today + +### The code + +| Component | LOC | Assessment | +|---|---|---| +| `core.py` faker generator | 382 | **Demo, not a product.** See below. | +| `fetchers/flights.py` | 336 | The only real engineering. BTS + carrier lookup → typed heterogeneous graph. | +| `fetchers/osm.py` | 120 | Thin pass-through to OSMnx. No moat — users would call OSMnx directly. | +| `fetchers/wiki.py` | 82 | Orphan. Emits raw text, no graph. Disconnected from the story. | +| `cli.py` | 123 | Works, with the bug below. | + +**The synthetic generator is not "realistic."** +[`core.py:generate_edges`](graphfaker/core.py) picks *both* endpoints with uniform +`random.choice` inside each type-pair. That is Erdős–Rényi attachment: Poisson degree +distribution, no preferential attachment, no clustering, no community structure, and no +correlation between attributes and topology (an organization's `industry` is literally +`fake.job()`). This matters beyond aesthetics — it is the first reason the ground-truth +thesis fails (§4, Front 2). + +**Three verified defects:** + +1. **No `seed` parameter anywhere** (`grep -rn seed graphfaker/` → no matches). Output is not + reproducible. Table stakes for anything calling itself a test-data library. +2. [`cli.py:5`](graphfaker/cli.py) is `from venv import logger` — the CLI logs through the + **standard library `venv` module's** logger by accident. +3. [`flights.py:128`](graphfaker/fetchers/flights.py) and + [`:187`](graphfaker/fetchers/flights.py) pass `verify=False` to `requests.get`. TLS + verification disabled — disqualifying inside any enterprise, and a genuine MITM exposure. + +### The traction, and the throughput constraint + +- **36 stars, 4 forks. 89 PyPI downloads in the last 30 days** (lifetime peak 194/mo). + Last commit 2025-08-26 — **11 months idle.** 11 open issues, 1 open PR. +- One effective maintainer (58 of 71 commits), 2 occasional contributors, no funding. +- GraphGeeks parent org: ~1,200 Discord members (2.2× YoY), YouTube 2k→24k views, flagship + in-person event **capped at 75 attendees**. + +**The uncomfortable observation, and it is the one that should drive planning:** the README +already publishes a roadmap — GraphML/JSON-LD/RDF export, Neo4j/Kuzu/TigerGraph +integration, million-node scale, LLM-powered fetching. **None of it shipped.** The binding +constraint is not strategy, it is throughput. A plan that is *more* ambitious than the +unshipped one is not a pivot; it is an escape from the real problem. Any credible plan must +be *smaller* than what has already failed to ship. + +### What is actually an asset + +1. **The name.** "Faker, but for graphs" explains itself and owns the PyPI slot. Worth more + than any code in the repo. +2. **The GraphGeeks channel.** ~1,200 graph practitioners reachable in a day. +3. **The flights pipeline** — real work, a good teaching dataset. +4. Everything else, including the synthetic generator, is replaceable. + +**The defensible asset is positioning plus audience, not code.** + +--- + +## 2. The four candidate ideas — and why each one dies + +### (a) Hosted notebook for schema-less graphs, "like DSPy" — **DROP** + +A distribution tactic mistaken for a product. The DSPy analogy doesn't parse (DSPy is a +programming model, not a notebook). No unique leverage, trivially cloneable, and hosting is +a recurring cost against a $0 budget. + +*Salvageable kernel:* "schema-less" really means *let users bring their own ontology instead +of hardcoded module-level constants.* A legitimate library feature — a Colab link, not a +service. + +### (b) A v2 that validates/lints someone's knowledge graph — **DROP** + +Validating *someone else's* graph requires *their* schema and expectations, which GraphFaker +does not have — so the generator asset provides **zero** leverage. And property graphs have +no dominant constraint standard (SHACL is RDF-only; a 2026 survey of 94 practitioners found +22 explicitly asking for property-graph support that does not exist). This means inventing a +spec language solo: a multi-year standards fight. + +### (c) Lightweight graph observability / "GraphOps" — **DROP, hardest** + +**The discipline does not exist.** "Graph observability," "graph data quality," and "graph +drift detection" are not terms — no vendor category, no analyst coverage, no conference +track. Great Expectations, Soda, Elementary, and Monte Carlo have **zero** graph connectors. +There is **no dbt adapter for Neo4j or Cypher** on dbt's community or trusted lists. Public +pain is single-digit forum threads. Zero job postings mention graph testing or observability. + +No startup died attempting this **because nobody tried.** And the closest analogue market is +a warning: Gartner sized *all* of data observability at **$346M in 2024 against $1B+ of VC +invested.** Great Expectations broke up (GX Cloud → FICO, shut down 2026-06-01; GX Core +handed to Fivetran). Monte Carlo has raised nothing since 2022. + +The vacuum is far more likely absence of demand than unmet demand. And it is a platform play +— integrations per graph DB, dashboards, on-call ergonomics — requiring exactly the +resources this project lacks. + +### (d) A graph-engineering workflow on LangGraph — **DROP** + +Maximally crowded (LightRAG 38.4k★, Microsoft GraphRAG 35.1k★, Cognee 29.6k★, Graphiti +29.4k★, LlamaIndex, neo4j-graphrag), tightly coupled to a framework that may not exist in +its current form in 18 months, discards the existing codebase, and the target user is +undefined. Trend-surfing, not asset-compounding. + +--- + +## 3. The thesis that looked right (and was tested to destruction) + +Stated fairly, because it is genuinely appealing: + +> A generator is the only artifact in the graph stack with **gold labels for free, because it +> created them.** GraphRAG quality, entity-resolution quality, and KG validation are all +> *measurement* problems blocked on ground truth that real graphs never carry. So GraphFaker +> should become the ground-truth substrate: `corrupt(G, spec) -> (G_dirty, gold_labels)`, +> then paired text↔graph corpora, then `score(pipeline_fn, corpus) -> report`. + +Supporting arguments that are still true: **contamination immunity** (a graph generated this +morning cannot be in any LLM's training data, unlike HotpotQA or Wikidata subsets), and +**difficulty as a dial** (sweep duplicate rate or hop depth and plot a degradation curve — +no fixed dataset allows that). + +There was also a real-looking gap. Every flavour of graph ground truth exists **in +isolation**: LFR/ABCD for communities (topology only), SynthKGQA for QA paths (derives from +Wikidata, doesn't synthesize the graph), TravelFraudBench for fraud rings (single vertical), +FEBRL/GeCo for duplicates (**flat record tables, not graphs**), PyGraft for schema +consistency (**713★, rigorous OWL reasoning, but no attribute values**), LDBC SNB for scale +(Scala/Spark, emits files not objects). Nothing combined typed graph + realistic attributes ++ gold labels + in-process NetworkX. + +And the go-to-market looked sharp: don't sell a benchmark to end users, sell fixtures to the +~6 framework maintainers with public, dated, unfixed dedup bugs — MS GraphRAG #1718 (closed +"not planned," outsiders still re-patching 16 months later via #2339/#2411), Neo4j shipping +`perform_entity_resolution` **off by default** while publicly saying it isn't happy with it, +Cognee's three-way dedup hackathon in Jun 2026 (#3627/#3628/#3629, all open), LangChain +`LLMGraphTransformer` producing a typed node on one run and an untyped node on the next +([#26614](https://github.com/langchain-ai/langchain/issues/26614)), `PropertyGraphIndex` +deduping against the store but not within a batch. + +**All of that is true, and the thesis still fails.** §4. + +--- + +## 4. Pressure test: what actually killed it + +Five attacks were run. Three land decisively. Two were conceded by the adversary itself and +are recorded as *not* killing anything — kept because a pressure test that only reports hits +is worthless. + +### KILL 1 — Splink built this and deleted it *(most lethal)* + +- `moj-analytical-services/splink_synthetic_data` → **404**, absent from the org's repo list. +- [#110 "Generate artificial linkage data to benchmark splink performance"](https://github.com/moj-analytical-services/splink/issues/110) + opened 2020-06-19, **closed**, punted to an external R package. #292, #316, #632: also closed. +- Splink today: **2,308★**, shipped 2026-07-31. It chose **static CSVs with a `cluster` + column** (`fake_1000`: `unique_id,first_name,surname,dob,city,email,cluster`). + +Corroborating: the canonical [Leipzig ER benchmark](https://dbs.uni-leipzig.de/research/projects/benchmark-datasets-for-entity-resolution) +**already used GeCo and DAPO to synthesize its duplicates** — then froze in 2019. Modern +Python GeCo ports: [geco3](https://github.com/T-Strojny/geco3) **2★**, +[GeCoWrapper](https://github.com/dobraczka/GeCoWrapper) **1★**, +[geco_data_generator_corruptor](https://github.com/sashaostr/geco_data_generator_corruptor) **0★**. + +Graph-shaped ER benchmark requests found across GitHub, Stack Overflow, r/dataengineering, +HN, five query variants: **zero.** + +> ⚠️ **One caveat worth checking before fully accepting this.** A deleted repo is not +> definitionally "rejected" — it could have been consolidated into `splink_datasets` or +> dropped for maintenance load rather than lack of demand. That specific inference is the +> load-bearing part of the strongest kill, so it deserves ten minutes of verification (ask in +> the Splink discussions) before anyone acts on it. The four independently closed issues make +> the inference reasonable, not certain. + +### KILL 2 — Ragas already ships Phase 2, at 400× the audience + +[Ragas](https://github.com/explodinggradients/ragas) — **15.1k★** — already +[builds a knowledge graph from documents](https://docs.ragas.io/en/stable/concepts/test_data_generation/rag/) +(hierarchical nodes, entity-similarity relations), then runs `QuerySynthesizer` to emit +**single-hop and multi-hop** queries as `EvalSample (Query, Context, Reference)` with +ground-truth `reference` answers. That *is* the proposed text↔graph paired corpus with gold +labels, shipping today. + +Phase 3 is occupied too: +[GraphRAG-Bench](https://github.com/GraphRAG-Bench/GraphRAG-Benchmark) — **471★**, active +2026-05-17, four difficulty tiers — and it uses **real** novel and medical corpora +*specifically because synthetic wouldn't be credible*. + +### KILL 3 — Nobody in pain asked for measurement + +[microsoft/graphrag#1718](https://github.com/microsoft/graphrag/issues/1718), the flagship +evidence: commenters discuss **the bug only**. No commenter mentions benchmarks, ground +truth, or evaluation. Solution-vs-measurement revealed preference: Splink **2,308★** and +Zingg **1.2k★** are solutions; the best synthetic-ER generator is **2★**. Cognee's dedup +hackathon signals they want an *implementation*, not a scoreboard. + +A benchmark is a vitamin sold to people with a broken leg. + +### KILL 4 — Synthetic corruption is ~100× too easy *(lands on fidelity, partially conceded)* + +[Lam et al., IJPDS 9:1:18 (2024)](https://discovery.ucl.ac.uk/id/eprint/10194216/1/ijpds-09-2389.pdf): +real ALSPAC linkage missed-match rates **4.59% / 2.61% / 2.40%**. Under naive corruption — +guessed error rates, errors independent of attributes, no co-occurrence, i.e. *exactly what a +`corrupt(G, spec)` API does* — the same methods score **0.12% / 0.03% / 0.02%**. + +[Alsadeeqi 2020](https://www.ros.hw.ac.uk/server/api/core/bitstreams/43d977dc-4acc-4b63-8f74-2e2cdbb50bec/content) +is worse: on real data all four string comparators returned **identical F=0.896091**, while +on corrupted synthetic data Jaro-Winkler "delivered the highest linkage." **The corruptor +manufactured a ranking the real data does not support.** + +The multi-hop QA half is independently dead. [Min et al., ACL 2019](https://arxiv.org/abs/1906.02900): +single-hop BERT hits **67 F1** on HotpotQA. [MuSiQue](https://arxiv.org/abs/2108.00573): +"existing multihop benchmarks are known to be largely solvable via shortcuts." *Human-curated* +multi-hop benchmarks are shortcut-riddled; questions mechanically generated by traversing a +graph you built will be strictly worse — the generator's traversal *is* the shortcut. + +Also: ER benchmarks are saturated — +[Primpeli & Bizer, CIKM 2020](https://www.uni-mannheim.de/media/Einrichtungen/dws/Files_Research/Web-based_Systems/pub/CIKM2020_Primpeli_Bizer.pdf): +"six benchmark tasks are perfectly solved by the baseline method (F1=1.00)." + +**Concession recorded:** no paper was found showing matcher *rankings invert* between +synthetic and real data; Lam et al. preserve the top-line ordering. So the lethal form of +this attack is unproven. The provable form — absolute numbers inflated ~100×, and +fine-grained comparator differences that are artifacts of the corruption model — is still +enough to gut the value proposition. Lam et al. also names the only fix: model **error +co-occurrence and error/attribute dependence** — which requires real labeled data you don't +have. That is circular, and it is a research program, not 200 lines. + +### KILL 5 — The unshipped roadmap *(execution, not strategy)* + +Covered in §1. The README's existing modest roadmap never shipped; 11 months idle. Proposing +something harder is an escape fantasy. **The corruption *model*, not the API, is the whole +product** — and §4 Kill 4 shows the model is the expensive part. + +### ATTACKS THAT FAILED — recorded honestly + +- **PyGraft will not eat this.** 713★, 53 forks, **2 open issues**; "Add support for + literals" is explicitly **LOW priority**; + [#6](https://github.com/nicolas-hbt/pygraft/issues) asking about the next release has been + open and unanswered since 2025-03-26. PyGraft is itself stalling. *But read the signal:* the + 713★ incumbent doesn't think attribute values are worth prioritizing, and its users aren't + pushing. +- **The base-rate argument is unsupported.** No data exists on solo-maintainer OSS pivot + success rates. The adversary withdrew it rather than invent a number. §1's throughput + argument replaces it and is stronger because it's specific to this repo. +- **There is one live thread of real demand.** Splink + [#2191 "Cluster evaluation — with ground truth data"](https://github.com/moj-analytical-services/splink/issues/2191) + and [#2274](https://github.com/moj-analytical-services/splink/issues/2274) are **both + open**. Note precisely what is wanted: better *scoring against labels users already have* — + a metrics API, **not** a generator to manufacture labels. §5. + +--- + +## 5. What survives + +Not a substrate, not a corpus, not a harness. Roughly 300 lines, under a different name: + +**Graph-aware cluster-evaluation metrics for NetworkX.** Given a predicted node clustering +and a gold clustering, compute pairwise and cluster-level precision/recall/F1, B-cubed, plus +split/merge diagnostics. + +It survives for three reasons: it is the only piece with a **verified public request** +(Splink #2191/#2274, both open — and Splink is Spark/DuckDB-native, so a NetworkX-side +implementation isn't redundant); it takes labels the user **already has** rather than +manufacturing labels whose realism is indefensible; and it is immune to Kill 4 entirely — +**a metric cannot be too easy, only the data can.** + +It is also not a business, not a pivot, and not an identity. It is a weekend PR, ideally +submitted to an existing library. **That the surviving fragment is too small to justify +repositioning a project around is itself the answer to the thesis.** + +--- + +## 6. The revised plan + +**Chosen objective: reputation and career.** Not revenue (§7 explains why that would be a +different project), and not stars. That choice sets the optimisation target: reputation comes +from artifacts *other people cite or use* — a merged PR in a large repo, a novel measurement, +a conference talk — not from a package you maintain alone. + +The sequencing inverts the discarded thesis: **solution first, measurement never as the +product.** + +### The forcing function + +The binding constraint is throughput, not strategy (§1). Strategy documents do not fix that; +deadlines do. Two real ones sit ~3.5 months out: + +- **NODES 2026** — Nov 12, free, virtual, CFP via sessionize +- **Connected Data London** — Nov 11–12, 10th anniversary + +Work backwards from a November talk. Every step below yields an artifact, and each feeds the +next. + +### Step 0 — Hygiene (½ day, unconditional) + +See Step 1. It comes first because the measurement in Step 1 is worthless if made with a tool +that has no `seed` and disables TLS. Reproducibility is the credibility of the result, not +housekeeping. + +### Step 1 — Publish the measurement nobody has published (2 weeks) + +Generate one controlled corpus. Run it through five frameworks — Microsoft GraphRAG, LightRAG, +Cognee, Graphiti, LlamaIndex `PropertyGraphIndex`. **Count how many nodes each creates for +entities known to be singular.** Publish the table. + +Why this specific experiment survives the §4 kills, which the discarded eval thesis did not: + +- **No corruption model is involved.** You feed *clean* input containing one + `Acme Corporation` and observe that a framework emitted five nodes. Kill 4 (Lam et al., + synthetic corruption ~100× too easy) is a critique of simulated *errors*; there are none + here. +- **It is not answer-quality evaluation.** Ragas and GraphRAG-Bench measure generation + quality. Per-framework duplicate rates on identical input are unoccupied, so Kill 2 does + not apply. +- **It is diagnosis, not a benchmark product.** Kill 3 says nobody wants to buy a scoreboard. + A published finding is not a product. + +This is the one place the ground-truth property is genuinely load-bearing and unattackable: +you know what you put in, so you can count what came out. + +### Step 2 — That post is the CFP abstract + +Submit it. "I measured entity duplication across five GraphRAG frameworks" is accepted because +it is a number, not an opinion. + +### Step 3 — Build the second act (weeks 3–8) + +A talk that is only a diagnosis is a complaint. Ship `resolve()` so the arc becomes +*problem → measurement → fix*, then take it to Cognee as a PR. A merged PR in a 29k-star repo +is the highest reputation-per-hour trade available, and it is the talk's climax. + +### Step 4 — November + +Give the talk, holding a novel measurement, a working tool, and ideally a merged upstream +contribution — from roughly eight weeks of work. + +### Step 1 detail — Hygiene (½ day, do regardless of everything else) + +1. Remove `verify=False` from [`flights.py`](graphfaker/fetchers/flights.py) or fence it + behind an explicit opt-in flag. +2. Fix `from venv import logger` at [`cli.py:5`](graphfaker/cli.py) and the stale + `"Use 'random' or 'osm'"` error message. +3. Add `seed` to `generate_graph`, threaded through `random` and `Faker`. + +These are unconditional. The library is published; #1 is a security defect and #3 is table +stakes. + +### Step 2 — Ship a dedup fix upstream (the actual play) + +1. **Target Cognee.** Its Jun 2026 dedup hackathon (#3627/#3628/#3629) is the only case where + *both* demand and maintainer receptivity are already verified. Second choice: LightRAG or + Graphiti (largest audiences). MS GraphRAG is the worst target — it closed its dedup issue + "not planned." +2. Write a working entity-resolution pass for their KG pipeline as a real PR. **Do not build + an ER engine from scratch** — Splink owns that (1.14M downloads/month). Wrap or borrow. +3. **Use GraphFaker inside that PR's tests** as the fixture proving the fix. The corruption + code gets written — as a means, never as the product, and never load-bearing on realism + claims. +4. If it merges, repeat on the next framework. Then, and only then, extract `graph-dedup` as + a standalone library — with verified users on day one instead of a hopeful launch. +5. Write the GraphGeeks post about it. At ~1,200 members that is the one distribution asset + that reliably works. + +**Why this beats the thesis on every constraint:** it borrows distribution from a +15k–30k-star ecosystem instead of hoping for adoption at 36★; it is immune to Kill 4 because +a dedup implementation is judged by whether it fixes the user's graph, so synthetic realism +never becomes load-bearing; it fits the throughput constraint (one PR, not an 18-month +research program); it sits on the 600×-larger side of the solution/measurement ratio; and a +merged PR in a 15k-star repo delivers the reputational payoff — the thesis's actual stated +monetization — far faster than a benchmark nobody cites. + +### Step 3 — Optional, opportunistic + +Contribute the §5 cluster-evaluation metrics as a PR to an existing library. One weekend. +Don't launch it as a project. + +### Standing kill list + +Hosted notebook service · KG linter · GraphOps platform · LangGraph workflow · "ground-truth +substrate" positioning · public `corrupt()` API as a product · multi-hop QA generation · +a GraphRAG eval harness. + +### What to do about `graphfaker` itself + +It becomes the **instrument**, which is a real role rather than a demotion. Keep the name and +the PyPI slot — the best assets here — and let it be a small, honest library that is credible +enough to make a measurement with. Fix the defects, keep the flights pipeline, add +`resolve()`, and delete the roadmap promises; the existing ones never shipped and unkept +promises read worse than a modest scope. + +The story stops being "I maintain a synthetic graph library" and becomes "I found and fixed a +bug class in the GraphRAG ecosystem, and built the tooling that proves it." For a career in a +field with ~245 open roles at $150–280K, the second sentence is worth an order of magnitude +more. + +### Shipped in this branch + +- Seeding (`GraphFaker(seed=...)`, `generate_graph(..., seed=...)`, `reseed()`, `--seed`), + without touching global `random` state. +- TLS verification restored by default; `GRAPHFAKER_INSECURE_TLS=1` opts out loudly. +- `tqdm` and `urllib3` declared — they were imported but missing from dependencies, so a clean + install could fail on `import graphfaker`. +- CLI logger no longer resolves to the stdlib `venv` module's logger. +- `graphfaker.resolve`: `resolve_entities()`, `merge_clusters()`, `evaluate_clusters()`, and + `GraphFaker.resolve()`. +- 45 new tests; the 3 CLI tests that were failing on `main` now pass. +- README stripped of unimplemented promises. + +--- + +## 7. The "lucrative" question, answered with arithmetic + +Where money actually flowed in graph/KG in 2025–26 (all verified): + +- **Applications, not tooling.** Glean: $150M Series F at $7.2B, ~$300M ARR, on a + permissions-aware knowledge graph. Palantir: FY2025 revenue $4.40B, US commercial +104%, + sold as "Ontology." Quantexa: $175M Series F at $2.6B, $100M+ ARR, entity-resolution-as-a-product. +- **Graph infrastructure is being harvested, not scaled.** Kuzu archived 2025-10-10, team + acqui-hired by Apple. Illumex → NVIDIA (~$60–75M on $13M raised). Linkurious → Nuix + (≤€20M). GraphAware/Hume → Neo4j. TigerGraph → PE. Memgraph, Stardog, FalkorDB, ArangoDB + financially frozen. Neo4j: $200M+ ARR at a $2B mark — but that was Nov 2024, no round + since, and secondary marks are $1.1–1.2B. +- **The eval category out-raised the entire graph long tail:** ~$230M in 12 months + (Braintrust $80M at $800M, LangChain $125M at $1.25B on $12–16M ARR, Arize $70M, Patronus + $50M). A solo project cannot enter there — and note this cuts *against* the eval thesis, + not for it. +- **Jobs are real, senior, and tiny.** Dice.com US, 2026-07-31: "knowledge graph" 245, + Neo4j 244, GraphRAG 50, ontology engineer 15 — versus data engineer 3,027. $150–280K. + LinkedIn's own *Skills on the Rise 2026* mentions none of: knowledge graph, ontology, + semantic, graph, taxonomy. + +| Outcome | Odds | +|---|---| +| Direct revenue > $50k/yr | **5–10%** | +| Vendor sponsorship of a neutral eval suite | **10–15%** (revised down from 20–30% — GraphRAG-Bench already occupies the neutral-benchmark slot with real data) | +| Reputational return → consulting / employment / talks | **~70%**, and Step 2 is the fastest route to it | + +**Nobody in 2025–26 paid meaningfully for graph dev tools or benchmarks.** Plan for the +reputational return; treat revenue as upside. If direct revenue is the goal, it is in the +application layer — a different project. + +--- + +## 8. GraphRAG deflation risk + +Treat GraphRAG as a **customer segment, not an identity.** Entity resolution predates it by +decades and will outlive it — which is a further argument for Step 2 over the eval thesis. + +The research does not support imminent deflation: arXiv GraphRAG papers went 13 (2024) → 91 +(2025) → **114 in 2026 YTD**, and Gartner places knowledge graphs on the Slope of +Enlightenment heading to the Plateau of Productivity. But the center of gravity has already +moved once — from batch GraphRAG to incremental agent memory (Graphiti's downloads went 311K +→ **1.59M/month** Feb→Jul 2026 while Microsoft GraphRAG's stayed flat and its Azure solution +accelerator was archived). It will move again. Anything coupled tightly to a single framework +inherits that risk; a dedup PR does not. + +--- + +## 9. Verify before acting + +1. **Why did Splink delete `splink_synthetic_data`?** This is the load-bearing inference in + the strongest kill (§4 Kill 1). Ten minutes in their discussions resolves it. If the answer + is "consolidated for maintenance reasons," Kill 1 weakens materially — though Kills 2, 3, + and 5 stand on their own and are individually sufficient. +2. **Confirm Cognee would accept an external dedup PR** before writing it. One comment on + #3627. This is the cheapest possible de-risking of Step 2 and should happen first. +3. **Do not fund the corruption-realism problem.** If Step 2 somehow creates pull for a public + `corrupt()` API, remember Lam et al.: credible realism needs error co-occurrence modeled + from real labeled data. Price that as research, not as a feature. diff --git a/README.md b/README.md index 80c836f..1bd9070 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,14 @@ GraphFaker is an open-source Python library designed to generate, load, and expo - `flights`: Flight/airline networks from Bureau of Transportation Statistics (airlines ↔ airports ↔ flight legs, complete with cancellation and delay flags) - **Unstructured Data Source:** - `WikiFetcher`: Raw Wikipedia page data (title, summary, content, sections, links, references) ready for custom graph or RAG pipelines +- **Entity Resolution:** + - `resolve()`: find and merge duplicate nodes using attribute similarity **and** neighbourhood overlap — the graph signal tabular record-linkage tools cannot see + - `evaluate_clusters()`: score a predicted clustering against gold labels you supply (pairwise and B-cubed) +- **Export Connectors:** + - CSV, `neo4j-admin` bulk-import CSV, Cypher, openCypher, and ISO GQL — file-based, so no driver or running database is needed +- **Measurement:** + - `generate_corpus()`: documents whose entities are known in advance, for counting how many nodes a graph builder creates per real entity +- **Reproducible:** every synthetic graph is seedable - **Easy CLI & Python Library** This removes friction around data acquisition, letting you focus on algorithms, teaching or rapid prototyping. @@ -134,31 +142,255 @@ You can also use `--date-range` for custom time spans (e.g., `--date-range "2024 --- -## Future Plans: Graph Export Formats +## Entity Resolution -- **GraphML**: General graph analysis/visualization (`--export graph.graphml`) -- **JSON/JSON-LD**: Knowledge graphs/web apps (`--export data.json`) -- **CSV**: Tabular analysis/database imports (`--export edges.csv`) -- **RDF**: Semantic web/linked data (`--export graph.ttl`) +LLM-built knowledge graphs routinely emit the same real-world entity as several +nodes, and every edge attached to a false node is a false edge. Tabular +record-linkage tools compare *rows*, so they cannot use the strongest signal a +graph offers: **two nodes that share most of their neighbours are probably the +same entity, however differently their names are spelled.** + +`resolve()` scores candidate pairs on attribute similarity *and* neighbourhood +overlap, clusters the survivors, and merges each cluster onto one canonical node +— rewiring its edges, dropping self-loops the merge creates, and recording what +was absorbed. + +```python +from graphfaker import GraphFaker + +gf = GraphFaker(seed=42) +gf.generate_graph(source="faker", total_nodes=500, total_edges=2000) + +result = gf.resolve(on=["name", "email"], threshold=0.85) +print(result.report()) +# candidate pairs scored : 1284 +# pairs above threshold : 12 +# clusters found : 5 +# duplicate nodes : 7 + +clean = result.apply() # merged copy; the original is untouched +``` + +`structural_weight` controls how much shared-neighbour evidence may lift a +pair's score. Structure can only *raise* a score, never lower it, so isolated +nodes are never penalised for having few neighbours — set it to `0` to fall back +to plain attribute matching: + +```python +gf.resolve(on=["name"], structural_weight=0.0) # attributes only +gf.resolve(on=["name"], structural_weight=0.8) # trust the graph structure +``` + +Shortened names get special handling. `"Hill"` against `"Allison Hill"` scores +only ~0.5 on character similarity, because most of the longer string is +unmatched — so it would be discarded before the structural signal was ever +consulted, even though referring back to an entity by a shorter form is one of +the commonest things a document does. When one name's tokens are contained in +the other's, the score is *floored* at `token_subset_floor` (0.75 by default) +rather than set to 1.0: containment is suggestive, not conclusive, so it lifts +the pair into consideration and leaves the decision to shared neighbours. Pass +`token_subset_floor=0.0` to switch it off. + +Already have labelled clusters? Score a prediction against them. This computes +metrics only — it does not manufacture ground truth: + +```python +from graphfaker import evaluate_clusters + +scores = evaluate_clusters(result.clusters, my_known_duplicates) +print(scores["pairwise_f1"], scores["b_cubed_f1"]) +``` --- -## Future Plans: Integration with Graph Tools +## Graph structure + +Synthetic graphs are built to look structurally like real ones. Until 0.5 both +endpoints of every edge were drawn uniformly at random, which produces an +Erdős–Rényi graph: a Poisson degree distribution with no hubs, effectively no +clustering, and no community structure. Edges are now formed by **preferential +attachment** (popular nodes attract more), **triadic closure** (friends of +friends become friends), and **homophily** over latent communities. + +Measured on 600 nodes / 2,400 edges, seed 1 — reproduce with +`graphfaker.metrics.compare_topology`: + +| metric | realistic | uniform (pre-0.5) | real graphs | +| --- | --- | --- | --- | +| degree Fano factor | **7.9** | 1.7 | ≫ 1 | +| max degree | **88** | 19 | hubs exist | +| degree Gini | **0.45** | 0.26 | unequal | +| average clustering | **0.180** | 0.013 | ≫ random baseline | +| community modularity | **0.72** | −0.004 | 0.4–0.7 | +| age homophily | **0.81** | 0.01 | positive | +| isolated nodes | **0** | 4 | giant component | + +The Fano factor — degree variance over mean — is the clearest single test: a +Poisson distribution has variance equal to its mean, so uniform attachment sits +near 1 by construction and cannot be made to look otherwise. + +The friendship layer alone (Person–Person edges) has clustering 0.51, modularity +0.91, and *positive* degree assortativity, which is what social networks look +like. The graph as a whole is mildly disassortative because it is multipartite — +people attach to hub cities and large employers — as real knowledge graphs are. + +```python +from graphfaker.metrics import compare_topology, graph_stats + +realistic = GraphFaker(seed=1).generate_graph(total_nodes=600, total_edges=2400) +uniform = GraphFaker(seed=1).generate_graph(total_nodes=600, total_edges=2400, + topology="uniform") +print(compare_topology({"realistic": realistic, "uniform": uniform})) +``` + +`topology="uniform"` is kept only for this comparison. It is not a supported way +to generate data. + +**Why it matters beyond looking right:** on a realistic graph, entity resolution +is measurably *harder*. Injecting the same known duplicates into both and running +`resolve()` gives precision 1.000 on the uniform graph but 0.88 on the realistic +one — because homophily means genuinely distinct people share attributes and +neighbours. Anything benchmarked against the old generator was flattered by it. -GraphFaker generates NetworkX graph objects that can be easily integrated with: -- **Graph databases**: Neo4j, Kuzu, TigerGraph -- **Analysis tools**: NetworkX, SNAP, graph-tool -- **ML frameworks**: PyTorch Geometric, DGL, TensorFlow GNN -- **Visualization**: G.V, Gephi, Cytoscape, D3.js +Attributes are also no longer independent of structure. `population` tracks a +place's connectivity, `employee_count` correlates 0.95 with the number of +`WORKS_AT` edges actually present, ages cluster by community, and `industry` +holds an industry rather than the job title `fake.job()` used to supply. +`LIVES_IN` and `BORN_IN` are singular — previously a person could live in four +cities at once. --- -## On the Horizon: +## Reproducibility + +Synthetic generation is seedable, per instance. The same seed and the same +arguments always produce an identical graph, and seeding does not disturb the +global `random` module: + +```python +GraphFaker(seed=42).generate_graph(source="faker", total_nodes=100) +# or per call: +gf.generate_graph(source="faker", total_nodes=100, seed=42) +``` + +```sh +python -m graphfaker.cli --fetcher faker --total-nodes 100 --seed 42 +``` + +--- + +## Getting the graph into a database + +Rather than shipping a driver per database — each needing credentials, a version +matrix, and a live service to test against — GraphFaker writes files that every +engine's own loader already understands. + +```python +from graphfaker.export import export_csv, export_neo4j_csv, export_cypher + +export_csv(G, "nodes.csv", "edges.csv") # pandas, Gephi, any bulk loader +export_neo4j_csv(G, "import/") # neo4j-admin bulk import headers +export_cypher(G, "load.cypher") # Neo4j, Memgraph, Kuzu +export_cypher(G, "load.gql", dialect="gql") # ISO GQL +``` + +Or from the CLI: + +```sh +python -m graphfaker.cli --fetcher faker --total-nodes 500 --format cypher --export load.cypher +python -m graphfaker.cli --fetcher faker --total-nodes 500 --format neo4j-csv --export import/ +``` + +| Format | Loads into | +| --- | --- | +| `graphml` | Gephi, Cytoscape, NetworkX, igraph | +| `csv` | pandas, TigerGraph `LOAD`, Amazon Neptune bulk loader, Spark/GraphFrames | +| `neo4j-csv` | `neo4j-admin database import` — the fast path for large graphs | +| `cypher` | Neo4j, Memgraph, Kuzu | +| `opencypher` | Amazon Neptune | +| `gql` | ISO GQL engines (`INSERT` in place of `CREATE`) | + +Node labels come from the `type` attribute and relationship types from +`relationship`, both configurable. Nodes of different types carry different +attributes, so CSV headers are the **union** of all keys seen — a node missing a +column gets an empty cell rather than having its values shifted into the wrong +one. Container values (coordinate tuples, merge provenance) are flattened, and +labels containing punctuation are sanitised. + +Once loaded, Neo4j Graph Data Science works directly on the result: + +```cypher +CALL gds.graph.project('g', '*', '*'); +CALL gds.pageRank.stream('g') YIELD nodeId, score +RETURN gds.util.asNode(nodeId).name AS name, score ORDER BY score DESC LIMIT 10; +``` + +--- + +## Measuring entity duplication + +`graphfaker.corpus` builds documents whose entities are known in advance, so you +can count how many nodes a graph builder creates for entities that are singular. + +```python +from graphfaker.corpus import generate_corpus, duplication_report + +corpus = generate_corpus(seed=42, n_entities=60, n_documents=80) +assert corpus.audit()["clean"] # no surface form belongs to two entities +corpus.write("corpus/") # documents + gold.json + +# ...run any graph builder over corpus/, then: +report = duplication_report(extracted_graph, corpus, framework="my-pipeline") +print(report.summary()) +``` + +Nothing is corrupted. The text is clean, well-formed English and every entity is +unambiguous to a human reader, so a correct pipeline scores zero. Entities are +referred to by the surface forms a normal writer uses — full name, surname +alone, an accepted abbreviation — which is ordinary prose, not injected noise. + +That restraint is deliberate. Synthetic *corruption* is far easier than +real-world error (Lam et al., IJPDS 2024, measured roughly a hundredfold gap), +so a benchmark built on guessed error rates mostly measures its own noise model. +Counting splits of entities a human would never split is a weaker claim, and one +a generator can actually support. + +`examples/duplication_experiment.py` runs this across several frameworks and +prints a comparison table. It refuses to run on an ambiguous corpus, includes a +perfect-extractor control that must score zero, and names any framework it +skipped rather than omitting it silently. + +**[`docs/notebooks/duplication_experiment.ipynb`](docs/notebooks/duplication_experiment.ipynb)** +walks the whole thing end to end — build the corpus, audit it, run +[Cognee](https://github.com/topoteretes/cognee), inspect what got split, repair +it with `resolve()`, sweep the threshold to see the precision/recall tradeoff, +and export the cleaned graph. Steps other than the Cognee run work without an +LLM key, using a clearly-labelled simulated extraction so the notebook is +runnable as a tutorial. + +On that simulated graph, `resolve()` takes duplication from **60% to 20%** and +node inflation from 2.15× to 1.20× at precision 1.000 — but read the +[caveats](docs/notebooks/duplication_experiment.ipynb) before quoting numbers +like that. Structural matching only helps when duplicate nodes share +neighbours; an extractor that *partitions* an entity's edges leaves almost no +overlap to find, which is the hard case. + +--- + +## Scope + +Anything not documented above is not implemented. Please open an issue if you +need something specific rather than assuming it is on the way. + +--- -- Handling large graph -> millions of nodes -- Using NLP/LLM to fetch graph data -> "Fetch flight data for Jan 2024" -- Connects to any graph database/engine of choice -> "Establish connections to graph database/engine of choice" +## Notes on network access +The `flights` fetcher downloads from BTS and OpenFlights with TLS verification +enabled. Some systems fail to validate the BTS certificate chain; if you hit +that, you can opt out with `GRAPHFAKER_INSECURE_TLS=1`, which logs a warning and +means the downloaded data is no longer authenticated. Verification is never +disabled silently. --- diff --git a/docs/notebooks/duplication_experiment.ipynb b/docs/notebooks/duplication_experiment.ipynb new file mode 100644 index 0000000..0a27f9c --- /dev/null +++ b/docs/notebooks/duplication_experiment.ipynb @@ -0,0 +1,852 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "cell-00", + "metadata": {}, + "source": [ + "# Measuring entity duplication in an LLM-built knowledge graph\n", + "\n", + "This notebook runs one experiment end to end:\n", + "\n", + "1. Build a corpus of documents whose entities are **known in advance**\n", + "2. Confirm a perfect extractor scores zero (the control)\n", + "3. Hand the same text to [Cognee](https://github.com/topoteretes/cognee) and count how many nodes it creates per real entity\n", + "4. Look at what actually got split\n", + "5. Repair the graph with `graphfaker.resolve()` and re-measure\n", + "6. Export the cleaned graph to Cypher / Neo4j\n", + "\n", + "## What this measures\n", + "\n", + "How often a pipeline emits **more than one node for one entity**, on clean input.\n", + "\n", + "## What it does not measure\n", + "\n", + "Answer quality, retrieval quality, or behaviour on messy real-world documents.\n", + "Those are different claims and this notebook does not support them.\n", + "\n", + "Nothing here is corrupted. The documents are clean, well-formed English and\n", + "every entity is unambiguous to a human reader, so a correct pipeline scores\n", + "**zero**. Entities are referred to by the surface forms a normal writer uses —\n", + "full name, surname alone, an accepted abbreviation — which is ordinary prose,\n", + "not injected noise.\n", + "\n", + "That restraint is deliberate. Synthetic *corruption* is far easier than\n", + "real-world error — Lam et al. ([IJPDS 2024](https://discovery.ucl.ac.uk/id/eprint/10194216/1/ijpds-09-2389.pdf))\n", + "measured real linkage error at 4.59% against 0.12% under naive corruption, a\n", + "~100x gap — so a benchmark built on guessed error rates mostly measures its own\n", + "noise model. Counting splits of entities a human would never split is a weaker\n", + "claim, and one a generator can actually support.\n", + "\n", + "---\n", + "\n", + "## Requirements\n", + "\n", + "- `pip install graphfaker`\n", + "- For step 3: `pip install cognee` **and an LLM API key**. Cognee calls an LLM\n", + " for every document, so this costs money. Start small (20 entities, 20\n", + " documents) before scaling up.\n", + "\n", + "Steps 1, 2, 4 and 6 run without cognee or a key." + ] + }, + { + "cell_type": "code", + "id": "cell-01", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import os\n", + "import sys\n", + "\n", + "import networkx as nx\n", + "\n", + "import graphfaker\n", + "from graphfaker import GraphFaker\n", + "from graphfaker.corpus import duplication_report, generate_corpus\n", + "from graphfaker.export import export_cypher, export_neo4j_csv\n", + "from graphfaker.resolve import resolve_entities\n", + "\n", + "print(\"graphfaker\", graphfaker.__version__)\n", + "print(\"networkx \", nx.__version__)\n", + "print(\"python \", sys.version.split()[0])" + ] + }, + { + "cell_type": "markdown", + "id": "cell-02", + "metadata": {}, + "source": [ + "---\n", + "## Step 1 — Build a corpus with a known answer key\n", + "\n", + "`generate_corpus` returns documents plus a registry of exactly which entities\n", + "they describe. Keep `SEED` fixed: a measurement made against an unreproducible\n", + "corpus cannot be checked by anyone else, including you next week.\n", + "\n", + "Start small. `ENTITIES = 20, DOCUMENTS = 20` is enough to see the effect and\n", + "cheap to run; scale up once the pipeline works." + ] + }, + { + "cell_type": "code", + "id": "cell-03", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "SEED = 42\n", + "ENTITIES = 20\n", + "DOCUMENTS = 20\n", + "\n", + "corpus = generate_corpus(seed=SEED, n_entities=ENTITIES, n_documents=DOCUMENTS)\n", + "\n", + "print(f\"{corpus.expected_entity_count} entities\")\n", + "print(f\"{len(corpus.documents)} documents\")\n", + "print(f\"{len(corpus.relations)} stated relations\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-04", + "metadata": {}, + "source": [ + "### Audit the corpus before trusting it\n", + "\n", + "This is the step that makes the result defensible. If any surface form could be\n", + "claimed by two entities, then some of the \"duplication\" you measure is the\n", + "**corpus's** fault, not the pipeline's.\n", + "\n", + "Two ways that happens, both of which bit me while building this:\n", + "\n", + "- **Confusable names.** Naming four people `Person 0`..`Person 3` produces\n", + " strings that are genuinely ~93% similar, so any name-based matcher merges\n", + " them.\n", + "- **Alias collisions.** Faker derives place names from surnames, so a person\n", + " aliased `Henderson` can appear beside an organization named\n", + " `Henderson, Ramirez and Lewis`. Both have a real claim on the string.\n", + "\n", + "`audit()` checks for both. The assertion below is not decoration — do not\n", + "publish a number from a corpus that fails it." + ] + }, + { + "cell_type": "code", + "id": "cell-05", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "audit = corpus.audit()\n", + "\n", + "print(f\"surface forms : {audit['surface_forms']}\")\n", + "print(f\"ambiguous forms : {len(audit['ambiguous_forms'])}\")\n", + "print(f\"containment pairs : {len(audit['containment_pairs'])}\")\n", + "print(f\"clean : {audit['clean']}\")\n", + "\n", + "if not audit[\"clean\"]:\n", + " print(\"\\nambiguous:\", audit[\"ambiguous_forms\"])\n", + " print(\"containment:\", audit[\"containment_pairs\"][:10])\n", + "\n", + "assert audit[\"clean\"], \"Corpus is ambiguous - try a different SEED before measuring anything\"\n", + "print(\"\\nOK: no surface form is claimable by two entities.\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-06", + "metadata": {}, + "source": [ + "### What the text and the answer key look like" + ] + }, + { + "cell_type": "code", + "id": "cell-07", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "document = corpus.documents[0]\n", + "print(\"DOCUMENT\", document.id)\n", + "print(document.text)\n", + "print()\n", + "print(\"ENTITIES IT MENTIONS\")\n", + "lookup = {e.id: e for e in corpus.entities}\n", + "for entity_id in document.entity_ids:\n", + " entity = lookup[entity_id]\n", + " print(f\" {entity.type:14s} {entity.name!r:34s} also written as {entity.aliases}\")" + ] + }, + { + "cell_type": "code", + "id": "cell-08", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Write the corpus to disk so any framework can read it, including ones that\n", + "# only accept a directory of files.\n", + "CORPUS_DIR = \"corpus\"\n", + "corpus.write(CORPUS_DIR)\n", + "print(f\"{len(os.listdir(CORPUS_DIR))} files in {CORPUS_DIR}/ (documents + gold.json)\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-09", + "metadata": {}, + "source": [ + "---\n", + "## Step 2 — The control\n", + "\n", + "A perfect extractor, one node per entity. **It must score zero.** If it does\n", + "not, the harness or the corpus is broken and every other number in this\n", + "notebook is meaningless.\n", + "\n", + "Always run a control. It is the cheapest possible guard against measuring your\n", + "own bug and publishing it." + ] + }, + { + "cell_type": "code", + "id": "cell-10", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def reference_extraction(corpus):\n", + " G = nx.DiGraph()\n", + " for entity in corpus.entities:\n", + " G.add_node(entity.id, name=entity.name, type=entity.type)\n", + " for relation in corpus.relations:\n", + " G.add_edge(relation.source, relation.target, relationship=relation.type)\n", + " return G\n", + "\n", + "\n", + "control = duplication_report(reference_extraction(corpus), corpus, framework=\"control\")\n", + "print(control.summary())\n", + "\n", + "assert control.duplication_rate == 0.0, \"control must score zero\"\n", + "assert control.missed == [], \"control must find every entity\"\n", + "print(\"\\nControl passes.\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-11", + "metadata": {}, + "source": [ + "---\n", + "## Step 3 — Run Cognee\n", + "\n", + "Verified against **cognee 1.4.1**: `add`, `cognify`, and `export` are all\n", + "coroutines, and `export` accepts `format=\"graphml\"`. (There is no\n", + "`get_graph_data` function — that was a guess in an earlier draft of this\n", + "harness, and it was wrong.)\n", + "\n", + "Two deliberate choices in the adapter:\n", + "\n", + "- **It writes to a run-specific dataset instead of calling `cognee.prune`.**\n", + " Pruning would wipe your local cognee store. `export` is scoped to a dataset,\n", + " so isolation is achieved without destroying anything.\n", + "- **`data_cache=False`,** so a re-run genuinely re-extracts. Left on, a second\n", + " run can replay the first one's output and look falsely stable.\n", + "\n", + "> **This costs money and takes minutes.** Cognee calls an LLM per document." + ] + }, + { + "cell_type": "code", + "id": "cell-12", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Cognee reads the key from its own config; LLM_API_KEY is the usual variable.\n", + "# Set it in your shell rather than pasting a key into a notebook you might share.\n", + "HAVE_KEY = bool(\n", + " os.environ.get(\"LLM_API_KEY\")\n", + " or os.environ.get(\"OPENAI_API_KEY\")\n", + " or os.environ.get(\"ANTHROPIC_API_KEY\")\n", + ")\n", + "\n", + "try:\n", + " import cognee\n", + "\n", + " HAVE_COGNEE = True\n", + " print(\"cognee\", cognee.get_cognee_version())\n", + "except ImportError:\n", + " HAVE_COGNEE = False\n", + " print(\"cognee not installed - pip install cognee\")\n", + "\n", + "print(\"API key present:\", HAVE_KEY)\n", + "if HAVE_COGNEE and not HAVE_KEY:\n", + " # cognee's add() runs a pipeline that tests the LLM connection before\n", + " # ingesting anything, so a missing key fails immediately rather than\n", + " # part-way through.\n", + " print(\"\\nCognee needs a key even for add() - step 3 will be skipped.\")" + ] + }, + { + "cell_type": "code", + "id": "cell-13", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import asyncio\n", + "\n", + "\n", + "async def _run_cognee(corpus, dataset):\n", + " # One add() call with the whole list; cognee batches internally and\n", + " # per-document calls are markedly slower.\n", + " await cognee.add([d.text for d in corpus.documents], dataset_name=dataset)\n", + " await cognee.cognify(datasets=[dataset], data_cache=False)\n", + " destination = os.path.abspath(f\"{dataset}.graphml\")\n", + " await cognee.export(dataset=dataset, format=\"graphml\", destination=destination)\n", + " return nx.read_graphml(destination)\n", + "\n", + "\n", + "cognee_raw = None\n", + "if HAVE_COGNEE and HAVE_KEY:\n", + " DATASET = f\"graphfaker_dup_{SEED}\"\n", + " cognee_raw = await _run_cognee(corpus, DATASET) # notebooks allow top-level await\n", + " print(f\"cognee produced {cognee_raw.number_of_nodes()} nodes, \"\n", + " f\"{cognee_raw.number_of_edges()} edges\")\n", + "else:\n", + " print(\"Skipped. Steps 4-6 will use a simulated graph so the notebook still runs.\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-14", + "metadata": {}, + "source": [ + "### Find out what cognee actually emitted, before filtering\n", + "\n", + "Cognee's graph holds bookkeeping nodes — `DocumentChunk`, `TextSummary`,\n", + "`TextDocument` — alongside extracted entities. Counting those as entities would\n", + "inflate the numbers badly.\n", + "\n", + "**Do not skip this cell.** The label names are not guaranteed stable across\n", + "versions, so inspect them rather than trusting a hard-coded filter. If you\n", + "filter on a label that no longer exists you get an empty graph, which reports as\n", + "\"100% of entities missed\" and looks like a catastrophic result instead of a\n", + "wrong filter." + ] + }, + { + "cell_type": "code", + "id": "cell-15", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def summarize_node_types(G, type_attr=\"type\"):\n", + " counts = {}\n", + " for _, data in G.nodes(data=True):\n", + " key = str(data.get(type_attr, \"\"))\n", + " counts[key] = counts.get(key, 0) + 1\n", + " return dict(sorted(counts.items(), key=lambda pair: -pair[1]))\n", + "\n", + "\n", + "if cognee_raw is not None:\n", + " # The export has been seen to use 'type' or 'label'; check which is present.\n", + " TYPE_ATTR = \"type\"\n", + " if not any(\"type\" in d for _, d in cognee_raw.nodes(data=True)):\n", + " TYPE_ATTR = \"label\"\n", + "\n", + " print(\"type attribute:\", TYPE_ATTR)\n", + " print(\"node labels :\", summarize_node_types(cognee_raw, TYPE_ATTR))\n", + " print()\n", + " print(\"sample node attributes:\")\n", + " for node, data in list(cognee_raw.nodes(data=True))[:3]:\n", + " print(\" \", node, \"->\", dict(list(data.items())[:6]))" + ] + }, + { + "cell_type": "markdown", + "id": "cell-16", + "metadata": {}, + "source": [ + "### Keep only the entity nodes\n", + "\n", + "Set `ENTITY_LABELS` from what you just saw. `{\"Entity\"}` is the expected value\n", + "for cognee 1.4.x; change it if the cell above says otherwise." + ] + }, + { + "cell_type": "code", + "id": "cell-17", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ENTITY_LABELS = {\"Entity\"}\n", + "\n", + "def simulate_llm_extraction(corpus, seed=0):\n", + " \"\"\"Stand-in so the rest of the notebook runs without an API key.\n", + "\n", + " Models the documented failure mode: an extractor is **consistent within a\n", + " document and inconsistent across documents**. Each document picks one\n", + " surface form per entity and links what that document states; a different\n", + " document may pick a different form for the same entity.\n", + "\n", + " Getting this shape right took three attempts, and the wrong versions are\n", + " instructive:\n", + "\n", + " 1. Orphan alias nodes with no edges. Not how extractors fail - duplicates\n", + " are harmful precisely *because* they carry edges.\n", + " 2. Edges scattered uniformly at random across variants. This made the\n", + " duplicates' neighbourhoods almost disjoint, which is the hardest possible\n", + " case for structural matching and not representative either.\n", + " 3. Per-document consistency (this version). Two variants of one entity share\n", + " neighbours to the extent their documents mention the same other entities,\n", + " which is what actually happens.\n", + "\n", + " Still INVENTED. Never report numbers from this - it exists only to exercise\n", + " steps 4-6 without an API key.\n", + " \"\"\"\n", + " import random\n", + "\n", + " rand = random.Random(seed)\n", + " G = nx.DiGraph()\n", + "\n", + " variants = {}\n", + " for entity in corpus.entities:\n", + " ids = [entity.id]\n", + " G.add_node(entity.id, name=entity.name, type=\"Entity\")\n", + " for index, alias in enumerate(entity.aliases):\n", + " node = f\"{entity.id}__{index}\"\n", + " G.add_node(node, name=alias, type=\"Entity\")\n", + " ids.append(node)\n", + " variants[entity.id] = ids\n", + "\n", + " relations_by_pair = {}\n", + " for relation in corpus.relations:\n", + " relations_by_pair.setdefault((relation.source, relation.target), relation.type)\n", + "\n", + " for document in corpus.documents:\n", + " mentioned = set(document.entity_ids)\n", + " # One choice per entity per document - the \"consistent within a chunk\"\n", + " # part. Seeded by document id so the corpus stays reproducible.\n", + " chosen = {\n", + " entity_id: rand.choice(variants[entity_id]) for entity_id in mentioned\n", + " }\n", + " for (source, target), rel_type in relations_by_pair.items():\n", + " if source in mentioned and target in mentioned:\n", + " G.add_edge(chosen[source], chosen[target], relationship=rel_type)\n", + " return G\n", + "\n", + "\n", + "if cognee_raw is not None:\n", + " keep = [n for n, d in cognee_raw.nodes(data=True)\n", + " if str(d.get(TYPE_ATTR, \"\")) in ENTITY_LABELS]\n", + " if not keep:\n", + " raise ValueError(\n", + " f\"No nodes matched {ENTITY_LABELS}. Set it from the labels printed above.\"\n", + " )\n", + " extracted = cognee_raw.subgraph(keep).copy()\n", + " for _, data in extracted.nodes(data=True):\n", + " data.setdefault(\"name\", data.get(\"text\", \"\"))\n", + " LABEL = \"cognee 1.4.1\"\n", + " IS_REAL = True\n", + "else:\n", + " extracted = simulate_llm_extraction(corpus, seed=SEED)\n", + " LABEL = \"SIMULATED (not a real result)\"\n", + " IS_REAL = False\n", + "\n", + "print(f\"{LABEL}: {extracted.number_of_nodes()} entity nodes\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-18", + "metadata": {}, + "source": [ + "---\n", + "## Step 4 — Measure\n", + "\n", + "`duplication_rate` is the headline: of the entities the pipeline found at all,\n", + "what share did it split across more than one node.\n", + "\n", + "Read `unmatched_nodes` as a caveat on everything else. A pipeline that invents\n", + "unrelated nodes, or labels them in a way the corpus cannot recognise, shows a\n", + "high count there — and its other numbers deserve proportionally less trust.\n", + "Attribution is itself a matching problem, so unmatched nodes are **reported\n", + "rather than forced** onto the nearest entity." + ] + }, + { + "cell_type": "code", + "id": "cell-19", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "before = duplication_report(extracted, corpus, framework=LABEL)\n", + "print(before.summary())\n", + "\n", + "if not IS_REAL:\n", + " print(\"\\n*** SIMULATED DATA - not a measurement of any real framework. ***\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-20", + "metadata": {}, + "source": [ + "### Look at what got split\n", + "\n", + "The aggregate number is what goes in a table; these examples are what make it\n", + "believable to a reader. Check them: if the \"duplicates\" are actually distinct\n", + "entities, the corpus is at fault and you should revisit step 1." + ] + }, + { + "cell_type": "code", + "id": "cell-21", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphfaker.corpus import attribute_nodes # noqa: F401\n", + "\n", + "attribution = attribute_nodes(extracted, corpus)\n", + "for entity_id, count in before.worst[:8]:\n", + " entity = lookup[entity_id]\n", + " nodes = attribution.matched.get(entity_id, [])\n", + " names = [extracted.nodes[n].get(\"name\", n) for n in nodes]\n", + " print(f\"{entity.name!r} ({entity.type}) -> {count} nodes\")\n", + " for name in names:\n", + " print(f\" {name!r}\")\n", + " print()\n", + "\n", + "if attribution.unmatched:\n", + " print(f\"{len(attribution.unmatched)} nodes matched no known entity, e.g.:\")\n", + " for node in attribution.unmatched[:5]:\n", + " print(\" \", repr(extracted.nodes[node].get(\"name\", node)))" + ] + }, + { + "cell_type": "markdown", + "id": "cell-22", + "metadata": {}, + "source": [ + "---\n", + "## Step 5 — Repair the graph, then re-measure\n", + "\n", + "A diagnosis on its own is a complaint. `resolve()` scores candidate pairs on\n", + "attribute similarity **and** neighbourhood overlap — two nodes sharing most of\n", + "their neighbours are probably the same entity, however differently their names\n", + "are spelled. That is the signal a tabular record-linkage tool structurally\n", + "cannot use, because it has no graph.\n", + "\n", + "The combination is `attr + w * structural * (1 - attr)`, so structure only ever\n", + "**raises** a score. An isolated node is never penalised for having few\n", + "neighbours.\n", + "\n", + "Then it merges each cluster onto one canonical node, rewiring incident edges,\n", + "dropping self-loops the merge creates, and recording what it absorbed in\n", + "`_merged_from` so the merge is not silently lossy." + ] + }, + { + "cell_type": "markdown", + "id": "cell-23", + "metadata": {}, + "source": [ + "### Where this gets hard, stated plainly\n", + "\n", + "Structural matching assumes duplicate nodes **share neighbours**. That holds when\n", + "two mentions of an entity appear alongside the same other entities. It fails when\n", + "an extractor *partitions* an entity's edges across its duplicates, because then\n", + "the duplicates' neighbourhoods are close to disjoint and there is no overlap to\n", + "find.\n", + "\n", + "So do not expect one setting to be right. Sweep the threshold and read precision\n", + "and recall together:\n", + "\n", + "- **High threshold** — merges only near-certain pairs. High precision, low recall.\n", + "- **Low threshold** — catches more real duplicates and starts merging distinct\n", + " entities. A resolver that merges everything scores a perfect duplication rate\n", + " and destroys the graph.\n", + "\n", + "Pick the operating point from these numbers, not from the duplication rate alone." + ] + }, + { + "cell_type": "code", + "id": "cell-24", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphfaker.resolve import evaluate_clusters\n", + "\n", + "gold_clusters_probe = [n for n in attribute_nodes(extracted, corpus).matched.values()\n", + " if len(n) > 1]\n", + "\n", + "print(f\"{'threshold':>10s} {'struct_w':>9s} {'clusters':>9s} {'merged':>7s} \"\n", + " f\"{'precision':>10s} {'recall':>7s} {'F1':>6s}\")\n", + "sweep = []\n", + "for threshold in (0.90, 0.85, 0.82, 0.75, 0.70, 0.65):\n", + " candidate = resolve_entities(\n", + " extracted, on=[\"name\"], threshold=threshold, structural_weight=0.6\n", + " )\n", + " s = evaluate_clusters(candidate.clusters, gold_clusters_probe)\n", + " sweep.append((threshold, candidate, s))\n", + " print(f\"{threshold:10.2f} {0.6:9.1f} {len(candidate.clusters):9d} \"\n", + " f\"{candidate.duplicate_nodes:7d} {s['pairwise_precision']:10.3f} \"\n", + " f\"{s['pairwise_recall']:7.3f} {s['pairwise_f1']:6.3f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-25", + "metadata": {}, + "source": [ + "Now pick the setting. `THRESHOLD` below is the one carried into the rest of the\n", + "notebook — change it based on the sweep, and say in any write-up which value you\n", + "used and why." + ] + }, + { + "cell_type": "code", + "id": "cell-26", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "THRESHOLD = 0.70 # from the sweep above: precision still 1.000, best F1\n", + "\n", + "result = resolve_entities(\n", + " extracted,\n", + " on=[\"name\"],\n", + " threshold=THRESHOLD,\n", + " structural_weight=0.6,\n", + ")\n", + "print(result.report())" + ] + }, + { + "cell_type": "code", + "id": "cell-27", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "repaired = result.apply()\n", + "after = duplication_report(repaired, corpus, framework=f\"{LABEL} + resolve()\")\n", + "\n", + "rows = [before, after]\n", + "print(f\"{'':38s} {'nodes':>7s} {'found':>7s} {'split':>7s} {'dup rate':>10s} {'inflation':>10s}\")\n", + "for report in rows:\n", + " print(f\"{report.framework[:38]:38s} {report.nodes_in_graph:7d} {report.found:7d} \"\n", + " f\"{report.duplicated:7d} {report.duplication_rate:9.1%} {report.node_inflation:9.2f}x\")\n", + "\n", + "print()\n", + "removed = before.nodes_in_graph - after.nodes_in_graph\n", + "if before.duplicated:\n", + " consolidated = before.duplicated - after.duplicated\n", + " print(f\"duplicate nodes removed : {removed}\")\n", + " print(f\"entities reduced to one node : {consolidated}/{before.duplicated}\")\n", + " if removed and not consolidated:\n", + " # Worth distinguishing: partial progress on an entity split three ways\n", + " # still leaves it counted as duplicated.\n", + " print(\"\\nNodes were merged but no entity was fully consolidated - each\")\n", + " print(\"still has 2+ nodes. Partial merges do not move the dup rate.\")\n", + "else:\n", + " print(\"Nothing was split, so there was nothing to repair.\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-28", + "metadata": {}, + "source": [ + "### Did it over-merge?\n", + "\n", + "Recall is easy; precision is the hard part. A resolver that merges everything\n", + "scores a perfect duplication rate and destroys the graph. `evaluate_clusters`\n", + "scores the clustering against the gold grouping, so over-merging shows up as\n", + "poor precision.\n", + "\n", + "This uses labels the corpus already holds — it does not manufacture ground\n", + "truth." + ] + }, + { + "cell_type": "code", + "id": "cell-29", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from graphfaker.resolve import evaluate_clusters\n", + "\n", + "gold_clusters = [\n", + " nodes for nodes in attribution.matched.values() if len(nodes) > 1\n", + "]\n", + "scores = evaluate_clusters(result.clusters, gold_clusters)\n", + "\n", + "for key in (\"pairwise_precision\", \"pairwise_recall\", \"pairwise_f1\",\n", + " \"b_cubed_precision\", \"b_cubed_recall\", \"b_cubed_f1\"):\n", + " print(f\" {key:20s} {scores[key]:.3f}\")\n", + "print()\n", + "print(f\" predicted pairs {scores['predicted_pairs']}, \"\n", + " f\"gold pairs {scores['gold_pairs']}, correct {scores['correct_pairs']}\")\n", + "print()\n", + "print(\"Low precision means resolve() merged entities that are genuinely distinct.\")\n", + "print(\"Low recall means it left real duplicates behind. Tune `threshold` and\")\n", + "print(\"`structural_weight` against these numbers, not against the dup rate alone.\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-30", + "metadata": {}, + "source": [ + "---\n", + "## Step 6 — Load the repaired graph into a database\n", + "\n", + "File-based, so no driver and no running database is needed. The CSV path also\n", + "covers TigerGraph `LOAD`, Amazon Neptune's bulk loader, and Spark/GraphFrames." + ] + }, + { + "cell_type": "code", + "id": "cell-31", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "export_cypher(repaired, \"repaired.cypher\", dialect=\"neo4j\")\n", + "export_neo4j_csv(repaired, \"neo4j_import\")\n", + "export_cypher(repaired, \"repaired.gql\", dialect=\"gql\")\n", + "\n", + "print(open(\"repaired.cypher\").read()[:600])" + ] + }, + { + "cell_type": "markdown", + "id": "cell-32", + "metadata": {}, + "source": [ + "Then, in Neo4j, graph data science runs directly on the result:\n", + "\n", + "```cypher\n", + "CALL gds.graph.project('g', '*', '*');\n", + "\n", + "CALL gds.pageRank.stream('g') YIELD nodeId, score\n", + "RETURN gds.util.asNode(nodeId).name AS name, score\n", + "ORDER BY score DESC LIMIT 10;\n", + "\n", + "CALL gds.louvain.stream('g') YIELD nodeId, communityId\n", + "RETURN communityId, count(*) AS size ORDER BY size DESC;\n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "cell-33", + "metadata": {}, + "source": [ + "---\n", + "## Step 7 — Record the result\n", + "\n", + "Save the numbers with the seed and versions attached. A result without the seed\n", + "is not reproducible, and a result without the corpus audit is not defensible." + ] + }, + { + "cell_type": "code", + "id": "cell-34", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "record = {\n", + " \"seed\": SEED,\n", + " \"graphfaker_version\": graphfaker.__version__,\n", + " \"corpus\": {\n", + " \"entities\": corpus.expected_entity_count,\n", + " \"documents\": len(corpus.documents),\n", + " \"audit_clean\": audit[\"clean\"],\n", + " },\n", + " \"is_real_measurement\": IS_REAL,\n", + " \"results\": [before.row(), after.row()],\n", + " \"resolve_quality\": scores,\n", + " \"per_entity_before\": before.per_entity,\n", + "}\n", + "with open(\"results.json\", \"w\") as handle:\n", + " json.dump(record, handle, indent=2)\n", + "\n", + "print(json.dumps(record[\"results\"], indent=2))\n", + "if not IS_REAL:\n", + " print(\"\\nis_real_measurement is false - simulated run.\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-35", + "metadata": {}, + "source": [ + "---\n", + "## Caveats to carry into any write-up\n", + "\n", + "State these alongside the number. They are what separate a result from a claim.\n", + "\n", + "1. **Clean input only.** Real documents are messier. This is a lower bound on\n", + " duplication, not an estimate of production behaviour.\n", + "2. **One corpus, one seed, one run.** LLM extraction is non-deterministic —\n", + " [langchain#26614](https://github.com/langchain-ai/langchain/issues/26614)\n", + " documents identical input yielding a typed node on one run and an untyped one\n", + " on the next. Run several seeds and report the spread, not one number.\n", + "3. **Attribution is itself a matching problem.** `unmatched_nodes` is the\n", + " honest measure of how much the harness could not account for. Report it.\n", + "4. **Entity-label filtering is a judgement call.** Which node types count as\n", + " entities is set by hand in step 3, and a reader may disagree. Say what you\n", + " filtered.\n", + "5. **Synthetic names are not real names.** Faker's distributions are not a\n", + " population. Cross-cultural naming, honorifics, and transliteration are all\n", + " places real extraction fails that this corpus does not probe.\n", + "6. **This does not measure answer quality.** A pipeline can duplicate entities\n", + " and still answer questions well. Do not extrapolate.\n", + "\n", + "### Reproducing\n", + "\n", + "```bash\n", + "pip install graphfaker cognee\n", + "export LLM_API_KEY=...\n", + "jupyter lab docs/notebooks/duplication_experiment.ipynb\n", + "```\n", + "\n", + "Or headless, across several seeds:\n", + "\n", + "```bash\n", + "for seed in 1 2 3 7 42; do\n", + " python examples/duplication_experiment.py --seed $seed --entities 60 --documents 80 \\\n", + " --results results_$seed.json\n", + "done\n", + "```" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/duplication_experiment.py b/examples/duplication_experiment.py new file mode 100644 index 0000000..f44173b --- /dev/null +++ b/examples/duplication_experiment.py @@ -0,0 +1,380 @@ +"""Measure entity duplication across graph-building frameworks. + +The experiment +-------------- +Build one corpus whose entities are known in advance, hand the *same* text to +several graph builders, and count how many nodes each one creates for entities +that a human reader would never split. + +Nothing is corrupted. The text is clean, well-formed English; every entity is +unambiguous; and the corpus is audited so that no surface form could be claimed +by two entities. A correct pipeline therefore scores zero. + +What this measures, and what it does not +---------------------------------------- +It measures node duplication on clean input. It does *not* measure answer +quality, retrieval quality, or how a pipeline behaves on messy real documents. +Those are different claims and this script does not support them. + +Running it +---------- +The adapters are stubs on purpose: each framework needs its own install and +usually an API key, and hard-wiring them here would rot within weeks. Fill in +the ones you have, run, and the table prints and writes to results.json. + + python examples/duplication_experiment.py --entities 60 --documents 80 + +Adding an adapter means writing one function: text in, NetworkX graph out, with +each node carrying a `name` attribute. +""" + +from __future__ import annotations + +import argparse +import json +import os +from collections.abc import Callable, Sequence + +import networkx as nx + +from graphfaker.corpus import ( + Corpus, + DuplicationReport, + duplication_report, + generate_corpus, +) + +# --------------------------------------------------------------------------- # +# adapters: text -> graph +# --------------------------------------------------------------------------- # +# +# Each returns a NetworkX graph whose nodes carry a `name` attribute, or None if +# the framework is not installed. Returning None is recorded as "not run" rather +# than as a zero, because a missing framework is not a result. + + +def build_with_reference(corpus: Corpus) -> nx.DiGraph: + """A perfect extractor, as a control. + + Include this in every run. If the control does not score zero, the harness + or the corpus is at fault and the other rows cannot be trusted. + """ + G = nx.DiGraph() + for entity in corpus.entities: + G.add_node(entity.id, name=entity.name, type=entity.type) + for relation in corpus.relations: + G.add_edge(relation.source, relation.target, relationship=relation.type) + return G + + +def build_with_langchain(corpus: Corpus) -> nx.DiGraph | None: + """LangChain `LLMGraphTransformer`.""" + try: + from langchain_core.documents import Document as LCDocument + from langchain_experimental.graph_transformers import LLMGraphTransformer + from langchain_openai import ChatOpenAI + except ImportError: + return None + + transformer = LLMGraphTransformer(llm=ChatOpenAI(temperature=0)) + documents = [LCDocument(page_content=d.text) for d in corpus.documents] + graph_documents = transformer.convert_to_graph_documents(documents) + + G = nx.DiGraph() + for graph_document in graph_documents: + for node in graph_document.nodes: + G.add_node(node.id, name=str(node.id), type=node.type) + for edge in graph_document.relationships: + G.add_edge(edge.source.id, edge.target.id, relationship=edge.type) + return G + + +def build_with_llamaindex(corpus: Corpus) -> nx.DiGraph | None: + """LlamaIndex `PropertyGraphIndex`.""" + try: + from llama_index.core import Document as LIDocument + from llama_index.core import PropertyGraphIndex + except ImportError: + return None + + index = PropertyGraphIndex.from_documents( + [LIDocument(text=d.text) for d in corpus.documents] + ) + store = index.property_graph_store + G = nx.DiGraph() + for node in store.get(): + name = getattr(node, "name", None) or getattr(node, "id", None) + G.add_node(str(name), name=str(name)) + for triplet in store.get_triplets(): + source, relation, target = triplet + G.add_edge( + str(getattr(source, "name", source)), + str(getattr(target, "name", target)), + relationship=str(getattr(relation, "label", relation)), + ) + return G + + +def build_with_graphrag(corpus: Corpus) -> nx.DiGraph | None: + """Microsoft GraphRAG. + + GraphRAG runs as a CLI over an input directory and writes parquet outputs, + so wire this to a completed run rather than calling it in-process. Point + GRAPHRAG_OUTPUT at the artifacts directory. + """ + output = os.environ.get("GRAPHRAG_OUTPUT") + if not output: + return None + try: + import pandas as pd + except ImportError: + return None + + entities = pd.read_parquet(os.path.join(output, "entities.parquet")) + G = nx.DiGraph() + for _, row in entities.iterrows(): + G.add_node(str(row.get("title") or row.get("id")), name=str(row.get("title"))) + relationships_path = os.path.join(output, "relationships.parquet") + if os.path.exists(relationships_path): + for _, row in pd.read_parquet(relationships_path).iterrows(): + G.add_edge(str(row["source"]), str(row["target"]), relationship="RELATED") + return G + + +def build_with_lightrag(corpus: Corpus) -> nx.DiGraph | None: + """LightRAG, which persists its graph as GraphML.""" + path = os.environ.get("LIGHTRAG_GRAPHML") + if not path or not os.path.exists(path): + return None + loaded = nx.read_graphml(path) + for node, data in loaded.nodes(data=True): + data.setdefault("name", node) + return loaded + + +#: Cognee node labels that represent extracted entities. Its graph also holds +#: DocumentChunk, TextSummary, and TextDocument nodes, which are bookkeeping +#: rather than entities and would otherwise be counted as unmatched noise. +COGNEE_ENTITY_TYPES = tuple( + filter(None, os.environ.get("COGNEE_ENTITY_TYPES", "Entity").split(",")) +) + + +def summarize_node_types(G: nx.Graph, type_attr: str = "type") -> dict[str, int]: + """Count node labels in an extracted graph. + + Worth printing on a first run against any framework: the filter below is + only correct if it names the labels the framework actually emits, and that + schema is not guaranteed stable across versions. + """ + counts: dict[str, int] = {} + for _, data in G.nodes(data=True): + key = str(data.get(type_attr, "")) + counts[key] = counts.get(key, 0) + 1 + return dict(sorted(counts.items(), key=lambda pair: -pair[1])) + + +def build_with_cognee( + corpus: Corpus, + dataset_name: str | None = None, + entity_types: Sequence[str] | None = COGNEE_ENTITY_TYPES, + keep_export: bool = False, +) -> nx.DiGraph | None: + """Cognee, via its GraphML export. + + Verified against cognee 1.4.1: `add`, `cognify`, and `export` are all + coroutines, and `export` accepts ``format="graphml"``. There is no + ``get_graph_data`` function. + + Requires an LLM API key — `add` itself runs a pipeline that tests the LLM + connection before ingesting, so nothing works without one. Set + ``LLM_API_KEY`` (or the provider variable cognee is configured for). + + Writes into a run-specific dataset rather than calling ``cognee.prune``, so + an existing local cognee store is left intact. `export` is scoped to that + dataset, so other data cannot contaminate the measurement. + + Args: + corpus: The corpus to ingest. + dataset_name: Dataset to write into. Defaults to one derived from the + corpus seed, so repeat runs of the same corpus land together and + different corpora stay separate. + entity_types: Node labels to keep. Pass None to keep every node, which + is the right choice when checking what cognee actually emits. + keep_export: Leave the intermediate .graphml on disk for inspection. + """ + try: + import asyncio + + import cognee + except ImportError: + return None + + dataset = dataset_name or f"graphfaker_dup_{corpus.seed}" + destination = os.path.abspath(f"{dataset}_export.graphml") + + async def run() -> nx.DiGraph: + # One add() call with the whole list: cognee batches internally, and + # per-document calls are markedly slower. + await cognee.add([d.text for d in corpus.documents], dataset_name=dataset) + # data_cache=False so a re-run genuinely re-extracts. Left on, a second + # run could replay the first run's output and look falsely stable. + await cognee.cognify(datasets=[dataset], data_cache=False) + await cognee.export(dataset=dataset, format="graphml", destination=destination) + return nx.read_graphml(destination) + + graph = asyncio.run(run()) + + if not keep_export and os.path.exists(destination): + os.remove(destination) + + # Cognee labels nodes by DataPoint class. Its GraphML export has been seen + # to use either 'type' or 'label', so check both before giving up. + type_attr = "type" + if not any("type" in data for _, data in graph.nodes(data=True)): + type_attr = "label" if any( + "label" in data for _, data in graph.nodes(data=True) + ) else "type" + + print(f" cognee node labels: {summarize_node_types(graph, type_attr)}") + + if entity_types: + keep = [ + node + for node, data in graph.nodes(data=True) + if str(data.get(type_attr, "")) in set(entity_types) + ] + if not keep: + # Filtering everything away would report 100% missed entities and + # look like a catastrophic result rather than a wrong filter. + print( + f" WARNING: no nodes matched {entity_types!r}; keeping all nodes. " + f"Set COGNEE_ENTITY_TYPES to one of the labels above." + ) + else: + graph = graph.subgraph(keep).copy() + + for node, data in graph.nodes(data=True): + data.setdefault("name", data.get("text") or node) + return graph + + +ADAPTERS: dict[str, Callable[[Corpus], nx.DiGraph | None]] = { + "reference (control)": build_with_reference, + "langchain-LLMGraphTransformer": build_with_langchain, + "llamaindex-PropertyGraphIndex": build_with_llamaindex, + "microsoft-graphrag": build_with_graphrag, + "lightrag": build_with_lightrag, + "cognee": build_with_cognee, +} + + +# --------------------------------------------------------------------------- # +# harness +# --------------------------------------------------------------------------- # + + +def render_table(reports: list[DuplicationReport]) -> str: + """Markdown table, ready to paste into a write-up.""" + headers = [ + ("framework", "Framework"), + ("expected_entities", "Entities"), + ("nodes_in_graph", "Nodes"), + ("entities_found", "Found"), + ("entities_missed", "Missed"), + ("entities_duplicated", "Split"), + ("duplication_rate", "Dup. rate"), + ("node_inflation", "Inflation"), + ("unmatched_nodes", "Unmatched"), + ] + lines = [ + "| " + " | ".join(title for _, title in headers) + " |", + "|" + "|".join("---" for _ in headers) + "|", + ] + for report in reports: + row = report.row() + cells = [] + for key, _ in headers: + value = row[key] + if key == "duplication_rate": + value = f"{value:.1%}" + elif key == "node_inflation": + value = f"{value:.2f}x" + cells.append(str(value)) + lines.append("| " + " | ".join(cells) + " |") + return "\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--entities", type=int, default=60) + parser.add_argument("--documents", type=int, default=80) + parser.add_argument("--corpus-dir", default="corpus") + parser.add_argument("--results", default="results.json") + arguments = parser.parse_args() + + corpus = generate_corpus( + seed=arguments.seed, + n_entities=arguments.entities, + n_documents=arguments.documents, + ) + + audit = corpus.audit() + print(f"Corpus: {audit['entities']} entities, {audit['documents']} documents, " + f"{audit['surface_forms']} surface forms") + if not audit["clean"]: + # Publishing a measurement from an ambiguous corpus would attribute the + # corpus's own defects to the frameworks. + print("\nREFUSING TO RUN: the corpus is ambiguous.") + print(f" ambiguous forms : {audit['ambiguous_forms']}") + print(f" containment : {audit['containment_pairs'][:10]}") + print("Try a different --seed or fewer --entities.") + raise SystemExit(1) + print("Audit clean: no surface form is claimable by two entities.\n") + + corpus.write(arguments.corpus_dir) + print(f"Corpus written to {arguments.corpus_dir}/\n") + + reports: list[DuplicationReport] = [] + skipped: list[str] = [] + for name, adapter in ADAPTERS.items(): + try: + graph = adapter(corpus) + except Exception as error: # noqa: BLE001 - one bad adapter must not end the run + print(f" {name}: FAILED ({type(error).__name__}: {error})") + skipped.append(f"{name} (error: {type(error).__name__})") + continue + if graph is None: + print(f" {name}: not installed / not configured, skipping") + skipped.append(f"{name} (not run)") + continue + report = duplication_report(graph, corpus, framework=name) + reports.append(report) + print(f" {name}: {report.duplication_rate:.1%} duplication, " + f"{report.node_inflation:.2f}x nodes") + + print("\n" + render_table(reports)) + + if skipped: + # Stated explicitly, because a table that silently omits frameworks + # reads as though they were tested and did well. + print(f"\nNot included in the table: {', '.join(skipped)}") + + with open(arguments.results, "w", encoding="utf-8") as handle: + json.dump( + { + "seed": arguments.seed, + "corpus": {k: v for k, v in audit.items() if k != "containment_pairs"}, + "results": [report.row() for report in reports], + "per_entity": {r.framework: r.per_entity for r in reports}, + "skipped": skipped, + }, + handle, + indent=2, + ) + print(f"\nResults written to {arguments.results}") + + +if __name__ == "__main__": + main() diff --git a/graphfaker/__init__.py b/graphfaker/__init__.py index d6a382d..6a2e7e5 100644 --- a/graphfaker/__init__.py +++ b/graphfaker/__init__.py @@ -2,10 +2,45 @@ __author__ = """Dennis Irorere""" __email__ = "denironyx@gmail.com" -__version__ = "0.2.0" +__version__ = "0.5.0" from .core import GraphFaker +from .corpus import ( + Corpus, + DuplicationReport, + attribute_nodes, + duplication_report, + generate_corpus, +) +from .export import export_csv, export_cypher, export_neo4j_csv from .fetchers.wiki import WikiFetcher -from .logger import configure_logging, logger +from .logger import add_file_logging, configure_logging, logger +from .metrics import compare_topology, graph_stats +from .resolve import ( + ResolutionResult, + evaluate_clusters, + merge_clusters, + resolve_entities, +) -__all__ = ["GraphFaker", "logger", "configure_logging", "add_file_logging"] +__all__ = [ + "Corpus", + "DuplicationReport", + "GraphFaker", + "ResolutionResult", + "WikiFetcher", + "add_file_logging", + "attribute_nodes", + "compare_topology", + "configure_logging", + "duplication_report", + "evaluate_clusters", + "export_csv", + "export_cypher", + "export_neo4j_csv", + "generate_corpus", + "graph_stats", + "logger", + "merge_clusters", + "resolve_entities", +] diff --git a/graphfaker/cli.py b/graphfaker/cli.py index d6966cc..7242d7e 100644 --- a/graphfaker/cli.py +++ b/graphfaker/cli.py @@ -2,14 +2,17 @@ Command-line interface for GraphFaker. """ -from venv import logger +import os + import typer + from graphfaker.core import GraphFaker -from graphfaker.enums import FetcherType -from graphfaker.fetchers.osm import OSMGraphFetcher +from graphfaker.enums import ExportFormat, FetcherType +from graphfaker.export import export_csv, export_cypher, export_neo4j_csv from graphfaker.fetchers.flights import FlightGraphFetcher +from graphfaker.fetchers.osm import OSMGraphFetcher +from graphfaker.logger import logger from graphfaker.utils import parse_date_range -import os app = typer.Typer() @@ -53,14 +56,29 @@ def gen( ), # common - export: str = typer.Option("graph.graphml", help="File path to export GraphML"), + seed: int = typer.Option( + None, help="Seed for reproducible synthetic generation (faker fetcher)." + ), + export: str = typer.Option( + "graph.graphml", + help="Output path. For csv/neo4j-csv this is used as a directory or stem.", + ), + export_format: ExportFormat = typer.Option( + ExportFormat.GRAPHML, + "--format", + help="Output format: graphml | csv | neo4j-csv | cypher | opencypher | gql.", + ), ): """Generate a graph using GraphFaker.""" - gf = GraphFaker() + gf = GraphFaker(seed=seed) if fetcher == FetcherType.FAKER: - g = gf.generate_graph(total_nodes=total_nodes, total_edges=total_edges) + g = gf.generate_graph( + source=FetcherType.FAKER.value, + total_nodes=total_nodes, + total_edges=total_edges, + ) logger.info( f"Generated random graph with {g.number_of_nodes()} nodes and {g.number_of_edges()} edges." ) @@ -114,9 +132,23 @@ def gen( abs_export_path = os.path.abspath(export) os.makedirs(os.path.dirname(abs_export_path) or ".", exist_ok=True) - - gf.export_graph(g, source=fetcher, path=abs_export_path) - logger.info(f"exported graph to {abs_export_path}, with {g.number_of_nodes()} nodes and {g.number_of_edges()} edges.") + + if export_format == ExportFormat.GRAPHML: + gf.export_graph(g, source=fetcher, path=abs_export_path) + elif export_format == ExportFormat.CSV: + stem, _ = os.path.splitext(abs_export_path) + export_csv(g, f"{stem}_nodes.csv", f"{stem}_edges.csv") + elif export_format == ExportFormat.NEO4J_CSV: + stem, _ = os.path.splitext(abs_export_path) + export_neo4j_csv(g, stem) + else: + dialect = "neo4j" if export_format == ExportFormat.CYPHER else export_format.value + export_cypher(g, abs_export_path, dialect=dialect) + + logger.info( + f"exported graph as {export_format.value} to {abs_export_path}, " + f"with {g.number_of_nodes()} nodes and {g.number_of_edges()} edges." + ) if __name__ == "__main__": diff --git a/graphfaker/core.py b/graphfaker/core.py index 9aa69a8..31863f0 100644 --- a/graphfaker/core.py +++ b/graphfaker/core.py @@ -3,14 +3,18 @@ """ import os -from typing import Optional +from typing import Any, Dict, Optional, Sequence import networkx as nx import random from faker import Faker from graphfaker.fetchers.osm import OSMGraphFetcher from graphfaker.fetchers.flights import FlightGraphFetcher from graphfaker.logger import logger +from graphfaker.resolve import ResolutionResult, resolve_entities +#: Module-level Faker kept for backwards compatibility only. Generation uses +#: the per-instance `GraphFaker._fake` so that seeding is reproducible without +#: reaching into global state. fake = Faker() # Define subtypes for each node category @@ -40,13 +44,101 @@ ("Person", "Product"): (REL_PERSON_PRODUCT, 0.02), } +#: Industries, replacing the `fake.job()` value that used to be stored here. A +#: job title is not an industry, and an organization whose industry was +#: "Chartered accountant" made attribute-based matching meaningless. +INDUSTRY_BY_SUBTYPE = { + "TechCompany": ["Software", "Semiconductors", "Cloud Infrastructure", "Robotics"], + "Hospital": ["Healthcare", "Medical Research", "Elder Care"], + "NGO": ["Humanitarian Aid", "Conservation", "Human Rights", "Education Access"], + "University": ["Higher Education", "Research"], + "RetailChain": ["Grocery", "Apparel Retail", "Consumer Electronics", "Home Goods"], +} + +#: Relationships a person can only sensibly have once. Nobody was born in four +#: cities, but the previous generator drew every Person-Place edge independently +#: and produced exactly that. +FUNCTIONAL_RELATIONSHIPS = {"LIVES_IN", "BORN_IN", "HEADQUARTERED_IN"} + +#: Fraction of Person-Person edges formed by closing a triangle rather than by +#: attaching to a stranger. Triadic closure is what gives social graphs their +#: clustering; uniform attachment has essentially none. +TRIADIC_CLOSURE_RATE = 0.55 + +#: Probability an attachment ignores degree and picks uniformly. Keeps the tail +#: from running away and guarantees low-degree nodes stay reachable. +UNIFORM_ATTACHMENT_RATE = 0.20 + +#: Candidates drawn per edge before scoring by affinity. Sampling rather than +#: scoring every node keeps edge formation linear in the number of edges instead +#: of quadratic in the number of nodes. +AFFINITY_SAMPLE_SIZE = 8 + +#: Chance that candidates are drawn from the source's own community rather than +#: the whole graph. Scoring alone is not enough to produce group structure: with +#: a dozen communities, a global sample of 8 often contains no same-community +#: candidate at all, so there is nothing for the affinity term to prefer. +SAME_COMMUNITY_RATE = 0.75 + + +class _AttachmentPool: + """Samples nodes with probability rising in their degree. + + Implemented with the standard repeated-entry trick: a node is appended each + time it gains an edge, so drawing uniformly from the list draws + proportionally to ``degree + 1``. That keeps sampling O(1) instead of + rebuilding a weight vector for every edge, which matters because the whole + point is to run this once per edge. + + Seeding the list with every node once supplies the ``+1``, so a node with no + edges yet can still be chosen. + """ + + def __init__(self, nodes: Sequence[Any], uniform_rate: float = UNIFORM_ATTACHMENT_RATE): + self._all = list(nodes) + self._weighted = list(nodes) + self._uniform_rate = uniform_rate + + def __bool__(self) -> bool: + return bool(self._all) + + def record(self, node: Any) -> None: + self._weighted.append(node) + + def sample(self, rand: random.Random) -> Any: + if rand.random() < self._uniform_rate: + return rand.choice(self._all) + return rand.choice(self._weighted) + class GraphFaker: - def __init__(self): + """Generate or load graphs into NetworkX. + + Args: + seed: Optional integer making synthetic generation reproducible. The + same seed and the same arguments always produce an identical graph. + Seeding is per-instance: two `GraphFaker(seed=1)` objects do not + interfere with each other or with the global `random` module. + """ + + def __init__(self, seed: Optional[int] = None): # We'll use a directed graph for directional relationships. self.G = nx.DiGraph() - - def generate_nodes(self, total_nodes=100): + self.seed = seed + self._rand = random.Random(seed) + self._fake = Faker() + if seed is not None: + self._fake.seed_instance(seed) + + def reseed(self, seed: Optional[int]) -> None: + """Reset the random state so the next generation call is reproducible.""" + self.seed = seed + self._rand = random.Random(seed) + self._fake = Faker() + if seed is not None: + self._fake.seed_instance(seed) + + def generate_nodes(self, total_nodes=100, communities=None): """ Generates nodes split into: - People (50%) @@ -54,6 +146,22 @@ def generate_nodes(self, total_nodes=100): - Organizations (15%) - Events (10%) - Products (5%) + + Every node is assigned a latent ``community``, and its attributes are + drawn from that community's distribution rather than independently. This + is what lets homophily appear during edge formation: without a hidden + variable shared between attributes and attachment, "like attaches to + like" has nothing to key on and the graph cannot help but look random. + + Places, organizations, and products also receive a heavy-tailed + ``prominence``, which edge formation uses to decide popularity. Real + cities, employers, and products differ in scale by orders of magnitude, + so drawing them from a uniform range made every one interchangeable. + + Args: + total_nodes: Total node count, split by the proportions above. + communities: How many latent groups to create. Defaults to roughly + one per 25 nodes, bounded to [2, 12]. """ counts = { "Person": int(total_nodes * 0.50), @@ -64,71 +172,147 @@ def generate_nodes(self, total_nodes=100): # Remaining nodes will be Products counts["Product"] = total_nodes - sum(counts.values()) - # Generate People + if communities is None: + communities = max(2, min(12, total_nodes // 25 or 2)) + self.communities = communities + + # Each community gets its own attribute distribution, so that community + # membership is inferable from attributes and vice versa. + profiles = {} + for index in range(communities): + profiles[index] = { + "mean_age": self._rand.randint(24, 62), + "education": self._rand.choices( + ["High School", "Bachelor", "Master", "PhD"], + weights=self._rand.sample([1, 2, 3, 4], 4), + k=1, + )[0], + "region": self._fake.city(), + } + self._community_profiles = profiles + + def pick_community(): + return self._rand.randrange(communities) + + def prominence(): + """Heavy-tailed scale factor. + + Lognormal rather than uniform: it is the difference between a graph + where every place is the same size and one with a handful of hubs. + """ + return round(self._rand.lognormvariate(0.0, 1.1), 4) + + # Places first, so that people can be given a home drawn from their own + # community's places. + places_by_community = {index: [] for index in range(communities)} + for i in range(counts["Place"]): + node_id = f"place_{i}" + subtype = self._rand.choice(PLACE_SUBTYPES) + community = pick_community() + scale = prominence() + self.G.add_node( + node_id, + type="Place", + name=self._fake.city(), + place_type=subtype, + # Population now tracks prominence, so the biggest cities are + # also the most connected ones. + population=int(8_000 * (1 + scale) ** 2.2) + self._rand.randint(0, 5_000), + coordinates=(self._fake.latitude(), self._fake.longitude()), + community=community, + prominence=scale, + ) + places_by_community[community].append(node_id) + for i in range(counts["Person"]): node_id = f"person_{i}" - subtype = random.choice(PERSON_SUBTYPES) + subtype = self._rand.choice(PERSON_SUBTYPES) + community = pick_community() + profile = profiles[community] + # Age clusters around the community mean instead of spanning the + # whole adult range uniformly, which is what produces measurable age + # homophily across edges. + age = int(min(80, max(18, self._rand.gauss(profile["mean_age"], 9)))) + # Small graphs can leave a community with no places at all, so fall + # back to any place and then to the empty string. Storing None here + # is not an option: GraphML cannot serialise it, which broke export + # for any graph under about 40 nodes. + local_places = places_by_community[community] + all_places = [ + node + for members in places_by_community.values() + for node in members + ] + home = self._rand.choice(local_places or all_places) if all_places else "" self.G.add_node( node_id, type="Person", - name=fake.name(), - age=random.randint(18, 80), - occupation=fake.job(), - email=fake.email(), - education_level=random.choice( - ["High School", "Bachelor", "Master", "PhD"] + name=self._fake.name(), + age=age, + occupation=self._fake.job(), + email=self._fake.email(), + education_level=( + profile["education"] + if self._rand.random() < 0.55 + else self._rand.choice( + ["High School", "Bachelor", "Master", "PhD"] + ) ), - skills=", ".join(fake.words(nb=3)), + skills=", ".join(self._fake.words(nb=3)), subtype=subtype, + community=community, + home_place=home, ) - # Generate Places - for i in range(counts["Place"]): - node_id = f"place_{i}" - subtype = random.choice(PLACE_SUBTYPES) - self.G.add_node( - node_id, - type="Place", - name=fake.city(), - place_type=subtype, - population=random.randint(10000, 1000000), - coordinates=(fake.latitude(), fake.longitude()), - ) - # Generate Organizations + for i in range(counts["Organization"]): node_id = f"org_{i}" - subtype = random.choice(ORG_SUBTYPES) + subtype = self._rand.choice(ORG_SUBTYPES) + community = pick_community() + scale = prominence() self.G.add_node( node_id, type="Organization", - name=fake.company(), - industry=fake.job(), - revenue=round(random.uniform(1e6, 1e9), 2), - employee_count=random.randint(50, 5000), + name=self._fake.company(), + # An industry, not a job title. + industry=self._rand.choice(INDUSTRY_BY_SUBTYPE[subtype]), + revenue=round(2.5e5 * (1 + scale) ** 3, 2), + # Filled in after edges exist, from the actual WORKS_AT count. + employee_count=0, subtype=subtype, + community=community, + prominence=scale, ) - # Generate Events + for i in range(counts["Event"]): node_id = f"event_{i}" - subtype = random.choice(EVENT_SUBTYPES) + subtype = self._rand.choice(EVENT_SUBTYPES) + community = pick_community() + scale = prominence() self.G.add_node( node_id, type="Event", - name=fake.catch_phrase(), + name=self._fake.catch_phrase(), event_type=subtype, - start_date=fake.date(), - duration=random.randint(1, 5), + start_date=self._fake.date(), + duration=self._rand.randint(1, 5), + community=community, + prominence=scale, ) # days - # Generate Products + for i in range(counts["Product"]): node_id = f"product_{i}" - subtype = random.choice(PRODUCT_SUBTYPES) + subtype = self._rand.choice(PRODUCT_SUBTYPES) + community = pick_community() + scale = prominence() self.G.add_node( node_id, type="Product", - name=fake.word().capitalize(), + name=self._fake.word().capitalize(), category=subtype, - price=round(random.uniform(10, 1000), 2), - release_date=fake.date(), + price=round(12 * (1 + scale) ** 2.5, 2), + release_date=self._fake.date(), + community=community, + prominence=scale, ) def add_relationship( @@ -144,12 +328,82 @@ def add_relationship( if bidirectional: self.G.add_edge(target, source, relationship=rel_type, **attributes) - def generate_edges(self, total_edges=1000): + def _affinity(self, source: str, target: str) -> float: + """How plausible an edge between two nodes is, higher being better. + + Only ever used to rank a small sample of candidates, so the scale is + arbitrary; what matters is that shared community and similar attributes + win. This is the homophily term. + """ + a = self.G.nodes[source] + b = self.G.nodes[target] + score = 0.0 + if a.get("community") == b.get("community"): + score += 1.0 + # Prominent targets are attractive regardless of similarity: a famous + # city draws visitors from everywhere. + score += 0.45 * float(b.get("prominence", 0.0)) + age_a, age_b = a.get("age"), b.get("age") + if isinstance(age_a, int) and isinstance(age_b, int): + score += 0.8 * max(0.0, 1.0 - abs(age_a - age_b) / 35.0) + if a.get("education_level") and a.get("education_level") == b.get( + "education_level" + ): + score += 0.25 + return score + + def _choose_affine( + self, + pool: "_AttachmentPool", + source: str, + exclude=(), + local_pool: "_AttachmentPool | None" = None, + ) -> Any: + """Draw several candidates by degree preference, keep the most plausible. + + Sampling then ranking gives homophily *and* a heavy tail while staying + linear in edges. Scoring every node for every edge would be quadratic and + is the reason a naive implementation reaches for uniform choice instead. + + When a community-local pool is supplied, most candidates are drawn from + it. That is what actually creates group structure — relying on the + affinity score alone leaves modularity near zero, because a global sample + usually contains no same-community candidate to prefer. + """ + best = None + best_score = float("-inf") + for _ in range(AFFINITY_SAMPLE_SIZE): + source_pool = pool + if local_pool and self._rand.random() < SAME_COMMUNITY_RATE: + source_pool = local_pool + candidate = source_pool.sample(self._rand) + if candidate == source or candidate in exclude: + continue + score = self._affinity(source, candidate) + if score > best_score: + best, best_score = candidate, score + return best + + def generate_edges(self, total_edges=1000, topology="realistic"): """ Generate edges based on the EDGE_DISTRIBUTION probabilities. The number of edges for each relationship category is determined by the weight. + + Args: + total_edges: Approximate number of edges to create. + topology: ``"realistic"`` (default) forms edges by preferential + attachment, triadic closure, and homophily, producing a + heavy-tailed degree distribution with clustering and community + structure. ``"uniform"`` restores the pre-0.5 behaviour of + drawing both endpoints uniformly at random, which yields an + Erdos-Renyi graph. Keep it only for comparison — see + :func:`graphfaker.metrics.compare_topology`. """ - # Get node lists by type + if topology not in ("realistic", "uniform"): + raise ValueError( + f"topology must be 'realistic' or 'uniform', got {topology!r}" + ) + nodes_by_type = { "Person": [], "Place": [], @@ -162,7 +416,25 @@ def generate_edges(self, total_edges=1000): if t in nodes_by_type: nodes_by_type[t].append(node) - # For each category in EDGE_DISTRIBUTION, calculate the number of edges + pools = { + node_type: _AttachmentPool(nodes) + for node_type, nodes in nodes_by_type.items() + if nodes + } + # One pool per (type, community), so candidates can be drawn from the + # source's own group. + local_pools: dict = {} + for node_type, nodes in nodes_by_type.items(): + grouped: dict = {} + for node in nodes: + grouped.setdefault(self.G.nodes[node].get("community"), []).append(node) + for community, members in grouped.items(): + local_pools[(node_type, community)] = _AttachmentPool(members) + # Tracks which functional relationships a node already has, so that + # LIVES_IN and BORN_IN stay singular. + assigned: dict = {} + neighbours: dict = {node: set() for node in self.G.nodes()} + for (src_type, tgt_type), (possible_rels, weight) in EDGE_DISTRIBUTION.items(): num_edges = int(total_edges * weight) src_nodes = nodes_by_type.get(src_type, []) @@ -170,24 +442,95 @@ def generate_edges(self, total_edges=1000): if not src_nodes or not tgt_nodes: continue - for _ in range(num_edges): - source = random.choice(src_nodes) - target = random.choice(tgt_nodes) - # Avoid self-loop in same category if not desired - if src_type == tgt_type and source == target: + # Pure preferential attachment leaves a large share of nodes never + # selected at all — measured at 24% isolated on a 600-node graph, + # against 0.5% for uniform attachment. Real graphs have a giant + # component, so the first pass over each category walks every source + # node once before preference takes over. + coverage = list(src_nodes) + self._rand.shuffle(coverage) + + # Budget counts edges actually added, not loop iterations. A + # bidirectional FRIENDS_WITH adds two directed edges while a skipped + # functional relationship adds none, so counting iterations made the + # result overshoot by ~14% or undershoot by ~12% depending on the + # mix. `total_edges` should mean what `number_of_edges()` reports. + added = 0 + index = 0 + attempts = 0 + while added < num_edges and attempts < num_edges * 8 + 32: + attempts += 1 + index += 1 + edges_before = self.G.number_of_edges() + rel = self._rand.choice(possible_rels) + + if topology == "uniform": + source = self._rand.choice(src_nodes) + target = self._rand.choice(tgt_nodes) + else: + source = ( + coverage[index - 1] + if index <= len(coverage) + else pools[src_type].sample(self._rand) + ) + local = local_pools.get( + (tgt_type, self.G.nodes[source].get("community")) + ) + + # A person already living somewhere does not acquire a + # second home; reuse the assignment instead of inventing one. + if rel in FUNCTIONAL_RELATIONSHIPS: + key = (source, rel) + if key in assigned: + continue + if rel == "LIVES_IN": + target = self.G.nodes[source].get("home_place") + else: + target = self._choose_affine(pools[tgt_type], source, local_pool=local) + if target is None: + continue + assigned[key] = target + elif ( + src_type == "Person" + and tgt_type == "Person" + and neighbours[source] + and self._rand.random() < TRIADIC_CLOSURE_RATE + ): + # Close a triangle: befriend a friend of a friend. This + # is the entire source of clustering; picking a stranger + # every time is why the old generator had none. + friend = self._rand.choice(tuple(neighbours[source])) + candidates = [ + candidate + for candidate in neighbours.get(friend, ()) + if candidate != source + and candidate not in neighbours[source] + and self.G.nodes[candidate].get("type") == "Person" + ] + target = ( + self._rand.choice(candidates) + if candidates + else self._choose_affine( + pools[tgt_type], source, local_pool=local + ) + ) + else: + target = self._choose_affine(pools[tgt_type], source, local_pool=local) + + if target is None or source == target: continue - rel = random.choice(possible_rels) + attr = {} # Add additional attributes for specific relationships if rel == "VISITED": - attr["visit_count"] = random.randint(1, 20) + attr["visit_count"] = self._rand.randint(1, 20) elif rel == "WORKS_AT": - attr["position"] = fake.job() + attr["position"] = self._fake.job() elif rel == "PURCHASED": - attr["date"] = fake.date() - attr["amount"] = round(random.uniform(1, 500), 2) + attr["date"] = self._fake.date() + attr["amount"] = round(self._rand.uniform(1, 500), 2) elif rel == "REVIEWED": - attr["rating"] = random.randint(1, 5) + attr["rating"] = self._rand.randint(1, 5) # Define directionality and bidirectionality # For Person-Person FRIENDS_WITH and COLLEAGUES, treat as bidirectional @@ -202,6 +545,125 @@ def generate_edges(self, total_edges=1000): self.add_relationship( source, target, rel, attributes=attr, bidirectional=bidir ) + added += self.G.number_of_edges() - edges_before + + if topology == "realistic": + pools[src_type].record(source) + pools[tgt_type].record(target) + neighbours[source].add(target) + neighbours[target].add(source) + + if topology == "realistic": + self._top_up_edges(total_edges, nodes_by_type, pools, local_pools, neighbours) + self._reconcile_attributes() + + def _top_up_edges( + self, + total_edges: int, + nodes_by_type: dict, + pools: dict, + local_pools: dict, + neighbours: dict, + ) -> None: + """Attach leftover nodes and make up the edge shortfall. + + Two problems are fixed here, both consequences of honest edge formation + rather than bugs in it: + + 1. **Isolated nodes.** Only source types get a coverage pass, so types + that appear mostly as targets — places, events, products — can be + missed entirely. A graph where an eighth of the nodes are unreachable + does not resemble anything real. + 2. **Edge shortfall.** Functional relationships are skipped once a node + already has one, so a run asks for 2400 edges and produces about 2000. + Silently returning 15% fewer edges than requested is a bad contract. + + Only non-functional relationships are used, so topping up cannot give + anyone a second birthplace. + """ + # Relationship options that may be repeated, per type pair. + repeatable = [] + for (src_type, tgt_type), (rels, weight) in EDGE_DISTRIBUTION.items(): + usable = [rel for rel in rels if rel not in FUNCTIONAL_RELATIONSHIPS] + if usable and nodes_by_type.get(src_type) and nodes_by_type.get(tgt_type): + repeatable.append((src_type, tgt_type, usable, weight)) + if not repeatable: + return + + # Which type pair can reach a given node as a target. + reachable_as_target: dict = {} + for src_type, tgt_type, usable, _ in repeatable: + reachable_as_target.setdefault(tgt_type, []).append((src_type, usable)) + + def connect(source, target, rel) -> bool: + if source is None or target is None or source == target: + return False + if self.G.has_edge(source, target): + return False + self.add_relationship(source, target, rel) + pools[self.G.nodes[source]["type"]].record(source) + pools[self.G.nodes[target]["type"]].record(target) + neighbours[source].add(target) + neighbours[target].add(source) + return True + + for node in list(self.G.nodes()): + if self.G.degree(node) > 0: + continue + options = reachable_as_target.get(self.G.nodes[node].get("type")) + if not options: + continue + src_type, usable = self._rand.choice(options) + community = self.G.nodes[node].get("community") + pool = local_pools.get((src_type, community)) or pools[src_type] + connect(pool.sample(self._rand), node, self._rand.choice(usable)) + + weights = [weight for _, _, _, weight in repeatable] + attempts = 0 + limit = max(200, 12 * total_edges) + while self.G.number_of_edges() < total_edges and attempts < limit: + attempts += 1 + src_type, tgt_type, usable, _ = self._rand.choices(repeatable, weights)[0] + source = pools[src_type].sample(self._rand) + local = local_pools.get((tgt_type, self.G.nodes[source].get("community"))) + target = self._choose_affine(pools[tgt_type], source, local_pool=local) + rel = self._rand.choice(usable) + if not connect(source, target, rel): + # Affinity concentrates candidates, so late in the fill the same + # plausible pairs keep coming back already connected. One uniform + # retry finds a fresh pair and stops the budget from stalling + # ~12% short. + connect(source, self._rand.choice(nodes_by_type[tgt_type]), rel) + + logger.debug( + "generate_edges: %d edges after top-up (%d requested, %d attempts)", + self.G.number_of_edges(), + total_edges, + attempts, + ) + + def _reconcile_attributes(self) -> None: + """Make scale attributes agree with the structure that was built. + + An organization whose ``employee_count`` is unrelated to how many people + actually work there is a trap for anyone using the graph to test + aggregation or validation logic, because the attribute and the topology + tell different stories. + """ + for node, data in self.G.nodes(data=True): + if data.get("type") != "Organization": + continue + employees = sum( + 1 + for source, _, edge in self.G.in_edges(node, data=True) + if edge.get("relationship") == "WORKS_AT" + ) + # Employees present in the graph are a sample, not the whole + # workforce, so scale up by prominence rather than reporting the + # raw count. + data["employee_count"] = max( + employees, int(employees * (12 + 40 * data.get("prominence", 0.0))) or 1 + ) def _generate_osm( self, @@ -282,11 +744,13 @@ def _generate_flights( raise - def _generate_faker(self, total_nodes=100, total_edges=1000): + def _generate_faker( + self, total_nodes=100, total_edges=1000, topology="realistic", communities=None + ): """Generates the complete Social Knowledge Graph.""" self.G = nx.DiGraph() # Reset the graph to a new instance - self.generate_nodes(total_nodes=total_nodes) - self.generate_edges(total_edges=total_edges) + self.generate_nodes(total_nodes=total_nodes, communities=communities) + self.generate_edges(total_edges=total_edges, topology=topology) return self.G def generate_graph( @@ -305,15 +769,34 @@ def generate_graph( year: int = 2024, month: int = 1, date_range: Optional[tuple] = None, + seed: Optional[int] = None, + topology: str = "realistic", + communities: Optional[int] = None, ) -> nx.DiGraph: """ - Unified entrypoint: choose 'random' or 'osm'. + Unified entrypoint: choose 'faker', 'osm', or 'flights'. Pass kwargs depending on source. + + Args: + seed: Reseed before generating, making the result reproducible. + Only affects the 'faker' source; 'osm' and 'flights' fetch real + data and are not randomised. + topology: For the 'faker' source. ``"realistic"`` (default) builds a + heavy-tailed, clustered, community-structured graph. + ``"uniform"`` reproduces the pre-0.5 Erdos-Renyi behaviour and + exists for comparison only. + communities: Number of latent groups. Defaults to about one per 25 + nodes. """ + if seed is not None: + self.reseed(seed) if source == "faker": return self._generate_faker( - total_nodes=total_nodes, total_edges=total_edges + total_nodes=total_nodes, + total_edges=total_edges, + topology=topology, + communities=communities, ) elif source == "osm": logger.info( @@ -339,7 +822,49 @@ def generate_graph( date_range=date_range, ) else: - raise ValueError(f"Unknown source '{source}'. Use 'random' or 'osm'.") + raise ValueError( + f"Unknown source '{source}'. Use 'faker', 'osm', or 'flights'." + ) + + def resolve( + self, + G: Optional[nx.Graph] = None, + on: Sequence[str] = ("name",), + threshold: float = 0.85, + structural_weight: float = 0.5, + **kwargs: Any, + ) -> ResolutionResult: + """Find nodes that look like duplicates of the same entity. + + Scores candidate pairs on attribute similarity *and* neighbourhood + overlap — the signal a tabular record-linkage tool cannot see. Nothing + is modified; call `.apply()` on the result to get a merged graph. + + Args: + G: Graph to resolve. Defaults to the graph held on this instance. + on: Attribute keys to compare, most identifying first. + threshold: Minimum combined score to link a pair, in [0, 1]. + structural_weight: How much shared-neighbour evidence may lift a + pair's score. 0 reduces this to plain attribute matching. + **kwargs: Passed through to + :func:`graphfaker.resolve.resolve_entities`. + + Example: + >>> gf = GraphFaker(seed=42) + >>> _ = gf.generate_graph(source="faker", total_nodes=100) + >>> result = gf.resolve(on=["name", "email"], threshold=0.9) + >>> clean = result.apply() + """ + target = self.G if G is None else G + if target is None: + raise ValueError("No graph available to resolve.") + return resolve_entities( + target, + on=on, + threshold=threshold, + structural_weight=structural_weight, + **kwargs, + ) def export_graph(self, G: nx.Graph = None, source: str = None, path: str = "graph.graphml"): """ @@ -364,11 +889,23 @@ def export_graph(self, G: nx.Graph = None, source: str = None, path: str = "grap if G is None: raise ValueError("No graph available to export.") - # Sanitize attributes that are not GraphML-friendly + # Sanitize attributes that are not GraphML-friendly. GraphML only + # accepts scalars, so containers (coordinate tuples, merge provenance + # from resolve(), anything a user attached) are flattened to strings. for _, data in G.nodes(data=True): if 'coordinates' in data and isinstance(data['coordinates'], tuple): lat, lon = data['coordinates'] data['coordinates'] = f"{lat},{lon}" + for key, value in list(data.items()): + if value is None: + # GraphML has no null; the writer raises on NoneType. + data[key] = "" + elif isinstance(value, (list, tuple, set)): + data[key] = ",".join(str(item) for item in value) + elif isinstance(value, dict): + data[key] = "; ".join( + f"{k}={v}" for k, v in sorted(value.items(), key=str) + ) if source == "osm": try: diff --git a/graphfaker/corpus.py b/graphfaker/corpus.py new file mode 100644 index 0000000..8ef4fd7 --- /dev/null +++ b/graphfaker/corpus.py @@ -0,0 +1,747 @@ +"""Paired text/entity corpora for measuring entity duplication. + +Purpose +------- +LLM-based graph builders frequently emit several nodes for one real-world +entity. Measuring how often requires a corpus where the correct answer is known +in advance: a set of documents plus a registry saying exactly which entities +they describe. + +What this deliberately does *not* do +------------------------------------ +It does not corrupt anything. No typos, no injected errors, no simulated OCR +noise. The documents are clean, well-formed English in which every entity is +unambiguous to a human reader. Entities are referred to by the surface forms a +normal writer would use — full name, surname alone, an accepted abbreviation — +because a document that says "Lovelace" once and "Ada Lovelace" twice is +ordinary prose, not a degraded signal. + +This matters for the validity of any result. Synthetic *corruption* is known to +be far easier than real-world error (Lam et al., IJPDS 2024, measured roughly a +hundredfold gap), so a benchmark built on guessed error rates measures its own +noise model. Counting how many nodes a pipeline creates for an entity that a +human would never split is a different and much weaker claim — and one the +generator can actually support. + +Names are checked for mutual distinctness, so a pipeline is never penalised for +conflating two entities the corpus itself made confusable. + +Example: + >>> corpus = generate_corpus(seed=42, n_entities=30, n_documents=40) + >>> corpus.write("corpus/") # NNN.txt files plus gold.json + >>> # ...run a graph builder over corpus/, then: + >>> report = duplication_report(extracted_graph, corpus) + >>> print(report.summary()) +""" + +from __future__ import annotations + +import json +import os +import random +from collections.abc import Hashable, Iterable, Sequence +from dataclasses import dataclass, field +from typing import Any + +import networkx as nx +from faker import Faker + +from graphfaker.logger import logger +from graphfaker.resolve import _string_similarity, normalize + +__all__ = [ + "AttributionResult", + "Corpus", + "Document", + "DuplicationReport", + "Entity", + "Relation", + "attribute_nodes", + "duplication_report", + "generate_corpus", +] + +#: Two entity names closer than this are considered confusable, and the +#: generator will draw a different name rather than emit them together. +DISTINCTNESS_CEILING = 0.62 + +ENTITY_TYPES = ("Person", "Organization", "Place", "Product") + + +# --------------------------------------------------------------------------- # +# data model +# --------------------------------------------------------------------------- # + + +@dataclass +class Entity: + """One real-world thing the corpus talks about.""" + + id: str + name: str + type: str + aliases: list[str] = field(default_factory=list) + + def surface_forms(self) -> list[str]: + """Every string the documents may use for this entity.""" + return [self.name] + list(self.aliases) + + +@dataclass +class Relation: + """A fact stated somewhere in the corpus.""" + + source: str + target: str + type: str + + +@dataclass +class Document: + """A single document, and the entities it mentions.""" + + id: str + text: str + entity_ids: list[str] = field(default_factory=list) + + +@dataclass +class Corpus: + """Documents plus the answer key describing them.""" + + entities: list[Entity] + relations: list[Relation] + documents: list[Document] + seed: int | None = None + + @property + def expected_entity_count(self) -> int: + """How many distinct entities a correct extraction would produce.""" + return len(self.entities) + + def alias_index(self) -> dict[str, str]: + """Normalised surface form -> entity id. + + Forms shared by more than one entity are dropped: an ambiguous string + cannot be used to judge a pipeline. + """ + owners: dict[str, set[str]] = {} + for entity in self.entities: + for form in entity.surface_forms(): + owners.setdefault(normalize(form), set()).add(entity.id) + return { + form: next(iter(ids)) for form, ids in owners.items() if len(ids) == 1 + } + + def text(self) -> str: + """The whole corpus as one string, for pipelines that want a blob.""" + return "\n\n".join(document.text for document in self.documents) + + def audit(self) -> dict[str, Any]: + """Report any surface form more than one entity could claim. + + Run this before publishing a measurement. If `ambiguous_forms` is + non-empty, some part of the observed "duplication" is the corpus's + fault rather than the pipeline's, and the result needs a caveat or a + different seed. + """ + owners: dict[str, set[str]] = {} + for entity in self.entities: + for form in entity.surface_forms(): + owners.setdefault(normalize(form), set()).add(entity.id) + ambiguous = { + form: sorted(ids) for form, ids in owners.items() if len(ids) > 1 + } + + # A form contained inside another entity's form is also ambiguous in + # running prose, even though the exact strings differ. + containment: list[tuple[str, str]] = [] + forms = sorted( + ((normalize(f), e.id) for e in self.entities for f in e.surface_forms()), + key=lambda pair: pair[0], + ) + for form, owner in forms: + for other_form, other_owner in forms: + if owner == other_owner or not form or not other_form: + continue + if form != other_form and form in other_form: + containment.append((form, other_form)) + + return { + "entities": len(self.entities), + "documents": len(self.documents), + "surface_forms": len(owners), + "ambiguous_forms": ambiguous, + "containment_pairs": sorted(set(containment)), + "clean": not ambiguous and not containment, + } + + def gold(self) -> dict[str, Any]: + """The answer key as a plain dict.""" + return { + "seed": self.seed, + "expected_entity_count": self.expected_entity_count, + "entities": [ + { + "id": entity.id, + "name": entity.name, + "type": entity.type, + "aliases": entity.aliases, + } + for entity in self.entities + ], + "relations": [ + {"source": r.source, "target": r.target, "type": r.type} + for r in self.relations + ], + "documents": [ + {"id": document.id, "entity_ids": document.entity_ids} + for document in self.documents + ], + } + + def write(self, directory: str, encoding: str = "utf-8") -> str: + """Write one .txt per document plus gold.json, and return the directory.""" + target = os.path.abspath(directory) + os.makedirs(target, exist_ok=True) + for document in self.documents: + path = os.path.join(target, f"{document.id}.txt") + with open(path, "w", encoding=encoding) as handle: + handle.write(document.text) + with open(os.path.join(target, "gold.json"), "w", encoding=encoding) as handle: + json.dump(self.gold(), handle, indent=2, ensure_ascii=False) + logger.info( + "corpus: wrote %d documents and gold.json to %s", + len(self.documents), + target, + ) + return target + + +# --------------------------------------------------------------------------- # +# generation +# --------------------------------------------------------------------------- # + +_PERSON_TEMPLATES = [ + "{person} joined {org} in {year} as a {role}.", + "{person} has worked at {org} since {year}.", + "Before {year}, {person} led the {role} team at {org}.", + "{person} was promoted to {role} at {org} last spring.", +] + +_PLACE_TEMPLATES = [ + "{org} operates its main site in {place}.", + "{org} opened a second office in {place} in {year}.", + "The {place} branch of {org} employs several dozen people.", +] + +_PRODUCT_TEMPLATES = [ + "{org} manufactures the {product}.", + "The {product} is {org}'s best-selling line.", + "{org} discontinued the {product} in {year}.", +] + +_RESIDENCE_TEMPLATES = [ + "{person} lives in {place}.", + "{person} relocated to {place} in {year}.", + "{person} grew up near {place}.", +] + +_REVIEW_TEMPLATES = [ + "{person} reviewed the {product} favourably.", + "According to {person}, the {product} is reliable.", + "{person} has used the {product} for years.", +] + +_ROLES = [ + "logistics", + "research", + "quality assurance", + "field operations", + "procurement", + "customer support", +] + + +def _distinct_enough(candidate: str, chosen: Iterable[str]) -> bool: + """Reject a name too similar to one already in use. + + Without this the corpus can defeat itself. Naming four people + "Person 0" through "Person 3" produces strings that are genuinely ~93% + similar, so any name-based matcher will merge them — and the resulting + "duplication" would be the corpus's fault, not the pipeline's. + """ + normalized = normalize(candidate) + if not normalized: + return False + for existing in chosen: + if _string_similarity(normalized, normalize(existing)) > DISTINCTNESS_CEILING: + return False + return True + + +def _draw_distinct( + make: Any, chosen: list[str], attempts: int = 200 +) -> str | None: + """Draw from `make` until the result is distinct from everything chosen.""" + for _ in range(attempts): + candidate = make() + if _distinct_enough(candidate, chosen): + return candidate + return None + + +def _other_aliases(entities: Sequence[Entity], exclude_id: str) -> list[str]: + """Aliases already granted to entities other than this one.""" + return [ + alias + for entity in entities + if entity.id != exclude_id + for alias in entity.aliases + ] + + +def _alias_is_safe(alias: str, others: Sequence[str]) -> bool: + """Reject an alias that any other entity could also plausibly claim. + + Similarity alone is not enough here. Faker derives place names from + surnames, so a corpus can end up with a person aliased "Henderson" beside an + organization named "Henderson, Ramirez and Lewis" — two entities with a + genuine claim on the same string. A document using it would be ambiguous to + a human too, which means it cannot be used to judge a pipeline. + + An alias is dropped when, against any other entity's surface form, it is + too similar, is a substring, or its tokens are a subset. + """ + normalized = normalize(alias) + if not normalized: + return False + tokens = set(normalized.split()) + for other in others: + other_normalized = normalize(other) + if not other_normalized: + continue + if _string_similarity(normalized, other_normalized) > DISTINCTNESS_CEILING: + return False + if normalized in other_normalized or other_normalized in normalized: + return False + if tokens and tokens <= set(other_normalized.split()): + return False + return True + + +def _person_aliases(name: str) -> list[str]: + """Surface forms a writer would naturally use for a person.""" + parts = name.split() + if len(parts) < 2: + return [] + first, last = parts[0], parts[-1] + return [last, f"{first[0]}. {last}", first] + + +def _org_aliases(name: str) -> list[str]: + """Drop legal suffixes and abbreviate, as ordinary prose does.""" + aliases: list[str] = [] + trimmed = name + for suffix in (" Inc", " Inc.", " LLC", " Ltd", " Ltd.", " and Sons", " Group"): + if trimmed.endswith(suffix): + trimmed = trimmed[: -len(suffix)].rstrip(",") + break + if trimmed != name: + aliases.append(trimmed) + # Strip trailing punctuation: the first word of "Barnes, Cole and Ramirez" + # is "Barnes," and an alias with a dangling comma appears verbatim in the + # generated prose. + head = trimmed.split()[0].strip(",.;:") + if len(head) > 3 and head.lower() not in {"the"}: + aliases.append(head) + words = [word for word in trimmed.split() if word[:1].isupper()] + if len(words) >= 2: + aliases.append("".join(word[0] for word in words).upper()) + return aliases + + +def generate_corpus( + seed: int | None = None, + n_entities: int = 30, + n_documents: int = 40, + sentences_per_document: int = 4, + alias_probability: float = 0.45, +) -> Corpus: + """Build a corpus of documents whose entities are known in advance. + + Args: + seed: Makes the corpus reproducible. Strongly recommended — a + measurement made against an unreproducible corpus cannot be checked + by anyone else. + n_entities: Approximate number of distinct entities. Fewer may be + produced if the distinctness check cannot find enough clearly + different names. + n_documents: How many documents to write. + sentences_per_document: Sentences per document. + alias_probability: Chance that a mention after the first uses an alias + rather than the canonical name. This is normal prose variation, not + corruption. + + Returns: + A `Corpus`. Every entity is mentioned at least once. + """ + rand = random.Random(seed) + faker = Faker() + if seed is not None: + faker.seed_instance(seed) + + per_type = max(1, n_entities // len(ENTITY_TYPES)) + entities: list[Entity] = [] + used_names: list[str] = [] + + makers = { + "Person": faker.name, + "Organization": faker.company, + "Place": faker.city, + "Product": lambda: f"{faker.word().capitalize()} {rand.randint(100, 900)}", + } + aliasers = { + "Person": _person_aliases, + "Organization": _org_aliases, + "Place": lambda name: [], + "Product": lambda name: [f"the {name.split()[0]}"], + } + + # First pass: draw every canonical name, each distinct from all the others. + for entity_type in ENTITY_TYPES: + for index in range(per_type): + name = _draw_distinct(makers[entity_type], used_names) + if name is None: + logger.warning( + "corpus: could only place %d distinct %s names", + index, + entity_type, + ) + break + used_names.append(name) + entities.append( + Entity(id=f"{entity_type.lower()}_{index}", name=name, type=entity_type) + ) + + # Second pass: assign aliases only now that every name is known. Doing this + # inside the first loop would vet each alias against a partial world and let + # a later entity collide with an alias already handed out. + dropped = 0 + for entity in entities: + others = [other.name for other in entities if other.id != entity.id] + # Trimming a legal suffix and taking the first word can both yield the + # same string ("Blake and Sons" -> "Blake" twice), so track what has + # already been granted, including the canonical name itself. + granted = {normalize(entity.name)} + for alias in aliasers[entity.type](entity.name): + key = normalize(alias) + if key in granted: + continue + if _alias_is_safe(alias, others + _other_aliases(entities, entity.id)): + entity.aliases.append(alias) + granted.add(key) + else: + dropped += 1 + if dropped: + logger.info( + "corpus: dropped %d ambiguous alias(es) that another entity could claim", + dropped, + ) + + by_type: dict[str, list[Entity]] = {t: [] for t in ENTITY_TYPES} + for entity in entities: + by_type[entity.type].append(entity) + + def mention(entity: Entity, first: bool) -> str: + if first or not entity.aliases or rand.random() > alias_probability: + return entity.name + return rand.choice(entity.aliases) + + relations: list[Relation] = [] + documents: list[Document] = [] + mentioned_once: set[str] = set() + + def build_sentence(seen: set[str]) -> tuple[str, list[str]]: + """Emit one factual sentence, returning it and the entities used.""" + options: list[str] = [] + if by_type["Person"] and by_type["Organization"]: + options += ["employment"] + if by_type["Organization"] and by_type["Place"]: + options += ["location"] + if by_type["Organization"] and by_type["Product"]: + options += ["product"] + if by_type["Person"] and by_type["Place"]: + options += ["residence"] + if by_type["Person"] and by_type["Product"]: + options += ["review"] + if not options: + return "", [] + + kind = rand.choice(options) + year = rand.randint(2015, 2025) + + if kind == "employment": + person = rand.choice(by_type["Person"]) + org = rand.choice(by_type["Organization"]) + relations.append(Relation(person.id, org.id, "WORKS_AT")) + text = rand.choice(_PERSON_TEMPLATES).format( + person=mention(person, person.id not in seen), + org=mention(org, org.id not in seen), + year=year, + role=rand.choice(_ROLES), + ) + return text, [person.id, org.id] + + if kind == "location": + org = rand.choice(by_type["Organization"]) + place = rand.choice(by_type["Place"]) + relations.append(Relation(org.id, place.id, "LOCATED_IN")) + text = rand.choice(_PLACE_TEMPLATES).format( + org=mention(org, org.id not in seen), + place=mention(place, place.id not in seen), + year=year, + ) + return text, [org.id, place.id] + + if kind == "product": + org = rand.choice(by_type["Organization"]) + product = rand.choice(by_type["Product"]) + relations.append(Relation(org.id, product.id, "MANUFACTURES")) + text = rand.choice(_PRODUCT_TEMPLATES).format( + org=mention(org, org.id not in seen), + product=mention(product, product.id not in seen), + year=year, + ) + return text, [org.id, product.id] + + if kind == "residence": + person = rand.choice(by_type["Person"]) + place = rand.choice(by_type["Place"]) + relations.append(Relation(person.id, place.id, "LIVES_IN")) + text = rand.choice(_RESIDENCE_TEMPLATES).format( + person=mention(person, person.id not in seen), + place=mention(place, place.id not in seen), + year=year, + ) + return text, [person.id, place.id] + + person = rand.choice(by_type["Person"]) + product = rand.choice(by_type["Product"]) + relations.append(Relation(person.id, product.id, "REVIEWED")) + text = rand.choice(_REVIEW_TEMPLATES).format( + person=mention(person, person.id not in seen), + product=mention(product, product.id not in seen), + ) + return text, [person.id, product.id] + + for index in range(n_documents): + seen: set[str] = set() + sentences: list[str] = [] + used: list[str] = [] + for _ in range(sentences_per_document): + sentence, ids = build_sentence(seen) + if not sentence: + continue + sentences.append(sentence) + for entity_id in ids: + seen.add(entity_id) + mentioned_once.add(entity_id) + if entity_id not in used: + used.append(entity_id) + documents.append( + Document(id=f"doc_{index:03d}", text=" ".join(sentences), entity_ids=used) + ) + + # Any entity the random walk never reached gets its own document, so the + # answer key never contains an entity the text does not mention. + unmentioned = [e for e in entities if e.id not in mentioned_once] + for offset, entity in enumerate(unmentioned): + documents.append( + Document( + id=f"doc_{n_documents + offset:03d}", + text=f"{entity.name} is a {entity.type.lower()} of record.", + entity_ids=[entity.id], + ) + ) + + logger.info( + "generate_corpus: %d entities, %d relations, %d documents (seed=%s)", + len(entities), + len(relations), + len(documents), + seed, + ) + return Corpus( + entities=entities, relations=relations, documents=documents, seed=seed + ) + + +# --------------------------------------------------------------------------- # +# measurement +# --------------------------------------------------------------------------- # + + +@dataclass +class AttributionResult: + """Which extracted nodes correspond to which known entities.""" + + matched: dict[str, list[Hashable]] + unmatched: list[Hashable] + + @property + def matched_node_count(self) -> int: + return sum(len(nodes) for nodes in self.matched.values()) + + +def attribute_nodes( + G: nx.Graph, + corpus: Corpus, + name_attr: str = "name", + fallback_to_id: bool = True, +) -> AttributionResult: + """Assign each node in an extracted graph to a known entity, if it matches. + + Matching is exact on a normalised surface form first, then falls back to + checking whether a canonical name appears inside a longer node label — LLM + extractors often emit "Ada Lovelace (engineer)". + + Nodes that match nothing are reported in `unmatched` rather than being + forced onto the nearest entity. Attribution is itself a matching problem, + and quietly guessing would contaminate the very number being measured. + """ + index = corpus.alias_index() + canonical = sorted( + ((normalize(e.name), e.id) for e in corpus.entities), + key=lambda pair: -len(pair[0]), + ) + + matched: dict[str, list[Hashable]] = {} + unmatched: list[Hashable] = [] + + for node, data in G.nodes(data=True): + label = data.get(name_attr) + if label in (None, "") and fallback_to_id: + label = node + key = normalize(label) + entity_id = index.get(key) + if entity_id is None: + # Longest canonical name contained in the label wins, so + # "Northwind Logistics" is preferred over "Northwind". + for name, candidate in canonical: + if name and name in key: + entity_id = candidate + break + if entity_id is None: + unmatched.append(node) + else: + matched.setdefault(entity_id, []).append(node) + + return AttributionResult(matched=matched, unmatched=unmatched) + + +@dataclass +class DuplicationReport: + """How badly a pipeline split the entities it was given.""" + + framework: str + expected_entities: int + nodes_in_graph: int + per_entity: dict[str, int] + missed: list[str] + unmatched_nodes: int + + @property + def found(self) -> int: + return sum(1 for count in self.per_entity.values() if count >= 1) + + @property + def duplicated(self) -> int: + return sum(1 for count in self.per_entity.values() if count > 1) + + @property + def duplication_rate(self) -> float: + """Share of the entities that were found which got split.""" + return self.duplicated / self.found if self.found else 0.0 + + @property + def node_inflation(self) -> float: + """Nodes produced per entity that actually exists.""" + return ( + self.nodes_in_graph / self.expected_entities + if self.expected_entities + else 0.0 + ) + + @property + def worst(self) -> list[tuple[str, int]]: + return sorted( + ((k, v) for k, v in self.per_entity.items() if v > 1), + key=lambda pair: (-pair[1], pair[0]), + ) + + def row(self) -> dict[str, Any]: + """One row for the comparison table.""" + return { + "framework": self.framework, + "expected_entities": self.expected_entities, + "nodes_in_graph": self.nodes_in_graph, + "entities_found": self.found, + "entities_missed": len(self.missed), + "entities_duplicated": self.duplicated, + "duplication_rate": round(self.duplication_rate, 4), + "node_inflation": round(self.node_inflation, 3), + "unmatched_nodes": self.unmatched_nodes, + } + + def summary(self) -> str: + lines = [ + f"{self.framework}", + f" entities in corpus : {self.expected_entities}", + f" nodes in graph : {self.nodes_in_graph}", + f" entities found : {self.found}", + f" entities missed : {len(self.missed)}", + ( + f" entities split in two+ : {self.duplicated}" + f" ({self.duplication_rate:.1%} of those found)" + ), + f" node inflation : {self.node_inflation:.2f}x", + f" nodes matching nothing : {self.unmatched_nodes}", + ] + for entity_id, count in self.worst[:10]: + lines.append(f" {entity_id}: {count} nodes") + return "\n".join(lines) + + +def duplication_report( + G: nx.Graph, + corpus: Corpus, + framework: str = "unknown", + name_attr: str = "name", +) -> DuplicationReport: + """Count how many nodes a pipeline created per known entity. + + The headline number is `duplication_rate`: of the entities the pipeline + found at all, what share did it split across more than one node. Because the + input text is clean and every entity is unambiguous to a human reader, a + correct pipeline scores zero. + + `unmatched_nodes` is reported alongside and should be read as a caveat on + the rest: a pipeline that invents unrelated nodes, or labels them in a way + the corpus cannot recognise, will show a high count there, and its other + numbers deserve less trust. + """ + attribution = attribute_nodes(G, corpus, name_attr=name_attr) + per_entity = { + entity.id: len(attribution.matched.get(entity.id, [])) + for entity in corpus.entities + } + missed = [entity_id for entity_id, count in per_entity.items() if count == 0] + return DuplicationReport( + framework=framework, + expected_entities=corpus.expected_entity_count, + nodes_in_graph=G.number_of_nodes(), + per_entity=per_entity, + missed=missed, + unmatched_nodes=len(attribution.unmatched), + ) diff --git a/graphfaker/enums.py b/graphfaker/enums.py index 4b3e845..c6a44fa 100644 --- a/graphfaker/enums.py +++ b/graphfaker/enums.py @@ -7,3 +7,14 @@ class FetcherType(str, Enum): OSM = "osm" FLIGHTS = "flights" FAKER = "faker" + + +class ExportFormat(str, Enum): + """Output formats a generated graph can be written to.""" + + GRAPHML = "graphml" + CSV = "csv" + NEO4J_CSV = "neo4j-csv" + CYPHER = "cypher" + OPENCYPHER = "opencypher" + GQL = "gql" diff --git a/graphfaker/export.py b/graphfaker/export.py new file mode 100644 index 0000000..9543800 --- /dev/null +++ b/graphfaker/export.py @@ -0,0 +1,372 @@ +"""Export graphs to formats other graph engines can load. + +Rather than maintaining a live driver per database — which means a service to +authenticate against, a version matrix, and a test environment for each — this +module writes files that every engine's own loader already understands: + + - ``export_csv`` generic node/edge CSV, for pandas, Gephi, or any + bulk loader + - ``export_neo4j_csv`` CSV with ``neo4j-admin database import`` headers + (``:ID``, ``:LABEL``, ``:START_ID``, ``:TYPE``) + - ``export_cypher`` a runnable script for Neo4j, Memgraph, Kuzu, + Amazon Neptune (openCypher), or ISO GQL + +The CSV output is also the path into TigerGraph (``LOAD``), Amazon Neptune +(bulk loader), and graph-data-science workflows, since all of them ingest +node/edge tables. + +Nodes are typed by their ``type`` attribute and edges by ``relationship``, +matching what the GraphFaker fetchers produce. Both are configurable. + +Example: + >>> from graphfaker import GraphFaker + >>> from graphfaker.export import export_cypher, export_csv + >>> gf = GraphFaker(seed=42) + >>> G = gf.generate_graph(source="faker", total_nodes=100) + >>> export_csv(G, "nodes.csv", "edges.csv") + >>> export_cypher(G, "load.cypher", dialect="neo4j") +""" + +from __future__ import annotations + +import csv +import json +import os +import re +from collections.abc import Hashable, Iterable, Sequence +from typing import Any + +import networkx as nx + +from graphfaker.logger import logger + +__all__ = [ + "DIALECTS", + "export_csv", + "export_cypher", + "export_neo4j_csv", + "flatten_value", +] + +#: Cypher-family dialects. ISO GQL (2024) uses INSERT where Cypher uses CREATE. +DIALECTS = ("neo4j", "opencypher", "gql") + +_UNSAFE_IDENT = re.compile(r"[^0-9a-zA-Z_]") +_DEFAULT_LABEL = "Node" +_DEFAULT_REL_TYPE = "RELATED_TO" + + +# --------------------------------------------------------------------------- # +# value handling +# --------------------------------------------------------------------------- # + + +def flatten_value(value: Any) -> Any: + """Reduce a value to something a tabular or query format can carry. + + Containers become strings — GraphML, CSV, and Cypher property values are all + scalar-only. Coordinate tuples and the provenance that `resolve()` attaches + to merged nodes both land here. + """ + if value is None: + return "" + if isinstance(value, bool): + return value + if isinstance(value, (int, float, str)): + return value + if isinstance(value, (list, tuple, set)): + return ",".join(str(flatten_value(item)) for item in value) + if isinstance(value, dict): + return "; ".join( + f"{key}={flatten_value(val)}" for key, val in sorted(value.items(), key=str) + ) + return str(value) + + +def _collect_fields(records: Iterable[dict[str, Any]], skip: Sequence[str] = ()) -> list[str]: + """Union of keys across all records, in first-seen order. + + Taking the first record's keys — the obvious shortcut — silently corrupts + heterogeneous graphs: a Person's header would be written and then a Place's + values would be filed under it. Graph nodes of different types rarely share + an attribute set, so the union is the only safe option. + """ + skipped = set(skip) + fields: list[str] = [] + seen = set() + for record in records: + for key in record: + if key not in seen and key not in skipped: + seen.add(key) + fields.append(key) + return fields + + +def _ensure_parent(path: str) -> str: + absolute = os.path.abspath(path) + parent = os.path.dirname(absolute) + if parent: + os.makedirs(parent, exist_ok=True) + return absolute + + +def _edge_records(G: nx.Graph) -> list[tuple[Hashable, Hashable, dict[str, Any]]]: + """Edges as (source, target, data), flattening multigraph keys away.""" + if G.is_multigraph(): + return [(u, v, data) for u, v, _, data in G.edges(keys=True, data=True)] + return list(G.edges(data=True)) + + +# --------------------------------------------------------------------------- # +# CSV +# --------------------------------------------------------------------------- # + + +def export_csv( + G: nx.Graph, + nodes_path: str = "nodes.csv", + edges_path: str = "edges.csv", + encoding: str = "utf-8", +) -> tuple[str, str]: + """Write the graph as two CSV files: one for nodes, one for edges. + + Every attribute seen on any node becomes a column; nodes missing it get an + empty cell. This is what makes the output safe for graphs whose node types + carry different attributes. + + Args: + G: Graph to export. + nodes_path: Destination for node rows. Gains an ``id`` column. + edges_path: Destination for edge rows. Gains ``source`` and ``target``. + encoding: Output encoding. UTF-8 by default, because generated names + frequently contain non-ASCII characters and the platform default + would fail on them. + + Returns: + The absolute paths written, as ``(nodes_path, edges_path)``. + """ + nodes_absolute = _ensure_parent(nodes_path) + edges_absolute = _ensure_parent(edges_path) + + node_data = [data for _, data in G.nodes(data=True)] + node_fields = _collect_fields(node_data, skip=("id",)) + with open(nodes_absolute, "w", newline="", encoding=encoding) as handle: + writer = csv.DictWriter( + handle, fieldnames=["id"] + node_fields, restval="", extrasaction="ignore" + ) + writer.writeheader() + for node, data in G.nodes(data=True): + row = {key: flatten_value(value) for key, value in data.items()} + row["id"] = node + writer.writerow(row) + + edges = _edge_records(G) + edge_fields = _collect_fields( + (data for _, _, data in edges), skip=("source", "target") + ) + with open(edges_absolute, "w", newline="", encoding=encoding) as handle: + writer = csv.DictWriter( + handle, + fieldnames=["source", "target"] + edge_fields, + restval="", + extrasaction="ignore", + ) + writer.writeheader() + for source, target, data in edges: + row = {key: flatten_value(value) for key, value in data.items()} + row["source"] = source + row["target"] = target + writer.writerow(row) + + logger.info( + "export_csv: %d nodes -> %s, %d edges -> %s", + G.number_of_nodes(), + nodes_absolute, + len(edges), + edges_absolute, + ) + return nodes_absolute, edges_absolute + + +def export_neo4j_csv( + G: nx.Graph, + directory: str = "neo4j_import", + type_attr: str = "type", + relationship_attr: str = "relationship", + encoding: str = "utf-8", +) -> dict[str, str]: + """Write CSV shaped for ``neo4j-admin database import``. + + Differs from `export_csv` in the header row, which uses Neo4j's typed + column syntax so the bulk importer assigns labels and relationship types + rather than treating them as ordinary properties. This is the fast path for + large graphs; `export_cypher` is friendlier for small ones. + + Returns: + A dict with ``nodes``, ``edges``, and ``command`` keys — the last being + the import command to run, which is also logged. + """ + target = os.path.abspath(directory) + os.makedirs(target, exist_ok=True) + nodes_path = os.path.join(target, "nodes.csv") + edges_path = os.path.join(target, "edges.csv") + + node_fields = _collect_fields( + (data for _, data in G.nodes(data=True)), skip=(type_attr,) + ) + with open(nodes_path, "w", newline="", encoding=encoding) as handle: + writer = csv.writer(handle) + writer.writerow(["id:ID"] + node_fields + [":LABEL"]) + for node, data in G.nodes(data=True): + label = _safe_identifier(str(data.get(type_attr) or _DEFAULT_LABEL)) + writer.writerow( + [node] + + [flatten_value(data.get(field, "")) for field in node_fields] + + [label] + ) + + edges = _edge_records(G) + edge_fields = _collect_fields( + (data for _, _, data in edges), skip=(relationship_attr,) + ) + with open(edges_path, "w", newline="", encoding=encoding) as handle: + writer = csv.writer(handle) + writer.writerow([":START_ID"] + edge_fields + [":TYPE", ":END_ID"]) + for source, target, data in edges: + rel = _safe_identifier( + str(data.get(relationship_attr) or _DEFAULT_REL_TYPE), upper=True + ) + writer.writerow( + [source] + + [flatten_value(data.get(field, "")) for field in edge_fields] + + [rel, target] + ) + + command = ( + "neo4j-admin database import full " + f"--nodes={nodes_path} --relationships={edges_path} " + "--overwrite-destination neo4j" + ) + logger.info("export_neo4j_csv: wrote %s and %s", nodes_path, edges_path) + logger.info("export_neo4j_csv: import with -> %s", command) + return {"nodes": nodes_path, "edges": edges_path, "command": command} + + +# --------------------------------------------------------------------------- # +# Cypher / GQL +# --------------------------------------------------------------------------- # + + +def _safe_identifier(raw: str, upper: bool = False, default: str = _DEFAULT_LABEL) -> str: + """Make a string usable as a label or relationship type. + + Labels cannot contain punctuation or start with a digit unescaped, and + generated data does produce such values. + """ + # Only trailing underscores are trimmed. A leading one is significant: + # stripping it rewrites the `_gf_id` property that the relationship MATCH + # clauses look up, and every relationship then silently fails to create. + cleaned = _UNSAFE_IDENT.sub("_", raw).rstrip("_") + if not cleaned.strip("_"): + cleaned = default + if cleaned[0].isdigit(): + cleaned = f"_{cleaned}" + return cleaned.upper() if upper else cleaned + + +def _cypher_literal(value: Any) -> str: + """Render a Python value as a Cypher/GQL literal.""" + flat = flatten_value(value) + if isinstance(flat, bool): + return "true" if flat else "false" + if isinstance(flat, (int, float)): + return repr(flat) + return json.dumps(str(flat), ensure_ascii=False) + + +def _properties(data: dict[str, Any], skip: Sequence[str] = ()) -> str: + skipped = set(skip) + items = [ + f"{_safe_identifier(key, default='prop')}: {_cypher_literal(value)}" + for key, value in data.items() + if key not in skipped + ] + return "{" + ", ".join(items) + "}" if items else "{}" + + +def export_cypher( + G: nx.Graph, + path: str = "graph.cypher", + dialect: str = "neo4j", + type_attr: str = "type", + relationship_attr: str = "relationship", + batch_size: int = 500, + encoding: str = "utf-8", +) -> str: + """Write a runnable script that recreates the graph. + + Args: + G: Graph to export. + path: Destination file. + dialect: One of ``neo4j``, ``opencypher``, or ``gql``. Neo4j gets a + uniqueness constraint so the relationship lookups are indexed; + ISO GQL uses ``INSERT`` in place of ``CREATE``. + type_attr: Node attribute supplying the label. + relationship_attr: Edge attribute supplying the relationship type. + batch_size: Statements between transaction boundaries. Neo4j Browser + and cypher-shell both dislike one enormous transaction. + + Returns: + The absolute path written. + + Note: + Node ids are written as a ``_gf_id`` property and used to match + endpoints when creating relationships, so importing into a database + that already holds ``_gf_id`` values will connect to those instead. + """ + if dialect not in DIALECTS: + raise ValueError(f"dialect must be one of {DIALECTS}, got {dialect!r}") + + absolute = _ensure_parent(path) + insert = "INSERT" if dialect == "gql" else "CREATE" + edges = _edge_records(G) + + with open(absolute, "w", encoding=encoding) as handle: + handle.write(f"// Generated by graphfaker ({dialect} dialect)\n") + handle.write(f"// {G.number_of_nodes()} nodes, {len(edges)} relationships\n\n") + + if dialect == "neo4j": + # Without this the MATCH pairs below degrade to full scans. + handle.write( + "CREATE CONSTRAINT gf_id IF NOT EXISTS\n" + "FOR (n:__GraphFaker__) REQUIRE n._gf_id IS UNIQUE;\n\n" + ) + + for written, (node, data) in enumerate(G.nodes(data=True), start=1): + label = _safe_identifier(str(data.get(type_attr) or _DEFAULT_LABEL)) + labels = f":{label}:__GraphFaker__" if dialect == "neo4j" else f":{label}" + properties = _properties({"_gf_id": node, **data}, skip=(type_attr,)) + handle.write(f"{insert} (n{labels} {properties});\n") + if batch_size and written % batch_size == 0: + handle.write(":commit\n:begin\n" if dialect == "neo4j" else "\n") + + handle.write("\n") + for source, target, data in edges: + rel = _safe_identifier( + str(data.get(relationship_attr) or _DEFAULT_REL_TYPE), upper=True + ) + properties = _properties(data, skip=(relationship_attr,)) + handle.write( + f"MATCH (a {{_gf_id: {_cypher_literal(source)}}}), " + f"(b {{_gf_id: {_cypher_literal(target)}}}) " + f"{insert} (a)-[r:{rel} {properties}]->(b);\n" + ) + + logger.info( + "export_cypher: %d nodes and %d relationships -> %s (%s)", + G.number_of_nodes(), + len(edges), + absolute, + dialect, + ) + return absolute diff --git a/graphfaker/fetchers/flights.py b/graphfaker/fetchers/flights.py index 1f93a71..be54491 100644 --- a/graphfaker/fetchers/flights.py +++ b/graphfaker/fetchers/flights.py @@ -35,11 +35,28 @@ import networkx as nx -# suppress only the single warning from unverified HTTPS +# TLS certificate verification is ON by default. The BTS host has historically +# served a chain that some systems fail to validate, which is why this was once +# disabled outright — but silently turning off verification for every user is a +# man-in-the-middle risk, so it is now opt-in and noisy. +# +# Set GRAPHFAKER_INSECURE_TLS=1 to skip verification, and understand that doing +# so means the downloaded data is no longer authenticated. import urllib3 from urllib3.exceptions import InsecureRequestWarning -urllib3.disable_warnings(InsecureRequestWarning) +VERIFY_TLS = os.environ.get("GRAPHFAKER_INSECURE_TLS", "").strip().lower() not in ( + "1", + "true", + "yes", +) + +if not VERIFY_TLS: + urllib3.disable_warnings(InsecureRequestWarning) + logger.warning( + "GRAPHFAKER_INSECURE_TLS is set: TLS certificate verification is " + "DISABLED for flight data downloads. Data integrity is not guaranteed." + ) # Data source URLs AIRLINE_LOOKUP_URL = ( @@ -125,7 +142,7 @@ def fetch_airlines() -> pd.DataFrame: HTTPError if download fails. """ logger.info("Fetching airlines lookup from BTS…") - resp = requests.get(AIRLINE_LOOKUP_URL, verify=False) + resp = requests.get(AIRLINE_LOOKUP_URL, verify=VERIFY_TLS) resp.raise_for_status() df = pd.read_csv(StringIO(resp.text)) return df.rename(columns={"Code": "carrier", "Description": "airline_name"}) @@ -184,7 +201,7 @@ def fetch_airports( @staticmethod def _download_extract_csv(url: str) -> io.BytesIO: """Stream-download a BTS zip file and return CSV data as BytesIO.""" - resp = requests.get(url, stream=True, verify=False) + resp = requests.get(url, stream=True, verify=VERIFY_TLS) resp.raise_for_status() buf = io.BytesIO() total = int(resp.headers.get("content-length", 0)) diff --git a/graphfaker/metrics.py b/graphfaker/metrics.py new file mode 100644 index 0000000..8c5fa1b --- /dev/null +++ b/graphfaker/metrics.py @@ -0,0 +1,233 @@ +"""Structural measurements for generated and extracted graphs. + +Claiming a generator is "realistic" means nothing without numbers attached, so +this module provides the numbers. It exists mainly to distinguish a graph with +real-world structure from one produced by uniform random attachment, which is +what GraphFaker's synthetic generator produced before version 0.5. + +The discriminating statistics: + +``degree_fano`` + Variance of degree divided by mean degree. Uniform random attachment gives a + Poisson degree distribution, whose variance equals its mean, so this sits + near 1.0. Real social, citation, and web graphs are heavy-tailed and land far + above it. This is the single clearest test of whether hubs exist. + +``degree_gini`` + Inequality of the degree distribution, 0 (every node identical) to 1 (one + node holds everything). Uniform attachment produces a low value. + +``average_clustering`` + How often two neighbours of a node are themselves connected. Uniform + attachment leaves this near ``mean_degree / n`` — effectively zero on a + sparse graph — while real graphs cluster strongly, because people who share a + friend tend to meet. + +``degree_assortativity`` + Whether high-degree nodes attach to other high-degree nodes. Social graphs + are usually mildly positive; technological ones negative. + +``community_modularity`` + How well the graph divides into groups, computed against the node + ``community`` attribute when present and otherwise against greedy + modularity communities. Uniform attachment has no group structure to find. + +Example: + >>> from graphfaker import GraphFaker + >>> from graphfaker.metrics import graph_stats, compare_topology + >>> realistic = GraphFaker(seed=1).generate_graph(total_nodes=500) + >>> uniform = GraphFaker(seed=1).generate_graph(total_nodes=500, + ... topology="uniform") + >>> print(compare_topology({"realistic": realistic, "uniform": uniform})) +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import networkx as nx + +from graphfaker.logger import logger + +__all__ = [ + "as_simple_undirected", + "compare_topology", + "degree_gini", + "graph_stats", +] + + +def as_simple_undirected(G: nx.Graph) -> nx.Graph: + """Collapse to a simple undirected graph. + + Clustering, modularity, and assortativity are defined on undirected simple + graphs. Reciprocal pairs (a FRIENDS_WITH written both ways) would otherwise + be double counted, and self-loops distort clustering. + """ + H = nx.Graph() + H.add_nodes_from(G.nodes(data=True)) + H.add_edges_from((u, v) for u, v in G.edges() if u != v) + return H + + +def degree_gini(degrees: list[int]) -> float: + """Gini coefficient of a degree sequence, in [0, 1].""" + if not degrees: + return 0.0 + ordered = sorted(degrees) + total = sum(ordered) + if total == 0: + return 0.0 + n = len(ordered) + weighted = sum((index + 1) * value for index, value in enumerate(ordered)) + return (2 * weighted) / (n * total) - (n + 1) / n + + +def _percentile(ordered: list[int], fraction: float) -> float: + if not ordered: + return 0.0 + position = max(0, min(len(ordered) - 1, round(fraction * (len(ordered) - 1)))) + return float(ordered[position]) + + +def graph_stats(G: nx.Graph, community_attr: str = "community") -> dict[str, Any]: + """Structural summary of a graph. + + Args: + G: Any NetworkX graph. Directed and multigraphs are collapsed first. + community_attr: Node attribute holding a ground-truth grouping. When + absent, modularity is computed against greedy communities instead, + which measures whether *any* group structure exists rather than + whether a known one was reproduced. + + Returns: + A dict of scalars, safe to serialise and put in a table. + """ + H = as_simple_undirected(G) + n = H.number_of_nodes() + m = H.number_of_edges() + if n == 0: + return {"nodes": 0, "edges": 0} + + degrees = [degree for _, degree in H.degree()] + ordered = sorted(degrees) + mean_degree = sum(degrees) / n + variance = sum((degree - mean_degree) ** 2 for degree in degrees) / n + median = _percentile(ordered, 0.5) + + stats: dict[str, Any] = { + "nodes": n, + "edges": m, + "density": nx.density(H), + "mean_degree": mean_degree, + "max_degree": max(degrees), + "median_degree": median, + # Poisson (uniform attachment) sits near 1.0; heavy tails far above. + "degree_fano": variance / mean_degree if mean_degree else 0.0, + "degree_gini": degree_gini(degrees), + "degree_p99_over_median": ( + _percentile(ordered, 0.99) / median if median else float("inf") + ), + "isolated_nodes": sum(1 for degree in degrees if degree == 0), + "average_clustering": nx.average_clustering(H) if m else 0.0, + # A sparse uniform-attachment graph clusters at roughly mean_degree / n, + # which is the baseline the measured value should be compared against. + "clustering_baseline": mean_degree / n if n else 0.0, + } + + try: + stats["degree_assortativity"] = nx.degree_assortativity_coefficient(H) + except (ZeroDivisionError, ValueError, nx.NetworkXError): + # Undefined when every node shares a degree. + stats["degree_assortativity"] = float("nan") + + groups = { + node: data.get(community_attr) + for node, data in H.nodes(data=True) + if data.get(community_attr) is not None + } + if len(groups) == n and len(set(groups.values())) > 1: + partition: dict[Any, set[Any]] = {} + for node, group in groups.items(): + partition.setdefault(group, set()).add(node) + stats["community_modularity"] = nx.community.modularity( + H, list(partition.values()) + ) + stats["community_source"] = community_attr + try: + stats["community_assortativity"] = nx.attribute_assortativity_coefficient( + H, community_attr + ) + except (ZeroDivisionError, ValueError, nx.NetworkXError): + stats["community_assortativity"] = float("nan") + elif m: + detected = list(nx.community.greedy_modularity_communities(H)) + stats["community_modularity"] = nx.community.modularity(H, detected) + stats["community_source"] = "greedy (detected)" + stats["detected_communities"] = len(detected) + + return stats + + +def numeric_assortativity(G: nx.Graph, attribute: str) -> float: + """Correlation of a numeric attribute across edges. + + Positive means like attaches to like — age homophily, for instance. Returns + NaN when the attribute is missing or constant. + """ + H = as_simple_undirected(G) + present = [ + node for node, data in H.nodes(data=True) if isinstance(data.get(attribute), (int, float)) + ] + if len(present) < 2: + return float("nan") + try: + return nx.numeric_assortativity_coefficient(H.subgraph(present), attribute) + except (ZeroDivisionError, ValueError, nx.NetworkXError): + return float("nan") + + +_TABLE_ROWS = [ + ("nodes", "{:.0f}"), + ("edges", "{:.0f}"), + ("mean_degree", "{:.2f}"), + ("max_degree", "{:.0f}"), + ("degree_fano", "{:.2f}"), + ("degree_gini", "{:.3f}"), + ("degree_p99_over_median", "{:.2f}"), + ("average_clustering", "{:.4f}"), + ("clustering_baseline", "{:.4f}"), + ("degree_assortativity", "{:+.3f}"), + ("community_modularity", "{:.3f}"), + ("isolated_nodes", "{:.0f}"), +] + + +def compare_topology( + graphs: Mapping[str, nx.Graph], community_attr: str = "community" +) -> str: + """Render a side-by-side table of `graph_stats` for several graphs. + + Useful for showing what a change to the generator actually did, rather than + asserting that it helped. + """ + computed = { + label: graph_stats(graph, community_attr) for label, graph in graphs.items() + } + labels = list(computed) + width = max([len("degree_p99_over_median")] + [len(label) for label in labels]) + 2 + + lines = ["metric".ljust(width) + "".join(label.rjust(14) for label in labels)] + lines.append("-" * (width + 14 * len(labels))) + for key, fmt in _TABLE_ROWS: + if not any(key in stats for stats in computed.values()): + continue + row = key.ljust(width) + for label in labels: + value = computed[label].get(key) + row += (fmt.format(value) if isinstance(value, (int, float)) else "-").rjust(14) + lines.append(row) + + logger.debug("compare_topology over %d graphs", len(graphs)) + return "\n".join(lines) diff --git a/graphfaker/resolve.py b/graphfaker/resolve.py new file mode 100644 index 0000000..65fd6ec --- /dev/null +++ b/graphfaker/resolve.py @@ -0,0 +1,722 @@ +"""Graph-native entity resolution for GraphFaker. + +Duplicate entities are the most commonly reported defect in LLM-extracted +knowledge graphs: the same real-world entity is emitted as several nodes, and +every edge attached to a false node is a false edge. Tabular record-linkage +tools compare *rows*, so they cannot use the strongest signal available in a +graph — two candidate nodes that share most of their neighbours are very likely +the same entity, however differently their names are spelled. + +This module scores candidate pairs on attribute similarity *and* neighbourhood +overlap, clusters the survivors, and merges each cluster while rewiring its +edges onto a single canonical node. + +Typical use:: + + from graphfaker import GraphFaker + + gf = GraphFaker(seed=42) + G = gf.generate_graph(source="faker", total_nodes=200, total_edges=800) + + result = gf.resolve(on=["name", "email"], threshold=0.85) + print(result.report()) + + G_clean = result.apply() + +`evaluate_clusters` is provided separately for the case where you already hold +labelled clusters and want to score a prediction against them. + +Design notes: + - Deterministic. Identical input yields identical output; every iteration + order is sorted. + - No new dependencies. String similarity uses `difflib` from the standard + library. + - Structural similarity can only *raise* a pair's score, never lower it, so + isolated nodes are never penalised for having few neighbours. +""" + +from __future__ import annotations + +import re +from collections.abc import Callable, Hashable, Iterable, Sequence +from dataclasses import dataclass, field +from difflib import SequenceMatcher +from typing import ( + Any, +) + +import networkx as nx + +from graphfaker.logger import logger + +__all__ = [ + "ResolutionResult", + "attribute_similarity", + "evaluate_clusters", + "merge_clusters", + "neighbor_overlap", + "normalize", + "resolve_entities", +] + +_NON_ALNUM = re.compile(r"[^0-9a-z]+") + +#: Attribute keys written onto merged nodes. +MERGED_FROM = "_merged_from" +MERGED_VARIANTS = "_merged_variants" +MERGE_COUNT = "_merge_count" + +#: Guard against accidentally running an O(n^2) comparison on a large graph. +MAX_UNBLOCKED_NODES = 2_000 + + +# --------------------------------------------------------------------------- # +# similarity primitives +# --------------------------------------------------------------------------- # + + +def normalize(value: Any) -> str: + """Lowercase, strip punctuation, and collapse whitespace. + + Returns an empty string for None so missing attributes compare as absent + rather than as the literal string "none". + """ + if value is None: + return "" + text = _NON_ALNUM.sub(" ", str(value).lower()) + return " ".join(text.split()) + + +def _tokens(value: Any) -> list[str]: + return normalize(value).split() + + +#: Floor applied when one name's tokens are a subset of the other's, e.g. +#: "Hill" inside "Allison Hill". Deliberately below the default threshold: a +#: containment alone must not merge anything, it only lifts the pair far enough +#: that shared-neighbour evidence can decide. Set `token_subset_floor=0.0` to +#: disable. +TOKEN_SUBSET_FLOOR = 0.75 + + +def _string_similarity(a: str, b: str, token_subset_floor: float = 0.0) -> float: + """Ratio in [0, 1]; token-set agreement is taken into account. + + Plain sequence matching alone treats "Acme Corporation" and + "Corporation Acme" as fairly different. Taking the better of the raw ratio + and the sorted-token ratio makes the comparison order-insensitive. + + `token_subset_floor` addresses a structural weakness of character ratios on + names. "Hill" against "Allison Hill" scores only 0.5, because most of the + longer string is unmatched — so with a 0.85 threshold the pair is discarded + before neighbourhood overlap is ever consulted, no matter how much context + the two nodes share. Shortening a name is one of the most common ways a + document refers to an entity it already introduced, so this case matters. + + When one token set is contained in the other, the score is floored rather + than set to 1.0. Containment is suggestive, not conclusive — "Smith" may + well be a different person from "John Smith" — so the floor sits below the + default threshold and leaves the decision to the structural signal. + """ + if not a or not b: + return 0.0 + if a == b: + return 1.0 + raw = SequenceMatcher(None, a, b).ratio() + sorted_a = " ".join(sorted(a.split())) + sorted_b = " ".join(sorted(b.split())) + if not (sorted_a == a and sorted_b == b): + raw = max(raw, SequenceMatcher(None, sorted_a, sorted_b).ratio()) + + if token_subset_floor > 0: + tokens_a, tokens_b = set(a.split()), set(b.split()) + # Single-letter tokens are initials; "a hill" vs "b hill" would + # otherwise look like containment of a shared surname. + meaningful_a = {t for t in tokens_a if len(t) > 1} + meaningful_b = {t for t in tokens_b if len(t) > 1} + if ( + meaningful_a + and meaningful_b + and (meaningful_a <= meaningful_b or meaningful_b <= meaningful_a) + ): + return max(raw, token_subset_floor) + return raw + + +def attribute_similarity( + a_data: dict[str, Any], + b_data: dict[str, Any], + on: Sequence[str], + weights: dict[str, float] | None = None, + token_subset_floor: float = 0.0, +) -> float: + """Weighted mean string similarity of two nodes' attributes. + + Fields absent from *both* nodes are skipped rather than scored as a match, + so comparing two nodes that each lack an `email` does not inflate their + similarity. If no field is present on either side the result is 0.0. + """ + weights = weights or {} + total_weight = 0.0 + accumulated = 0.0 + for key in on: + left = normalize(a_data.get(key)) + right = normalize(b_data.get(key)) + if not left and not right: + continue + weight = float(weights.get(key, 1.0)) + if weight <= 0: + continue + accumulated += weight * _string_similarity(left, right, token_subset_floor) + total_weight += weight + if total_weight == 0: + return 0.0 + return accumulated / total_weight + + +def _neighbors(G: nx.Graph, node: Hashable, relationship_aware: bool) -> set[Any]: + """Undirected neighbour set, optionally qualified by relationship type. + + Direction is deliberately ignored: an entity duplicated by an extraction + pipeline often ends up with its edges pointing the opposite way, and we + still want to recognise the shared context. + """ + seen: set[Any] = set() + if G.is_directed(): + incident = list(G.out_edges(node, data=True)) + [ + (v, u, d) for u, v, d in G.in_edges(node, data=True) + ] + else: + incident = [(node, other, data) for other, data in G[node].items()] + if G.is_multigraph(): + incident = [ + (node, other, data) + for other, keyed in G[node].items() + for data in keyed.values() + ] + + for _, other, data in incident: + if relationship_aware: + seen.add((data.get("relationship"), other)) + else: + seen.add(other) + return seen + + +def neighbor_overlap( + G: nx.Graph, + a: Hashable, + b: Hashable, + relationship_aware: bool = False, +) -> float: + """Jaccard overlap of two nodes' neighbourhoods, in [0, 1]. + + This is the signal a tabular linkage tool cannot compute. The two nodes are + removed from each other's neighbour sets first, so a direct edge between + them neither helps nor hurts. + """ + left = _neighbors(G, a, relationship_aware) + right = _neighbors(G, b, relationship_aware) + if relationship_aware: + left = {item for item in left if item[1] != b} + right = {item for item in right if item[1] != a} + else: + left = left - {b} + right = right - {a} + if not left or not right: + return 0.0 + union = left | right + if not union: + return 0.0 + return len(left & right) / len(union) + + +# --------------------------------------------------------------------------- # +# blocking +# --------------------------------------------------------------------------- # + + +def _block_keys(data: dict[str, Any], fields: Sequence[str], prefix: int) -> set[str]: + """Cheap keys used to propose candidate pairs. + + A node gets one key per token prefix of each blocking field, so + "Acme Corporation" and "Corporation Acme Ltd" share a key even though their + full strings differ. + """ + keys: set[str] = set() + for field_name in fields: + for token in _tokens(data.get(field_name)): + if token: + keys.add(token[:prefix]) + return keys + + +def _candidate_pairs( + G: nx.Graph, + nodes: list[Hashable], + on: Sequence[str], + block_on: Sequence[str] | None, + block_prefix: int, + respect_type: bool, +) -> set[tuple[Hashable, Hashable]]: + """Propose pairs worth scoring. + + With `block_on` set (the default is the first field of `on`) this is + near-linear. With `block_on=[]` every same-type pair is compared, which is + quadratic and therefore guarded by `MAX_UNBLOCKED_NODES`. + """ + if block_on is None: + block_on = on[:1] + + def type_of(node: Hashable) -> Any: + return G.nodes[node].get("type") if respect_type else None + + pairs: set[tuple[Hashable, Hashable]] = set() + + if not block_on: + if len(nodes) > MAX_UNBLOCKED_NODES: + raise ValueError( + f"Refusing to compare {len(nodes)} nodes without blocking " + f"({len(nodes) * (len(nodes) - 1) // 2} pairs). Pass block_on=[...] " + f"with at least one field, or raise " + f"graphfaker.resolve.MAX_UNBLOCKED_NODES if you mean it." + ) + for i, left in enumerate(nodes): + for right in nodes[i + 1 :]: + if type_of(left) == type_of(right): + pairs.add((left, right)) + return pairs + + buckets: dict[tuple[Any, str], list[Hashable]] = {} + for node in nodes: + for key in _block_keys(G.nodes[node], block_on, block_prefix): + buckets.setdefault((type_of(node), key), []).append(node) + + for bucket in buckets.values(): + if len(bucket) < 2: + continue + ordered = sorted(bucket, key=repr) + for i, left in enumerate(ordered): + for right in ordered[i + 1 :]: + pairs.add((left, right) if repr(left) <= repr(right) else (right, left)) + return pairs + + +# --------------------------------------------------------------------------- # +# result object +# --------------------------------------------------------------------------- # + + +@dataclass +class ResolutionResult: + """Outcome of a resolution pass. + + Nothing has been changed on the graph yet — call `apply()` to get a merged + copy, or read `clusters` and decide for yourself. + """ + + graph: nx.Graph + clusters: list[list[Hashable]] + scores: dict[tuple[Hashable, Hashable], float] + threshold: float + candidates_considered: int + on: list[str] = field(default_factory=list) + + @property + def duplicate_nodes(self) -> int: + """How many nodes would disappear if every cluster were merged.""" + return sum(len(cluster) - 1 for cluster in self.clusters) + + def mapping(self) -> dict[Hashable, Hashable]: + """Absorbed node id -> canonical node id.""" + result: dict[Hashable, Hashable] = {} + for cluster in self.clusters: + canonical = pick_canonical(self.graph, cluster) + for node in cluster: + if node != canonical: + result[node] = canonical + return result + + def apply(self, inplace: bool = False) -> nx.Graph: + """Merge every cluster and return the resulting graph.""" + return merge_clusters(self.graph, self.clusters, inplace=inplace) + + def report(self) -> str: + """Human-readable summary, safe to print in a notebook.""" + lines = [ + ( + f"Resolution over {self.graph.number_of_nodes()} nodes " + f"on {self.on or '[]'} at threshold {self.threshold}" + ), + f" candidate pairs scored : {self.candidates_considered}", + f" pairs above threshold : {len(self.scores)}", + f" clusters found : {len(self.clusters)}", + f" duplicate nodes : {self.duplicate_nodes}", + ] + for cluster in self.clusters[:10]: + canonical = pick_canonical(self.graph, cluster) + names = ", ".join( + f"{node}({self.graph.nodes[node].get('name', '?')})" for node in cluster + ) + lines.append(f" -> {canonical}: {names}") + if len(self.clusters) > 10: + lines.append(f" ... and {len(self.clusters) - 10} more") + return "\n".join(lines) + + +# --------------------------------------------------------------------------- # +# resolution +# --------------------------------------------------------------------------- # + + +def resolve_entities( + G: nx.Graph, + on: Sequence[str] = ("name",), + threshold: float = 0.85, + structural_weight: float = 0.5, + node_types: Iterable[str] | None = None, + block_on: Sequence[str] | None = None, + block_prefix: int = 4, + respect_type: bool = True, + relationship_aware: bool = False, + weights: dict[str, float] | None = None, + token_subset_floor: float = TOKEN_SUBSET_FLOOR, +) -> ResolutionResult: + """Find clusters of nodes that appear to be the same entity. + + Args: + G: The graph to examine. Never modified. + on: Attribute keys compared for similarity, most identifying first. + threshold: Minimum combined score for a pair to be linked, in [0, 1]. + structural_weight: How much neighbourhood overlap may lift a pair's + score, in [0, 1]. 0 disables the graph signal entirely, reducing + this to ordinary attribute matching. The combination is + ``attr + w * structural * (1 - attr)``, so structure corroborates + attribute evidence but never contradicts it — a pair with no shared + neighbours simply scores its attribute similarity. + node_types: Restrict to nodes whose `type` attribute is in this set. + block_on: Fields used to propose candidates. Defaults to the first + entry of `on`. Pass `[]` to compare all same-type pairs. + block_prefix: Token prefix length for blocking keys. + respect_type: Only compare nodes sharing the same `type` attribute. + relationship_aware: Qualify neighbours by edge `relationship`, so + structural overlap requires the same relationship to the same + target. Stricter, and useful when edge types are trustworthy. + weights: Per-field weights for attribute similarity. + token_subset_floor: Minimum attribute score for a pair where one name's + tokens are contained in the other's ("Hill" inside "Allison Hill"). + Character ratios score that pair around 0.5, so without this it is + discarded before the structural signal is ever consulted. The floor + sits below the default threshold on purpose, so containment alone + never merges anything — shared neighbours still have to agree. Pass + 0.0 to disable. + + Returns: + A `ResolutionResult`. The graph is untouched until you call `apply()`. + """ + if not 0.0 <= threshold <= 1.0: + raise ValueError(f"threshold must be in [0, 1], got {threshold}") + if not 0.0 <= structural_weight <= 1.0: + raise ValueError( + f"structural_weight must be in [0, 1], got {structural_weight}" + ) + on = list(on) + if not on: + raise ValueError("`on` must name at least one attribute to compare") + + wanted = set(node_types) if node_types is not None else None + nodes = sorted( + ( + node + for node, data in G.nodes(data=True) + if wanted is None or data.get("type") in wanted + ), + key=repr, + ) + if len(nodes) < 2: + return ResolutionResult(G, [], {}, threshold, 0, on) + + pairs = _candidate_pairs(G, nodes, on, block_on, block_prefix, respect_type) + + accepted: dict[tuple[Hashable, Hashable], float] = {} + for left, right in sorted(pairs, key=repr): + attr = attribute_similarity( + G.nodes[left], G.nodes[right], on, weights, token_subset_floor + ) + if attr <= 0: + continue + # Skip the structural computation when attributes alone already decide + # the outcome: structure can only raise the score, so a pair that + # already clears the threshold is accepted, and one that cannot be + # lifted to the threshold even by perfect overlap is rejected. + best_possible = attr + structural_weight * (1.0 - attr) + if best_possible < threshold: + continue + if attr >= threshold: + accepted[(left, right)] = attr + continue + structural = neighbor_overlap(G, left, right, relationship_aware) + score = attr + structural_weight * structural * (1.0 - attr) + if score >= threshold: + accepted[(left, right)] = score + + link_graph = nx.Graph() + link_graph.add_nodes_from(node for pair in accepted for node in pair) + link_graph.add_edges_from(accepted.keys()) + + clusters = [ + sorted(component, key=repr) + for component in nx.connected_components(link_graph) + if len(component) > 1 + ] + clusters.sort(key=lambda cluster: (-len(cluster), repr(cluster[0]))) + + logger.info( + "resolve_entities: %d candidate pairs, %d accepted, %d clusters, " + "%d duplicate nodes", + len(pairs), + len(accepted), + len(clusters), + sum(len(c) - 1 for c in clusters), + ) + return ResolutionResult(G, clusters, accepted, threshold, len(pairs), on) + + +# --------------------------------------------------------------------------- # +# merging +# --------------------------------------------------------------------------- # + + +def pick_canonical(G: nx.Graph, cluster: Sequence[Hashable]) -> Hashable: + """Choose which node in a cluster survives the merge. + + The best-connected node wins, since it carries the most context. Ties break + on the number of populated attributes and then on the node id, which keeps + the choice deterministic. + """ + return max( + cluster, + key=lambda node: ( + G.degree(node), + len([v for v in G.nodes[node].values() if v not in (None, "")]), + repr(node), + ), + ) + + +def _merge_node_attributes( + target: dict[str, Any], source: dict[str, Any], source_id: Hashable +) -> None: + """Fold `source` attributes into `target`, recording disagreements. + + The canonical node's own values win. Anything it is missing gets filled in; + anything it disagrees with is preserved under `_merged_variants` so the + merge is never silently lossy. + """ + variants: dict[str, list[Any]] = target.setdefault(MERGED_VARIANTS, {}) + for key, value in source.items(): + if key in (MERGED_FROM, MERGED_VARIANTS, MERGE_COUNT): + continue + current = target.get(key) + if key not in target or current in (None, ""): + target[key] = value + elif current != value and value not in (None, ""): + seen = variants.setdefault(key, []) + if value not in seen: + seen.append(value) + target.setdefault(MERGED_FROM, []).append(source_id) + + +def _add_or_reinforce_edge( + H: nx.Graph, u: Hashable, v: Hashable, data: dict[str, Any] +) -> None: + """Attach a rewired edge, collapsing exact duplicates into a count. + + Multigraphs keep parallel edges as-is. For simple graphs, an edge that + already exists with the same relationship gets its `_merge_count` + incremented rather than being dropped without trace. + """ + if H.is_multigraph(): + H.add_edge(u, v, **data) + return + if H.has_edge(u, v): + existing = H.edges[u, v] + if existing.get("relationship") == data.get("relationship"): + existing[MERGE_COUNT] = existing.get(MERGE_COUNT, 1) + 1 + return + H.add_edge(u, v, **data) + + +def merge_clusters( + G: nx.Graph, + clusters: Iterable[Sequence[Hashable]], + inplace: bool = False, + canonical: Callable[[nx.Graph, Sequence[Hashable]], Hashable] | None = None, +) -> nx.Graph: + """Collapse each cluster into one node, rewiring its edges. + + This is the part that has no tabular equivalent. Every edge touching an + absorbed node is re-pointed at the canonical node; self-loops created by + the merge are dropped; and edges that become duplicates are counted rather + than discarded. + + Args: + G: Source graph. + clusters: Groups of node ids to merge. Groups of one are ignored. + inplace: Mutate `G` instead of working on a copy. + canonical: Optional `(graph, cluster) -> node` override for choosing + the surviving node. + + Returns: + The merged graph — `G` itself when `inplace` is true, otherwise a copy. + """ + H = G if inplace else G.copy() + chooser = canonical or pick_canonical + + mapping: dict[Hashable, Hashable] = {} + absorbed_by: dict[Hashable, list[Hashable]] = {} + for cluster in clusters: + members = sorted({node for node in cluster if node in H}, key=repr) + if len(members) < 2: + continue + survivor = chooser(H, members) + for node in members: + if node != survivor: + mapping[node] = survivor + absorbed_by.setdefault(survivor, []).append(node) + + if not mapping: + return H + + # Rewire first, using the complete mapping so that a cluster member + # pointing at another cluster's member resolves correctly in one pass. + for u, v, data in list(H.edges(data=True)): + new_u = mapping.get(u, u) + new_v = mapping.get(v, v) + if (new_u, new_v) == (u, v): + continue + if new_u == new_v: + continue # the merge turned this into a self-loop + _add_or_reinforce_edge(H, new_u, new_v, dict(data)) + + # Then fold attributes and drop the absorbed nodes, which removes their + # now-redundant edges along with them. + for survivor, absorbed in absorbed_by.items(): + target = H.nodes[survivor] + for node in sorted(absorbed, key=repr): + _merge_node_attributes(target, dict(H.nodes[node]), node) + H.remove_nodes_from(absorbed) + if not target.get(MERGED_VARIANTS): + target.pop(MERGED_VARIANTS, None) + + logger.info( + "merge_clusters: %d nodes merged into %d canonical nodes (%d -> %d)", + len(mapping), + len(absorbed_by), + G.number_of_nodes(), + H.number_of_nodes(), + ) + return H + + +# --------------------------------------------------------------------------- # +# scoring against known clusters +# --------------------------------------------------------------------------- # + + +def _as_cluster_map( + clustering: Any, universe: set[Hashable] | None = None +) -> dict[Hashable, Any]: + """Accept either a list of clusters or a node -> cluster-id mapping.""" + if isinstance(clustering, dict): + return dict(clustering) + result: dict[Hashable, Any] = {} + for index, cluster in enumerate(clustering): + for node in cluster: + result[node] = index + return result + + +def _pairs_within(cluster_map: dict[Hashable, Any]) -> set[tuple[Any, Any]]: + groups: dict[Any, list[Hashable]] = {} + for node, cluster_id in cluster_map.items(): + groups.setdefault(cluster_id, []).append(node) + pairs: set[tuple[Any, Any]] = set() + for members in groups.values(): + ordered = sorted(members, key=repr) + for i, left in enumerate(ordered): + for right in ordered[i + 1 :]: + pairs.add((left, right)) + return pairs + + +def evaluate_clusters(predicted: Any, gold: Any) -> dict[str, float]: + """Score a predicted clustering against a known-correct one. + + Both arguments may be a list of clusters or a `{node: cluster_id}` mapping. + Any node named on one side but not the other is treated as a singleton on + the missing side, so you can pass only the non-trivial clusters. + + Returns pairwise and B-cubed precision/recall/F1. Pairwise is the stricter + and more familiar measure; B-cubed is less punishing about one large cluster + being split and is the usual choice for entity resolution. + + This function makes no assumptions about where the gold labels came from and + does not manufacture them — supply your own. + """ + predicted_map = _as_cluster_map(predicted) + gold_map = _as_cluster_map(gold) + + universe = set(predicted_map) | set(gold_map) + for index, node in enumerate(sorted(universe, key=repr)): + predicted_map.setdefault(node, f"__pred_singleton_{index}") + gold_map.setdefault(node, f"__gold_singleton_{index}") + + predicted_pairs = _pairs_within(predicted_map) + gold_pairs = _pairs_within(gold_map) + shared = predicted_pairs & gold_pairs + + def ratio(numerator: int, denominator: int) -> float: + return numerator / denominator if denominator else 1.0 + + pairwise_precision = ratio(len(shared), len(predicted_pairs)) + pairwise_recall = ratio(len(shared), len(gold_pairs)) + + predicted_groups: dict[Any, set[Hashable]] = {} + gold_groups: dict[Any, set[Hashable]] = {} + for node, cluster_id in predicted_map.items(): + predicted_groups.setdefault(cluster_id, set()).add(node) + for node, cluster_id in gold_map.items(): + gold_groups.setdefault(cluster_id, set()).add(node) + + b3_precision_total = 0.0 + b3_recall_total = 0.0 + for node in universe: + in_predicted = predicted_groups[predicted_map[node]] + in_gold = gold_groups[gold_map[node]] + overlap = len(in_predicted & in_gold) + b3_precision_total += overlap / len(in_predicted) + b3_recall_total += overlap / len(in_gold) + count = len(universe) or 1 + + def f1(precision: float, recall: float) -> float: + return ( + 2 * precision * recall / (precision + recall) + if (precision + recall) + else 0.0 + ) + + b3_precision = b3_precision_total / count + b3_recall = b3_recall_total / count + + return { + "pairwise_precision": pairwise_precision, + "pairwise_recall": pairwise_recall, + "pairwise_f1": f1(pairwise_precision, pairwise_recall), + "b_cubed_precision": b3_precision, + "b_cubed_recall": b3_recall, + "b_cubed_f1": f1(b3_precision, b3_recall), + "predicted_pairs": len(predicted_pairs), + "gold_pairs": len(gold_pairs), + "correct_pairs": len(shared), + } diff --git a/pyproject.toml b/pyproject.toml index fb93a8e..cd52b50 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "graphfaker" -version = "0.3.0" +version = "0.5.0" description = "an open-source python library for generating, and loading both synthetic and real-world graph datasets" readme = "README.md" authors = [ @@ -17,14 +17,29 @@ classifiers = [ ] license = "MIT" -keywords = ["faker", "graph-data", "flights", "osmnx", "graphs", "graphfaker"] +requires-python = ">=3.10" +keywords = [ + "faker", + "graph-data", + "flights", + "osmnx", + "graphs", + "graphfaker", + "entity-resolution", + "knowledge-graph", + "graphrag", + "cypher", + "neo4j", +] dependencies = [ "faker>=37.1.0", "networkx>=3.4.2", "osmnx==2.0.2", "pandas>=2.2.2", "requests>=2.32.3", + "tqdm>=4.66.0", "typer", + "urllib3>=2.0.0", "wikipedia>=1.4.0", ] @@ -63,3 +78,9 @@ warn_unreachable = true warn_no_return = true + + +[tool.pytest.ini_options] +markers = [ + "network: test reaches the internet; deselect with -m 'not network'", +] diff --git a/tests/test_cli.py b/tests/test_cli.py index e528ebd..5541354 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,31 +1,101 @@ -from typer.testing import CliRunner from unittest.mock import patch + +import networkx as nx +from typer.testing import CliRunner + from graphfaker.cli import app runner = CliRunner() -def test_faker_mode_generates_graph(): +def _tiny_graph(): + """A minimal real graph. + + The fetchers must be mocked with an actual graph rather than a string: the + CLI reports node and edge counts and then exports, so a stand-in has to + support the NetworkX interface. + """ + G = nx.DiGraph() + G.add_node("a", type="Test", name="A") + G.add_node("b", type="Test", name="B") + G.add_edge("a", "b", relationship="LINKS_TO") + return G + + +def _tiny_osm_graph(): + """A graph shaped enough like an OSMnx result to survive osmnx export.""" + G = nx.MultiDiGraph() + G.graph["crs"] = "epsg:4326" + G.graph["simplified"] = True + G.add_node(1, x=-0.13, y=51.51, street_count=1) + G.add_node(2, x=-0.14, y=51.52, street_count=1) + G.add_edge(1, 2, osmid=1001, length=120.5, oneway=False) + return G + + +def test_faker_mode_generates_graph(tmp_path): result = runner.invoke( app, - ["--fetcher", "faker", "--total-nodes", "10", "--total-edges", "20"], + [ + "--fetcher", + "faker", + "--total-nodes", + "10", + "--total-edges", + "20", + "--export", + str(tmp_path / "out.graphml"), + ], ) assert result.exit_code == 0 assert "Graph" in result.output or "nodes" in result.output + assert (tmp_path / "out.graphml").exists() + + +def test_faker_mode_is_reproducible_with_seed(tmp_path): + def run(target): + return runner.invoke( + app, + [ + "--fetcher", + "faker", + "--total-nodes", + "20", + "--total-edges", + "40", + "--seed", + "42", + "--export", + str(tmp_path / target), + ], + ) + + assert run("first.graphml").exit_code == 0 + assert run("second.graphml").exit_code == 0 + assert (tmp_path / "first.graphml").read_bytes() == ( + tmp_path / "second.graphml" + ).read_bytes() @patch("graphfaker.fetchers.osm.OSMGraphFetcher.fetch_network") -def test_osm_mode_with_place(mock_fetch): - mock_fetch.return_value = "Mocked OSM Graph" +def test_osm_mode_with_place(mock_fetch, tmp_path): + mock_fetch.return_value = _tiny_osm_graph() result = runner.invoke( app, - ["--fetcher", "osm", "--place", "Soho Square, London, UK"], + [ + "--fetcher", + "osm", + "--place", + "Soho Square, London, UK", + "--export", + str(tmp_path / "osm.graphml"), + ], ) assert result.exit_code == 0 - assert "Mocked OSM Graph" in result.output mock_fetch.assert_called_once() + assert (tmp_path / "osm.graphml").exists() @patch("graphfaker.fetchers.flights.FlightGraphFetcher.fetch_airlines") @@ -33,23 +103,37 @@ def test_osm_mode_with_place(mock_fetch): @patch("graphfaker.fetchers.flights.FlightGraphFetcher.fetch_flights") @patch("graphfaker.fetchers.flights.FlightGraphFetcher.build_graph") def test_flight_mode_valid_inputs( - mock_build_graph, mock_fetch_flights, mock_fetch_airports, mock_fetch_airlines + mock_build_graph, + mock_fetch_flights, + mock_fetch_airports, + mock_fetch_airlines, + tmp_path, ): mock_fetch_airlines.return_value = ["airline1"] mock_fetch_airports.return_value = ["airport1"] mock_fetch_flights.return_value = ["flight1"] - mock_build_graph.return_value = "Mocked Flight Graph" + mock_build_graph.return_value = _tiny_graph() result = runner.invoke( - app, ["--fetcher", "flights", "--year", "2024", "--month", "1"] + app, + [ + "--fetcher", + "flights", + "--year", + "2024", + "--month", + "1", + "--export", + str(tmp_path / "flights.graphml"), + ], ) assert result.exit_code == 0 - assert "Mocked Flight Graph" in result.output mock_fetch_airlines.assert_called_once() mock_fetch_airports.assert_called_once() mock_fetch_flights.assert_called_once() mock_build_graph.assert_called_once() + assert (tmp_path / "flights.graphml").exists() def test_invalid_month(): @@ -75,7 +159,7 @@ def test_invalid_daterange(): @patch("graphfaker.fetchers.flights.FlightGraphFetcher.fetch_flights") -def test_flight_mode_with_date_range(mock_fetch_flights): +def test_flight_mode_with_date_range(mock_fetch_flights, tmp_path): mock_fetch_flights.return_value = ["flightX"] with patch( "graphfaker.fetchers.flights.FlightGraphFetcher.fetch_airlines", return_value=[] @@ -83,7 +167,7 @@ def test_flight_mode_with_date_range(mock_fetch_flights): "graphfaker.fetchers.flights.FlightGraphFetcher.fetch_airports", return_value=[] ), patch( "graphfaker.fetchers.flights.FlightGraphFetcher.build_graph", - return_value="Mocked Graph", + return_value=_tiny_graph(), ): result = runner.invoke( app, @@ -96,7 +180,10 @@ def test_flight_mode_with_date_range(mock_fetch_flights): "1", "--date-range", "2024-01-01,2024-01-10", + "--export", + str(tmp_path / "range.graphml"), ], ) assert result.exit_code == 0 - assert "Mocked Graph" in result.output + mock_fetch_flights.assert_called_once() + assert (tmp_path / "range.graphml").exists() diff --git a/tests/test_corpus.py b/tests/test_corpus.py new file mode 100644 index 0000000..f9f8ab4 --- /dev/null +++ b/tests/test_corpus.py @@ -0,0 +1,240 @@ +# tests/test_corpus.py +"""Paired text/entity corpora and the duplication measurement.""" + +import json + +import networkx as nx +import pytest + +from graphfaker.corpus import ( + Corpus, + Entity, + attribute_nodes, + duplication_report, + generate_corpus, +) + + +def _perfect_extraction(corpus: Corpus) -> nx.DiGraph: + """What a pipeline that never splits an entity would produce.""" + G = nx.DiGraph() + for entity in corpus.entities: + G.add_node(entity.id, name=entity.name, type=entity.type) + return G + + +def _splits_every_alias(corpus: Corpus) -> nx.DiGraph: + """A pipeline that emits a separate node per surface form.""" + G = nx.DiGraph() + for entity in corpus.entities: + G.add_node(entity.id, name=entity.name, type=entity.type) + for index, alias in enumerate(entity.aliases): + G.add_node(f"{entity.id}_a{index}", name=alias, type=entity.type) + return G + + +# --------------------------------------------------------------------------- # +# generation +# --------------------------------------------------------------------------- # + + +def test_corpus_is_reproducible(): + a = generate_corpus(seed=42, n_entities=16, n_documents=10) + b = generate_corpus(seed=42, n_entities=16, n_documents=10) + assert [e.name for e in a.entities] == [e.name for e in b.entities] + assert [d.text for d in a.documents] == [d.text for d in b.documents] + + +def test_different_seeds_differ(): + a = generate_corpus(seed=1, n_entities=16, n_documents=10) + b = generate_corpus(seed=2, n_entities=16, n_documents=10) + assert [d.text for d in a.documents] != [d.text for d in b.documents] + + +def test_every_entity_is_mentioned_somewhere(): + corpus = generate_corpus(seed=42, n_entities=24, n_documents=15) + mentioned = {eid for doc in corpus.documents for eid in doc.entity_ids} + assert mentioned == {entity.id for entity in corpus.entities} + + +def test_entity_names_appear_verbatim_in_the_text(): + corpus = generate_corpus(seed=42, n_entities=16, n_documents=20) + blob = corpus.text() + for entity in corpus.entities: + assert any( + form in blob for form in entity.surface_forms() + ), f"{entity.id} never appears" + + +@pytest.mark.parametrize("seed", [1, 2, 7, 42, 99]) +def test_generated_corpora_are_unambiguous(seed): + """No surface form may be claimable by two entities. + + This is the property the whole measurement rests on: if the corpus is + ambiguous, observed "duplication" is partly the corpus's fault. + """ + audit = generate_corpus(seed=seed, n_entities=24, n_documents=15).audit() + assert audit["ambiguous_forms"] == {} + assert audit["containment_pairs"] == [] + assert audit["clean"] is True + + +def test_audit_detects_a_deliberately_ambiguous_corpus(): + """The audit must actually be capable of failing.""" + corpus = Corpus( + entities=[ + Entity(id="e1", name="Acme Corporation", type="Organization", aliases=["Acme"]), + Entity(id="e2", name="Acme", type="Organization"), + ], + relations=[], + documents=[], + ) + audit = corpus.audit() + assert audit["clean"] is False + assert audit["ambiguous_forms"] + + +@pytest.mark.parametrize("seed", [1, 2, 7, 42, 99]) +def test_aliases_carry_no_dangling_punctuation(seed): + """The first word of "Barnes, Cole and Ramirez" is "Barnes," with a comma. + + Aliases appear verbatim in the generated prose, so trailing punctuation + shows up as visible garbage in the corpus. + """ + corpus = generate_corpus(seed=seed, n_entities=24, n_documents=10) + for entity in corpus.entities: + for alias in entity.aliases: + assert alias == alias.strip(" ,.;:"), (entity.id, alias) + assert ",," not in alias + + +def test_aliases_never_duplicate_the_canonical_name(): + corpus = generate_corpus(seed=42, n_entities=24, n_documents=10) + for entity in corpus.entities: + assert len(entity.aliases) == len(set(entity.aliases)), entity.id + assert entity.name not in entity.aliases + + +def test_alias_index_maps_every_unambiguous_form(): + from graphfaker.resolve import normalize + + corpus = generate_corpus(seed=42, n_entities=16, n_documents=10) + index = corpus.alias_index() + for entity in corpus.entities: + # Every surface form must resolve back to its own entity. + for form in entity.surface_forms(): + assert index[normalize(form)] == entity.id + + +def test_write_produces_documents_and_gold(tmp_path): + corpus = generate_corpus(seed=42, n_entities=16, n_documents=8) + target = corpus.write(str(tmp_path / "corpus")) + + for document in corpus.documents: + assert (tmp_path / "corpus" / f"{document.id}.txt").exists() + + gold = json.loads((tmp_path / "corpus" / "gold.json").read_text(encoding="utf-8")) + assert gold["seed"] == 42 + assert gold["expected_entity_count"] == corpus.expected_entity_count + assert len(gold["entities"]) == len(corpus.entities) + assert target.endswith("corpus") + + +# --------------------------------------------------------------------------- # +# attribution +# --------------------------------------------------------------------------- # + + +def test_attribution_matches_canonical_names(): + corpus = generate_corpus(seed=42, n_entities=16, n_documents=10) + result = attribute_nodes(_perfect_extraction(corpus), corpus) + assert result.unmatched == [] + assert len(result.matched) == len(corpus.entities) + + +def test_attribution_matches_aliases(): + corpus = generate_corpus(seed=42, n_entities=16, n_documents=10) + result = attribute_nodes(_splits_every_alias(corpus), corpus) + assert result.unmatched == [] + + +def test_attribution_handles_decorated_labels(): + """Extractors often emit 'Ada Lovelace (engineer)'.""" + corpus = generate_corpus(seed=42, n_entities=16, n_documents=10) + entity = corpus.entities[0] + G = nx.DiGraph() + G.add_node("n1", name=f"{entity.name} (mentioned in report)") + result = attribute_nodes(G, corpus) + assert result.matched.get(entity.id) == ["n1"] + + +def test_attribution_reports_rather_than_guesses_unknown_nodes(): + corpus = generate_corpus(seed=42, n_entities=16, n_documents=10) + G = nx.DiGraph() + G.add_node("junk", name="Completely Unrelated Thing 12345") + result = attribute_nodes(G, corpus) + assert result.unmatched == ["junk"] + assert result.matched == {} + + +def test_attribution_falls_back_to_the_node_id(): + corpus = generate_corpus(seed=42, n_entities=16, n_documents=10) + entity = corpus.entities[0] + G = nx.DiGraph() + G.add_node(entity.name) # id is the name, no 'name' attribute + result = attribute_nodes(G, corpus) + assert result.matched.get(entity.id) == [entity.name] + + +# --------------------------------------------------------------------------- # +# the measurement +# --------------------------------------------------------------------------- # + + +def test_perfect_pipeline_scores_zero_duplication(): + corpus = generate_corpus(seed=42, n_entities=24, n_documents=15) + report = duplication_report(_perfect_extraction(corpus), corpus, framework="perfect") + assert report.duplication_rate == 0.0 + assert report.missed == [] + assert report.unmatched_nodes == 0 + assert report.node_inflation == pytest.approx(1.0) + + +def test_alias_splitting_pipeline_is_penalised(): + corpus = generate_corpus(seed=42, n_entities=24, n_documents=15) + report = duplication_report( + _splits_every_alias(corpus), corpus, framework="splitter" + ) + assert report.duplication_rate > 0.5 + assert report.node_inflation > 1.5 + assert report.worst[0][1] > 1 + + +def test_missing_entities_are_counted(): + corpus = generate_corpus(seed=42, n_entities=24, n_documents=15) + G = _perfect_extraction(corpus) + G.remove_node(corpus.entities[0].id) + report = duplication_report(G, corpus) + assert corpus.entities[0].id in report.missed + assert report.found == corpus.expected_entity_count - 1 + + +def test_report_row_is_serialisable(): + corpus = generate_corpus(seed=42, n_entities=16, n_documents=10) + row = duplication_report(_perfect_extraction(corpus), corpus, framework="x").row() + assert json.loads(json.dumps(row))["framework"] == "x" + + +def test_report_summary_is_printable(): + corpus = generate_corpus(seed=42, n_entities=16, n_documents=10) + text = duplication_report(_splits_every_alias(corpus), corpus, "s").summary() + assert "duplication" in text.lower() or "split" in text.lower() + assert "node inflation" in text + + +def test_empty_graph_reports_everything_missed(): + corpus = generate_corpus(seed=42, n_entities=16, n_documents=10) + report = duplication_report(nx.DiGraph(), corpus) + assert report.found == 0 + assert report.duplication_rate == 0.0 # nothing found, so nothing split + assert len(report.missed) == corpus.expected_entity_count diff --git a/tests/test_export.py b/tests/test_export.py new file mode 100644 index 0000000..3e89722 --- /dev/null +++ b/tests/test_export.py @@ -0,0 +1,308 @@ +# tests/test_export.py +"""Export connectors.""" + +import csv + +import networkx as nx +import pytest + +from graphfaker.core import GraphFaker +from graphfaker.export import ( + export_csv, + export_cypher, + export_neo4j_csv, + flatten_value, +) + + +def _read(path): + with open(path, encoding="utf-8") as handle: + return handle.read() + + +def _heterogeneous_graph(): + """Node types with genuinely different attribute sets. + + This is the case that breaks a naive exporter which takes the first node's + keys as the header for every row. + """ + G = nx.DiGraph() + G.add_node("p1", type="Person", name="Ada Lovelace", age=36, email="ada@x.com") + G.add_node("pl1", type="Place", name="London", population=9000000) + G.add_node("o1", type="Organization", name="Acme Ltd", revenue=1000.5) + G.add_edge("p1", "o1", relationship="WORKS_AT", position="Engineer") + G.add_edge("p1", "pl1", relationship="LIVES_IN") + return G + + +# --------------------------------------------------------------------------- # +# value flattening +# --------------------------------------------------------------------------- # + + +def test_flatten_value_handles_containers_and_scalars(): + assert flatten_value(None) == "" + assert flatten_value(True) is True + assert flatten_value(3) == 3 + assert flatten_value("x") == "x" + assert flatten_value([1, 2]) == "1,2" + assert flatten_value((1.5, 2.5)) == "1.5,2.5" + assert flatten_value({"b": 2, "a": 1}) == "a=1; b=2" + + +# --------------------------------------------------------------------------- # +# generic CSV +# --------------------------------------------------------------------------- # + + +def test_csv_header_is_the_union_of_all_node_keys(tmp_path): + G = _heterogeneous_graph() + nodes_path, _ = export_csv( + G, str(tmp_path / "nodes.csv"), str(tmp_path / "edges.csv") + ) + with open(nodes_path, encoding="utf-8") as handle: + header = next(csv.reader(handle)) + for column in ("id", "type", "name", "age", "email", "population", "revenue"): + assert column in header, column + + +def test_csv_values_land_in_the_right_columns(tmp_path): + """The bug this guards against writes Place values under Person columns.""" + G = _heterogeneous_graph() + nodes_path, _ = export_csv( + G, str(tmp_path / "nodes.csv"), str(tmp_path / "edges.csv") + ) + with open(nodes_path, encoding="utf-8") as handle: + rows = {row["id"]: row for row in csv.DictReader(handle)} + + assert rows["pl1"]["name"] == "London" + assert rows["pl1"]["population"] == "9000000" + assert rows["pl1"]["age"] == "" # Place has no age + assert rows["p1"]["age"] == "36" + assert rows["p1"]["population"] == "" # Person has no population + assert rows["o1"]["revenue"] == "1000.5" + + +def test_csv_edges_carry_source_target_and_attributes(tmp_path): + G = _heterogeneous_graph() + _, edges_path = export_csv( + G, str(tmp_path / "nodes.csv"), str(tmp_path / "edges.csv") + ) + with open(edges_path, encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + assert len(rows) == 2 + works = next(r for r in rows if r["relationship"] == "WORKS_AT") + assert works["source"] == "p1" + assert works["target"] == "o1" + assert works["position"] == "Engineer" + + +def test_csv_handles_an_empty_graph(tmp_path): + """A naive implementation raises StopIteration here.""" + nodes_path, edge_path = export_csv( + nx.DiGraph(), str(tmp_path / "n.csv"), str(tmp_path / "e.csv") + ) + with open(nodes_path, encoding="utf-8") as handle: + assert next(csv.reader(handle)) == ["id"] + with open(edge_path, encoding="utf-8") as handle: + assert next(csv.reader(handle)) == ["source", "target"] + + +def test_csv_writes_non_ascii(tmp_path): + G = nx.DiGraph() + G.add_node("n1", type="Person", name="Zoë Ünicode 北京") + nodes_path, _ = export_csv(G, str(tmp_path / "n.csv"), str(tmp_path / "e.csv")) + with open(nodes_path, encoding="utf-8") as handle: + assert "Zoë Ünicode 北京" in handle.read() + + +def test_csv_flattens_container_attributes(tmp_path): + G = nx.DiGraph() + G.add_node("n1", type="Place", coordinates=(51.5, -0.13), tags=["a", "b"]) + nodes_path, _ = export_csv(G, str(tmp_path / "n.csv"), str(tmp_path / "e.csv")) + with open(nodes_path, encoding="utf-8") as handle: + row = next(csv.DictReader(handle)) + assert row["coordinates"] == "51.5,-0.13" + assert row["tags"] == "a,b" + + +def test_csv_handles_multigraph_edges(tmp_path): + G = nx.MultiDiGraph() + G.add_node("a", type="X") + G.add_node("b", type="X") + G.add_edge("a", "b", relationship="R1") + G.add_edge("a", "b", relationship="R2") + _, edges_path = export_csv(G, str(tmp_path / "n.csv"), str(tmp_path / "e.csv")) + with open(edges_path, encoding="utf-8") as handle: + assert len(list(csv.DictReader(handle))) == 2 + + +# --------------------------------------------------------------------------- # +# neo4j-admin CSV +# --------------------------------------------------------------------------- # + + +def test_neo4j_csv_uses_typed_headers(tmp_path): + G = _heterogeneous_graph() + result = export_neo4j_csv(G, str(tmp_path / "import")) + with open(result["nodes"], encoding="utf-8") as handle: + header = next(csv.reader(handle)) + assert header[0] == "id:ID" + assert header[-1] == ":LABEL" + + with open(result["edges"], encoding="utf-8") as handle: + header = next(csv.reader(handle)) + assert header[0] == ":START_ID" + assert ":TYPE" in header + assert header[-1] == ":END_ID" + assert "neo4j-admin database import" in result["command"] + + +def test_neo4j_csv_labels_come_from_the_type_attribute(tmp_path): + G = _heterogeneous_graph() + result = export_neo4j_csv(G, str(tmp_path / "import")) + with open(result["nodes"], encoding="utf-8") as handle: + rows = list(csv.reader(handle))[1:] + labels = {row[0]: row[-1] for row in rows} + assert labels["p1"] == "Person" + assert labels["pl1"] == "Place" + + +# --------------------------------------------------------------------------- # +# Cypher / GQL +# --------------------------------------------------------------------------- # + + +def test_cypher_emits_one_statement_per_node_and_edge(tmp_path): + G = _heterogeneous_graph() + body = _read(export_cypher(G, str(tmp_path / "g.cypher"))) + assert body.count("CREATE (n:") == 3 + assert body.count("MATCH (a {_gf_id:") == 2 + assert ":Person" in body + assert "WORKS_AT" in body + + +def test_gql_dialect_uses_insert(tmp_path): + G = _heterogeneous_graph() + with open(export_cypher(G, str(tmp_path / "g.gql"), dialect="gql"), encoding="utf-8") as handle: + body = handle.read() + assert "INSERT (n:" in body + assert "CREATE (n:" not in body + + +def test_neo4j_dialect_emits_a_constraint_but_opencypher_does_not(tmp_path): + G = _heterogeneous_graph() + neo4j = _read(export_cypher(G, str(tmp_path / "a.cypher"), dialect="neo4j")) + opencypher = _read( + export_cypher(G, str(tmp_path / "b.cypher"), dialect="opencypher") + ) + assert "CREATE CONSTRAINT" in neo4j + assert "CREATE CONSTRAINT" not in opencypher + + +def test_cypher_rejects_an_unknown_dialect(tmp_path): + with pytest.raises(ValueError, match="dialect must be one of"): + export_cypher(_heterogeneous_graph(), str(tmp_path / "g.cypher"), dialect="sql") + + +def test_cypher_escapes_quotes_and_newlines(tmp_path): + G = nx.DiGraph() + G.add_node("n1", type="Person", name='He said "hi"\nthen left', note="back\\slash") + with open(export_cypher(G, str(tmp_path / "g.cypher")), encoding="utf-8") as handle: + body = handle.read() + # The literal must not break out of its quoting, so no raw newline may + # appear inside the statement. + statement = next(line for line in body.splitlines() if line.startswith("CREATE (n:")) + assert statement.endswith(");") + assert '\\"hi\\"' in statement + assert "\\n" in statement + + +def test_cypher_node_id_property_matches_what_relationships_look_up(tmp_path): + """Regression: identifier sanitising must not rewrite `_gf_id`. + + Stripping the leading underscore made nodes carry `gf_id` while the + relationship MATCH clauses searched for `_gf_id`, so every relationship + silently failed to create while the script still ran without error. + """ + G = _heterogeneous_graph() + with open(export_cypher(G, str(tmp_path / "g.cypher")), encoding="utf-8") as handle: + body = handle.read() + + creates = [line for line in body.splitlines() if line.startswith("CREATE (n:")] + matches = [line for line in body.splitlines() if line.startswith("MATCH (a {")] + assert creates and matches + for line in creates: + assert "_gf_id:" in line, line + assert " gf_id:" not in line, line + for line in matches: + assert "_gf_id:" in line, line + + +def test_cypher_preserves_underscore_prefixed_properties(tmp_path): + """`resolve()` writes `_merged_from`; it must survive export intact.""" + G = nx.DiGraph() + G.add_node("a", type="Person", name="X", _merged_from=["b", "c"]) + with open(export_cypher(G, str(tmp_path / "g.cypher")), encoding="utf-8") as handle: + body = handle.read() + assert "_merged_from:" in body + + +def test_cypher_sanitises_labels_and_relationship_types(tmp_path): + G = nx.DiGraph() + G.add_node("a", type="Weird Type-With Punctuation!") + G.add_node("b", type="123numeric") + G.add_edge("a", "b", relationship="has spaces & symbols") + with open(export_cypher(G, str(tmp_path / "g.cypher")), encoding="utf-8") as handle: + body = handle.read() + assert "Weird_Type_With_Punctuation" in body + assert ":_123numeric" in body # a label may not begin with a digit + assert "HAS_SPACES___SYMBOLS" in body + + +def test_cypher_falls_back_to_defaults_for_untyped_elements(tmp_path): + G = nx.DiGraph() + G.add_node("a") + G.add_node("b") + G.add_edge("a", "b") + with open(export_cypher(G, str(tmp_path / "g.cypher")), encoding="utf-8") as handle: + body = handle.read() + assert ":Node" in body + assert "RELATED_TO" in body + + +# --------------------------------------------------------------------------- # +# integration +# --------------------------------------------------------------------------- # + + +def test_all_exporters_run_on_a_generated_graph(tmp_path): + gf = GraphFaker(seed=42) + G = gf.generate_graph(source="faker", total_nodes=60, total_edges=180) + + nodes_path, edges_path = export_csv( + G, str(tmp_path / "n.csv"), str(tmp_path / "e.csv") + ) + with open(nodes_path, encoding="utf-8") as handle: + assert len(list(csv.DictReader(handle))) == 60 + + export_neo4j_csv(G, str(tmp_path / "import")) + for dialect in ("neo4j", "opencypher", "gql"): + assert export_cypher(G, str(tmp_path / f"{dialect}.cypher"), dialect=dialect) + + +def test_export_survives_resolve_provenance(tmp_path): + """Merged nodes carry list and dict attributes; CSV must handle them.""" + from graphfaker.resolve import merge_clusters + + gf = GraphFaker(seed=3) + G = gf.generate_graph(source="faker", total_nodes=40, total_edges=100) + people = sorted(n for n, d in G.nodes(data=True) if d.get("type") == "Person") + merged = merge_clusters(G, [[people[0], people[1]]]) + + nodes_path, _ = export_csv( + merged, str(tmp_path / "n.csv"), str(tmp_path / "e.csv") + ) + with open(nodes_path, encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + assert any(row.get("_merged_from") for row in rows) diff --git a/tests/test_fetchers_wiki.py b/tests/test_fetchers_wiki.py index 3fcf370..58053ee 100644 --- a/tests/test_fetchers_wiki.py +++ b/tests/test_fetchers_wiki.py @@ -1,16 +1,69 @@ -# tests/test_graph_from_source_osm.py +# tests/test_fetchers_wiki.py + +from typing import ClassVar +from unittest.mock import patch import pytest + from graphfaker.fetchers.wiki import WikiFetcher wiki = WikiFetcher() -def test_wiki_fetch_page(): - page = wiki.fetch_page("Graph Theory") - - # check url - assert page['url'] == 'https://en.wikipedia.org/wiki/Graph_theory' - # check title - assert page['title'] == "Graph theory" - +class _FakePage: + """Stands in for `wikipedia.page`'s return value. + + The previous version of this test called the live Wikipedia API, so the suite + failed intermittently on any network hiccup and could not run offline or in a + sandboxed CI job. Mocking keeps the field-mapping assertions — which is what + `fetch_page` is actually responsible for — without the flakiness. + """ + + title = "Graph theory" + url = "https://en.wikipedia.org/wiki/Graph_theory" + summary = "Graph theory is the study of graphs." + content = "Graph theory is the study of graphs, which are mathematical structures." + images: ClassVar[list[str]] = ["https://upload.wikimedia.org/example.png"] + links: ClassVar[list[str]] = [ + "Vertex (graph theory)", + "Edge (graph theory)", + "Adjacency matrix", + ] + references: ClassVar[list[str]] = ["https://example.org/reference"] + sections: ClassVar[list[str]] = ["History", "Definitions"] + + def section(self, name): + return f"Body of {name}" + + +def test_wiki_fetch_page_maps_core_fields(): + with patch("wikipedia.page", return_value=_FakePage()) as mock_page: + page = wiki.fetch_page("Graph Theory") + + mock_page.assert_called_once_with("Graph Theory") + assert page["url"] == "https://en.wikipedia.org/wiki/Graph_theory" + assert page["title"] == "Graph theory" + assert "study of graphs" in page["summary"] + assert page["links"][0] == "Vertex (graph theory)" + + +def test_wiki_export_page_json(tmp_path): + import json + + with patch("wikipedia.page", return_value=_FakePage()): + page = wiki.fetch_page("Graph Theory") + + target = tmp_path / "graph_theory.json" + WikiFetcher.export_page_json(page, str(target)) + + assert target.exists() + reloaded = json.loads(target.read_text(encoding="utf-8")) + assert reloaded["title"] == "Graph theory" + + +@pytest.mark.network +def test_wiki_fetch_page_live(): + """Hits the real API. Deselect with `-m "not network"`.""" + page = wiki.fetch_page("Graph Theory") + assert page["url"] == "https://en.wikipedia.org/wiki/Graph_theory" + assert page["title"] == "Graph theory" diff --git a/tests/test_resolve.py b/tests/test_resolve.py new file mode 100644 index 0000000..856354b --- /dev/null +++ b/tests/test_resolve.py @@ -0,0 +1,514 @@ +# tests/test_resolve.py +"""Graph-native entity resolution.""" + +import networkx as nx +import pytest + +from graphfaker.core import GraphFaker +from graphfaker.resolve import ( + MERGED_FROM, + MERGED_VARIANTS, + attribute_similarity, + evaluate_clusters, + merge_clusters, + neighbor_overlap, + normalize, + pick_canonical, + resolve_entities, +) + + +def _duplicate_pair_graph(): + """Two spellings of one company, sharing most of their neighbourhood. + + Attribute similarity alone is mediocre; the shared neighbours are what make + this pair recognisable. + """ + G = nx.DiGraph() + G.add_node("org_a", type="Organization", name="Acme Corporation", industry="Tools") + G.add_node("org_b", type="Organization", name="ACME Corp.", industry="Tools") + G.add_node("org_far", type="Organization", name="Zenith Industries") + # Deliberately dissimilar names: "Person 0" and "Person 1" would themselves + # score as near-duplicates, which would make this fixture test the wrong + # thing. + employees = ["Ada Lovelace", "Grace Hopper", "Alan Turing", "Edsger Dijkstra"] + for i, full_name in enumerate(employees): + person = f"person_{i}" + G.add_node(person, type="Person", name=full_name) + G.add_edge(person, "org_a", relationship="WORKS_AT") + G.add_edge(person, "org_b", relationship="WORKS_AT") + G.add_node("person_x", type="Person", name="Katherine Johnson") + G.add_edge("person_x", "org_far", relationship="WORKS_AT") + return G + + +# --------------------------------------------------------------------------- # +# primitives +# --------------------------------------------------------------------------- # + + +def test_normalize_strips_punctuation_and_case(): + assert normalize("ACME Corp.") == "acme corp" + assert normalize(None) == "" + assert normalize(" multiple spaces ") == "multiple spaces" + + +def test_attribute_similarity_is_token_order_insensitive(): + left = {"name": "Acme Corporation"} + right = {"name": "Corporation Acme"} + assert attribute_similarity(left, right, ["name"]) > 0.9 + + +def test_attribute_similarity_skips_fields_absent_on_both_sides(): + # A shared missing email must not be scored as agreement. + both_present = attribute_similarity( + {"name": "Ann", "email": "a@x.com"}, {"name": "Ann", "email": "a@x.com"}, ["name", "email"] + ) + email_missing = attribute_similarity({"name": "Ann"}, {"name": "Ann"}, ["name", "email"]) + assert both_present == pytest.approx(1.0) + assert email_missing == pytest.approx(1.0) + + # And a field present on only one side counts as disagreement. + half = attribute_similarity( + {"name": "Ann", "email": "a@x.com"}, {"name": "Ann"}, ["name", "email"] + ) + assert half < 1.0 + + +def test_attribute_similarity_with_no_comparable_fields_is_zero(): + assert attribute_similarity({}, {}, ["name"]) == 0.0 + + +def test_token_subset_lifts_a_shortened_name(): + """"Hill" vs "Allison Hill" scores ~0.5 on character ratio alone.""" + plain = attribute_similarity({"name": "Hill"}, {"name": "Allison Hill"}, ["name"]) + lifted = attribute_similarity( + {"name": "Hill"}, {"name": "Allison Hill"}, ["name"], token_subset_floor=0.75 + ) + assert plain < 0.6 + assert lifted == pytest.approx(0.75) + + +def test_token_subset_floor_never_reaches_certainty(): + """Containment is suggestive, not conclusive, so it must stay below 1.0.""" + from graphfaker.resolve import TOKEN_SUBSET_FLOOR + + assert TOKEN_SUBSET_FLOOR < 1.0 + score = attribute_similarity( + {"name": "Smith"}, {"name": "John Smith"}, ["name"], + token_subset_floor=TOKEN_SUBSET_FLOOR, + ) + assert score < 1.0 + + +def test_initials_are_dropped_so_a_shortened_name_still_matches(): + """"A. Hill" reduces to its meaningful token, "hill". + + That token is contained in "allison hill", so the pair is lifted — which is + the intent: an initialised form is one of the commonest ways a document + refers back to a person it already named. + """ + score = attribute_similarity( + {"name": "A. Hill"}, {"name": "Allison Hill"}, ["name"], token_subset_floor=0.75 + ) + assert score == pytest.approx(0.75) + assert ( + attribute_similarity( + {"name": "A. Hill"}, {"name": "Allison Hill"}, ["name"], + token_subset_floor=0.0, + ) + < 0.75 + ) + + +def test_same_surname_different_initial_is_a_known_precision_hazard(): + """Documents a real weakness rather than asserting it is absent. + + "A. Hill" and "B. Hill" are different people, but character ratio scores + them ~0.83 because five of six characters agree. Nothing in this module + fixes that; it is why `structural_weight` exists and why a resolution result + should be reviewed rather than applied blindly. + """ + score = attribute_similarity({"name": "A. Hill"}, {"name": "B. Hill"}, ["name"]) + assert score > 0.8 + + # Structure is what separates them: with no shared neighbours they are not + # merged at the default threshold. + G = nx.Graph() + G.add_node("a", type="Person", name="A. Hill") + G.add_node("b", type="Person", name="B. Hill") + G.add_edge("a", "acme") + G.add_edge("b", "zenith") + assert resolve_entities(G, on=["name"], threshold=0.9).clusters == [] + + +def test_containment_alone_does_not_merge_without_structural_support(): + """The floor must lift a pair into consideration, not decide it.""" + G = nx.Graph() + G.add_node("a", type="Person", name="Hill") + G.add_node("b", type="Person", name="Allison Hill") + # No shared neighbours at all. + result = resolve_entities(G, on=["name"], threshold=0.85, structural_weight=0.5) + assert result.clusters == [] + + +def test_containment_plus_shared_neighbours_does_merge(): + G = nx.Graph() + G.add_node("a", type="Person", name="Hill") + G.add_node("b", type="Person", name="Allison Hill") + for shared in ("acme", "london", "project_x", "team_y"): + G.add_edge("a", shared) + G.add_edge("b", shared) + result = resolve_entities(G, on=["name"], threshold=0.85, structural_weight=0.5) + assert result.clusters == [["a", "b"]] + + +def test_token_subset_can_be_disabled(): + G = nx.Graph() + G.add_node("a", type="Person", name="Hill") + G.add_node("b", type="Person", name="Allison Hill") + for shared in ("acme", "london", "project_x", "team_y"): + G.add_edge("a", shared) + G.add_edge("b", shared) + off = resolve_entities( + G, on=["name"], threshold=0.85, structural_weight=0.5, token_subset_floor=0.0 + ) + assert off.clusters == [] + + +def test_neighbor_overlap_ignores_the_direct_edge_between_candidates(): + G = nx.DiGraph() + G.add_edge("a", "b") # only connection is to each other + assert neighbor_overlap(G, "a", "b") == 0.0 + + +def test_neighbor_overlap_counts_shared_context(): + G = nx.DiGraph() + for shared in ("x", "y", "z"): + G.add_edge("a", shared) + G.add_edge("b", shared) + assert neighbor_overlap(G, "a", "b") == pytest.approx(1.0) + + +def test_neighbor_overlap_ignores_direction(): + G = nx.DiGraph() + G.add_edge("a", "shared") + G.add_edge("shared", "b") # opposite direction, same context + assert neighbor_overlap(G, "a", "b") == pytest.approx(1.0) + + +def test_neighbor_overlap_works_on_undirected_and_multigraphs(): + for graph_type in (nx.Graph, nx.MultiGraph, nx.MultiDiGraph): + G = graph_type() + G.add_edge("a", "shared") + G.add_edge("b", "shared") + assert neighbor_overlap(G, "a", "b") == pytest.approx(1.0), graph_type + + +# --------------------------------------------------------------------------- # +# resolution +# --------------------------------------------------------------------------- # + + +def test_structural_signal_finds_a_pair_attributes_alone_would_miss(): + G = _duplicate_pair_graph() + attributes_only = resolve_entities( + G, on=["name"], threshold=0.85, structural_weight=0.0 + ) + with_structure = resolve_entities( + G, on=["name"], threshold=0.85, structural_weight=0.8 + ) + assert attributes_only.clusters == [] + assert with_structure.clusters == [["org_a", "org_b"]] + + +def test_structure_never_lowers_a_confident_attribute_match(): + G = nx.DiGraph() + G.add_node("a", type="Person", name="Jane Doe") + G.add_node("b", type="Person", name="Jane Doe") + # No edges at all, so structural overlap is 0. + result = resolve_entities(G, on=["name"], threshold=0.95, structural_weight=0.9) + assert result.clusters == [["a", "b"]] + + +def test_unrelated_nodes_are_not_clustered(): + G = _duplicate_pair_graph() + result = resolve_entities(G, on=["name"], threshold=0.85, structural_weight=0.8) + clustered = {node for cluster in result.clusters for node in cluster} + assert "org_far" not in clustered + assert "person_x" not in clustered + + +def test_respect_type_prevents_cross_type_merges(): + G = nx.DiGraph() + G.add_node("p", type="Person", name="Acme") + G.add_node("o", type="Organization", name="Acme") + assert resolve_entities(G, on=["name"], respect_type=True).clusters == [] + assert resolve_entities(G, on=["name"], respect_type=False).clusters == [["o", "p"]] + + +def test_node_types_filter_restricts_the_search(): + G = _duplicate_pair_graph() + result = resolve_entities( + G, on=["name"], threshold=0.85, structural_weight=0.8, node_types=["Person"] + ) + assert result.clusters == [] + + +def test_resolution_is_deterministic(): + G = _duplicate_pair_graph() + first = resolve_entities(G, on=["name"], structural_weight=0.8) + second = resolve_entities(G, on=["name"], structural_weight=0.8) + assert first.clusters == second.clusters + assert first.scores == second.scores + + +def test_resolve_does_not_mutate_the_input_graph(): + G = _duplicate_pair_graph() + before = (G.number_of_nodes(), G.number_of_edges()) + resolve_entities(G, on=["name"], structural_weight=0.8).apply() + assert (G.number_of_nodes(), G.number_of_edges()) == before + + +def test_invalid_arguments_are_rejected(): + G = _duplicate_pair_graph() + with pytest.raises(ValueError): + resolve_entities(G, on=["name"], threshold=1.5) + with pytest.raises(ValueError): + resolve_entities(G, on=["name"], structural_weight=-0.1) + with pytest.raises(ValueError): + resolve_entities(G, on=[]) + + +def test_unblocked_comparison_is_guarded_on_large_graphs(): + from graphfaker import resolve as resolve_module + + G = nx.Graph() + for i in range(resolve_module.MAX_UNBLOCKED_NODES + 1): + G.add_node(i, type="Person", name=f"Person {i}") + with pytest.raises(ValueError, match="Refusing to compare"): + resolve_entities(G, on=["name"], block_on=[]) + + +def test_report_is_printable(): + G = _duplicate_pair_graph() + text = resolve_entities(G, on=["name"], structural_weight=0.8).report() + assert "clusters found" in text + assert "duplicate nodes" in text + + +# --------------------------------------------------------------------------- # +# merging +# --------------------------------------------------------------------------- # + + +def test_merge_rewires_edges_onto_the_canonical_node(): + G = _duplicate_pair_graph() + merged = merge_clusters(G, [["org_a", "org_b"]]) + + assert merged.number_of_nodes() == G.number_of_nodes() - 1 + survivor = "org_a" if "org_a" in merged else "org_b" + assert ("org_a" in merged) != ("org_b" in merged) + + # All four employees must still reach the surviving organization. + for i in range(4): + assert merged.has_edge(f"person_{i}", survivor) + # And no edge may still point at the absorbed node. + absorbed = "org_b" if survivor == "org_a" else "org_a" + assert absorbed not in merged + assert all(absorbed not in (u, v) for u, v in merged.edges()) + + +def test_merge_drops_self_loops_created_by_the_merge(): + G = nx.DiGraph() + G.add_node("a", type="Person", name="Dup") + G.add_node("b", type="Person", name="Dup") + G.add_edge("a", "b", relationship="FRIENDS_WITH") + merged = merge_clusters(G, [["a", "b"]]) + assert merged.number_of_nodes() == 1 + assert merged.number_of_edges() == 0 + + +def test_merge_records_provenance_and_conflicting_values(): + G = nx.DiGraph() + G.add_node("a", type="Person", name="Jane Doe", email="jane@x.com") + G.add_node("b", type="Person", name="Jane Doe", email="j.doe@x.com") + G.add_edge("a", "peer") + G.add_edge("b", "peer") + + merged = merge_clusters(G, [["a", "b"]]) + survivor = next(n for n in ("a", "b") if n in merged) + absorbed = "b" if survivor == "a" else "a" + data = merged.nodes[survivor] + + assert data[MERGED_FROM] == [absorbed] + assert "email" in data[MERGED_VARIANTS] + + +def test_merge_fills_gaps_from_absorbed_nodes(): + G = nx.Graph() + G.add_node("a", type="Person", name="Jane", email=None) + G.add_node("b", type="Person", name="Jane", email="jane@x.com", phone="555") + merged = merge_clusters(G, [["a", "b"]]) + survivor = next(n for n in ("a", "b") if n in merged) + assert merged.nodes[survivor]["email"] == "jane@x.com" + assert merged.nodes[survivor]["phone"] == "555" + + +def test_merge_counts_collapsed_duplicate_edges(): + G = nx.DiGraph() + G.add_node("a", type="Organization", name="Dup") + G.add_node("b", type="Organization", name="Dup") + G.add_node("p", type="Person", name="Person") + G.add_edge("p", "a", relationship="WORKS_AT") + G.add_edge("p", "b", relationship="WORKS_AT") + merged = merge_clusters(G, [["a", "b"]]) + survivor = next(n for n in ("a", "b") if n in merged) + assert merged.edges["p", survivor]["_merge_count"] == 2 + + +def test_merge_handles_a_cluster_of_three(): + G = nx.Graph() + for name in ("a", "b", "c"): + G.add_node(name, type="Person", name="Same Person") + G.add_edge(name, "anchor") + merged = merge_clusters(G, [["a", "b", "c"]]) + assert merged.number_of_nodes() == 2 # one survivor plus the anchor + + +def test_merge_ignores_singleton_and_missing_clusters(): + G = _duplicate_pair_graph() + merged = merge_clusters(G, [["org_a"], ["ghost_1", "ghost_2"]]) + assert merged.number_of_nodes() == G.number_of_nodes() + + +def test_merge_inplace_mutates_the_original(): + G = _duplicate_pair_graph() + original_count = G.number_of_nodes() + returned = merge_clusters(G, [["org_a", "org_b"]], inplace=True) + assert returned is G + assert G.number_of_nodes() == original_count - 1 + + +def test_pick_canonical_prefers_the_better_connected_node(): + G = nx.Graph() + G.add_node("sparse", name="X") + G.add_node("rich", name="X") + G.add_edge("rich", "n1") + G.add_edge("rich", "n2") + assert pick_canonical(G, ["sparse", "rich"]) == "rich" + + +def test_canonical_override_is_respected(): + G = _duplicate_pair_graph() + merged = merge_clusters( + G, [["org_a", "org_b"]], canonical=lambda graph, cluster: "org_b" + ) + assert "org_b" in merged + assert "org_a" not in merged + + +# --------------------------------------------------------------------------- # +# evaluation +# --------------------------------------------------------------------------- # + + +def test_evaluate_perfect_prediction(): + gold = [["a", "b"], ["c", "d"]] + scores = evaluate_clusters(gold, gold) + assert scores["pairwise_f1"] == pytest.approx(1.0) + assert scores["b_cubed_f1"] == pytest.approx(1.0) + + +def test_evaluate_penalises_a_missed_cluster(): + gold = [["a", "b"]] + predicted = [] + scores = evaluate_clusters(predicted, gold) + assert scores["pairwise_recall"] == pytest.approx(0.0) + assert scores["gold_pairs"] == 1 + assert scores["correct_pairs"] == 0 + + +def test_evaluate_penalises_an_over_merge(): + gold = [["a", "b"]] + predicted = [["a", "b", "c"]] + scores = evaluate_clusters(predicted, gold) + assert scores["pairwise_recall"] == pytest.approx(1.0) + assert scores["pairwise_precision"] < 1.0 + + +def test_evaluate_accepts_mapping_form(): + as_clusters = evaluate_clusters([["a", "b"]], [["a", "b"]]) + as_mapping = evaluate_clusters({"a": 0, "b": 0}, {"a": "x", "b": "x"}) + assert as_clusters["pairwise_f1"] == as_mapping["pairwise_f1"] + + +def test_evaluate_treats_unlisted_nodes_as_singletons(): + # 'c' appears only in the prediction, and must not be assumed to match. + scores = evaluate_clusters([["a", "b"], ["c", "d"]], [["a", "b"]]) + assert scores["pairwise_precision"] == pytest.approx(0.5) + + +def test_evaluate_empty_inputs_do_not_divide_by_zero(): + scores = evaluate_clusters([], []) + assert scores["pairwise_f1"] == pytest.approx(1.0) + + +# --------------------------------------------------------------------------- # +# integration with the generator +# --------------------------------------------------------------------------- # + + +def test_resolve_runs_on_a_generated_graph(): + gf = GraphFaker(seed=42) + gf.generate_graph(source="faker", total_nodes=100, total_edges=300) + result = gf.resolve(on=["name"], threshold=0.9) + # Distinct faker names should not collapse; the point is that it completes + # and reports cleanly on a realistic graph. + assert result.candidates_considered >= 0 + assert isinstance(result.report(), str) + assert result.apply().number_of_nodes() <= 100 + + +def test_resolve_round_trip_recovers_injected_duplicates(): + """End-to-end: duplicate a node, then check resolve finds it back. + + This is the only claim being made — that a node copied with a perturbed + name and a shared neighbourhood is recoverable. It is not a claim that real + extraction errors look like this one. + """ + gf = GraphFaker(seed=7) + G = gf.generate_graph(source="faker", total_nodes=60, total_edges=200) + + people = [n for n, d in G.nodes(data=True) if d.get("type") == "Person"] + original = min(people) + twin = f"{original}__twin" + data = dict(G.nodes[original]) + data["name"] = data["name"].upper() + "." + G.add_node(twin, **data) + for neighbor in list(G.successors(original)): + G.add_edge(twin, neighbor, **G.edges[original, neighbor]) + + result = resolve_entities( + G, on=["name", "email"], threshold=0.85, structural_weight=0.6 + ) + found = [set(cluster) for cluster in result.clusters] + assert {original, twin} in found + + scores = evaluate_clusters(result.clusters, [[original, twin]]) + assert scores["pairwise_recall"] == pytest.approx(1.0) + + +def test_export_survives_merge_provenance(tmp_path): + """Merged nodes carry list/dict attributes; GraphML export must not break.""" + gf = GraphFaker(seed=3) + G = gf.generate_graph(source="faker", total_nodes=40, total_edges=100) + people = sorted(n for n, d in G.nodes(data=True) if d.get("type") == "Person") + a, b = people[0], people[1] + merged = merge_clusters(G, [[a, b]]) + + out = tmp_path / "merged.graphml" + gf.export_graph(merged, path=str(out)) + assert out.exists() + reloaded = nx.read_graphml(str(out)) + assert reloaded.number_of_nodes() == merged.number_of_nodes() diff --git a/tests/test_seed.py b/tests/test_seed.py new file mode 100644 index 0000000..9aa9b0a --- /dev/null +++ b/tests/test_seed.py @@ -0,0 +1,83 @@ +# tests/test_seed.py +"""Reproducibility of synthetic generation.""" + +import networkx as nx + +from graphfaker.core import GraphFaker + + +def _fingerprint(G: nx.Graph): + """Everything that should be identical between two seeded runs.""" + return ( + sorted(G.nodes()), + sorted( + (n, tuple(sorted((k, str(v)) for k, v in d.items()))) + for n, d in G.nodes(data=True) + ), + sorted((u, v, d.get("relationship")) for u, v, d in G.edges(data=True)), + ) + + +def test_same_seed_produces_identical_graph(): + a = GraphFaker(seed=42).generate_graph( + source="faker", total_nodes=40, total_edges=120 + ) + b = GraphFaker(seed=42).generate_graph( + source="faker", total_nodes=40, total_edges=120 + ) + assert _fingerprint(a) == _fingerprint(b) + + +def test_different_seeds_produce_different_graphs(): + a = GraphFaker(seed=1).generate_graph(source="faker", total_nodes=40, total_edges=120) + b = GraphFaker(seed=2).generate_graph(source="faker", total_nodes=40, total_edges=120) + assert _fingerprint(a) != _fingerprint(b) + + +def test_seed_via_generate_graph_argument(): + gf = GraphFaker() + a = gf.generate_graph(source="faker", total_nodes=30, total_edges=90, seed=7) + b = gf.generate_graph(source="faker", total_nodes=30, total_edges=90, seed=7) + assert _fingerprint(a) == _fingerprint(b) + + +def test_reseed_resets_state(): + gf = GraphFaker(seed=99) + first = gf.generate_graph(source="faker", total_nodes=30, total_edges=90) + fingerprint = _fingerprint(first) + gf.reseed(99) + second = gf.generate_graph(source="faker", total_nodes=30, total_edges=90) + assert _fingerprint(second) == fingerprint + + +def test_seeding_does_not_disturb_global_random(): + import random + + random.seed(123) + expected = [random.random() for _ in range(3)] + + random.seed(123) + GraphFaker(seed=42).generate_graph(source="faker", total_nodes=20, total_edges=40) + actual = [random.random() for _ in range(3)] + + assert actual == expected + + +def test_unseeded_instances_still_differ(): + a = GraphFaker().generate_graph(source="faker", total_nodes=40, total_edges=120) + b = GraphFaker().generate_graph(source="faker", total_nodes=40, total_edges=120) + assert _fingerprint(a) != _fingerprint(b) + + +def test_unknown_source_names_the_valid_options(): + gf = GraphFaker() + try: + gf.generate_graph(source="nonsense") + except ValueError as exc: + message = str(exc) + assert "faker" in message + assert "osm" in message + assert "flights" in message + assert "random" not in message + else: + raise AssertionError("expected ValueError for an unknown source") diff --git a/tests/test_topology.py b/tests/test_topology.py new file mode 100644 index 0000000..c22799d --- /dev/null +++ b/tests/test_topology.py @@ -0,0 +1,304 @@ +# tests/test_topology.py +"""Structural realism of the synthetic generator. + +These are differential tests. Rather than asserting an absolute threshold for +"realistic" — which would be arbitrary — each one compares the default generator +against ``topology="uniform"``, the pre-0.5 behaviour of drawing both endpoints +uniformly at random. Uniform attachment is a known quantity (an Erdos-Renyi +graph), so it makes a principled baseline for every property that was supposed to +have been fixed. +""" + +import networkx as nx +import pytest + +from graphfaker.core import GraphFaker +from graphfaker.metrics import ( + as_simple_undirected, + degree_gini, + graph_stats, + numeric_assortativity, +) + +NODES = 400 +EDGES = 1600 + + +@pytest.fixture(scope="module") +def realistic(): + return GraphFaker(seed=11).generate_graph(total_nodes=NODES, total_edges=EDGES) + + +@pytest.fixture(scope="module") +def uniform(): + return GraphFaker(seed=11).generate_graph( + total_nodes=NODES, total_edges=EDGES, topology="uniform" + ) + + +def _friendship_subgraph(G): + """Person-Person edges only. + + The whole graph is multipartite — people attach to hub cities and employers — + so it is legitimately disassortative. The social layer is where positive + assortativity is expected, so mixing them hides both effects. + """ + H = nx.Graph() + H.add_nodes_from( + (node, data) for node, data in G.nodes(data=True) if data.get("type") == "Person" + ) + H.add_edges_from( + (u, v) + for u, v, data in G.edges(data=True) + if u != v + and data.get("relationship") in {"FRIENDS_WITH", "COLLEAGUES", "MENTORS"} + and u in H + and v in H + ) + return H + + +# --------------------------------------------------------------------------- # +# heavy-tailed degree distribution +# --------------------------------------------------------------------------- # + + +def test_degree_distribution_is_heavy_tailed(realistic, uniform): + """Fano factor separates a Poisson degree distribution from a heavy tail. + + Uniform attachment gives variance approximately equal to mean, so its Fano + factor sits near 1. A hub-bearing graph is far above it. + """ + fano_realistic = graph_stats(realistic)["degree_fano"] + fano_uniform = graph_stats(uniform)["degree_fano"] + assert fano_uniform < 3.0, "uniform attachment should be roughly Poisson" + assert fano_realistic > 3 * fano_uniform + + +def test_hubs_exist(realistic, uniform): + realistic_max = graph_stats(realistic)["max_degree"] + uniform_max = graph_stats(uniform)["max_degree"] + assert realistic_max > 2 * uniform_max + + +def test_degree_inequality_is_higher(realistic, uniform): + assert graph_stats(realistic)["degree_gini"] > graph_stats(uniform)["degree_gini"] + 0.1 + + +def test_degree_gini_bounds(): + assert degree_gini([]) == 0.0 + assert degree_gini([0, 0, 0]) == 0.0 + assert degree_gini([5, 5, 5, 5]) == pytest.approx(0.0, abs=0.01) + # One node holding every edge is maximally unequal. + assert degree_gini([0, 0, 0, 9]) > 0.6 + + +# --------------------------------------------------------------------------- # +# clustering +# --------------------------------------------------------------------------- # + + +def test_clustering_far_exceeds_random_baseline(realistic): + """A sparse random graph clusters at about mean_degree / n.""" + stats = graph_stats(realistic) + assert stats["average_clustering"] > 5 * stats["clustering_baseline"] + + +def test_clustering_beats_uniform_attachment(realistic, uniform): + assert ( + graph_stats(realistic)["average_clustering"] + > 3 * graph_stats(uniform)["average_clustering"] + ) + + +def test_friendship_layer_is_strongly_clustered(realistic, uniform): + """Triadic closure should be most visible among people.""" + assert ( + graph_stats(_friendship_subgraph(realistic))["average_clustering"] + > 5 * graph_stats(_friendship_subgraph(uniform))["average_clustering"] + ) + + +# --------------------------------------------------------------------------- # +# communities and homophily +# --------------------------------------------------------------------------- # + + +def test_community_structure_is_recoverable(realistic, uniform): + """Modularity against the latent `community` label. + + Uniform attachment ignores the label entirely, so its modularity sits near + zero — there is no group structure to find. + """ + assert graph_stats(uniform)["community_modularity"] < 0.05 + assert graph_stats(realistic)["community_modularity"] > 0.35 + + +def test_attributes_correlate_with_topology(realistic, uniform): + """Age homophily: neighbours should be closer in age than chance allows.""" + assert numeric_assortativity(realistic, "age") > 0.3 + assert abs(numeric_assortativity(uniform, "age")) < 0.15 + + +def test_every_node_carries_a_community(realistic): + for _, data in realistic.nodes(data=True): + assert isinstance(data.get("community"), int) + + +def test_community_count_is_configurable(): + G = GraphFaker(seed=3).generate_graph( + total_nodes=200, total_edges=600, communities=4 + ) + assert len({data["community"] for _, data in G.nodes(data=True)}) <= 4 + + +# --------------------------------------------------------------------------- # +# connectivity +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("nodes,edges", [(40, 120), (100, 400), (300, 1200)]) +def test_no_isolated_nodes(nodes, edges): + """Preferential attachment alone left 24% of a 600-node graph isolated. + + Real graphs have a giant component, so coverage and top-up passes exist to + guarantee every node is reachable. + """ + G = GraphFaker(seed=5).generate_graph(total_nodes=nodes, total_edges=edges) + assert [node for node, degree in G.degree() if degree == 0] == [] + + +def test_graph_is_connected_as_one_component(): + G = GraphFaker(seed=5).generate_graph(total_nodes=300, total_edges=1200) + components = list(nx.connected_components(as_simple_undirected(G))) + largest = max(len(component) for component in components) + assert largest / G.number_of_nodes() > 0.9 + + +@pytest.mark.parametrize("nodes,edges", [(20, 50), (100, 1000), (300, 1200)]) +def test_edge_budget_is_respected(nodes, edges): + """`total_edges` should mean what `number_of_edges()` reports. + + Counting loop iterations rather than edges made this overshoot by ~14% when + bidirectional friendships dominated, and undershoot by ~12% when skipped + functional relationships did. + """ + G = GraphFaker(seed=5).generate_graph(total_nodes=nodes, total_edges=edges) + assert edges * 0.9 <= G.number_of_edges() <= edges * 1.1 + + +# --------------------------------------------------------------------------- # +# semantic realism +# --------------------------------------------------------------------------- # + + +def test_functional_relationships_are_singular(): + """Nobody lives in three cities or was born in two.""" + G = GraphFaker(seed=9).generate_graph(total_nodes=400, total_edges=1600) + for relationship in ("LIVES_IN", "BORN_IN"): + counts = {} + for source, _, data in G.edges(data=True): + if data.get("relationship") == relationship: + counts[source] = counts.get(source, 0) + 1 + assert not counts or max(counts.values()) == 1, relationship + + +def test_uniform_topology_still_produces_the_old_defect(): + """Guards the baseline: the comparison is meaningless if it also got fixed.""" + G = GraphFaker(seed=9).generate_graph( + total_nodes=300, total_edges=1200, topology="uniform" + ) + counts = {} + for source, _, data in G.edges(data=True): + if data.get("relationship") == "LIVES_IN": + counts[source] = counts.get(source, 0) + 1 + assert counts and max(counts.values()) > 1 + + +def test_employee_count_matches_the_actual_workforce(): + """An attribute that contradicts the topology is a trap for users.""" + G = GraphFaker(seed=7).generate_graph(total_nodes=400, total_edges=1600) + pairs = [] + for node, data in G.nodes(data=True): + if data.get("type") != "Organization": + continue + employed = sum( + 1 + for _, _, edge in G.in_edges(node, data=True) + if edge.get("relationship") == "WORKS_AT" + ) + pairs.append((data["employee_count"], employed)) + assert data["employee_count"] >= employed + + assert len(pairs) > 5 + import statistics + + assert statistics.correlation(*zip(*pairs)) > 0.5 + + +def test_place_population_tracks_prominence(): + G = GraphFaker(seed=7).generate_graph(total_nodes=400, total_edges=1600) + pairs = [ + (data["population"], data["prominence"]) + for _, data in G.nodes(data=True) + if data.get("type") == "Place" + ] + import statistics + + assert statistics.correlation(*zip(*pairs)) > 0.8 + + +def test_industry_is_an_industry_not_a_job_title(): + """`industry=fake.job()` produced values like "Chartered accountant".""" + from graphfaker.core import INDUSTRY_BY_SUBTYPE + + allowed = {value for values in INDUSTRY_BY_SUBTYPE.values() for value in values} + G = GraphFaker(seed=7).generate_graph(total_nodes=200, total_edges=600) + industries = { + data["industry"] + for _, data in G.nodes(data=True) + if data.get("type") == "Organization" + } + assert industries + assert industries <= allowed + + +# --------------------------------------------------------------------------- # +# contract +# --------------------------------------------------------------------------- # + + +def test_realistic_topology_is_reproducible(): + a = GraphFaker(seed=42).generate_graph(total_nodes=200, total_edges=600) + b = GraphFaker(seed=42).generate_graph(total_nodes=200, total_edges=600) + assert sorted(a.edges()) == sorted(b.edges()) + assert [d.get("community") for _, d in a.nodes(data=True)] == [ + d.get("community") for _, d in b.nodes(data=True) + ] + + +def test_node_count_is_exact(): + G = GraphFaker(seed=1).generate_graph(total_nodes=137, total_edges=400) + assert G.number_of_nodes() == 137 + + +def test_unknown_topology_is_rejected(): + with pytest.raises(ValueError, match="topology must be"): + GraphFaker(seed=1).generate_graph(total_nodes=50, topology="scale-free") + + +def test_stats_on_an_empty_graph(): + assert graph_stats(nx.DiGraph()) == {"nodes": 0, "edges": 0} + + +def test_compare_topology_renders_a_table(): + from graphfaker.metrics import compare_topology + + a = GraphFaker(seed=1).generate_graph(total_nodes=80, total_edges=240) + b = GraphFaker(seed=1).generate_graph( + total_nodes=80, total_edges=240, topology="uniform" + ) + table = compare_topology({"realistic": a, "uniform": b}) + assert "degree_fano" in table + assert "realistic" in table and "uniform" in table + assert len(table.splitlines()) > 5